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

1 Introduction

• KV-cache memory grows with context length, and reading the full cache at each decode step adds latency; existing selection methods trade accuracy for speed.For Llama3.1-8B, a 16K context gives a 2GB KV-cache footprint and takes twice as long per token as a 1K context on an RTX 4090.Source: §1 · p. 2
• Existing decoding methods face three challenges: serialized layer-wise retrieval, complete retention of newly generated KV pairs, and offloading policies that cannot adapt as sequences grow.Layer-wise retrieval overhead can reach 60% of latency; fixed offloading strategies can cause over 80% performance degradation after a tiny sequence-length increase.Source: §1 · p. 2
• SpeContext uses a distilled language model as a retrieval algorithm because its information focus is similar to the original LLM’s for the same inputs.The paper analyzes this similarity using mutual information and the data processing inequality.Source: §1 · p. 3
• At the algorithm level, a lightweight retrieval head uses head-level DLM attention weights and prunes redundant operations, reducing parameters by more than 90%.The paper reports higher similarity between focused tokens at the head level than at the batch level.Source: §1 · p. 3
• At the system level, asynchronous prefetch and elastic loading overlap KV-cache retrieval with LLM inference and load only the KV pairs that differ from the prior generation.Elastic loading reduces data transfer by up to 90%.Source: §1 · p. 3
• At the compilation level, a theoretical memory model guides adaptive memory allocation to maximize GPU utilization as reasoning sequences grow.The system adapts memory allocation to increase inference speed in resource-constrained environments.Source: §1 · p. 4
• 24.89× throughput improvement in cloud and 10.06× speedup in edge, with negligible accuracy loss, are reported for SpeContext across long-context input and reasoning scenarios.The results push the Pareto frontier of accuracy and throughput.Source: §1 · p. 4

2.1 Large Language Model

• The KV cache stores past key-value pairs for autoregressive decoding, reducing computation while creating memory overhead that scales with context length.The prefill phase creates the cache; decoding uses and extends it.Source: §2.1 · p. 4, §2.2 · p. 4

2.2 KV Cache Optimization

• Permanent eviction bounds KV memory by retaining recent pairs, but can lose historical context; StreamingLLM also retains initial “attention sink” tokens.StreamingLLM combines retained initial tokens with a sliding window.Source: §2.2 · p. 4
• Dynamic selection retrieves input-relevant KV pairs during inference, but commonly preprocesses prompt KV and retains newly generated pairs during decoding.Quest uses pages, ClusterKV uses clusters, and ShadowKV quantizes keys; repeated preprocessing is costly as reasoning KV grows.Source: §2.2 · p. 5

2.3 Knowledge Distillation in LLMs

• Knowledge distillation trains a smaller student to mimic a teacher’s output distribution by minimizing KL divergence.In speculative decoding, EAGLE uses a distilled model to draft tokens for parallel LLM verification.Source: §2.3 · p. 5

3.1 Two Core Questions of KV Selection

• Retrieval algorithms aim to align efficiently with the LLM’s contextual focus, and their alignment degree determines retrieval accuracy.Permanent eviction uses input-agnostic properties; dynamic selection uses query-aware retrieval based on content-dependent focus.Source: §3.1 · p. 5
• Preprocessing reduces retrieval computation, which otherwise entails per-layer query–key multiplication and Top-K selection over candidate keys.The paper models single-inference retrieval overhead as layers × batch size × heads × dimension × key length × multiplication cost; existing preprocessing reduces key length or multiplication cost.Source: §3.1 · p. 5

3.2 Key Insight and Theoretical Analysis

• The paper argues that distillation’s output-distribution alignment requires the student to capture contextual information similarly to the teacher.It invokes mutual information and the data processing inequality to argue that the student’s internal representation must capture information the teacher deems important, motivating the use of a distilled model as the retrieval algorithm.Source: §3.2 · pp. 5–6

