In che modo Flash Attention riduce il traffico di memoria durante l’inferenza locale?

Flash Attention riduce il traffico di memoria suddividendo l’attenzione esatta in blocchi, così che i punteggi intermedi possano essere elaborati direttamente sul chip invece di essere scritti e letti ripetutamente dalla HBM della GPU.

Su un server IA locale, l’effetto diventa più importante all’aumentare della lunghezza del prompt, perché l’attenzione standard crea una quantità di lavoro intermedio sempre maggiore. Flash Attention non modifica le conoscenze del modello né elimina la sua cache KV; riorganizza il calcolo dell’attenzione per spostare meno dati attraverso i livelli più lenti della gerarchia di memoria della GPU.

L’attenzione standard sposta grandi tensori intermedi attraverso la HBM

L’attenzione inizia con tensori di query, key e value e deve calcolare i punteggi, normalizzarli e combinarli con i valori. Un’implementazione diretta materializza grandi matrici intermedie nella memoria ad alta larghezza di banda tra questi passaggi.

L’articolo originale su FlashAttention sostiene che le prestazioni dell’attenzione sono limitate non solo dai calcoli, ma anche dalle letture e scritture tra la HBM e la SRAM sul chip. Le sequenze lunghe amplificano questo traffico perché la matrice dei punteggi cresce in base alle coppie di token.

Una GPU domestica può quindi trascorrere molto tempo a spostare dati invece di eseguire calcoli matriciali utili, soprattutto quando aumenta la lunghezza del contesto.

Flash Attention suddivide il problema in insiemi di lavoro sul chip

Flash Attention elabora blocchi di query, key e value dimensionati per adattarsi alla memoria veloce sul chip. Ogni blocco esegue diversi passaggi dell’attenzione prima che i dati intermedi lascino il chip.

FlashAttention-3 descrive l’approccio fondamentale della famiglia come una strategia di suddivisione in blocchi che evita il traffico intermedio nella memoria globale. L’algoritmo continua a produrre un’attenzione esatta per la formulazione supportata.

La dimensione dei blocchi dipende dalle risorse hardware e dall’implementazione del kernel. Il meccanismo importante è la località dei dati: riutilizzare i valori caricati mentre si trovano vicino alle unità di calcolo.

La softmax online evita di memorizzare l’intera matrice dei punteggi

La softmax ordinaria sembra richiedere tutti i punteggi di una riga prima della normalizzazione, costringendo l’implementazione a materializzare una matrice di grandi dimensioni. Flash Attention mantiene invece statistiche di normalizzazione progressive tra i vari blocchi.

Il tutorial di Triton sull’attenzione fusa mostra un’attenzione fusa con stato softmax organizzato per blocchi. In questo modo, i blocchi parziali dei punteggi possono essere utilizzati e scartati senza riscrivere ogni valore intermedio nella HBM.

È qui che la riduzione del traffico di memoria e quella della memoria temporanea si incontrano: l’algoritmo riorganizza il calcolo in modo che un grande oggetto intermedio non debba mai esistere nella memoria globale.

Featured

ZimaCube 2 NAS personale cloud domestico

ZimaCube2

$799.00 - $2,499.00

Buy Now

La fusione dei kernel mantiene una parte maggiore della catena dell’attenzione in un unico percorso di esecuzione

Il calcolo dei punteggi, il ridimensionamento, l’applicazione delle maschere, la normalizzazione e l’aggregazione dei valori possono essere fusi in un kernel coordinato, invece di essere eseguiti da kernel separati che scambiano dati intermedi completi attraverso la memoria globale.

L’interfaccia di PyTorch per l’attenzione con prodotto scalare può selezionare implementazioni di attenzione fuse quando sono supportati il dispositivo, il tipo di dati, la forma e le funzionalità richieste.

Il risultato può essere un’elaborazione più rapida dei prompt e una minore quantità massima di memoria temporanea, ma il guadagno effettivo dipende dal fatto che il runtime utilizzi realmente un kernel compatibile con Flash Attention.

La riduzione dell’I/O non equivale all’eliminazione della cache KV

Durante l’inferenza autoregressiva, le key e i value generati in precedenza devono comunque restare disponibili per i token successivi, a meno che il modello o il runtime non utilizzi uno schema separato di eliminazione, compressione, finestratura o trasferimento su memoria esterna.

La documentazione ROCm di AMD descrive Flash Attention come un miglioramento della località della memoria tra l’archiviazione sul chip e la HBM. Questa ottimizzazione riguarda il modo in cui l’attenzione legge ed elabora lo stato attivo, non l’esistenza dello stato KV storico.

Questa distinzione evita un errore comune nella pianificazione: Flash Attention può ridurre la memoria di lavoro e l’I/O, mentre un server con contesto lungo può comunque esaurire la capacità perché la cache KV conservata continua a crescere.

Il supporto hardware e la forma dei dati determinano l’uso del percorso rapido

I kernel Flash sono specializzati. La dimensione della testa, il tipo di dati, la maschera di attenzione, l’architettura della GPU, la versione del runtime e altre funzionalità possono imporre il ricorso a un’implementazione diversa.

L’articolo di ZimaSpace sulla memoria dell’attenzione rispetto ai parametri del modello tratta l’aspetto della capacità. Flash Attention deve essere misurato separatamente come ottimizzazione del percorso di esecuzione.

Esegui benchmark con prompt brevi e lunghi, verifica il backend selezionato e registra la latenza di prefill, la memoria allocata massima, la larghezza di banda della memoria e la velocità di generazione dei token. La guida di ZimaSpace all’hardware per l’IA locale fornisce il quadro generale della capacità che l’ottimizzazione dei kernel non può sostituire.

Transformer Engine di NVIDIA tratta i backend di attenzione ottimizzati e sull’esecuzione dipendente dall’hardware come parte di uno stack più ampio per la selezione dei kernel.