Joint Subspace Sparse Attention per una decodifica efficiente con contesti lunghi

Abstract degli autori

L’attenzione sparsa riduce il traffico della cache KV che domina la decodifica degli LLM con contesti lunghi, ma i metodi utilizzabili nella pratica operano per lo più a granularità grossolana. La selezione a livello di token offre un controllo più fine, ma pone due problemi: il calcolo dei punteggi in uno spazio a bassa dimensionalità può perdere la struttura query–key necessaria per classificare accuratamente i token, mentre il recupero dei token selezionati, sparsi nella cache, può annullare il vantaggio in termini di latenza. Presentiamo Joint Subspace Sparse Attention (JSSA) per affrontare entrambi i problemi. Per migliorare la qualità della selezione, JSSA apprende uno spazio condiviso a bassa dimensionalità per query e key minimizzando l’errore nella conservazione dei punteggi dei token. Stabiliamo un limite superiore per la funzione obiettivo non convessa che ne risulta tramite una funzione obiettivo marginale trattabile, che ammette una soluzione spettrale in forma chiusa; il raffinamento diretto della funzione obiettivo esatta non produce miglioramenti misurabili nelle configurazioni esaminate. Per rendere più efficiente l’esecuzione, sviluppiamo un kernel di attenzione indicizzata a livello di token che non richiede operazioni di gather e supporta selezioni distinte per ciascuna testa KV nella disposizione standard della cache KV dell’attenzione grouped-query (GQA). Nei modelli con GQA standard, con attenzione ibrida a finestra scorrevole e globale e con Multi-head Latent Attention (MLA), JSSA mantiene sistematicamente un’accuratezza vicina a quella dell’attenzione full-KV, evitando al tempo stesso il notevole aumento della lunghezza delle generazioni e i problemi dipendenti dall’architettura osservati con i metodi concorrenti. Su una GPU NVIDIA B200, JSSA raggiunge un throughput di decodifica fino a $6.43\times$ quello dell’attenzione full-KV con un contesto di 128K. JSSA raggiunge un’accuratezza nel ragionamento paragonabile a quella del DSA nativo del modello senza addestrare il modello e arriva al 96--99\% del suo throughput di decodifica anche con modelli di grandi dimensioni su scala 700B.