提出跨度查询统一多种推理场景,显著提升缓存与注意力效率。
Using Span Queries to Optimize for Cache and Attention Locality
- 用带交换性约束的表达式树抽象推理任务,支持多类工作负载。
- 在非对话场景下实现TTFT降低10-20倍,性能大幅提升。
- 兼容现有系统改造小(仅492行),适合部署在vLLM等框架中。
客户端正从简单对话扩展到多样化的推理时缩放和深度推理技术,但推理服务器仍以对话生成为核心优化。已有研究显示,若服务器适配非对话场景,可显著提升键值缓存命中率。然而这些方案通常只针对单一场景(如RAG)。本文提出跨度查询(span query),将对话、RAG、推理时缩放及代理型工作负载统一为同一接口。关键发现是:此前被忽视的本质区别在于输入顺序是否可交换——对话不可交换,而RAG常可交换。跨度查询以表达式树形式组织推理调用,并加入交换性约束。我们定义其语法与语义,展示如何自动优化以提升键值缓存局部性。仅修改vLLM中492行代码,即可高效执行跨度查询。实验表明,在两个独立非对话场景中,跨度查询使首次输出时间(TTFT)降低10-20倍。进一步地,通过优化注意力局部性,可缓解“中间丢失”问题。在20亿参数模型上,经注意力优化的跨度查询,其准确率远超使用80亿参数模型的标准推理服务器。
原文摘要 · Abstract (English)
Clients are evolving beyond chat completion, and now include a variety of innovative inference-time scaling and deep reasoning techniques. At the same time, inference servers remain heavily optimized for chat completion. Prior work has shown that large improvements to KV cache hit rate are possible if inference servers evolve towards these non-chat use cases. However, they offer solutions that are also optimized for a single use case, RAG. In this paper, we introduce the span query to generalize the interface to the inference server. We demonstrate that chat, RAG, inference-time scaling, and agentic workloads can all be expressed as span queries. We show how the critical distinction that had been assumed by prior work lies in whether the order of the inputs matter -- do they commute? In chat, they do not. In RAG, they often do. This paper introduces span queries, which are expression trees of inference calls, linked together with commutativity constraints. We describe span query syntax and semantics. We show how they can be automatically optimized to improve KV cache locality. We show how a small change to vLLM (affecting only 492 lines) can enable high-performance execution of span queries. Using this stack, we demonstrate that span queries can achieve 10-20x reductions in TTFT for two distinct non-chat use cases. Finally, we show that span queries can also be optimized to improve attention locality, so as to avoid the so-called lost-in-the-middle problem. We demonstrate that an attention-optimized span query on a 2b parameter model vastly outperforms the accuracy of a stock inference server using an 8b model.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。