QUOKA erklärt: Welche KV-Historie die Attention lesen sollte

Wer einem Coding-Agent bei der Arbeit zusieht, merkt schnell: Er versteht eine Aufgabe selten anhand einer einzigen Datei. Er öffnet mehrere Dateien quer durchs Repository, verfolgt, wie die Teile zusammenhängen, geht Abhängigkeiten nach und ändert erst dann etwas. Manchmal heißt das, dem Modell erstaunlich viel Kontext zu schicken.

Das wirft eine Frage auf. Anbieter werben mit Kontextfenstern von Hunderttausenden, ja einer Million Tokens. Was bedeutet es wirklich, so viel Text zu verarbeiten, und woher kommen die Kosten? Kein einzelner Trick beantwortet das. Inferenz mit langem Kontext berührt die Kodierung von Positionen, die Berechnung der Attention, die Speicherverwaltung und die Planung von Workloads.

Ein Paper, das sich einen schmalen Ausschnitt des Problems vornimmt, ist QUOKA, kurz für Query-Oriented KV Selection for Efficient LLM Prefill. Es geht darum, Attention günstiger zu machen, während das Modell einen langen Eingabeprompt liest. Die Kernidee ist intuitiv, aber um zu sehen, warum sie funktioniert, braucht man drei Bausteine Hintergrund: den KV-Cache, die Trennung von Prefill und Decode und Chunked Prefill.

Der KV-Cache in Kürze

In einem Transformer erzeugt jedes Token für die Attention eine Query (Q), einen Key (K) und einen Value (V). Grob kann man sie so verstehen:

• Die Query fragt, welche Information das Token sucht.
• Der Key zeigt an, welche Art von Information das Token enthält.
• Der Value ist das, was übergeben wird, wenn sich das Token als relevant erweist.

Bei autoregressiver Generierung lassen sich früher berechnete Keys und Values in einem KV-Cache speichern und später wiederverwenden, statt sie neu zu berechnen. Im Umriss:

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

Je länger die Sequenz wird, desto mehr K- und V-Einträge sind gespeichert. Das ist für QUOKA wichtig, und um zu verstehen warum, braucht man zuerst die zwei Phasen der Inferenz.

Prefill: Ihren Prompt lesen

Wenn Sie einem Modell einen Prompt schicken, ob eine einzeilige Frage oder einen Berg Code, müssen die Eingabetokens verarbeitet werden, bevor irgendeine Ausgabe beginnt. Diese Phase heißt Prefill.

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

Decode: die Antwort schreiben

Ist die Eingabe verarbeitet, erzeugt das Modell seine Antwort Token für Token. Das ist Decode.

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

Die Unterscheidung zählt, weil es ziemlich leicht ist zu entscheiden, welche Teile einer langen Historie Aufmerksamkeit verdienen, wenn es nur eine neue Query gibt. Ganz anders sieht es aus, wenn ein ganzer Chunk Hunderte oder Tausende Queries auf einmal enthält. Vorher braucht es noch ein Konzept.

Chunked Prefill

Einen sehr langen Prompt in einem einzigen Prefill-Durchlauf zu verarbeiten ist teuer. Inferenzsysteme nutzen deshalb oft Chunked Prefill: Die Eingabe wird in kleinere Chunks geteilt, die der Reihe nach verarbeitet werden. Bei einem Prompt mit 3.000 Tokens:

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

Hinter jedem Chunk nach dem ersten steht mehr Historie:

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 macht ein langes Prefill erträglich, doch jeder spätere Chunk kann auf einen immer größeren Berg gecachter Historie zugreifen. Genau diesen Berg nimmt sich QUOKA vor.

Die Beobachtung hinter der Methode

Stellen Sie sich vor, ein später Chunk wird verarbeitet, während bereits 100.000 historische KV-Positionen gespeichert sind. Dichte Attention lässt die aktuellen Queries jede einzelne davon betrachten.

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

