Move the Query, Not the Cache: Characterizing Cross-Instance Latent Attention Redistribution Across GPU Fabrics

Sintesi Tecnica: Sposta la Query, Non la Cache

Definizione del Problema
Con la scalabilità dei Large Language Models (LLM), i carichi di lavoro di serving dipendono sempre più da meccanismi di Multi-head Latent Attention (MLA) e attention sparsa. In queste architetture, la cache Key-Value (KV) viene compressa in vettori latenti stretti o selezionata tramite indizzatori sparsi, creando una "asimmetria di byte": una singola riga di query (es. $\approx 1$ KB) è ordini di grandezza più piccola del chunk di cache KV a cui presta attenzione (es. megabyte).

Quando il corpus di interesse (un codebase canonico, documenti legali o una base di conoscenza aziendale condivisa) eccede l'HBM di una singola GPU, deve essere partizionato tra più istanze. Ciò richiede l'attenzione cross-istanza, dove una query di decoding su una GPU deve accedere a blocchi KV residenti su un'altra. Il riflesso prevalente nell'industria è quello di spostare la cache (fetch): prelevare i blocchi KV remoti, portarli al richiedente, inserirli nella cache locale ed eseguire l'attenzione localmente. Tuttavia, questo approccio comporta un costo significativo di "adattamento della posizione" (ri-rotazione delle codifiche posizionali memorizzate per il nuovo offset) e un enorme trasferimento di dati. Il paper postula che, sotto MLA e selezione sparsa, l'aritmetica possa invertirsi: potrebbe essere più economico instradare la query verso il detentore, calcolare l'attenzione parziale lì e fondere il risultato ridotto di ritorno.

Metodologia
Gli autori caratterizzano questo trade-off su un cluster reale di H100 multi-nodo connesso tramite NDR-200 InfiniBand con IBGDA (device-initiated RDMA) abilitato. La loro metodologia prevede:

1. Decomposizione dei Costi: Modellano la ridistribuzione cross-istanza come una funzione di costo additiva: $T_{redist} = T_{probe} + T_{transfer} + T_{compute} + T_{return} + T_{merge}$.
2. Confronto tra Primitivi: Confrontano tre primitivi:
   • Fetch: Prelevare l'intero chunk KV, adattare le posizioni (splice) ed eseguire l'attenzione localmente.
   • Route: Inviare la query al detentore, calcolare l'attenzione parziale e fondere il risultato tramite softmax online.
   • Local: Ricalcolare il chunk da zero (re-prefill).
3. Caratterizzazione dell'Hardware: Misurano la latenza e la larghezza di banda di IBGDA per payload di forma attention (batch di query su scala kilobyte) e le confrontano con i percorsi proxy-CPU. Analizzano inoltre l'impatto della topologia (same-leaf vs. spine-traversing), del tipo di fabric (NVLink, PCIe, InfiniBand) e della concorrenza.
4. Verifica della Correttezza: Verificano che l'attenzione parziale instradata, quando fusa tramite softmax online, riproduca i risultati dell'attenzione a singola istanza entro i limiti di rumore di bf16, anche per scenari di selezione sparsa.

Contributi Chiave

• Modello di Costo Topology-Aware: Gli autori derivano un modello di costo a forma chiusa per l'attenzione cross-istanza che si adatta alle misurazioni IBGDA su H100 con un errore assoluto medio percentuale (MAPE) di $\approx 7\%$ per il dispatch batch ($M_q \ge 512$). Il modello identifica che il costo è dominato da una latenza di probe fissa ($\approx 16 \mu s$) e una larghezza di banda effettiva indipendente dal payload ($\approx 25$ GB/s per il dispatch di un singolo blocco), piuttosto che dalla larghezza di banda di picco del link.
• Caratterizzazione del Regime IBGDA (Device-Initiated RDMA): A differenza del lavoro precedente sul dispatch MoE (dove i percorsi proxy-CPU a volte vincono per messaggi minuscoli), questo paper stabilisce che per i payload di attenzione ($\approx 2$ KB), l'RDMA avviato dal dispositivo è superiore, riducendo la latenza per fetch del 40% rispetto ai percorsi mediati dall'host.
• Predicato di Selezione a Forma Chiusa: Forniscono una regola decisionale per gli scheduler per scegliere tra route, fetch o local basandosi sulla forma della richiesta ($M_q$), la dimensione del chunk ($c_t$) e le caratteristiche del fabric.
  • Route vince in fase di decode (piccolo $M_q$) perché evita lo splice di adattamento della posizione di $\approx 3$ ms richiesto dal fetch e evita il costo computazionale del local.
  • Fetch diventa vitale solo quando si ammortizza il costo dello splice su molti passaggi locali successivi.
