🤖 machine learning

RKSC: Reasoning-Aware KV Cache Sharing and Confident Early Exit for Multi-Step LLM Inference

RKSC is a training-free inference framework that accelerates multi-step LLM reasoning by eliminating structural redundancies through attention-similarity-based KV cache sharing, confidence-gated early exits, and a selective block cache manager, achieving a mean 3x speedup with negligible error rates across diverse models and benchmarks.

Original authors: Anirudh Sekar

Published 2026-06-10

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Anirudh Sekar

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a brilliant detective trying to solve a complex mystery. Instead of just guessing one answer, you decide to ask eight different detectives (branches) to write down their own theories simultaneously. This is how modern "reasoning" AI models work: they generate multiple possible paths to a solution to find the best one.

The paper introduces RKSC, a new "manager" for this team of detectives. Its goal is simple: stop the detectives from wasting time doing the exact same work over and over, and stop them from reading the entire case file if they already know the answer.

Here is how RKSC works, broken down into three simple tricks:

1. The "Shared Notebook" Trick (ASKS)

The Problem:
Usually, when eight detectives start working, they all read the first 1,000 pages of the case file (the "prefix") independently. Even if they are all reading the exact same text, the computer calculates the meaning of those words eight separate times. It's like eight people in a library each copying the first chapter of a book by hand, even though they could just share one copy.

The RKSC Solution:
RKSC introduces a Shared Notebook.

The system reads the common part of the case file once.
It then hands this single, pre-calculated "notebook" to all eight detectives.
The Magic: Unlike older systems that only share the notebook if the detectives are using the exact same words, RKSC is smart enough to share it even if they phrase things slightly differently, as long as they are thinking about the same thing. It checks their "thoughts" (hidden states) rather than just their "words."

The Result: The team saves a massive amount of time because they stop re-reading the beginning of the case file eight times.

2. The "Confident Exit" Trick (CGEE)

The Problem:
After the detectives write their theories, a supervisor usually checks every single theory from start to finish to see which one is right. This "verification" step is slow. But sometimes, one detective is so obviously confident and correct that checking the rest of the file is a waste of time.

The RKSC Solution:
RKSC acts like a smart supervisor with two rules:

Rule A (The Skip): If one detective is 95% sure of their answer and the others are clearly unsure, the supervisor says, "Great job, you're done!" and skips the rest of the verification entirely.
Rule B (The Early Stop): If the supervisor must check, they don't need to read the whole file. They can stop halfway through the book. Why? Because the paper found that once a detective gets to a certain point in their reasoning, their confidence stabilizes. Reading the last 30% of the book doesn't change the conclusion; it just wastes time.

The Result: The system often skips the verification step entirely or cuts it short, saving huge amounts of time.

3. The "Memory Janitor" (RSBCM)

The Problem:
If the detectives keep branching out into deeper and deeper theories (like a tree with many roots), the "Shared Notebook" could get too big to fit in the computer's memory.

The RKSC Solution:
RKSC has a janitor who watches the notebook. If it gets too full, the janitor throws away the notes from the detectives who are deep in a dead-end theory or who seem less confident. This keeps the memory clean and ensures the system doesn't crash, even during very long reasoning sessions.

What Did They Find?

The authors tested this system on five different AI models (ranging from small to medium size) across four different types of difficult puzzles (science, math, and logic).

Speed: The system made the AI 3 times faster on average compared to standard methods. In the best cases, it was nearly 4 times faster.
Accuracy: It was incredibly accurate. Out of over 1,600 verification checks, the "Early Exit" trick made a mistake only 6 times (an error rate of 0.37%).
No Training Needed: The best part is that this doesn't require re-training the AI. It's like putting a new, smarter engine in a car without having to rebuild the car itself. It works with existing models right out of the box.

In Summary

RKSC is like giving a team of AI detectives a shared notebook so they don't re-read the same pages, a smart supervisor who stops checking work when the answer is obvious, and a janitor to keep their workspace from getting cluttered. The result is a reasoning system that thinks much faster without losing its sharpness.

Technical Summary: RKSC – Reasoning-Aware KV Cache Sharing and Confident Early Exit

1. Problem Statement

Modern large language models (LLMs) designed for complex reasoning (e.g., DeepSeek-R1, o1-style models) often employ multi-branch inference strategies, such as Tree of Thoughts or self-consistency, where multiple reasoning trajectories are generated for a single problem. The paper identifies two structural redundancies in current production inference systems (like vLLM and SGLang) that degrade efficiency in these pipelines:

Cross-Branch KV Waste: Parallel branches typically share a significant portion of their prefix (the problem statement and initial context). However, existing systems only reuse Key-Value (KV) caches when token sequences are byte-identical. If branches diverge semantically but share the same prefix tokens, or if they share a prefix but use different phrasing for hints, the system recomputes the expensive $O(n^2)$ attention for the shared prefix independently for each branch.
Over-Deep Verification: In process-reward verification pipelines, the model performs a full-depth forward pass to score each generated branch, even when the generation confidence is already decisive or when the model's internal state (logit distribution) has stabilized at an intermediate layer. Current systems lack mechanisms to skip the verification pass entirely or terminate it early based on per-layer entropy.

