SpeContext: Enabling Efficient Long-context Reasoning with Speculative Context Sparsity in LLMs

Jiaming Xu, Jiayi Pan, Hanzhen Wang, Yongkang Zhou, Jiancai Ye, Yu Wang and 1 other

When an LLM reasons step by step for thousands of tokens, its Key-Value cache balloons and every new token must read all of it, so long-context reasoning chokes on memory and latency — and the usual trick of keeping only the important cache entries backfires here because selecting them layer by layer is serial and slow; this paper (SJTU, Tsinghua, Infinigence — Jiaming Xu, Yu Wang, Guohao Dai, and colleagues) presents SpeContext, which notices that picking the right cache entries is the same alignment goal as knowledge distillation, so it runs a tiny distilled draft model ahead of the big one to speculatively predict which context matters — prefetching it asynchronously and reaching up to 24.89 times higher throughput in the cloud and 10.06 times on the edge with almost no accuracy loss.

The argument in one glance

Problem

Long step-by-step reasoning drowns the model in its own Key-Value cache

Modern LLMs get much better at hard problems through test-time scaling: instead of answering in one shot, they reason step by step, generating a long chain of thought — sometimes thousands of tokens for planning or math, and far more inside AI agents. To generate each new token efficiently, an LLM keeps a Key-Value (KV) cache: the key and value vectors for every earlier token, so it can attend to the past without recomputing it. The catch is that this cache grows with the context, and every new token must read the entire cache to compute attention. So long-context reasoning is doubly punishing: the cache balloons in memory, and per-token latency climbs as it grows. The paper gives a concrete feel — on an NVIDIA RTX 4090, generating one token with a 16K context takes twice as long as with a 1K context, and the KV cache alone reaches about 2GB, so only three requests fit on the card at once. For cloud vendors this is lost throughput and higher cost; on edge devices it is a hard wall. The KV cache, not the model weights, becomes the bottleneck that stops long reasoning from being deployed cheaply.
What people did

Keep only the important KV entries, selecting them layer by layer

A large body of work attacks the KV cache by making attention sparse: rather than attend to every past token, retrieve just the important KV pairs and compute over those. Representative systems (Quest, ClusterKV, ShadowKV) preprocess the cache after the prompt is read — paging it, clustering it, or quantizing it — and then, during generation, score and select the most relevant entries for each layer, loading only those. This works well for long-context input, where a big prompt is processed once and then reused, and it lets these systems offload the rest of the cache to slower memory to save space.
Why it fell short

Input-tuned KV tricks break down when the long part is generated, not given

Applied to long-context reasoning, this paradigm fails three ways. First, the per-layer retrieval is serial: because each layer must retrieve and load its important KV before it can compute attention, the selection sits on the critical path and its overhead scales with model depth, adding up to 60 percent latency and destroying the natural overlap between computation and memory access. Second, these methods were built assuming the long part is the input prompt, so they preprocess the prompt once and then simply keep every newly generated KV pair to avoid re-preprocessing — but in reasoning the generated part is the long part and never stops growing, so retention plus per-layer retrieval overhead makes them slower than plain full attention. Third, in memory-constrained settings they offload KV to slower memory using a strategy fixed before inference (fully offload or never), which cannot adapt to a sequence that keeps lengthening, so a tiny increase in length can trigger a full offload and over 80 percent slowdown. The methods are tuned for a static long input, and reasoning is a dynamic, growing one.
This paper's idea

Let a tiny distilled model speculatively predict the context that matters

