Grouped Query Attention (GQA)

La Grouped Query Attention (GQA) è un design di attenzione in cui più teste di query condividono un insieme più piccolo di teste di chiave e valore. Preserva le diverse prospettive delle query a teste multiple riducendo al contempo la memoria e lo spostamento di dati richiesti per chiavi e valori. La GQA è particolarmente preziosa durante l'inferenza autoregressiva, in cui un modello genera un token alla volta e legge ripetutamente gli stati di attenzione memorizzati.

Come funziona la GQA#

All'interno di un attention mechanism, ogni token viene proiettato in tre rappresentazioni:

• Query: Descrive le informazioni che il token corrente sta cercando.
• Chiave: Descrive ciò che ciascun token disponibile rappresenta.
• Valore: Contiene le informazioni recuperate quando una query corrisponde a una chiave.

Nella self-attention convenzionale, queste proiezioni sono divise in teste in modo che il modello possa apprendere diverse relazioni in parallelo. La GQA mantiene molte teste di query ma assegna gruppi di esse a teste di chiave e valore condivise. Ad esempio, un livello di attenzione può utilizzare otto teste di query e due teste di chiave-valore. Ciascuna testa di chiave-valore serve quindi quattro teste di query.

Il calcolo dell'attenzione in sé rimane l'attention scaled dot-product. Ciò che cambia è la disposizione delle teste e la quantità di dati chiave-valore prodotti, memorizzati e letti. Framework come PyTorch scaled dot-product attention richiedono pertanto che il numero di teste di query sia divisibile per il numero di teste di chiave-valore quando la GQA è abilitata. (docs.pytorch.org)

GQA vs. Multi-Head e Multi-Query Attention#

La GQA si posiziona tra la multi-head attention e la multi-query attention:

• Multi-head attention: Ogni testa di query ha le sue teste di chiave e valore dedicate. Ciò offre la massima indipendenza delle teste, ma crea il più alto requisito di memoria per chiave e valore.
• Grouped query attention: Diverse teste di query condividono ciascuna testa di chiave-valore. Bilancia la capacità di rappresentazione con l'efficienza della memoria.
• Multi-query attention: Tutte le teste di query condividono una testa di chiave e una testa di valore. Questo riduce al minimo l'archiviazione di chiavi e valori, ma fornisce una minore diversità di chiavi e valori.

Se un livello ha 32 teste di query, la multi-head attention può utilizzare anche 32 teste di chiave-valore, la GQA potrebbe usarne 8 e la multi-query attention ne usa 1. La GQA non sostituisce la più ampia Transformer architecture; è una configurazione possibile all'interno di un livello di attenzione Transformer. La documentazione di NVIDIA sulla multi-head, multi-query, and grouped-query attention documentation descrive queste varianti come differenti principalmente nel numero di teste di chiave-valore che servono le teste di query. (nvidia.github.io)

Perché la GQA è importante#

Durante la generazione autoregressiva, un modello memorizza chiavi e valori precedenti in una KV cache. Con la standard multi-head attention, ogni testa di attenzione contribuisce con chiavi e valori memorizzati separatamente. Sequenze lunghe, batch ampi e molti livelli del modello possono quindi consumare una notevole memoria GPU.

Poiché la GQA utilizza meno teste di chiave-valore, la sua cache è più piccola in proporzione alla riduzione di tali teste. Questo può migliorare:

• Capacità di memoria: Prompt più lunghi o richieste simultanee maggiori possono rientrare nella memoria disponibile.
• Throughput di generazione: È necessario leggere una minore quantità di dati chiave-valore per ogni token generato.
• Inference latency: Un traffico di memoria ridotto può abbreviare il tempo di elaborazione da token a token.
• Scalabilità del Context window: Gestire sequenze lunghe diventa più pratico.

Questi guadagni dipendono dall'hardware, dalla lunghezza della sequenza, dalla dimensione del batch e dal supporto del kernel. I runtime di produzione richiedono comunque un'allocazione efficiente della cache; il sistema TensorRT-LLM KV cache system di NVIDIA, ad esempio, combina il supporto GQA con caching basato su blocchi, riutilizzo e offloading. (nvidia.github.io)

Applicazioni nel mondo reale#

Assistenti a contesto lungo: Un assistente di codifica può elaborare migliaia di token di codice sorgente prima di generare una risposta. La GQA riduce lo stato chiave-valore memorizzato nella cache associato a tale cronologia, consentendo a un server di gestire file più lunghi o utenti più concorrenti senza aumentare proporzionalmente la memoria GPU.

Assistenti multimodali per immagini e video: Un vision-language model può rappresentare patch di immagini o fotogrammi video come lunghe sequenze di token. La GQA riduce la pressione sulla cache di attenzione dopo che questi token visivi entrano nel decoder del linguaggio, lasciando potenzialmente più memoria per ulteriori immagini, fotogrammi o richieste. Questo differisce da un Vision Transformer standard, in cui l'attenzione può elaborare tutte le patch di immagini in parallelo senza caching KV autoregressivo.

Implementazione pratica e compromessi#

PyTorch espone la GQA tramite enable_gqa=True. Questo esempio documentato in CUDA utilizza otto teste di query e due teste di chiave-valore condivise:

import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

assert torch.cuda.is_available(), "A CUDA device is required."

query = torch.randn(1, 8, 16, 32, device="cuda")
key = torch.randn(1, 2, 16, 32, device="cuda")
value = torch.randn(1, 2, 16, 32, device="cuda")

# Four query heads share each key-value head.
with sdpa_kernel(SDPBackend.MATH):
    output = F.scaled_dot_product_attention(
        query,
        key,
        value,
        enable_gqa=True,
    )

print(output.shape)

L'output mantiene otto teste di query, mentre chiavi e valori utilizzano solo due teste. Il selettore del PyTorch attention backend selector può controllare quale kernel supportato esegue l'operazione, mentre la guida PyTorch Transformer building-block guide fornisce un contesto di implementazione più ampio. (docs.pytorch.org)

La GQA deve essere scelta durante la progettazione o l'adattamento di un modello; non è generalmente un interruttore attivabile in fase di inferenza per un checkpoint esistente. Gli sviluppatori dovrebbero verificare il supporto del framework, la divisibilità delle teste, il comportamento numerico, l'uso della memoria e la qualità dell'attività. Per il deployment della computer vision, flussi di lavoro complementari come la Ultralytics TensorRT integration e la Ultralytics benchmark mode aiutano a misurare se le ottimizzazioni dell'architettura e del runtime offrono miglioramenti significativi sull'hardware di destinazione.