Riepilogo Tecnico: Richieste della stessa piuma devono radunarsi insieme

Enunciato del Problema

L'inferenza dei Modelli Linguistici di Grande Dimensione (LLM) è fondamentalmente limitata dalla memoria durante la fase di decodifica auto-regressiva. Ogni token generato richiede di prestare attenzione alla cache Key-Value (KV) di tutti i token precedenti, rendendo necessario un passaggio completo della cache KV nella memoria GPU per ogni step. Mentre il lavoro precedente si è concentrato sulla massimizzazione della dimensione del batch per migliorare l'utilizzo della GPU, questo documento identifica un compromesso critico e trascurato: l'omogeneità dei prefissi.

Nei carichi di lavoro reali (ad esempio, prompting few-shot, prompt di sistema, RAG), le richieste spesso condividono prefissi significativi. I motori di inferenza esistenti (come vLLM e SGLang) utilizzano alberi radix o strutture hash per rilevare questi prefissi condivisi e riutilizzare le cache KV, principalmente per ridurre l'impronta di memoria. Tuttavia, gli attuali scheduler danno priorità alla massimizzazione della dimensione del batch o alla minimizzazione delle espulsioni di memoria senza considerare l'impatto dell'eterogeneità dei prefissi sui pattern di accesso alla memoria.