SpeContext's key insight reframes the whole problem: the job of a KV-retrieval algorithm is to align with the LLM — to pick exactly the tokens the LLM itself would focus on — and that is the same objective as knowledge distillation, which trains a small model to match a large one. The paper argues from information theory (mutual information and the data processing inequality) that if a small distilled language model (DLM) reproduces the LLM's output distribution, its internal information focus — which tokens matter — must be highly similar to the LLM's. So a DLM can serve as the retrieval algorithm, speculatively predicting the important context, much as a draft model speculatively predicts tokens in speculative decoding. On this insight SpeContext builds an algorithm-and-system co-design. At the algorithm level, it prunes the DLM down to a lightweight retrieval head — keeping only the parts needed to compute attention weights (the embedding and query-key projection), a parameter reduction of over 90 percent. At the system level, because that head depends only on the input and not on per-layer results, SpeContext runs it ahead of the LLM and prefetches KV asynchronously, overlapping loading with computation and reloading only the roughly 20 percent of entries that change between steps (elastic loading). At the compilation level, a theoretical memory model drives adaptive management that splits KV between GPU and CPU as the sequence grows. Across cloud and edge, SpeContext reaches up to 24.89 times and 2.19 times the throughput of Huggingface and FlashInfer in the cloud, and 10.06 times and 8.02 times on the edge, with negligible accuracy loss.

Why this problem is worth solving

The value of long reasoning is exactly what makes its cost unavoidable, so the KV cache bottleneck stands directly in the way of deploying the models people most want.

Test-time scaling — letting a model think longer, in more steps — is one of the most effective ways to raise LLM capability on genuinely hard tasks, and it is what powers agents that plan over long horizons. But thinking longer means generating more tokens, and every generated token becomes part of the KV cache that all subsequent tokens must read. Two costs grow together: memory, because the cache is proportional to sequence length, and latency, because attention for each new token scans the whole cache. Unlike the model weights, which are fixed, the cache is unbounded — it expands with exactly the reasoning that makes the model good. That is why the paper frames it as the barrier to practical deployment: a 4090 that can hold only three long-context requests, or an edge device that cannot hold one, is not limited by compute or by model size but by the cache the reasoning produces. Making long reasoning affordable therefore requires attacking the KV cache directly, which is what every method in this space, including SpeContext, is really about.

The deeper reason the prior tricks fail is a mismatch between a static assumption and a dynamic reality, and naming that mismatch is what unlocks the fix.

It is tempting to think a long sequence is a long sequence — that a method which handles a 100K-token prompt should handle 100K tokens of generated reasoning. The paper shows why not. When the long part is the input, it arrives all at once: you can afford an expensive one-time preprocessing pass (clustering, quantization, paging) over the prompt, then reuse that structure cheaply for every generated token, and simply keep the comparatively few new tokens as they trickle out. When the long part is the output, none of that holds. There is no single moment to preprocess, because the sequence keeps growing token by token; the newly generated tokens are not a trickle but the bulk of the context; and the memory footprint is a moving target that a strategy fixed before inference cannot track. So methods optimized for a static long input inherit assumptions that are actively wrong for a dynamic long output, which is why they can end up slower than doing nothing clever at all. Recognizing that reasoning demands a fundamentally different design — cheap per-step selection with no preprocessing, and memory management that adapts as the sequence grows — is the conceptual step that the rest of SpeContext delivers on.

What people did before, and where it broke

Per-layer dynamic KV retrieval (Quest, ClusterKV). Score and select the important KV entries in each layer during decoding, loading only those. Where it broke. The selection is serial and data-dependent, so it sits on the critical path and its overhead scales with model depth — up to 60 percent latency — breaking the overlap of compute and memory access.
Preprocess the prompt, retain all new KV. Preprocess the input cache once (clustering, quantization) and keep every newly generated KV pair to avoid redoing it. Where it broke. In reasoning the generated part is the long part and grows without bound, so retention plus per-layer overhead makes these methods slower than full attention.
Quantize or cluster keys to cut retrieval cost (ShadowKV). Shrink the per-comparison cost of retrieval by working on compressed keys. Where it broke. It still relies on preprocessing the cache, which is ill-suited to a cache that continuously grows during generation, and still pays the per-layer serialization.
Fixed offloading of KV to slower memory. Decide before inference whether to fully offload the cache to CPU memory or never offload. Where it broke. A predetermined split cannot adapt to a lengthening sequence, so a tiny length increase can trigger a full offload and over 80 percent slowdown.
The gap. No method selected important context cheaply, without preprocessing, off the critical path, while adapting its memory placement to a sequence that keeps growing — the regime that long-context reasoning actually lives in.