4.1 Challenge: Time-consuming DLM

• The DLM processes the same inputs with a full KV cache, adding approximately 20% overhead and motivating a lightweight retrieval algorithm.The paper uses the EAGLE-3 DLM, which has a single Transformer decoder layer, and extends its 2K context support with training-free YaRN.Source: §4.1 · p. 6, §4.3 · p. 6

4.2 Insight and Analysis: Redundant Operation

• Head-level retrieval has higher similarity to the original LLM’s important tokens and a higher generated-token hit rate than batch-level retrieval.The analysis concludes that only attention-weight-related operations, such as Query and Key generation, are needed; other DLM operations are redundant.Source: §4.2 · p. 6

4.3 Approach: Lightweight Retrieval Head

• The lightweight head retains the DLM embedding and QK projections, computes attention weights, and selects important tokens per head for the original LLM.Figure 5 depicts redundancy pruning and head-level selection, alongside implementations for four attention mechanisms.Source: §4.3 · p. 6, Figure 5 · p. 7
• For MHA, selected KV pairs map to corresponding heads; for GQA and MQA, group-level attention weights guide selection for shared KV heads.GQA takes an element-wise maximum across query heads in a group; MQA treats all query heads as one group.Source: §4.3 · pp. 6–7
• For MLA, retrieval remains head-level, and only selected latent-cache entries are expanded for attention computation.MLA caches a lower-dimensional latent representation rather than full key-value pairs.Source: §4.3 · p. 7

5.1 Motivation: Data independence

• Deploying the retrieval head before LLM inference removes the retrieval–loading data dependency and allows computation and KV prefetching to run concurrently on multiple CUDA streams.The retrieval head depends only on LLM input, unlike prior retrieval that loads KV for attention in each layer.Source: §5.1 · p. 7

5.2 Challenge: Heavy Data transfer

• With a large KV budget, data-transfer latency exceeds LLM layer-inference latency, making KV loading a major end-to-end latency component.The challenge is reducing transferred KV volume without sacrificing accuracy.Source: §5.2 · p. 7, Figure 6 · p. 7

5.3 Insight: Contextual Similarity

• More than 80% of selected important tokens overlap between adjacent generations, so about 20% of GPU KV cache needs updating for the subsequent generation.The paper states that loading only the updating KV cache can maintain accuracy while reducing data transfer.Source: §5.3 · pp. 7–8

5.4 Approach: Elastic Loading

• Elastic loading reuses KV cache already on the GPU and fetches only entries absent from that cache for the current generation.With a fixed KV budget, the strategy uses set differences between the prior, current, and resident selections to determine updates and loads, performed in place.Source: §5.4 · p. 8, Figure 7 · p. 8

6.1 Motivation: Performance Degradation

• A small workload increase can trigger full KV-cache offloading to CPU and more than 80% performance degradation.The motivating problem is that reasoning sequence lengths are dynamic and unpredictable.Source: §6.1 · p. 8

6.2 Approach: Adaptive Memory Management

• The system models memory overhead from LLM architecture, hardware specifications, and inference workload, then uses that model for adaptive memory management.Source: §6.2 · p. 8
• As context grows, SpeContext progressively offloads KV-cache layers to CPU memory, freeing GPU memory for cache in other layers.It precalculates sequence-length thresholds during compilation, using sequence length as the primary factor influencing memory overhead.Source: §6.2.1 · p. 9

6.2.1 Implementation Details.

• At each threshold, inference offloads another layer’s KV cache; between the second and third thresholds, the example keeps the remaining layers on GPU and offloads the last two to CPU DRAM.The system uses the thresholds to maintain GPU memory utilization as sequence length increases.Source: §6.2.1 · p. 9

