The Pulse
I kernel cuTile di NVIDIA sono fino a 1,82 volte più veloci su Mamba-2
NVIDIA e Jane Street hanno sviluppato kernel per il passaggio in avanti di Mamba-2 con NVIDIA CUDA Tile, registrando velocità da 1,14 a 1,82 volte superiori rispetto a un’implementazione in Triton. Il confronto è stato condotto su una sola
AI.info Team · Scritto con l’IA · verificato ·
Secondo NVIDIA, un passaggio in avanti di Mamba-2 scritto con il suo modello di programmazione CUDA Tile è da 1,14 a 1,82 volte più veloce di un’implementazione open source in Triton. L’azienda ha pubblicato i risultati il 19 ottobre, descrivendo il lavoro svolto dagli ingegneri di machine learning di Jane Street in collaborazione con NVIDIA. Il confronto misura i kernel del passaggio in avanti, non cicli di addestramento completi.
Tre kernel gestiscono il passaggio in avanti di Mamba-2
L’implementazione usa cuTile, il modello di programmazione basato su Python di NVIDIA per scrivere kernel GPU mediante operazioni orientate ai tile. Invece di calcolare lo stato nascosto di ciascun token in un lungo ciclo seriale, Mamba-2 può elaborare segmenti di sequenza con moltiplicazioni di matrici adatte ai Tensor Cores delle GPU. Questo approccio rende l’implementazione dei kernel, e il modo in cui sposta i dati nella GPU, determinanti per il benchmark.
Gli autori suddividono il passaggio in avanti tra tre kernel. chunk_cumsum calcola i valori cumulativi di decadimento all’interno di ciascun segmento; fwd_pass_state fa avanzare lo stato nascosto tra i segmenti; e compute_out combina quello stato con gli input del segmento per produrre gli output. La suddivisione consente all’implementazione di riutilizzare i calcoli del decadimento, affidando al tempo stesso le principali operazioni matriciali ai Tensor Cores.
Il risultato di 1,82 volte proviene da un test su B300
NVIDIA ha confrontato la sua implementazione cuTile con un kernel open source per Mamba-2 scritto in Triton. I test sono stati eseguiti su una NVIDIA B300 SXM6, indicata nel post come GPU Blackwell Ultra, con 268 GiB di HBM. La configurazione software comprendeva CUDA 13.3, PyTorch 2.13.0, cuTile 1.5.0 e Triton 3.7.1.
Le velocità superiori registrate vanno da 1,14 a 1,82 volte, con vantaggi maggiori all’aumentare della lunghezza della sequenza e delle dimensioni dello stato. Il benchmark presentato da NVIDIA come esempio usa una sequenza di 16.384 token e uno stato di dimensione 128, con batch di dimensione quattro e otto head. Questi valori descrivono la configurazione dell’esempio nel post; l’intervallo di velocità riportato nel titolo riassume i risultati ottenuti nelle configurazioni testate.
«I risultati mostrano che i kernel cuTile ottimizzati per il passaggio in avanti sono da 1,14 a 1,82 volte più veloci, con vantaggi che aumentano al crescere della lunghezza delle sequenze e delle dimensioni dello stato.»
Sylvain Gugger, Axel Feldmann, Yinuo Liu e Siddharth Samsi, NVIDIA Technical Blog
La disposizione dei dati in memoria e l’ottimizzazione incidono sui vantaggi
Il post descrive diverse ottimizzazioni oltre alla struttura dei kernel. Gli autori usano indicazioni sulla latenza per orientare la pianificazione delle operazioni di memoria da parte del compilatore, ridispongono alcuni operandi matriciali per ridurre la contesa nella memoria condivisa e traspongono la disposizione dei dati per consentire accessi più contigui. Un sistema di ottimizzazione automatica cerca le impostazioni relative alla dimensione dei blocchi e all’occupazione per ciascuna forma dell’input, anziché affidarsi a un’unica configurazione per ogni carico di lavoro.
Questa ottimizzazione conta perché il margine registrato varia con la lunghezza della sequenza e le dimensioni dello stato. Il confronto di NVIDIA mostra quindi le prestazioni di una specifica implementazione su una determinata GPU e configurazione software, non un vantaggio di velocità universale di cuTile rispetto a Triton su ogni carico di lavoro Mamba-2.
Il benchmark si ferma al passaggio in avanti
Il post esplicita il limite principale: si tratta di misurazioni a livello di kernel. L’addestramento end-to-end dipende anche dal passaggio all’indietro, che il benchmark non valuta. Un kernel più veloce per il passaggio in avanti può ridurre i tempi di una parte del carico di lavoro, ma il valore di 1,82 volte riportato non dimostra un miglioramento equivalente nel tempo totale di addestramento.
NVIDIA afferma che l’implementazione è disponibile tramite TileGym. Per ora, il risultato concreto è un confronto del passaggio in avanti su una B300: nelle configurazioni testate, i kernel cuTile risultano da 1,14 a 1,82 volte più veloci del riferimento in Triton.