How SpeContext works

From serial per-layer retrieval to a speculative, prefetched context selection

↓ →
new
↓ →
new
↓ →
new
↓ →
new
↓ →
new

The idea

Let a small model speculate about what the big model needs. SpeContext runs a lightweight distilled model ahead of the LLM to guess which context entries matter, then acts on that guess — the same speculative pattern as running a cheap draft model to propose what an expensive model will do. Using a small distilled draft model to run ahead and predict the large model’s behavior, turning an expensive step into a cheap prediction plus a check, is the same speculative-decoding strategy the work on a heterogeneous architecture for speculative decoding with a draft model (DFVG) covered elsewhere on this site is built on. SpeContext transplants that draft-model idea from predicting tokens to predicting which parts of the context the LLM should attend to.

Prune the helper down to only what the task needs. The DLM, run in full, would add real overhead, so SpeContext strips it to a retrieval head that keeps only the attention-weight computation and discards everything else, a parameter cut of over 90 percent. Removing the parts of a neural network that a specific task does not need so that only a lean, purpose-built remainder runs is the same dynamic-pruning instinct the work on hardware-accelerated dynamic neural network pruning (DPACS) covered elsewhere on this site pursues, spending compute only on the parts that matter. SpeContext prunes the draft model to the single function it plays here — identifying important tokens from attention weights.

Move data across the memory hierarchy ahead of when it is needed. Because the retrieval head predicts the needed KV up front, SpeContext prefetches it asynchronously from slower memory and reloads only the fraction that changed, keeping the transfer off the critical path. Prefetching and migrating tensors between fast and slow memory tiers so the data arrives before the computation that consumes it is the same overlap-and-migrate discipline the work on tensor migration and prefetching in unified memory (DeepUM) covered elsewhere on this site is built on, hiding transfer latency behind compute. SpeContext applies it to the KV cache, using the speculative head’s early prediction to know what to prefetch.

Following one generation step through SpeContext

From a new reasoning token to its answer, without serial per-layer retrieval

A trace, drawn from the paper, of how SpeContext selects and loads context for a single generated token during long-context reasoning.

Step 1 of 5

Run the lightweight retrieval head on the input

Before the LLM computes anything, SpeContext runs its pruned retrieval head — the DLM reduced to embedding and query-key projection — over the current input. It maintains a full key cache and computes attention weights to score every past token's importance.

This head is over 90 percent smaller than the full distilled model.

Why it works

Retrieval is distillation, so a distilled model already knows the answer. Selecting the right KV means aligning with the LLM’s focus → but a distilled model is trained to reproduce that focus → so it predicts the important tokens without any new training.
Input-only dependence removes the serial bottleneck. Per-layer retrieval must wait for each layer’s results → but the retrieval head depends only on the input → so all important tokens are known up front and can be prefetched off the critical path.
Pruning to the attention path makes the helper nearly free. Running the full draft model adds real overhead → but only its query-key attention computation is needed for retrieval → so a head over 90 percent smaller does the job cheaply.
Step-to-step overlap makes loading incremental. Reloading the full selected cache each step is wasteful → but adjacent steps share over 80 percent of their tokens → so elastic loading transfers only the roughly 20 percent that changed.
A memory model keeps offloading from collapsing. A fixed offload split breaks as the sequence grows → but a theoretical model adapts the GPU-CPU placement to length → so a small increase no longer triggers an over-80-percent slowdown.

What it costs (trade-offs and scope)