Brauchen die aktuellen Queries wirklich alle 100.000? Wahrscheinlich nicht. Manche Keys werden für das, was dieser Chunk tut, sehr wichtig sein, viele tragen fast nichts bei. Es gibt also eine Gelegenheit:

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

Das ist die Grundintuition. Statt teure Attention über die gesamte Historie laufen zu lassen, wählt man eine kleinere Menge historischer KVs aus, die für den aktuellen Chunk wahrscheinlich wichtig sind. Einfach genug, aber es führt zu einer schwierigeren Frage: Wer entscheidet, welche KVs wichtig sind?

Eine Query kann selbst wählen, ein Chunk hat Tausende

Bei der Generierung Token für Token gibt es jeweils eine neue Query-Position, und diese Query kann im Grunde selbst fragen, welche historischen Keys für sie relevant sind.

one query
    ↓
Which historical keys
are relevant to me?

Beim Chunked Prefill sind viele Query-Positionen gleichzeitig vorhanden:

q1
q2
q3
...
q4096

Verschiedene Queries können sich für völlig unterschiedliche Teile der Historie interessieren. Man könnte sie zu einer Durchschnitts-Query zusammenfassen, riskiert dann aber, ungewöhnliche, doch wichtige Bedürfnisse zu verlieren. Angenommen, die meisten Queries zeigen in dieselbe Richtung:

q1  ↗
q2  ↗
q3  ↗
q4  ↗
q5  →

Der Durchschnitt wird von q1 bis q4 dominiert. Ist q5 die einzige Query, die einen bestimmten Key braucht, begräbt die Mittelung diese Interaktion.

Die Queries wählen, die der Masse widersprechen

QUOKA berechnet zuerst die mittlere Query des aktuellen Chunks und misst dann die Kosinus-Ähnlichkeit jeder Query zu diesem Mittelwert. Zum Beispiel:

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

Statt die Queries zu behalten, die dem Mittelwert am nächsten liegen, priorisiert es die am wenigsten ähnlichen:

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

Die Überlegung ist einfach. Sehr ähnliche Queries zeigen ungefähr in dieselbe Richtung, viele davon zu behalten liefert also redundante Information. Die Queries mit geringer Ähnlichkeit sind diejenigen, die sagen, dass sie etwas anderes suchen, und genau sie dürften Keys aufspüren, die die Mehrheit übersehen würde. So werden aus Tausenden Queries eine kleine, vielfältige Menge:

thousands of queries
        ↓
a small set of diverse queries

Diese Queries bewerten die Historie

Jetzt durchsuchen die vielfältigen Queries die gespeicherten Keys, und jede fragt, welche Teile der Vergangenheit für sie nützlich aussehen. QUOKA nutzt wieder die Kosinus-Ähnlichkeit, diesmal zwischen den ausgewählten Queries und den historischen Keys. Angenommen, drei Queries wurden gewählt:

qA
qB
qC

und es gibt Tausende Keys:

K1, K2, K3, ... K10000

Jede ausgewählte Query bewertet jeden Key per Kosinus-Ähnlichkeit:

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

Das ist nützlich, weil verschiedene Queries verschiedene Teile der Vergangenheit finden. Hier interessiert sich qA stark für K1, qB für K2 und qC für K3.

Nun hat aber jeder Key mehrere Wichtigkeitswerte, einen pro ausgewählter Query, während QUOKA eine einzige Zahl braucht, um zu entscheiden, ob der Key bleibt.

Warum das Maximum

Betrachten wir K1:

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

Mitteln wäre die naheliegende Option:

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

Damit wirkt K1 mäßig wichtig, und etwas Wertvolles geht verloren: qA hält ihn für extrem wichtig. Ein Key muss nicht für jede Query zählen, um einen Platz zu verdienen. Er kann genau das enthalten, was eine Query braucht. Deshalb nimmt QUOKA das Maximum über die ausgewählten 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

Jeder Key erhält einen finalen Wert, und starke Query-Key-Interaktionen bleiben erhalten. Wenn auch nur eine repräsentative Query einen Key dringend braucht, dürfen die anderen dieses Signal nicht verwässern.

