QUOKA: Query-Oriented KV Selection For Efficient LLM Prefill
This addresses the computational bottleneck in large language model inference for applications requiring fast response times, though it is an incremental improvement on existing sparse attention methods.
The paper tackles the problem of inefficient attention computation in transformer inference during prefill by introducing QUOKA, a training-free sparse attention algorithm that selects key-value pairs based on query similarity, achieving a 3x reduction in time-to-first-token and up to 7x speedup while maintaining near-baseline accuracy with 88% fewer key-value pairs.
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.