QUOKA: Query-Oriented KV Selection For Efficient LLM Prefill

TL;DR

QUOKA is a training-free, hardware-agnostic sparse attention method that accelerates transformer inference by selectively focusing on low similarity queries and their aligned keys, significantly reducing computation while maintaining accuracy.

Contribution

It introduces a novel query-oriented key-value selection technique that approximates full attention efficiently without training, applicable across hardware platforms.

Findings

3x reduction in time-to-first-token

5x speedup on Nvidia GPUs

up to 7x speedup on Intel CPUs

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.