• SpeContext is evaluated on an NVIDIA A100-80GB in the cloud and an RTX 4060 Laptop GPU with 8GB memory and an Intel i7-13650HX CPU at the edge.Source: §7.1 · p. 9
• Baselines include Huggingface and FlashInfer for full attention, plus Quest, ClusterKV, and ShadowKV for sparse attention; accuracy tasks are four LongBench tasks and LongWriter.Source: §7.1 · pp. 9–10
• The evaluation uses Llama3.1-8B, DeepSeek-R1-Distill-Llama-8B, and Qwen3-8B in cloud, and Reasoning-Llama-3.2-1B at the edge.Source: §7.1 · p. 9

7.2 Evaluation on Accuracy

• At a 512-token budget, SpeContext accuracy is slightly below ClusterKV; at 1k, it surpasses the baselines and reaches full-attention accuracy.Source: §7.2.1 · p. 10

7.2.1 Results in Long-context Input Scenario.

• For long-context input, outputs from Quest, ClusterKV, and ShadowKV score close to full attention because their roughly 100-token input fits within every KV budget, so all input content is selected.The passage reports that these methods consequently produce the same scores across KV budgets, with poor throughput due to invalid KV optimization.Source: §7.2.2 · p. 10

7.2.2 Results in Long-context Reasoning Scenario.

• Figure 10 covers end-to-end throughput for a single request in both cloud and edge environments.Source: Figure 10 · p. 10

7.3 Evaluation on Speedup and Throughput

• The speedup evaluation uses a 2048 KV budget and DeepSeek-Distill-Llama-8B and Qwen3-8B.Source: §7.3 · p. 10
• 24.89× throughput improvement over full-attention Eager and 2.20× over FlashInfer in cloud; edge speedups are 10.06× over full-attention Eager and 1.17× over ShadowKV.Source: §7.3.1 · p. 10, §7.3.2 · p. 10

7.3.2 Low-end GPU with limited memory in edge.

• With edge GPU memory limited to 4GB, SpeContext is compared against full attention and ShadowKV with offloading.Source: §7.3.2 · p. 10

7.4 Overhead Evaluation

• The retrieval head for Llama3-8B or Qwen3-8B is about 60MB, and EAGLE-3 training for either model takes 24 hours on an RTX 3090 GPU.Source: §7.4 · p. 11

7.5 Ablation Study

• The ablation study uses DeepSeek-Distill-Llama-8B results from Table 3, presented in Figure 11.Source: §7.5 · p. 11

7.5.1 C1: Lightweight retrieval head.

• C1’s speedup comes from both FlashInfer’s backend and sparse attention, which permits more parallel requests through lower GPU memory requirements.Source: §7.5.1 · p. 11
• The asynchronous dataflow’s speedup improvement is mainly from reducing data volume because C1 inference is bound by memory access.The passage attributes the memory-access bound to heavy KV transfer in C1.Source: §7.5.2 · p. 11

7.5.3 C3: Adaptive memory management.

• Adaptive memory management aims to place more KV cache on the GPU to minimize data-transfer overhead and improve speedup over HuggingFace with complete offloading.The strategy is guided by a theoretical model.Source: §7.5.3 · p. 11

8 Conclusion

• The paper analyzes information-focus similarity between distilled and original language models using information theory, then proposes a DLM as the retrieval algorithm.Source: §8 · p. 11
• SpeContext achieves 24.89× throughput improvement in cloud and 10.06× speedup in edge with negligible accuracy loss, pushing the accuracy–speedup Pareto frontier.Source: §8 · p. 11
• SpeContext is an algorithm and system co-design for speculative context sparsity in long-context reasoning.Source: §8 · p. 11

A LongWriter Benchmark

• The detailed LongWriter Benchmark scores are cited as support for the description in Section 7.2.2.Source: A LongWriter Benchmark · p. 12
• Table 4 provides detailed accuracy results on the LongWriter Benchmark.Source: A LongWriter Benchmark · p. 12, Table 4 · p. 13
• The LongWriter Benchmark is the scope of the detailed results presented in Table 4.Source: Table 4 · p. 13