LLM Serving Optimization with Variable Prefill and Decode Lengths

Sintesi Tecnica: Ottimizzazione del Servizio LLM con Lunghezze di Prefill e Decode Variabili

Formulazione del Problema
Questo articolo affronta il problema dello scheduling offline per l'inferenza di Large Language Model (LLM) su un singolo worker computazionale con un budget di memoria per la KV cache fisso. Il sistema gestisce un backlog di $n$ richieste che arrivano simultaneamente al tempo $t=0$. Ogni richiesta $i$ è caratterizzata da una lunghezza del prompt eterogenea (dimensione di prefill, $s_i$) e una lunghezza della risposta (dimensione di decode, $o_i$).

La sfida principale deriva dalla natura dinamica del consumo di memoria durante il decoding autoregressivo. Mentre la fase di prefill consuma memoria proporzionalmente a $s_i$, la fase di decode consuma memoria che cresce linearmente con ogni token generato ($s_i + j$ per il $j$-esimo token). Di conseguenza, il numero di richieste concorrenti che un worker può gestire fluttua dinamicamente in base al progresso della generazione delle richieste attive. L'obiettivo è minimizzare la Latenza Totale End-to-End (TEL), definita come la somma dei tempi di completamento per tutte le richieste, soggetto al vincolo che l'utilizzo aggregato della memoria di qualsiasi batch in ogni step temporale non superi la capacità $M$.

Gli autori rilassano l'assunzione di dimensioni di input uniformi presente nel lavoro precedente (Jaillet et al., 2025), considerando uno scenario realistico in cui i carichi di lavoro mescolano brevi prompt conversazionali con compiti di riassunto di documenti lunghi.

Risultati Teorici e Sfide
L'articolo stabilisce che il rilassamento dell'assunzione di dimensione di input uniforme altera fondamentalmente la complessità del problema di scheduling:

1. NP-Hardness: Il problema di scheduling generalizzato con $s_i$ e $o_i$ eterogenei è dimostrato essere NP-hard tramite riduzione dal problema ristretto della 3-Partizione.
2. Fallimento delle Euristiche Standard: Le policy standard, inclusi First-Come-First-Served (FCFS), Shortest-Output-First (MC-SF) e la prioritizzazione basata sulla dimensione totale (Total-Size-Based), mostrano rapporti di approssimazione illimitati in questo contesto eterogeneo. Nello specifico, una richiesta con un input grande ma un output piccolo può consumare molta memoria brevemente, mentre una richiesta con un input piccolo ma un output lungo può occupare memoria per un periodo prolungato, creando trade-off che le semplici euristiche non possono risolvere.

Metodologia: L'Algoritmo Sorted-F
Per affrontare queste sfide, gli autori propongono Sorted-F, un algoritmo di scheduling progettato per ottenere una garanzia di approssimazione a fattore costante. L'algoritmo opera in due fasi:

1. Costruzione del Batch (Fase 1): L'algoritmo seleziona iterativamente batch fattibili di richieste per formare una lista ordinata. La selezione è guidata da una nuova metrica di qualità, $F(X)$, definita per un insieme di richieste $X$ come:
   $$F(X) = \frac{\sum_{r_i \in X} o_i}{|X|^2}$$
   Questa metrica bilancia la dimensione del batch ($|X|$) rispetto al costo di decode aggregato a valle ($\sum o_i$). Minimizzare $F(X)$ equivale a massimizzare la "densità di throughput" (richieste completate per unità di tempo medio di decode). All'interno di ogni batch selezionato, le richieste sono ordinate per lunghezza di output non decrescente.
2. Esecuzione (Fase 2): Il sistema esegue lo scheduling su step temporali discreti, ammettendo le richieste dalla lista ordinata nel batch corrente finché i vincoli di fattibilità della memoria sono soddisfatti, utilizzando il continuous batching e il mixing di prefill/decode.

Garanzia di Approssimazione
Gli autori dimostrano che Sorted-F raggiunge un rapporto di approssimazione costante di al massimo 48 rispetto allo scheduling offline ottimale. La prova prevede una serie di trasformazioni strutturali (separazione, raggruppamento e allineamento) che convertono lo schedule di Sorted-F e lo schedule ottimale in forme comparabili, dimostrando che l'incremento di latenza ad ogni step è limitato da un fattore costante.

Implementazioni Pratiche
Riconoscendo che l'ottimizzazione combinatoria richiesta nella Fase 1 è computazionalmente costosa per $n$ elevati, l'articolo propone tre solver pratici:

• Dynamic Programming Esatto: Fornisce una soluzione esatta per istanze piccole ($n \le 100$) con complessità $O(n^2 M)$.
• Local Swap Search: Un'euristica per scale moderate ($n \le 500$) che migliora iterativamente una soluzione greedy iniziale tramite scambi a coppie ($O(n^2)$).
• Selezione Quantile Greedy: Un'euristica in tempo lineare ($O(n)$) per sistemi su larga scala ($n \ge 1000$) che utilizza il campionamento statistico per filtrare gli outlier e riempire greedy la capacità.

Inoltre, l'articolo introduce Sorted-LP, un'euristica che sfrutta la soluzione di una rilassazione di Programmazione Lineare (LP) del problema per ordinare le richieste, e Sorted-F+Amin, una variante adattiva che gestisce l'incertezza della predizione delle lunghezze di output raffinando dinamicamente le stime durante l'inferenza.

Risultati Sperimentali
Esperimenti numerici su carichi di lavoro misti (combinando il dataset LMSYS-Chat e la riassunzione di documenti lunghi di Arxiv) dimostrano:

• Riduzione della Latenza: Sorted-F (utilizzando Local Swap Search) supera costantemente FCFS, MC-SF e le baseline guidate da LP. In un mix 80/20 di breve/lungo con 2000 richieste, Sorted-F ha ridotto la latenza media di 4.87$\times$ rispetto a FCFS e di 2.09$\times$ rispetto a MC-SF.
• Gap di Ottimalità: Su istanze trattabili, il gap empirico tra Sorted-F e il limite inferiore della rilassazione LP è piccolo (circa 1.03–1.09), suggerendo che l'algoritmo opera vicino all'ottimo nella pratica nonostante il limite teorico ampio.
• Sensibilità al Carico di Lavoro: I benefici di Sorted-F sono più pronunciati in regimi eterogenei e sovraccarichi, dove la memoria è un vincolo stringente. In regimi omogenei o con carico leggero, il vantaggio diminuisce.
• Performance Online: Una variante a orizzonte recedente (Sorted-F-Online) testata sotto arrivi di tipo Poisson mostra miglioramenti significativi rispetto a MC-SF e FCFS in regimi online sovraccarichi, sebbene offra meno benefici in regimi stabili a bassa latenza.

Significatività e Rivendicazioni
L'articolo sostiene che le lunghezze variabili del prefill non sono un dettaglio di modellazione minore, ma un fattore fondamentale che rende inefficaci le euristiche di scheduling standard e rende il problema NP-hard. Il contributo primario è l'algoritmo Sorted-F, che fornisce la prima garanzia di approssimazione a fattore costante per lo scheduling di LLM sotto vincoli eterogenei di input/output. Il lavoro colma il divario tra i limiti teorici dello scheduling e il servizio pratico di LLM, offrendo un framework scalabile e consapevole della memoria che riduce significativamente la latenza in scenari di carichi di lavoro misti senza richiedere la predizione perfetta delle lunghezze di output. Gli autori sottolineano che, sebbene il limite teorico sia 48, i risultati empirici suggeriscono che l'algoritmo sia molto più vicino alla soluzione ottimale in contesti realistici.