Technische Samenvatting: RKSC – Reasoning-Aware KV Cache Sharing en Confident Early Exit

1. Probleemstelling

Moderne Large Language Models (LLMs) die zijn ontworpen voor complexe redenering (bijv. DeepSeek-R1, o1-stijl modellen) maken vaak gebruik van multi-branch inferentiestrategieën, zoals Tree of Thoughts of self-consistency, waarbij meerdere redeneertrajecten worden gegenereerd voor één enkele probleemstelling. Het artikel identificeert twee structurele redundanties in huidige productie-inferentiesystemen (zoals vLLM en SGLang) die de efficiëntie in deze pipelines degraderen:

1. Cross-Branch KV Waste: Parallelle takken delen doorgaans een aanzienlijk deel van hun prefix (de probleemstelling en de initiële context). Bestaande systemen hergebruiken Key-Value (KV) caches echter alleen wanneer tokensequenties byte-identiek zijn. Als takken semantisch divergeren maar dezelfde prefix-tokens delen, of als ze een prefix delen maar verschillende formuleringen gebruiken voor hints, herberekent het systeem de dure $O(n^2)$ aandacht voor de gedeelde prefix onafhankelijk voor elke tak.
2. Over-Deep Verification: In process-reward verificatie-pipelines voert het model een volledige forward pass uit om elke gegenereerde tak te scoren, zelfs wanneer de generatie-betrouwbaarheid al beslissend is of wanneer de interne staat van het model (logit distributie) op een intermediaire laag is gestabiliseerd. Huidige systemen missen mechanismen om de verificatiepass volledig over te slaan of deze vroegtijdig te beëindigen op basis van per-layer entropie.

2. Methodologie: Het RKSC-framework

De auteurs stellen RKSC (Reasoning-Aware KV Cache Sharing) voor, een training-vrij inferentieframework dat deze redundanties elimineert door middel van drie complementaire, schakelbare mechanismen. RKSC vereist geen fine-tuning, geen architecturale wijzigingen en geen aangepaste CUDA-kernels.

A. ASKS (Attention-Similarity KV Sharing)

ASKS generaliseert token-exacte prefix caching naar een semantisch regime.

• Mechanisme: In plaats van byte-identieke tokens te vereisen, berekent ASKS een gewogen cosinus-gelijkenis tussen de hidden states van de suffix van een tak en de hidden states van de root.
• Weging: De gelijkenis-score ($\sigma_b$) past exponentieel toenemende gewichten toe op latere transformer-lagen ($w_l = \exp(\alpha l/L)$), aangezien deze lagen meer taakspecifieke informatie bevatten en beter zijn in het detecteren van semantische divergentie.
• Actie: Als $\sigma_b \geq \tau$, wordt de tak als semantisch consistent met de prefix beschouwd en ontvangt deze de gedeelde KV cache via een zero-copy repeat_interleave operatie. Takken onder de drempelwaarde vallen terug op onafhankelijke berekening.
• Generalisatie: Dit generaliseert strikt token-exacte caching (waar $\sigma_b=1$ voor identieke tokens) en maakt hergebruik mogelijk voor takken met verschillende oppervlakkige formuleringen maar vergelijkbare semantische representaties.

B. CGEE (Confidence-Gated Early Exit)

CGEE vermindert de verificatiekosten door een dual-level exit-strategie:

1. Verify-Skip (Niveau 1): Voordat de verificatiepass wordt uitgevoerd, controleert het systeem de generatie-betrouwbaarheid. Als de winnende tak een hoge absolute betrouwbaarheid heeft ($p(b) \geq \tau_{conf}$) en de concurrenten met een duidelijke relatieve marge domineert, wordt de gehele verificatie forward pass overgeslagen.
2. Layer-Level Entropy Exit (Niveau 2): Indien de verificatie wel plaatsvindt, monitoren lightweight forward hooks de logit-entropie bij elke laag. Als de entropie onder een drempelwaarde ($\theta$) zakt en stabiliseert (verandering $< \epsilon$) bij een intermediaire laag $l^*$, wordt de forward pass vroegtijdig beëindigd. De gedeeltelijke output wordt gebruikt als het verificatieresultaat, waarbij de resterende $L - l^*$ lagen worden overgeslagen.

C. RSBCM (Reasoning-Selective Block Cache Manager)