Gli autori dimostrano che i batch eterogenei (contenenti richieste con prefissi diversi) interrompono la località spaziale e temporale negli accessi alla DRAM della GPU, portando a una scarsa utilizzazione della larghezza di banda e a una ridotta capacità di elaborazione. Al contrario, i batch omogenei per prefisso (dove tutte le richieste condividono un prefisso comune) consentono un efficiente prefetching hardware e il riutilizzo della cache. Il documento sostiene che gli scheduler esistenti non riescono a riconoscere che un batch moderatamente piccolo e omogeneo può superare le prestazioni di uno più grande ed eterogeneo. Inoltre, i meccanismi utilizzati per rilevare la condivisione dei prefissi (ad esempio, le traversate dell'albero radix) comportano un sovraccarico della CPU proibitivo, spesso paragonabile o superiore al tempo di esecuzione della GPU.

Metodologia

Il documento introduce Feather, uno scheduler consapevole dei prefissi progettato per ottimizzare il compromesso tra dimensione del batch e omogeneità dei prefissi. Feather integra due componenti principali:

1. Albero Hash a Blocchi (CHT)

Per affrontare l'elevato sovraccarico della CPU delle traversate tradizionali degli alberi radix, Feather impiega una struttura dati leggera chiamata Albero Hash a Blocchi (Chunked Hash Tree).

• Meccanismo: Invece di confronti token per token, le richieste sono rappresentate come vettori di hash cumulativi su blocchi di token di dimensione fissa.
• Efficienza: Il CHT mantiene un "insieme di lavoro" di hash per il batch attivo e una min-heap di richieste in attesa ordinate in base al loro "conteggio mancante" (il numero di blocchi non ancora coperti dal batch attivo).
• Complessità: Ciò consente allo scheduler di identificare la migliore richiesta candidata da aggiungere a un batch in tempo $O(\log W)$ (dove $W$ è il numero di richieste in attesa), un miglioramento significativo rispetto alla complessità $O(W \times T)$ o $O(E)$ delle traversate in profondità degli alberi radix.
• Operazioni: Il CHT supporta un'inserzione efficiente, la ricerca della migliore richiesta, l'aggiunta di una richiesta al batch attivo e la rimozione delle richieste completate, utilizzando aggiornamenti lazy della heap e manutenzione incrementale delle punte per minimizzare il sovraccarico.

2. Politica di Batching basata su Apprendimento per Rinforzo (RL)

Feather utilizza una politica basata su RL per determinare il punto di arresto ottimale per la formazione del batch.

• Il Compromesso: Aggiungere più richieste aumenta l'utilizzo della GPU (dimensione del batch più grande) ma può diluire l'omogeneità dei prefissi (riducendo la località della memoria).
• Stato e Azione: Lo scheduler osserva la dimensione attuale del batch, la perdita nella profondità del prefisso condiviso ($\Delta$) e il numero di richieste in attesa che condividono la nuova profondità del prefisso ($w$). Decide se Aggiungere un'altra richiesta o Fermarsi e inviare il batch.
• Politiche: Gli autori valutano tre politiche: un'euristica statica, un Bandit Contestuale (UCB) e Q-Learning. L'obiettivo è massimizzare la ricompensa, definita come capacità di elaborazione nella decodifica.
• Adattabilità: La politica RL impara a formare batch "giusti": abbastanza grandi da utilizzare la GPU ma abbastanza omogenei da massimizzare la località della memoria. Si adatta dinamicamente all'intensità del carico di lavoro e ai pattern di condivisione dei prefissi.

Contributi Chiave

1. Identificazione di un Compromesso Critico: Il documento stabilisce che l'omogeneità dei prefissi è un determinante più forte della capacità di elaborazione nella decodifica rispetto alla sola dimensione del batch. Dimostra che anche una singola richiesta divergente in un batch può interrompere la località e ridurre la capacità di elaborazione di quasi 2 volte.
2. Quantificazione del Sovraccarico della CPU: Gli autori rivelano che i meccanismi di rilevamento dei prefissi all'avanguardia (come le traversate degli alberi radix in SGLang) comportano sovraccarichi della CPU che possono raggiungere il 50–90% della latenza totale, annullando di fatto i benefici del riutilizzo dei prefissi.
3. Scheduler Feather: La progettazione e l'implementazione di Feather, che integra la struttura dati CHT con una politica di batching basata su RL. Feather è integrato sia in vLLM che in SGLang.
4. Indipendenza dall'Hardware: A differenza dei kernel di attenzione consapevoli dei prefissi che richiedono un'ottimizzazione specifica dell'hardware, Feather opera a livello di schedulazione, rendendolo compatibile con i kernel di attenzione esistenti e le configurazioni hardware.

Risultati

La valutazione è stata condotta su GPU NVIDIA RTX 6000 Ada e A100 utilizzando modelli come Llama 3 8B, Qwen e LongChat 13B su carichi di lavoro diversificati (L-Eval, LongBench).

• Guadagni di Capacità di Elaborazione: Feather raggiunge una capacità di elaborazione end-to-end da 2 a 10 volte superiore rispetto agli scheduler esistenti (FCFS di vLLM, LPM/DFS-W di SGLang) in carichi di lavoro con prefissi condivisi che vanno da 1K a 10K token.
• Sensibilità ai Prefissi: Negli esperimenti con 800 richieste divise in 8 gruppi di prefissi, l'elaborazione in 8 batch omogenei da 100 richieste ha prodotto una capacità di elaborazione quasi doppia rispetto all'elaborazione in batch più grandi ed eterogenei.
• Riduzione del Sovraccarico: Il CHT di Feather riduce il sovraccarico di rilevamento dei prefissi fino a 1000 volte rispetto agli approcci basati su DFS per sequenze lunghe. Il sovraccarico di schedulazione rimane inferiore all'1% del tempo di esecuzione della GPU, mentre gli approcci di base possono superare il tempo di esecuzione della GPU.
• Robustezza: Feather performa alla pari con lo status quo (FCFS) quando i carichi di lavoro mancano di condivisione dei prefissi, garantendo nessun degrado delle prestazioni in scenari non condivisi.
• Confronto con i Kernel: Feather supera o eguaglia i kernel di attenzione specializzati consapevoli dei prefissi (come PAT) nonostante non modifichi il kernel stesso, dimostrando che le ottimizzazioni a livello di schedulazione possono essere efficaci quanto quelle a livello di kernel.

Significato

Il documento afferma che massimizzare la dimensione del batch non è più sufficiente per un servizio efficiente degli LLM. Il collo di bottiglia principale nelle fasi di decodifica moderne è la larghezza di banda della memoria, che è fortemente influenzata dalla località degli accessi. Spostando l'attenzione dal "massimizzare la dimensione del batch" all'"ottimizzare l'omogeneità dei prefissi", Feather sblocca significativi guadagni di prestazioni senza richiedere modifiche specifiche dell'hardware ai kernel.

Il lavoro evidenzia che il costo del trovare prefissi condivisi (sovraccarico della CPU) è stato un collo di bottiglia nascosto nei sistemi attuali. Il leggero CHT di Feather e la politica RL adattiva dimostrano che è possibile bilanciare dinamicamente l'utilizzo computazionale e la località della memoria, portando a miglioramenti sostanziali sia nella capacità di elaborazione che nella latenza per i carichi di lavoro reali di inferenza LLM. Gli autori intendono rendere open-source Feather contribuendolo ai repository vLLM e SGLang.