QUOKA: Query-Oriented KV Selection For Efficient LLM Prefill

TL;DR

QUOKA tackles prefill latency in long-context LLM inference by introducing a training-free, hardware-agnostic sparse attention method that is tailored for chunked prefill. By selecting a small set of representative queries via cosine-dissimilarity to a mean query and scoring their interactions with keys using cosine similarity, QUOKA subselects KV tokens and aggregates across queries and KV groups to closely approximate full attention while reducing compute and memory. Across Needle-in-a-Haystack, RULER, LongBench, and Math500, QUOKA achieves near-baseline accuracy with substantial latency reductions, delivering up to 5x module-level speedups on GPUs and 3x TTFT improvements on CPUs, and reducing KV usage by about 88%. The method generalizes across decoder-only LLM families and remains robust to hyperparameter choices, highlighting its practicality for deploying long-context models on resource-constrained hardware. Overall, QUOKA provides a portable, efficient approach to sparse attention that preserves performance while enabling faster prefill in diverse deployment environments, facilitating more scalable and cost-effective LLM inference.

Abstract

We present QUOKA: Query-oriented KV selection for efficient attention, a training-free and hardware agnostic sparse attention algorithm for accelerating transformer inference under chunked prefill. While many queries focus on a smaller group of keys in the attention operator, we observe that queries with low cosine similarity with respect to the mean query interact more strongly with more keys and have the greatest contribution to final attention logits. By prioritizing these low cosine similarity queries, the behavior of full attention during the prefill stage can be closely approximated. QUOKA leverages this observation, accelerating attention by (1) first retaining a small set of representative queries and (2) then subselectin the keys most aligned with those queries. Through experiments on Needle-In-A-Haystack, LongBench, RULER, and Math500, we show that, while realizing a 3x reduction in time-to-first-token, 5x speedup in attention on Nvidia GPUs and up to nearly a 7x speedup on Intel Xeon CPUs, QUOKA achieves near-baseline accuracy, utilizing 88% fewer key-value pairs per attention evaluation.

Figures (7)

• Figure 1: Overview of chunked prefill with QuoKA:(a) A prompt of 6 tokens is divided into three chunks of 2 tokens each. For each chunk, QuoKA subselects the KV cache and feeds the reduced cache into a dense attention kernel. (b) Subselection is performed by applying query subselection based on cosine dissimilarity, followed by key subselection using query–key cosine similarity.
• Figure 2: Empirical observations from Llama 3.2-3B-Instruct, layer 0 head 11. (a) Attention map $A$. (b) PCA visualization of $Q$ and $K$, showing that queries with higher $S_q$ lie closer to the keys. (c) Correlation between $S_q$ and $\max_k(A)$, indicating stronger key interactions for higher $S_q$ queries.
• Figure 3: Distribution of attention score max deviation from mean along query and head dimension.
• Figure 4: Accuracy across document length and needle depth for NIAH with $B_\text{SA}=2048$ and $B_{\text{CP}}=128$. Results for additional sparse attention methods are shown in \ref{['fig:needle_haystack_extended']}.
• Figure 5: Relative speedup of attention and TTFT compared to dense attention baseline using $B_\text{CP}=128$ on different hardware.

Theorems & Definitions (2)

• Theorem 1
• proof