Technische Samenvatting: Sparse Prefix Caching voor Hybride en Recurrente LLM-serving

1. Probleemstelling

Bestaande Large Language Model (LLM) serving-systemen vertrouwen sterk op dense prefix caching, waarbij de Key/Value (KV) toestanden voor elk token in een gedeelde prefix worden opgeslagen om redundante prefill-berekeningen te vermijden. Deze aanpak is standaard voor Transformer-architecturen. Echter, State-Space Models (SSM's) en hybride architecturen (die attention combineren met recurrente lagen) vertonen een structurele asymmetrie: een recurrente laag evolueert haar toestand als $h_t = F(h_{t-1}, x_t)$, wat betekent dat ze de berekening kan hervatten vanuit een enkele opgeslagen toestand $h_c$ in plaats van dat de volledige geschiedenis van tussenliggende toestanden vereist is.

Dense per-token caching van recurrente toestanden is vaak onnodig en onbetaalbaar duur. Voor modellen zoals Qwen-3.5 schaalt de recurrente toestand per kop als $d_{head}^2$, terwijl attention KV-opslag lineair schaalt. Bijgevolg verbruikt het opslaan van elke token's recurrente toestand aanzienlijk meer geheugen dan attention KV-caches, maar biedt het geen extra voordeel ten opzichte van het opslaan van een schaarse set checkpoints.

Het kernprobleem dat wordt aangepakt is: Gegeven een geheugenbudget van $M$ checkpoint-slots voor een prefix van lengte $N$, waar moeten checkpoints worden geplaatst om de verwachte herberekeningkosten te minimaliseren onder een specifieke verdeling van toekomstige overlapdieptes? Deze optimalisatie verschilt van cache-toelatings-/verwijderingsbeleid; het richt zich op welke tokens binnen een enkele behouden entiteit te checkpointen om de bruikbaarheid voor gedeeltelijke prefix-overlappen te maximaliseren.

2. Methodologie

Theoretische Formulering

De auteurs formaliseren het probleem als een eenzijdig gewogen k-median-probleem op een lijn.

• Laat $C = \{c_1, \dots, c_M\}$ de set checkpoint-posities zijn.
• Voor een verzoek dat de prefix overlapt tot diepte $t$, is de diepst herbruikbare checkpoint $\ell(t; C) = \max(\{0\} \cup \{c \in C : c \le t\})$.
• De kosten zijn het aantal tokens om opnieuw te berekenen: $r(t; C) = t - \ell(t; C)$.
• Het doel is om de verwachte herberekeningkosten $E[r(T; C)]$ te minimaliseren, gegeven een overlapdiepte-verdeling $T$ met kansen $p_t$.

Algorithmische Aanpak

1. Optimale Gebalanceerde Plaatsing (Uniform/Worst-Case): De auteurs bewijzen dat onder een uniforme overlapverdeling of in het worst-case scenario, gebalanceerde spacing (gelijkmatig verdeelde checkpoints) optimaal is.
2. Distributie-bewuste Dynamische Programmering (DP): Voor niet-uniforme, real-world overlapverdelingen leiden de auteurs een exact $O(NM)$ dynamisch programmeringsalgoritme af.
   • De recurrente relatie is: $dp[m, j] = \min_{1 \le s \le j} (dp[m-1, s-1] + w(s, j))$, waarbij $w(s, j)$ de kosten zijn van het gebruik van $s$ als diepste checkpoint voor dieptes $s \dots j$.
   • Dit wordt efficiënt opgelost met de monotone convex hull-truc, gebruikmakend van de monotonie van prefix-sommen.
3. Omgaan met Distributiedrift: Aangezien de ware overlapverdeling onbekend is en moet worden geschat op basis van historische data, biedt het artikel theoretische garanties (via Lemma 1 en Theorema 3) dat de waarde-functie Lipschitz-continu is met betrekking tot de overlapverdeling. Dit zorgt ervoor dat het gebruik van een empirisch histogram (of een exponentieel gewogen histogram voor driftende verdelingen) een bijna-optimale planning oplevert met begrenste suboptimaliteit.

Implementatiestrategie

• Checkpointing: Het systeem slaat exacte recurrente toestanden op op schaarse posities. Bij een cache-hit hervat het systeem vanaf de diepst opgeslagen checkpoint en berekent het suffix exact opnieuw.
• Exactheid: Omdat de recurrente update deterministisch is, levert het hervatten vanaf een checkpoint en het opnieuw afspelen van het suffix bit-identieke outputs op ten opzichte van een volledige prefill.
• Hybride Integratie: Voor hybride modellen slaat de methode de volledige attention KV-cache op (dense), maar past ze alleen schaarse checkpointing toe op de recurrente lagen.

3. Belangrijkste Bijdragen

1. Theoretische Optimaliteit: Bewezen dat gebalanceerde spacing optimaal is onder uniforme overlap en worst-case omstandigheden (Theorema 1).
2. Exact Algoritme: Geformuleerd sparse prefix caching als een gewogen k-median-probleem en een exacte $O(NM)$ dynamisch programmeringsoplossing verstrekt (Theorema 2).
3. Robuustheidsgaranties: Aangetoond dat de optimale planning stabiel is onder distributiefouten. De suboptimaliteit wordt begrensd door de totale variatie-afstand tussen de empirische en ware overlapverdelingen (Theorema 3), en uitgebreid naar scenario's met distributiedrift met behulp van exponentiële herweging (Theorema 4).
4. Empirische Validatie: De methode gevalideerd op real-world datasets (QuALITY, NarrativeQA, System Prompts), waarbij consistente verbeteringen werden aangetoond ten opzichte van baselines met vast budget.

4. Experimentele Resultaten

Datasets en Setup

• Datasets:
  • QuALITY: Eén document met meerdere vragen (hoge overlap nabij volledige lengte).
  • NarrativeQA: Vergelijkbare structuur met brede multimodale overlap.
  • System Prompts: Real system prompts gecombineerd met gebruikersvragen (overlap geconcentreerd bij korte lengtes).
• Baselines: Vergelijken met "Geen Cache", "KV Alleen", "Gebalanceerd" (gelijkmatige spacing), "Blok" (vast blokformaat), "Logaritmisch", en "$\sqrt{L}$-gespatieerde" strategieën.
• Hardware: Prototypen uitgevoerd op NVIDIA RTX 2080 Super met Qwen-3.5-0.8B-model.

Bevindingen

• Tokenbesparing: De DP-optimale planning domineert consistent deterministische baselines (Gebalanceerd, Blok, Logaritmisch) in termen van factoren voor reductie van recurrent werk.
• Voordeel bij Laag Budget: Het prestatieverschil is het grootst bij lage checkpoint-budgetten waar de overlapverdeling het meest niet-uniform is. Bijvoorbeeld, op System Prompts (waaroverlappen kort zijn), plaatst de DP-methode checkpoints dicht bij de divergentiepiek, waardoor dezelfde runtime wordt bereikt als bij block caching met minder checkpoints.
• Wall-Clock Tijd: Op QuALITY en System Prompts levert de afgeleide DP-planning meetbare wall-clock-snelheidswinsten op voor representatieve laaggroepen. De reductie in runtime correleert lineair met het aantal bespaarde tokens.
• Overhead: De overhead van het laden van checkpoints van CPU naar GPU is klein ten opzichte van de besparingen die worden behaald door verminderde herberekening, wat resulteert in een netto verbetering van de runtime.

5. Betekenis en Claims

Het artikel claimt dat sparse prefix caching een nieuw ontwerppunt ontsluit voor hybride en recurrente LLM-serving dat eerder niet beschikbaar was.

• Geheugen/Berekening Trade-off: Het biedt een trade-off tussen geheugengebruik en herberekening die dense KV-caching niet kan bieden, specifiek voor architecturen waarbij toestanden exact kunnen worden geëxtraheerd en hersteld.
• Exactheid: In tegenstelling tot benaderende KV-compressietechnieken, behoudt deze methode exacte outputs zonder de recurrente berekeningkernen te wijzigen.
• Toepasbaarheid: Het is het meest voordelig voor werklasten waarbij verzoeken een aanmerkelijke maar niet identieke prefix delen (bijv. RAG met verschillende vragen, variërende system prompts). Het biedt geen extra voordelen voor strikt append-only chat-werklasten (waar de laatste toestand voldoende is), maar veroorzaakt in dergelijke gevallen geen overhead.
• Toekomstig Werk: De auteurs erkennen engineering-uitdagingen, specifiek de behoefte aan efficiënte APIs voor staatsextractie/herstel in productieruntimes en de integratie van deze plaatsingsstrategie met toelatings-/verwijderingsbeleid (bijv. Marconi et al.). Ze merken ook op dat de huidige theorie voor single-prefix kan worden uitgebreid naar algemene prefixbomen.