Von Werten zu einer kleineren KV-Menge

Mit den Werten in der Hand sortiert QUOKA die Keys und behält nur die besten, gemäß seinem Budget für Sparse Attention.

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  ✓

Wird eine Key-Position ausgewählt, kommt der passende Value mit. Behält QUOKA diese Keys:

K3, K1, K2

dann behält es auch diese Values:

V3, V1, V2

Das Ergebnis ist eine kleinere Menge von K,V-Paaren, die an der Attention teilnimmt. Im Kern zwingt QUOKA die Attention nicht, jede historische Position zu betrachten, sondern baut zuerst eine viel kleinere Kandidatenmenge auf, je nachdem, was die aktuellen Queries tatsächlich zu brauchen scheinen.

Ein Haken: Grouped-Query Attention

Moderne Modelle haben nicht unbedingt so viele Query-Heads wie KV-Heads, und hier kommt Grouped-Query Attention (GQA) ins Spiel. Angenommen, ein Modell hat vier Query-Heads und zwei KV-Heads:

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

Statt dass jeder Query-Head einen eigenen Key- und Value-Head hat, teilen sich mehrere Query-Heads einen. Das betrifft QUOKA, denn in einem GQA-Modell hat jedes ausgewählte Token mehrere Query-Head-Repräsentationen. Für ein ausgewähltes Token qA könnte das so aussehen:

qA_head1
qA_head2
qA_head3
qA_head4

Die ersten beiden interagieren mit KV-Head 1, die anderen beiden mit KV-Head 2. Jeden Head einzeln zu bewerten und danach zu mitteln sähe so aus:

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

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

Das wiederholt teure Query-Key-Vergleiche gegen denselben KV-Head. QUOKA nutzt eine einfache algebraische Tatsache und mittelt zuerst die normalisierten Query-Repräsentationen:

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

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

Das funktioniert, weil (Q1 K^T + Q2 K^T) / 2 gleich ((Q1 + Q2) / 2) K^T ist. Statt zwei Bewertungs-Multiplikationen mit anschließender Mittelung mittelt man zuerst die Queries und multipliziert einmal. Im QUOKA-Paper heißt dieser Kniff Pre-Aggregation (Voraggregation).

MEAN und MAX haben unterschiedliche Aufgaben

Man verwechselt dieses MEAN leicht mit dem MAX von vorhin, aber sie arbeiten entlang verschiedener Dimensionen.

MEAN arbeitet über die Query-Heads, die sich einen KV-Head teilen, und liefert eine günstigere Ersatz-Query, um diesen Head zu bewerten:

qA_head1 ─┐
          ├── MEAN
qA_head2 ─┘

MAX läuft über verschiedene ausgewählte Query-Tokens und schützt einen Key, wenn auch nur eine abweichende Query ihn für hoch relevant hält:

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

Zur Unterscheidung:

MEAN across related GQA heads
          ↓
make KV selection cheaper

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

Zwei weitere Details zählen. Gemittelte Query-Heads dienen ausschließlich dazu, KV-Positionen auszuwählen. Die ursprünglichen Query-Heads des Modells bleiben bestehen und führen die eigentliche Sparse Attention über die Auswahl aus, die Voraggregation ist also nur eine Abkürzung für die Auswahl.

Beachten Sie außerdem, dass nichts verdrängt wird. Der vollständige historische KV-Cache bleibt erhalten, die ausgewählten KVs sind nur die Teilmenge für den aktuellen Attention-Schritt, und ein neuer Chunk kann anhand seiner eigenen Queries eine andere Teilmenge wählen. Die Einsparung liegt darin, wie viel Historie in die Attention-Berechnung eingeht, nicht in der Größe des gespeicherten Caches.

Wo die Zahlen stehen

Benchmark-Ergebnisse stehen in drei Tabellen im Paper, und dort sollte man nachsehen, wie die Methode im Vergleich zu Alternativen abschneidet. Lesen Sie es auf arXiv.