QuoKA: Query-Oriented KV Selection for Efficient LLM Prefill

Dalton Jones, Junyoung Park, Matthew J Morse, Mingu Lee, Matthew Harper Langston, Christopher Lott

Keywords: Efficient LLM Inference, LLM Prefill Acceleration, Sparse Attention, KV Cache Subselection, Training-Free

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 subselecting the keys most aligned with those queries. Through experiments on Needle-In-A-Haystack, LongBench, RULER, and Math500, we show that, while realizing a 3× reduction in time-to-first-token, 5× speedup in attention on Nvidia GPUs and up to nearly a 7× speedup on Intel Xeon CPUs, QuoKA achieves near-baseline accuracy, utilizing 88% fewer key-value pairs per attention evaluation.

Primary Area: foundation or frontier models, including LLMs

Submission Number: 14170