The retrieval head is an approximation, not the LLM. It predicts important tokens from a distilled model’s attention rather than the LLM’s own, so selection can differ from truly optimal; the paper reports negligible accuracy loss, but this is an approximation with a (small) accuracy budget.
It depends on a suitable distilled model. SpeContext builds on a DLM (from the EAGLE-3 speculative-decoding family) and extends its short context with a training-free method; a model without such a draft companion needs one obtained first.
The head still keeps a full key cache. To score all past tokens the retrieval head maintains a full key cache and computes attention over it, so some memory and compute for selection remain, even though the value path and language-model head are pruned.
Gains depend on step-to-step overlap holding. Elastic loading’s benefit rests on adjacent steps sharing most of their important tokens; workloads where the focus shifts sharply each step would reload more and gain less.
Results are on specific models, hardware, and baselines. The up-to-24.89-times and 10.06-times figures are against Huggingface and FlashInfer on particular cloud and edge setups and reasoning benchmarks; other models, budgets, or hardware may differ.

The comparison

Serving long-context LLM reasoning under a memory budget

Property	Per-layer KV retrieval	Full attention	SpeContext
When context is selected
Handles a growing generated cache
Overlaps loading with compute
Adapts offloading to length
Throughput vs baselines

A comparison of ways to serve long-context reasoning. Per-layer KV retrieval methods (Quest, ClusterKV, ShadowKV) were tuned for long inputs and, on long generation, pay serial per-layer selection, retain all new KV, and use fixed offloading that collapses as the sequence grows. Full attention (Huggingface, FlashInfer) reads the whole cache and slows with length. SpeContext runs a pruned distilled retrieval head once, off the critical path, to predict important context, prefetches it asynchronously with elastic loading, and adapts KV placement with a memory model — reaching up to 24.89 and 2.19 times the throughput of Huggingface and FlashInfer in the cloud and 10.06 and 8.02 times on the edge with negligible accuracy loss. Exact figures should be read from the paper.Underlined values explain themselves when clicked.

Where this stops working

What to keep in mind

This explainer is based on the paper’s full text (the authors’ arXiv copy, 2512.00722, of the ASPLOS 2026 paper). The described system is the paper’s: SpeContext is an algorithm-and-system co-design for efficient long-context LLM reasoning built on the insight that KV-retrieval’s objective — aligning with the LLM’s contextual focus — matches the objective of knowledge distillation, justified via an information-theory argument (mutual information and the data processing inequality) that a distilled model’s information focus resembles the original LLM’s. Its three contributions are a lightweight retrieval head (the distilled model pruned to embedding and query-key projection, over 90 percent parameter reduction, with head-level token selection and support for MHA, GQA, MQA, and MLA attention, extended to long context via a training-free method), an asynchronous prefetch dataflow with elastic loading (exploiting that the head depends only on the input and that adjacent steps share over 80 percent of selected tokens, so only about 20 percent are reloaded), and an adaptive memory-management system driven by a theoretical memory model that splits KV between GPU and CPU as the sequence grows. The headline results — up to 24.89 times and 2.19 times throughput over Huggingface and FlashInfer in the cloud, and 10.06 times and 8.02 times on the edge, with negligible accuracy loss — are from the paper’s evaluation on specific models, benchmarks, and hardware. Background on KV cache, sparse attention, speculative decoding, and knowledge distillation is standard context. The step-by-step walkthrough illustrates the described mechanism. Exact formulations, the memory model, and full results should be read from the paper.

Background

Why efficient LLM serving keeps reusing small models and the memory hierarchy

It helps to place SpeContext in the wider effort to serve large models cheaply, especially on constrained hardware. Two recurring tools show up again and again. The first is a small model that stands in for the big one where full precision is not needed: speculative decoding uses a draft model to propose tokens, early-exit skips later layers, and here a distilled model proposes which context matters — all cases of spending a little cheap computation to save a lot of expensive computation. The second is the memory hierarchy: because the data an LLM needs rarely all fits in fast memory, systems shuttle it between GPU and CPU, and the winning designs prefetch and overlap that movement so it hides behind compute rather than stalling it. SpeContext combines both — a pruned distilled head to predict the sparse context, and asynchronous prefetching with adaptive offloading to move only what is needed, when it is needed. This drive to run demanding language-model inference within tight compute and memory limits by being clever about what to compute and what to move is the same terrain the work on turbocharging NLP inference at the edge via elastic pipelining (STI) covered elsewhere on this site occupies, fitting large models into small budgets. SpeContext’s particular contribution is to unify the small-model and memory-hierarchy tricks specifically for the growing, dynamic cache of long-context reasoning.

