I kernel di attenzione rapida e la fusione degli operatori riorganizzano il calcolo dei transformer per limitare gli spostamenti dei dati e i lanci dei kernel sull’acceleratore. Queste ottimizzazioni lasciano in principio invariati i pesi del modello e la sua funzione matematica, anche se l’ordine delle operazioni e la precisione possono produrre piccole differenze numeriche. Il tutorial pubblicato da Machine Learning Mastery dettaglia queste tecniche, con esempi di codice.
Un layer di transformer si basa su operazioni comuni: proiezioni lineari, attenzione, normalizzazione, funzioni di attivazione e addizioni residue. Un’implementazione diretta lancia un kernel di calcolo distinto, chiamato «kernel», per ogni operazione, e scrive ogni tensore intermedio nella memoria dell’acceleratore. Il calcolo rimane corretto, ma il tempo di esecuzione può essere dedicato più allo spostamento dei valori e al lancio dei kernel che ai calcoli utili.
Fusione degli operatori
Il blocco SwiGLU, utilizzato da molti modelli a solo decodificatore, illustra questo problema. Un’implementazione diretta lancia moltiplicazioni di matrici per le proiezioni, un kernel per la funzione SiLU, un altro per la moltiplicazione elemento per elemento, poi un’ultima moltiplicazione di matrici. Ad ogni confine, un tensore intermedio può essere scritto nella memoria ad alta larghezza di banda (HBM) e poi riletto. La fusione consiste nel combinare questi passaggi in un unico kernel. Fa guadagnare tempo quando gli spostamenti di memoria e i lanci pesano molto, e molto meno quando una grande moltiplicazione di matrici domina già l’esecuzione. La scelta del kernel rimane empirica: bisogna profilare l’intero carico di lavoro.
Altre fusioni sono comuni nei transformer, come l’aggiunta di una connessione residua seguita da una normalizzazione RMSNorm o LayerNorm, o la combinazione della dequantizzazione con una moltiplicazione di matrici. La fusione ha tuttavia i suoi limiti: operazioni con organizzazioni dei dati incompatibili o uno stato intermedio molto voluminoso si combinano male. Un kernel troppo grande può aumentare la pressione sui registri, ridurre il tasso di occupazione e diventare più lento di diversi kernel ben regolati.
Attenzione di tipo FlashAttention e attenzione paginata
Un’attenzione classica calcola una matrice di punteggi di dimensione T² per testa, T essendo la lunghezza della sequenza. Su una GPU, produrla, scriverla nella memoria globale e poi rileggerla consuma molta larghezza di banda, il che può diventare un collo di bottiglia durante i lunghi pre-riempimenti («prefill»). FlashAttention elabora query, chiavi e valori per blocchi. Per ciascuno, calcola una parte dei punteggi, aggiorna le statistiche del softmax e accumula l’output corrispondente nella memoria situata sul chip, senza memorizzare la matrice completa. PyTorch espone questa ottimizzazione tramite la funzione scaled_dot_product_attention(). Il tutorial precisa che imporre un backend in modo rigido può far fallire il codice se nessun backend supporta gli input o l’hardware.
La cache KV, che conserva le chiavi e i valori già calcolati, cresce di un token alla volta, per una lunghezza finale imprevedibile. Un’allocazione contigua per richiesta lascia quindi spazio riservato e inutilizzato. La cache paginata suddivide la memorizzazione in blocchi di dimensione fissa e si appoggia a una tabella di blocchi che collega le posizioni logiche ai blocchi fisici. Il kernel di attenzione deve quindi tradurre ogni posizione attraverso questa tabella. La paginazione non riduce il numero di elementi logici, ma limita la frammentazione e la capacità riservata.
Nell’esempio fornito, vLLM utilizza il modello Qwen/Qwen2.5-1.5B-Instruct con blocchi di 16 token (block_size=16) e il backend «TRITON_ATTN». Secondo il tutorial, vLLM verifica la compatibilità del backend con il modello e l’hardware durante l’inizializzazione, e segnala un errore in caso di incompatibilità.
Fusione e CUDA Graphs con torch.compile
La funzione torch.compile() cattura una sequenza di operazioni con TorchDynamo, poi la compila con TorchInductor, il suo backend predefinito. TorchInductor può generare un unico kernel Triton per operazioni compatibili che la modalità standard lancerebbe separatamente. Le interruzioni del grafo, dovute ad esempio a un controllo di flusso dipendente dai dati o a conversioni di tensori in scalari Python, limitano le zone che il compilatore può ottimizzare.
La modalità «reduce-overhead» chiede a TorchInductor di utilizzare CUDA Graphs quando il grafo catturato è idoneo. Un CUDA Graph registra la sequenza delle operazioni GPU e la riproduce con una sola sottomissione alla CPU. Non crea fusione aggiuntiva, completa quella di TorchInductor. Richiede indirizzi di memoria stabili e forme compatibili, e può consumare più memoria. Il tutorial raccomanda di verificare i risultati con tolleranze adeguate al tipo di dati, di profilare l’esecuzione e di misurare le prestazioni end-to-end.
