QUOKA Explained: Choosing Which KV History Attention Should Read

Watch a coding agent work and you'll notice it rarely understands a task from one file. It opens several files across a repository, follows how the pieces connect, traces dependencies, and only then edits something. Sometimes that means shipping a startling amount of context to the underlying model.

Which raises a question. Vendors advertise context windows of hundreds of thousands, even a million tokens. What does processing that much text really involve, and where does the cost come from? No single trick answers it. Long-context inference touches how positions are encoded, how attention is computed, how memory is managed and how workloads are scheduled.

One paper that tackles a narrow slice of the problem is QUOKA, short for Query-Oriented KV Selection for Efficient LLM Prefill. It's about making attention cheaper while the model is reading a long input prompt. The core idea is intuitive, but seeing why it works takes three pieces of background: the KV cache, the split between prefill and decode, and chunked prefill.

The KV cache in brief

Inside a Transformer, every token produces a query (Q), a key (K) and a value (V) for attention. A rough way to think about them:

• The query asks what information the token is looking for.
• The key advertises what kind of information the token holds.
• The value is what gets handed over if the token turns out to be relevant.

When generating autoregressively, keys and values computed earlier can be stored in a KV cache and reused later instead of being recomputed. In outline:

Previously processed tokens
          ↓
     K,V representations
          ↓
        KV cache
          ↓
reused by later attention

The longer the sequence gets, the more stored K and V entries there are. That matters for QUOKA, and to see why you first need the two phases of inference.

Prefill: reading your prompt

When you send a model a prompt, whether a one-line question or a mountain of code, the input tokens have to be processed before any output starts. That phase is prefill.

Input prompt
     ↓
   PREFILL
     ↓
process input tokens
     ↓
produce their K,V representations

Decode: writing the answer

Once the input is processed, the model produces its reply one token at a time. That's decode.

Processed input
      ↓
generate token 1
      ↓
generate token 2
      ↓
generate token 3
      ↓
     ...

The distinction matters because choosing which parts of a long history deserve attention is fairly easy when there's a single new query. It's a different problem when a whole chunk contains hundreds or thousands of queries at once. One more concept is needed before getting there.

Chunked prefill

Handling a very long prompt in a single prefill pass is expensive. Inference systems often use chunked prefill instead, which splits the input into smaller chunks processed in order. For a 3,000-token prompt:

3,000-token prompt
Chunk 0 → tokens    0–999
Chunk 1 → tokens 1000–1999
Chunk 2 → tokens 2000–2999

Each chunk after the first has more history behind it:

Chunk 0 → no previous chunk
Chunk 1 → 1,000 historical KV positions
Chunk 2 → 2,000 historical KV positions
Chunk 3 → 3,000 historical KV positions
...

Chunking makes a long prefill tolerable, yet every later chunk can attend to an ever larger pile of cached history. That pile is what QUOKA goes after.

The observation behind the method

Picture a later chunk being processed with 100,000 historical KV positions already stored. Dense attention lets the current queries look at every one of them.

Current chunk
     ↓
K1 K2 K3 K4 K5 .............. K100000

Do the current queries really need all 100,000? Probably not. Some keys will matter a lot for what this chunk is doing, and many will contribute almost nothing. So there's an opening:

100,000 historical KV
          ↓
find the important ones
          ↓
much smaller KV subset
          ↓
attention

That's the basic intuition. Instead of running costly attention across the whole history, pick out a smaller set of historical KVs likely to matter for the current chunk. Simple enough, but it prompts a harder question: who decides which KVs matter?

One query can choose for itself, but a chunk has thousands

In token-by-token generation there's one new query position at a time, and that query can effectively ask which historical keys are relevant to it.

one query
    ↓
Which historical keys
are relevant to me?

During chunked prefill, many query positions are present at once:

q1
q2
q3
...
q4096

Different queries may care about entirely different parts of the history. You could collapse them into one average query, but that risks losing unusual yet important needs. Suppose most queries point the same way:

q1  ↗
q2  ↗
q3  ↗
q4  ↗
q5  →

The average will be dominated by q1 through q4. If q5 is the only query that needs some particular key, the averaging buries that interaction.

Pick the queries that disagree with the crowd

QUOKA first computes the mean query of the current chunk, then measures each query's cosine similarity to that mean. For example:

Similarity to mean
q1 → 0.96
q2 → 0.94
q3 → 0.91
q4 → 0.35
q5 → 0.18

Rather than keeping the queries closest to the mean, it prioritizes the ones least similar to it:

q1 → 0.96
q2 → 0.94
q3 → 0.91
q4 → 0.35  ← keep
q5 → 0.18  ← keep

The reasoning is straightforward. Highly similar queries point in roughly the same direction, so keeping many of them gives redundant information. The low-similarity queries are the ones saying they're looking for something different, and they're exactly the ones likely to uncover keys the majority would overlook. So thousands of queries become a small, diverse set:

thousands of queries
        ↓
a small set of diverse queries

Letting those queries score the history