2. Methodology: The RKSC Framework

The authors propose RKSC (Reasoning-Aware KV Cache Sharing), a training-free inference framework that eliminates these redundancies through three complementary, switchable mechanisms. RKSC requires no fine-tuning, no architecture changes, and no custom CUDA kernels.

A. ASKS (Attention-Similarity KV Sharing)

ASKS generalizes token-exact prefix caching to a semantic regime.

Mechanism: Instead of requiring byte-identical tokens, ASKS computes a weighted cosine similarity between the hidden states of a branch's suffix and the root's hidden states.
Weighting: The similarity score ( $\sigma_b$ ) applies exponentially increasing weights to later transformer layers ( $w_l = \exp(\alpha l/L)$ ), as these layers carry more task-specific information and are better at detecting semantic divergence.
Action: If $\sigma_b \geq \tau$ , the branch is deemed semantically consistent with the prefix and receives the shared KV cache via a zero-copy repeat_interleave operation. Branches below the threshold fall back to independent computation.
Generalization: This strictly generalizes token-exact caching (where $\sigma_b=1$ for identical tokens) and enables reuse for branches with different surface-level phrasing but similar semantic representations.

B. CGEE (Confidence-Gated Early Exit)

CGEE reduces verification costs through a dual-level exit strategy:

Verify-Skip (Level 1): Before running the verification pass, the system checks the generation confidence. If the winning branch has high absolute confidence ( $p(b) \geq \tau_{conf}$ ) and dominates competitors by a clear relative margin, the entire verification forward pass is skipped.
Layer-Level Entropy Exit (Level 2): If verification proceeds, lightweight forward hooks monitor the logit entropy at each layer. If the entropy drops below a threshold ( $\theta$ ) and stabilizes (change $< \epsilon$ ) at an intermediate layer $l^*$ , the forward pass is terminated early. The partial output is used as the verification result, skipping the remaining $L - l^*$ layers.

C. RSBCM (Reasoning-Selective Block Cache Manager)

To prevent unbounded memory growth during deep tree searches, RSBCM manages the shared cache.

Eviction Policy: Blocks are evicted based on an importance score $\omega = \text{branch\_score} / (\text{depth} + 1)$ . This prioritizes keeping blocks from high-scoring branches at shallower depths while evicting blocks from deep, low-confidence trajectories first.

3. Key Contributions

ASKS: Introduces a semantic similarity gate that achieves 28.6% additional KV reuse over token-exact matching in diverse-phrasing scenarios, strictly generalizing existing prefix caching.
CGEE: Proposes a dual-level exit mechanism that skips verification entirely when confidence is decisive and truncates the pass when per-layer entropy stabilizes, reducing verification latency without retraining.
RSBCM: Provides a depth-priority eviction strategy to maintain cache stability in deep search trees.
Zero-Overhead Design: The framework adds no learnable parameters and requires no architectural modifications, acting as a drop-in wrapper for existing reasoning pipelines.

4. Experimental Results

The framework was evaluated across five model families (7B–10B parameters, including Qwen2.5, Mistral, Falcon3, and Llama-3) and four benchmarks (GPQA Diamond, MMLU-STEM, ARC-Challenge, GSM8K) on 1,000 problems.

Speedup: RKSC achieves a mean speedup of 3.008× over the No-KV baseline (which recomputes prefixes) and a 1.66× improvement over vLLM-equivalent token-exact prefix caching. The peak speedup observed was 3.990× (Llama-3-8B on MMLU-STEM).
Accuracy: The framework maintains high accuracy with a CGEE-induced error rate of only 0.37% (6 errors out of 1,616 verify calls). Errors were concentrated in multiple-choice benchmarks where generation fluency occasionally diverged from correctness signals; free-form generation tasks showed near-zero error rates.
Component Breakdown:
- KV sharing alone provided a 1.343× speedup (25.5% latency reduction).
- The combined RKSC system (KV + CGEE) provided 1.622× speedup in throughput decomposition and 3.008× in extended evaluation.
- CGEE's verify-skip rate was adaptive to task difficulty, ranging from ~42% on hard GPQA problems to ~92% on easier MMLU-STEM problems.
Architecture Agnosticism: The KV sharing speedup was consistent across all five models (std 0.02×), confirming that the gains are driven by the inference workload structure (prefix length and branching factor) rather than specific model architectures.

5. Significance and Claims

The paper claims that RKSC addresses fundamental structural inefficiencies in multi-step LLM reasoning that current production systems ignore. By decoupling KV reuse from lexical identity (ASKS) and exploiting the "entropy collapse" phenomenon in later transformer layers (CGEE), RKSC significantly lowers the energy and time costs of inference-time scaling.

The authors emphasize that these gains are achieved without the computational overhead of training or the architectural complexity of speculative decoding methods. The framework is presented as a practical, drop-in optimization that makes state-of-the-art reasoning more accessible on constrained hardware by reducing redundant computation in multi-branch pipelines. The results suggest that for multi-step reasoning tasks, a significant portion of the verification and prefill computation is structurally redundant and can be safely eliminated or shortened based on semantic similarity and confidence metrics.