Riepilogo Tecnico: PersistentKV

Problema

L'elaborazione di modelli linguistici autoregressivi (LLM) è sempre più limitata dal movimento della cache Key-Value (KV) piuttosto che dalla moltiplicazione di matrici dense. Sebbene i moderni sistemi di paged-attention (es. vLLM, FlashInfer) abbiano ridotto con successo la frammentazione della cache KV e ottimizzato i kernel di decode nativi-paged, non esiste un singolo kernel "ottimale" per tutti gli scenari di serving.

Nello specifico, esiste una tensione tra:

1. Decode a lungo contesto con bassa attività: dove un singolo kernel ottimizzato può sotto-utilizzare le GPU commodity a causa di un'insufficiente parallelizzazione del lavoro.
2. Lunghezze di sequenza miste: dove l'esecuzione di molti lanci a lunghezza esatta genera un elevato overhead, mentre i batch grossolani con padding introducono inefficienze.

Il documento investiga se un motore di decode basato su tabella di blocchi nativa, combinato con uno scheduling consapevole delle pagine (page-aware), possa migliorare la latenza di serving sulle GPU commodity nei regimi in cui i kernel di paged attention di qualità produttiva sono già performanti.

Metodologia

Gli autori propongono PersistentKV, un sistema composto da un motore di decode basato su tabella di blocchi nativa e una politica di scheduling consapevole delle pagine.

1. Motore di Decode Nativo su Tabella di Blocchi

PersistentKV è progettato per la Grouped-Query Attention (GQA), dove $G = H_q/H_{kv}$ teste di query condividono una testa KV.

• Layout: Esegue direttamente sui layout nativi della cache KV a pagine (tabelle di blocchi) senza riorganizzare (repacking) le pagine K,V in tensori contigui.
• Suddivisione della Sequenza (Sequence Splitting): Per esporre una parallelizzazione sufficiente per carichi di lavoro a bassa attività, la dimensione della sequenza viene divisa in $S$ intervalli disgiunti. Ogni divisione calcola uno stato di online-softmax parziale. Un kernel di merge separato combina questi stati esattamente.
• Limiti Locali per Riga (Row-Local Bounds): A differenza degli approcci a bucket che utilizzano lunghezze con padding, PersistentKV utilizza le vere lunghezze di sequenza per riga ($N^*_i$) per determinare il conteggio dei tile e gli intervalli di divisione. Le divisioni vuote scrivono stati neutri ed escono anticipatamente, evitando letture KV con padding.
• Coda di Lavoro Compatta (Compact Workqueue): Per gestire batch irregolari (richieste con lunghezze variabili) senza eccessivi lanci di kernel, lo scheduler emette una griglia 1D di task solo per gli intervalli non vuoti di riga–testa-KV–divisione. Ciò evita il "fan-out di lancio" del bucketing a lunghezza esatta, preservando al contempo il parallelismo di sequenza.

2. Politica Adattiva Calibrata

Il sistema impiega una politica calibrata manualmente per selezionare la rotta di esecuzione in base alla dimensione del batch ($B$) e alla lunghezza massima del contesto ($N_{max}$):

• $B=1$ (Lungo Contesto): Utilizza PersistentKV con bucket di lunghezza e suddivisione della sequenza per esporre un numero sufficiente di Array di Thread di Calcolo (CTA).
• $B \ge 8$ (Lungo Contesto): Utilizza PersistentKV con la coda di lavoro compatta per ammortizzare l'overhead del doppio kernel (split + merge).
• $B=4$ (Confine): Instrada verso FlashInfer. A questa dimensione di batch, l'overhead di splitting e merging in PersistentKV supera i benefici, causando una regressione nel tempo di esecuzione (wall-clock time).
• Altrimenti: Utilizza FlashInfer.

3. Protocollo di Valutazione

• Hardware: NVIDIA RTX 3060 (28 SM, 12GB VRAM).
• Baseline: FlashInfer 0.2.5 (primaria), vLLM, TensorRT-LLM.
• Tracce: Tracce sintetiche (Bucketed, Bimodal, Uniform, Zipf) progettate per isolare la lunghezza del contesto, la dimensione del batch attivo e l'irregolarità della lunghezza.
• Metriche: Correttezza (tolleranza FP16 rispetto a FlashInfer), temporizzazione degli eventi CUDA (lavoro GPU in coda) e tempo di parete sincronizzato (latenza end-to-end inclusa la pianificazione dell'host).

Contributi Chiave

1. Motore GQA su Tabella di Blocchi Nativa: Un motore di decode che espone il lavoro di parallelismo di sequenza per il decode a lungo contesto con bassa attività senza riorganizzare le cache KV.
2. Scheduling Consapevole delle Pagine: Implementazione di uno scheduling della lunghezza di sequenza locale per riga e di una coda di lavoro compatta che allinea i confini di split e prefetch con le vere lunghezze delle richieste, evitando griglie rettangolari vuote.
3. Metodologia Adattiva: Un framework di valutazione rigoroso che utilizza seed di calibrazione/held-out, controlli di correttezza e contatori di carico derivati dallo scheduler per dimostrare che l'assegnazione del lavoro è una variabile decisiva insieme alla matematica dell'attenzione.

Risultati

Sul seed held-out (20260623) con soglie fisse derivate da un seed di calibrazione (20260622):

• Singola Richiesta ($B=1$): PersistentKV con bucket di lunghezza ottiene un miglioramento di 1.399× nel throughput di parete sincronizzato rispetto a FlashInfer su tracce a lungo contesto con bucket.
• Alto Batch ($B=8$):
  • Bimodal: 1.265× di miglioramento.
  • Uniform: 1.082× di miglioramento.
  • Zipf-like: 1.063× di miglioramento.
• Caso di Confine ($B=4$): La politica adattiva seleziona correttamente FlashInfer, evitando una regressione delle prestazioni dove PersistentKV altrimenti sottoperformerebbe.
• Overhead di Scheduling: La coda di lavoro compatta riduce il fan-out di lancio da 16.00 lanci/step (bucket esatti) a 2.00 lanci/step su tracce bimodal $B=8$, riducendo al contempo il traffico dello stato di merge di circa il 37%.

Significato e Rivendicazioni

Il documento esprime una rivendicazione modesta e focalizzata sui sistemi: PersistentKV identifica una nicchia concreta per lo scheduling adattivo di decode consapevole delle pagine.

• Non è un Kernel Universale: Gli autori dichiarano esplicitamente che PersistentKV non è universalmente più veloce; FlashInfer rimane superiore per piccoli batch attivi dove l'overhead di split-merge domina.
• L'Assegnazione del Lavoro è Importante: I risultati dimostrano che "l'assegnazione del lavoro, e non solo la matematica dell'attenzione, è una variabile decisiva per il sistema di serving".
• Contributo di livello Workshop: Il documento si posiziona come un contributo di livello workshop che fornisce una metodologia riproducibile e analisi di ablazione. Evidenzia che, nei regimi a lungo contesto e bassa attività su GPU commodity, suddividere la sequenza e utilizzare una coda di lavoro compatta può superare significativamente un singolo kernel altamente ottimizzato.
• Direzioni Future: Gli autori osservano che un sistema di produzione richiederebbe un modello di costo appreso per le soglie, l'integrazione in un runtime di serving completo (inclusi MLP e sampling) e test su hardware di generazione Hopper con diverse caratteristiche di performance.