Now the diverse queries scan the stored keys, each asking which parts of the past look useful to it. QUOKA again uses cosine similarity, this time between the selected queries and the historical keys. Say three queries were chosen:

qA
qB
qC

and there are thousands of keys:

K1, K2, K3, ... K10000

Each selected query scores every key by cosine similarity:

            K1      K2     K3     K4
qA          .90    .10    .20    .15
qB          .10    .85    .15    .20
qC          .20    .10    .95    .10

That's useful because different queries find different parts of the past. Here qA cares strongly about K1, qB about K2 and qC about K3.

But now each key has several importance scores, one per selected query, while QUOKA needs a single number to decide whether the key survives.

Why take the maximum

Look at K1:

qA → K1 = .90
qB → K1 = .10
qC → K1 = .20

Averaging would be the obvious option:

(.90 + .10 + .20) / 3 = .40

That makes K1 look moderately important and hides something valuable, namely that qA considers it extremely important. A key doesn't have to matter to every query to deserve a place. It may hold exactly what one query needs. So QUOKA takes the maximum across the selected queries:

K1 → max(.90, .10, .20) = .90
K2 → max(.10, .85, .10) = .85
K3 → max(.20, .15, .95) = .95
K4 → max(.15, .20, .10) = .20

Every key ends up with one final score, and strong query-key interactions are preserved. If even one representative query strongly needs a key, the others shouldn't be allowed to wash that signal away.

From scores to a smaller KV set

With scores in hand, QUOKA ranks the keys and keeps only the top ones, according to its sparse-attention budget.

Historical keys
K1 → .90  ✓
K2 → .85  ✓
K3 → .95  ✓
K4 → .20
K5 → .12
...
K10000 → .08

        ↓
      Top-K
        ↓

small set of selected keys
K1 → .90  ✓
K2 → .85  ✓
K3 → .95  ✓

Whenever a key position is picked, its matching value comes along. If QUOKA keeps these keys:

K3, K1, K2

it also keeps these values:

V3, V1, V2

The result is a smaller set of K,V pairs that takes part in attention. In essence, rather than forcing attention to consider every historical position, QUOKA first builds a much smaller candidate set based on what the current queries actually seem to need.

A wrinkle: grouped-query attention

Modern models don't necessarily have as many query heads as KV heads, and that's where grouped-query attention (GQA) comes in. Suppose a model has four query heads and two KV heads:

Q head 1 ─┐
          ├── KV head 1
Q head 2 ─┘
Q head 3 ─┐
          ├── KV head 2
Q head 4 ─┘

Instead of every query head owning its own key and value head, several query heads share one. That affects QUOKA, because in a GQA model each selected token has several query-head representations. For one selected token qA you might have:

qA_head1
qA_head2
qA_head3
qA_head4

The first two interact with KV head 1 and the other two with KV head 2. Scoring each head separately and averaging afterward would look like this:

qA_head1 × (K head 1)ᵀ ─┐
                         ├── average scores
qA_head2 × (K head 1)ᵀ ─┘

qA_head3 × (K head 2)ᵀ ─┐
                         ├── average scores
qA_head4 × (K head 2)ᵀ ─┘

That repeats expensive query-key comparisons against the same KV head. QUOKA exploits a simple algebraic fact and averages the normalized query representations first:

qA_head1 ─┐
          ├── MEAN → q̄A_1 → × (K head 1)ᵀ
qA_head2 ─┘

qA_head3 ─┐
          ├── MEAN → q̄A_2 → × (K head 2)ᵀ
qA_head4 ─┘

This works because (Q1 K^T + Q2 K^T) / 2 equals ((Q1 + Q2) / 2) K^T. Rather than two scoring multiplications followed by an average, you average the queries first and multiply once. That trick goes by the name pre-aggregation in the QUOKA paper.

MEAN and MAX do different jobs

It's easy to mix up this MEAN with the MAX from earlier, but they work along different dimensions.

MEAN works across the query heads sharing one KV head, producing a cheaper stand-in query for scoring that head:

qA_head1 ─┐
          ├── MEAN
qA_head2 ─┘

MAX runs across different selected query tokens, protecting a key when even one diverse query finds it highly relevant:

qA → K50 = .20
qB → K50 = .91
qC → K50 = .30
                 ↓
                MAX
                 ↓
          K50 score = .91

To keep them straight:

MEAN across related GQA heads
          ↓
make KV selection cheaper

MAX across diverse query tokens
          ↓
preserve important query–key interactions

Two more details matter. Averaged query heads serve strictly to choose KV positions. The model's original query heads stay in place and carry out the real sparse attention over whatever was chosen, which makes pre-aggregation a shortcut for selection alone.

Note too that nothing is evicted. The complete historical KV cache remains, the selected KVs are just the subset used in the current attention step, and a new chunk can choose a different subset based on its own queries. The saving is in how much history joins the attention computation, not in the size of the stored cache.

Where to find the numbers

Benchmark results are given in three tables in the paper, which is the place to check how the method stacks up against alternatives. Read it on arXiv.