If you remember only three things

1. Long-context reasoning — letting an LLM think step by step for thousands of tokens — sharply improves capability but generates an ever-growing Key-Value cache that must be read in full for every new token, making memory and latency, not model size, the deployment bottleneck (on an RTX 4090, a 16K-context token takes twice as long as a 1K one, and only three long requests fit at once). Existing KV-cache-sparsity methods were designed for long inputs and break on long generation: their per-layer retrieval is serial and adds up to 60 percent latency, they wastefully retain all newly generated cache, and their offloading strategy is fixed before inference so it collapses (over 80 percent slowdown) as the sequence grows.
2. SpeContext's key insight is that selecting the right KV entries means aligning with what the LLM would attend to, which is exactly the objective of knowledge distillation — so a small distilled model can speculatively predict the important context, just as a draft model predicts tokens in speculative decoding. On this it builds an algorithm-and-system co-design: a lightweight retrieval head (the distilled model pruned to its attention-weight computation, over 90 percent smaller) that selects important tokens per attention head; an asynchronous prefetch dataflow that, because the head depends only on the input, predicts all needed KV up front and overlaps loading with compute, reloading only the roughly 20 percent of entries that change between steps; and an adaptive memory manager, driven by a theoretical memory model, that splits KV between GPU and CPU as the sequence lengthens.
3. By moving context selection off the critical path and reducing both what is computed and what is moved, SpeContext reaches up to 24.89 times and 2.19 times the throughput of Huggingface and FlashInfer in the cloud, and 10.06 times and 8.02 times on the edge, with negligible accuracy loss, pushing the accuracy-throughput Pareto frontier. The trade-offs: the retrieval head is a distilled approximation with a small accuracy budget, it depends on having a suitable draft model, it still keeps a full key cache to score tokens, and its elastic-loading gains rest on adjacent steps sharing most of their important tokens — with results measured on specific models, benchmarks, and hardware.

Check yourself

Check your understanding

1. Why does long-context reasoning create a severe efficiency problem?
2. Why do existing KV-cache-sparsity methods, built for long inputs, fail on long-context reasoning?
3. What is SpeContext's key insight, and how does it use it?
4. What are SpeContext's three main techniques for acting on that insight?

Sources

The paper is “SpeContext: Enabling Efficient Long-context Reasoning with Speculative Context Sparsity in LLMs” by Jiaming Xu, Jiayi Pan, Hanzhen Wang, Yongkang Zhou, Jiancai Ye, Yu Wang, and Guohao Dai (Shanghai Jiao Tong University, Tsinghua University, and Infinigence-AI), published at ASPLOS 2026, DOI 10.1145/3779212.3790224. It is available on arXiv as 2512.00722.

This explainer is written from the paper’s full text (the authors’ arXiv copy). SpeContext is an algorithm-and-system co-design for long-context LLM reasoning built on the insight that KV-cache retrieval aligns with the LLM’s contextual focus in the same way knowledge distillation aligns a small model with a large one, so a distilled language model can speculatively select important context. It contributes a lightweight retrieval head (the distilled model pruned to its attention-weight path, over 90 percent smaller, with head-level selection), an asynchronous prefetch dataflow with elastic loading (reloading only about 20 percent of entries per step because adjacent steps overlap by over 80 percent), and an adaptive memory-management system driven by a theoretical memory model that splits KV between GPU and CPU as the sequence grows. Evaluated on cloud and edge, the paper reports up to 24.89 times and 2.19 times throughput over Huggingface and FlashInfer in the cloud and 10.06 times and 8.02 times on the edge, with negligible accuracy loss. Background on KV cache, sparse attention, speculative decoding, and distillation is standard context. The numbers are the paper’s own; exact mechanisms and full results should be confirmed in the paper.