• Analisi di Scatter e Selezione: Per l'attention sparsa (es. DeepSeek-V3.2, GLM-5.1), il paper dimostra che il fetch diventa un gather sparso attraverso più detentori, incorrendo in una penalità che cresce con il numero di detentori. Il costo di Route rimane piatto perché la query viene inviata una sola volta e il detentore esegue l'attenzione sui blocchi sparsi in loco senza ri-rotazione.
• Elbow di Capacità del Detentore: Identificano colli di bottiglia specifici per il lato detentore:
  • Copy Elbow: Un pool di $K=8$ stream CUDA è ottimale per il caricamento dei dati per richieste concorrenti.
  • Compute Elbow: Un detentore può servire circa 8 richiedenti instradati concorrenti con un aumento trascurabile della latenza prima della saturazione computazionale.

Risultati

• Dominanza della Latenza: Alle tipiche dimensioni di batch di decode ($M_q = 256$), un round trip instradato costa 31–48 $\mu s$ attraverso vari fabric (NVLink, PCIe, IBGDA). Questo è 60–100$\times$ più veloce del costo di splice di $\approx 3$ ms del fetch di un chunk.
• Efficienza dei Byte sul Filo: Il routing muove $\ge 76\%$ meno byte sul filo rispetto al fetch durante il decode batch.
• Indipendenza dal Fabric: Alle dimensioni di decode, il dispatch instradato è limitato dal dispatch (limitato dalla capacità di un singolo thread block di emettere trasferimenti, $\approx 18\text{--}25$ GB/s) piuttosto che dalla larghezza di banda. Di conseguenza, un link NVLink 4.0 da 900 GB/s si comporta in modo simile a un link InfiniBand cross-nodo da 25 GB/s per query a singolo blocco; la larghezza di banda di picco conta solo per batch di grandi dimensioni.
• Resilienza alla Congestione: La primitiva di routing è robusta alla congestione di rete. La latenza rimane piatta finché il link è completamente saturo ($K=3$ flussi concorrenti, e anche in quel caso rimane un ordine di grandezza più veloce del fetch.
• Avvertenza sull'Overhead dell'Host: Sebbene il routing vinca sui byte sul filo e sulla latenza di trasporto, gli autori notano che nel loro prototipo attuale, l'overhead dell'host (client Python, logica di dispatch) domina il tempo di "time-to-first-token" end-to-end. Per batch di piccole dimensioni, un semplice fetch "bytes-back" può attualmente vincere l'end-to-end grazie al minor overhead dell'host, sebbene siano state identificate tre specifiche riduzioni implementative (risposta collassata, ammortamento del calcolo, batching) per colmare questo divario.

Significatività e Rivendicazioni
Il paper sostiene che il riflesso "move-the-cache" sia subottimale per il regime emergente di attention compressa e sparsa. Invertendo l'aritmetica, instradare la query diventa la primitiva dominante per l'attenzione cross-istanza, particolarmente per i carichi di lavoro agentici dove molti sub-agenti interrogano un corpus condiviso e immutabile.

Gli autori sottolineano che il loro contributo non è un nuovo algoritmo di routing (che esiste già in lavori come DistAttention), ma una caratterizzazione e un modello di costo che valida il caso economico per il routing sotto specifiche condizioni hardware e architettoniche (MLA, IBGDA, selezione sparsa). Argomentano che questa asimmetria di byte è la proprietà chiave che motiva l'uso di RDMA avviato dal dispositivo per l'attention. Il paper conclude che, sebbene il routing sia teoricamente superiore, realizzare i suoi pieni benefici di latenza end-to-end in uno stack di serving di produzione richiede di affrontare gli overhead lato host, un compito lasciato agli sforzi ingegneristici futuri. Il modello di costo e il predicato forniti servono come artefatti riutilizzabili per i progettisti di sistemi per prendere decisioni di scheduling topology-aware senza dover profilare ogni nuova configurazione.