QUOKA: Query-Oriented KV Selection For Efficient LLM Prefill

Paper: 2602.08722 Authors: Dalton Jones, Junyoung Park, Matthew Morse, Mingu Lee, Chris Lott, Harper Langston Categories: cs.LG

Abstract

QUOKA introduces a training-free, hardware-agnostic sparse attention mechanism for accelerating transformer inference under chunked prefill. The method is grounded in a key observation: queries with low cosine similarity to the mean query vector interact more strongly with a broader set of keys and contribute most significantly to attention logits. By retaining a small set of representative queries and subselecting the most aligned keys, QUOKA achieves a 3x reduction in time-to-first-token (TTFT), 5x attention speedup on Nvidia GPUs, 7x on Intel Xeon CPUs, and near-baseline accuracy while using 88% fewer KV pairs.

Key Contributions

• Query-oriented KV selection: A novel perspective that selects KV pairs based on query characteristics rather than key/value statistics alone
• Cosine similarity insight: Discovery that queries deviating from the mean query are the most informationally rich and interact with more keys
• Training-free deployment: No fine-tuning or retraining required, works as a drop-in replacement
• Hardware-agnostic speedups: Demonstrated gains on both GPU (Nvidia) and CPU (Intel Xeon) platforms
• Extreme KV reduction: Near-baseline accuracy with only 12% of original KV pairs

Background: The Prefill Bottleneck

In autoregressive LLM inference, the prefill phase processes the entire input prompt to populate the KV cache before generating the first output token. For long prompts, this phase dominates latency and is bounded by the quadratic cost of self-attention:

$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

Chunked prefill splits the input into manageable chunks, but each chunk still computes full attention over all accumulated KV pairs. As context grows, this becomes the primary bottleneck.

The QUOKA Method

Step 1: Query Analysis

QUOKA first computes the mean query vector $\bar{q}$ across all queries in a chunk:

$\bar{q} = \frac{1}{n}\sum_{i=1}^{n} q_i$

Then it computes the cosine similarity of each query to this mean:

$s_i = \frac{q_i \cdot \bar{q}}{\|q_i\| \|\bar{q}\|}$

Step 2: Representative Query Selection

Queries with low cosine similarity to the mean (i.e., outlier queries) are selected as representatives. These queries are the most “demanding” in terms of attention — they attend to a wider range of keys and contribute disproportionately to the final attention output.

Step 3: Key Subselection

For each representative query, QUOKA identifies the most aligned keys using dot-product scoring. Only these high-alignment KV pairs are retained for the full attention computation. The remaining queries (those close to the mean) can be adequately served by this reduced KV set.

Why Outlier Queries Matter

The core insight is counterintuitive: queries that are most similar to each other (close to the mean) tend to produce diffuse, low-entropy attention distributions. Outlier queries produce sharper, more focused attention patterns that capture the critical information flow. By ensuring these outlier queries have access to their most relevant keys, QUOKA preserves the essential attention computation while discarding redundant KV interactions.

Results

Metric	Value
TTFT Reduction	3x
Attention Speedup (Nvidia GPU)	5x
Attention Speedup (Intel Xeon CPU)	7x
KV Pairs Used	12% (88% reduction)
Accuracy Degradation	Near-baseline

• Consistent improvements across multiple model sizes and architectures
• Benefits compound with longer context lengths
• Compatible with existing KV cache compression techniques

Significance

1. Practical inference acceleration: Directly addresses the prefill bottleneck that limits real-world LLM deployment
2. No training cost: ike learned sparse attention methods, QUOKA requires zero additional training
3. Cross-platform: Rare to see a method that accelerates both GPU and CPU inference
4. Composable: Can be combined with other efficiency techniques (quantization, pruning, etc.)
5. Theoretical insight: The query-centric view of attention sparsity opens new research directions for understanding which computations in attention are truly necessary