Om ongecontroleerde geheugengroei tijdens diepe boomzoekopdrachten te voorkomen, beheert RSBCM de gedeelde cache.

• Eviction Policy: Blocks worden verwijderd op basis van een belangrijkheidsscore $\omega = \text{branch\_score} / (\text{depth} + 1)$. Dit geeft prioriteit aan het behouden van blocks van hoog scorende takken op ondiepe dieptes, terwijl blocks van diepe, laag-betrouwbare trajecten als eerste worden verwijderd.

3. Belangrijkste Bijdragen

• ASKS: Introduceert een semantische gelijkenis-gate die 28,6% extra KV-hergebruik realiseert ten opzichte van token-exacte matching in scenario's met diverse formuleringen, wat de bestaande prefix caching strikt generaliseert.
• CGEE: Stelt een dual-level exit-mechanisme voor dat de verificatie volledig overslaat wanneer de betrouwbaarheid beslissend is en de pass afkapt wanneer de per-layer entropie stabiliseert, waardoor de verificatie-latentie wordt verminderd zonder hertraining.
• RSBCM: Biedt een diepte-prioritaire eviction-strategie om cache-stabiliteit te handhaven in diepe zoekbomen.
• Zero-Overhead Design: Het framework voegt geen leerbare parameters toe en vereist geen architecturale wijzigingen, waardoor het fungeert als een drop-in wrapper voor bestaande reasoning pipelines.

4. Experimentele Resultaten

Het framework werd geëvalueerd over vijf modelfamilies (7B–10B parameters, inclusief Qwen2.5, Mistral, Falcon3, en Llama-3) en vier benchmarks (GPQA Diamond, MMLU-STEM, ARC-Challenge, GSM8K) op 1.000 problemen.

• Speedup: RKSC bereikt een gemiddelde versnelling van 3,008× ten opzichte van de No-KV baseline (die prefixes opnieuw berekent) en een verbetering van 1,66× ten opzichte van de vLLM-equivalente token-exacte prefix caching. De piekversnelling die werd waargenomen was 3,990× (Llama-3-8B op MMLU-STEM).
• Nauwkeurigheid: Het framework behoudt een hoge nauwkeurigheid met een door CGEE veroorzaakte foutenmarge van slechts 0,37% (6 fouten op 1.616 verify calls). Fouten concentreerden zich in multiple-choice benchmarks waar generatieve vloeiendheid soms afweek van correctheidssignalen; bij taken met vrije tekst-generatie was de foutenmarge nagenoeg nul.
• Componentenoverzicht:
  • Alleen KV sharing zorgde voor een 1,343× versnelling (25,5% latentiereductie).
  • Het gecombineerde RKSC-systeem (KV + CGEE) leverde een 1,622× versnelling in throughput decompositie en 3,008× in uitgebreide evaluatie.
  • De CGEE verify-skip rate was adaptief aan de taak moeilijkheid, variërend van ~42% op moeilijke GPQA-problemen tot ~92% op eenvoudigere MMLU-STEM problemen.
• Architectuur-agnosticisme: De versnelling door KV sharing was consistent over alle vijf de modellen (std 0,02×), wat bevestigt dat de winsten worden gedreven door de structuur van de inferentie-workload (prefix lengte en branching factor) in plaats van specifieke modelarchitecturen.

5. Betekenis en Claims

Het artikel claimt dat RKSC fundamentele structurele inefficiënties in multi-step LLM reasoning aanpakt die huidige productie-systemen negeren. Door KV-hergebruik los te koppelen van lexicale identiteit (ASKS) en gebruik te maken van het "entropy collapse" fenomeen in latere transformer-lagen (CGEE), verlaagt RKSC de energie- en tijdkosten van inferentie-scaling aanzienlijk.

De auteurs benadrukken dat deze winsten worden behaald zonder de computationele overhead van training of de architecturale complexiteit van speculative decoding methoden. Het framework wordt gepresenteerd als een praktische, drop-in optimalisatie die state-of-the-art reasoning toegankelijker maakt op beperkte hardware door redundante berekeningen in multi-branch pipelines te verminderen. De resultaten suggereren dat voor multi-step reasoning taken een aanzienlijk deel van de verificatie- en prefill-berekening structureel redundant is en veilig kan worden geëlimineerd of verkort op basis van semantische gelijkenis en betrouwbaarheidsmetrieken.