一句话:提出SpeContext,以蒸馏语言模型作为检索算法,协同算法、系统与编译优化,大幅提升长上下文推理吞吐量。
长上下文推理中,LLM逐步生成导致KV缓存激增,成为部署瓶颈。现有KV缓存优化直接用于推理时存在三个问题:逐层检索同步开销随深度增长(高达60%延迟)、新生成KV完全保留导致无效、序列长度微增即引发>80%性能下降。本文指出检索目标与知识蒸馏目标一致,从信息论角度分析蒸馏语言模型(DLM)与原始LLM的信息焦点相似性,提出以DLM作为检索算法。基于此设计SpeContext,进行算法与系统协同:算法层面提出基于DLM注意力权重的轻量检索头,剪枝冗余后参数减少>90%;系统层面设计弹性加载的异步预取数据流,重叠KV检索与LLM计算;编译层面构建理论内存模型并实现自适应内存管理,最大化GPU利用率。在云和边缘环境评估,相比Huggingface和FlashInfer,云上吞吐量提升最高24.89倍和2.19倍,边缘加速10.06倍和8.02倍,精度损失可忽略,推动了精度-吞吐量帕累托前沿。
As test-time scaling in large language model(LLM) reasoning has been proven effective in enhancing the model performance through step-by-step generation, this long-context generation incurs substantial Key-Value(KV) cache, posing a critical bottleneck for practical applications deployment(e.g., Agents). While recent KV cache optimizations perform well in the long-context input scenario, the following problems remain unsolved if directly applied to long-context reasoning. (1) Time-consuming layer-wise retrieval operation. The retrieval operation, which selects the important KV pairs in each layer, brings the synchronization overhead that scales with model depth due to the data dependency, resulting in up to 60% latency overhead. (2) Complete retention of the newly generated KV cache. Existing works designed for long-context input choose to retain the KV pair of newly generated tokens to avoid repeated, time-consuming processing on the KV cache, rendering them ineffective in long-context reasoning. (3) Performance degradation with a tiny increase in sequence length. Existing offloading strategies determined before inference cannot adapt to the increasing sequence length, resulting in >80% performance degradation with a tiny increase in sequence length. In this paper, we point out that the objective of the retrieval algorithms is to align with the LLM, which is similar to the objective of knowledge distillation in LLMs. We analyze the similarity in information focus between the distilled language model(DLM) and the original LLM from the perspective of information theory, and thus propose a novel paradigm that leverages a DLM as the retrieval algorithm. Based on the insight, we present SpeContext, an algorithm and system co-design for long-context reasoning. (1) At the algorithm level, SpeContext proposes lightweight retrieval head based on the head-level attention weights of DLM, achieving >90% parameters reduction by pruning the redundancy. (2) At the system level, SpeContext designs an asynchronous prefetch dataflow via the elastic loading strategy, effectively overlapping KV cache retrieval with the LLM computation. (3) At the compilation level, SpeContext constructs the theoretical memory model and implements an adaptive memory management system to achieve acceleration by maximizing GPU memory utilization. We deploy and evaluate SpeContext in two resource-constrained environments, cloud and edge. Extensive experiments show that, compared with the Huggingface and FlashInfer framework, SpeContext achieves up to 24.89× and 2.19× throughput improvement in cloud and 10.06× and 8.02× speedup in edge with negligible accuracy loss, pushing the Pareto frontier of accuracy and throughput.
DOI 原文 · arXiv · PDF(开放获取) · DBLP
@inproceedings{abs-2512-00722,
title = {<i>SpeContext:</i>
Enabling Efficient Long-context Reasoning with Speculative Context Sparsity in LLMs},
author = {Jiaming Xu and Jiayi Pan and Hanzhen Wang and Yongkang Zhou and Jiancai Ye and Yu Wang and Guohao Dai},
booktitle = {ASPLOS},
year = {2026},
doi = {10.1145/3779212.3790224}
}