Using Span Queries to Optimize for Cache and Attention Locality

TL;DR

This paper introduces span queries as a flexible interface for inference servers, enabling optimization for various workloads like chat, RAG, and reasoning, leading to significant performance improvements and better cache and attention locality.

Contribution

The paper proposes span queries as a unified, expressive interface for inference workloads, with automatic optimization for cache and attention locality, and demonstrates practical implementation and performance gains.

Findings

Achieved 10-20x reductions in TTFT for non-chat workloads.

Enabled high-performance execution with a small code change in vLLM.

Improved attention locality to outperform larger models on accuracy.

Abstract

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.