TLDR: Span queries introduce a new declarative interface for AI inference servers, allowing them to understand if the order of inputs matters (commutativity). This enables automatic optimization for KV cache locality, reducing time-to-first-token by 10-20x for non-chat workloads like RAG and agentic systems, and improving attention locality to overcome the “lost-in-the-middle” problem, even allowing smaller models to outperform larger ones.
The landscape of artificial intelligence is rapidly evolving beyond simple chat interactions. Modern AI applications, from sophisticated retrieval-augmented generation (RAG) systems to complex agentic workflows, demand more from their underlying infrastructure. However, current AI inference servers, like vLLM, are primarily designed and optimized for traditional chat completion tasks. This mismatch often leads to inefficiencies, particularly in how these servers manage their Key-Value (KV) cache, a critical component for accelerating transformer-based language models.
A new research paper, “Using Span Queries to Optimize for Cache and Attention Locality”, introduces a novel concept called “span queries” to address these challenges. Authored by Paul Castro, Nick Mitchell, Nathan Ordonez, Thomas Parnell, Mudhakar Srivatsa, and Antoni Viros i Martin, the paper proposes a generalized interface for inference servers that can accommodate a wide array of non-chat AI workloads, significantly boosting performance.
Understanding the Problem: Beyond Chat
In chat applications, the conversation history builds up linearly, meaning previous parts of the conversation (the “prefix”) remain constant and can be efficiently reused from the KV cache. This is what vLLM excels at. However, RAG and agentic workloads often combine information from various sources or generate multiple outputs that are then processed further. In these scenarios, the order of inputs can change, or new information might be inserted in the middle, breaking the simple prefix-based reuse pattern that vLLM relies on. This leads to a drastically reduced KV cache hit rate, essentially forcing the system to recompute information it has already processed, wasting valuable resources and time.
Introducing Span Queries: A Flexible Approach
The core idea behind span queries is to provide a declarative way to tell the inference server how different parts of an AI task relate to each other, especially regarding whether the order of inputs matters. This concept is called “commutativity.” For instance, in a chat, the order of messages is crucial (non-commutative). But in a RAG system, if you retrieve several fragments of information, their exact order might not always be critical for the final answer (commutative). Span queries allow developers to express these commutativity constraints, enabling the server to make smarter optimization decisions.
The paper demonstrates that chat, RAG, inference-time scaling, and agentic workloads can all be expressed using this unified “span query” framework. These queries are essentially “expression trees” of inference calls, linked by these commutativity constraints.
Optimizing for KV Cache Locality
One of the primary benefits of span queries is their ability to dramatically improve KV cache locality. By understanding which parts of the input can be reordered without affecting the outcome, the system can rearrange computations to maximize the reuse of cached data. The researchers implemented a small change to vLLM (affecting only 492 lines of code) to enable high-performance execution of span queries.
A key challenge addressed is the “dual output paradox,” where the model server sends one output to the client but caches a slightly different version internally. Span queries, through a high-level optimizer, help resolve this by aligning the query structure with the server’s caching behavior. Furthermore, a new algorithm called CIDRA (Concurrent In-place Duplicating ReROPE Algorithm) was developed to efficiently reposition cached KV vectors when their context changes, ensuring that reused data is correctly aligned for new requests.
The results are impressive: span queries achieved 10-20x reductions in “Time to First Token” (TTFT) for non-chat use cases. Even when there’s a cache miss, span queries showed a 3x reduction in prefill load due to optimized attention sparsity, significantly outperforming existing solutions like CacheBlend.
Improving Attention Locality: Solving “Lost in the Middle”
Beyond cache optimization, span queries offer a surprising benefit: they can also improve “attention locality,” addressing the “lost-in-the-middle” problem. This phenomenon describes how smaller models often struggle to focus on important information buried in the middle of long inputs, performing better when key data is at the beginning or end.
By treating the span query’s expression tree like a “divide-and-conquer” or “map/reduce” structure, the high-level optimizer can transform complex queries into a series of smaller, more manageable steps. For example, an 8-way judge/generator task can be broken down into three 2-way judge steps. This optimization helps models attend more effectively to relevant information. Remarkably, the paper shows that a 2-billion parameter model using an attention-optimized span query vastly outperformed the accuracy of an 8-billion parameter model running on a stock inference server, highlighting the power of this approach without requiring any changes to the model server itself.
Also Read:
- Boosting LLM Performance: How Processing-Near-Memory Redefines KV-Cache Management
- GLM: Enhancing LLM Reasoning on Graphs with Multi-Agent Collaboration and Optimized Serving
A New Era for AI Inference
Span queries represent a significant step forward in optimizing AI inference for the diverse and complex workloads emerging today. By providing a declarative framework that expresses commutativity constraints, they enable intelligent optimizations for both KV cache and attention locality, leading to substantial performance gains and even improved accuracy. This work paves the way for more efficient and capable AI systems, moving beyond the limitations of chat-centric server designs.
