QUOKA explicado: cómo elegir qué historial KV debe leer la atención

Si observas trabajar a un agente de programación, verás que casi nunca entiende una tarea con un solo archivo. Abre varios archivos por todo el repositorio, sigue cómo se conectan las piezas, rastrea dependencias y solo entonces edita algo. A veces eso supone enviar al modelo una cantidad sorprendente de contexto.

Lo que plantea una pregunta. Los proveedores anuncian ventanas de contexto de cientos de miles de tokens, incluso un millón. ¿Qué implica de verdad procesar tanto texto y de dónde sale el coste? Ningún truco aislado lo resuelve. La inferencia con contexto largo toca cómo se codifican las posiciones, cómo se calcula la atención, cómo se gestiona la memoria y cómo se planifican las cargas.

Un artículo que aborda una porción estrecha del problema es QUOKA, siglas de Query-Oriented KV Selection for Efficient LLM Prefill. Trata de abaratar la atención mientras el modelo lee un prompt de entrada largo. La idea central es intuitiva, pero para ver por qué funciona hacen falta tres piezas de contexto: la caché KV, la separación entre prefill y decode, y el chunked prefill.

La caché KV en pocas palabras

Dentro de un Transformer, cada token produce para la atención una consulta (Q), una clave (K) y un valor (V). Una forma aproximada de verlos:

• La consulta pregunta qué información busca el token.
• La clave anuncia qué tipo de información contiene el token.
• El valor es lo que se entrega si el token resulta relevante.

Al generar de forma autorregresiva, las claves y valores calculados antes pueden guardarse en una caché KV y reutilizarse después en lugar de recalcularse. A grandes rasgos:

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

Cuanto más larga es la secuencia, más entradas K y V hay guardadas. Eso importa para QUOKA, y para ver por qué primero hay que conocer las dos fases de la inferencia.

Prefill: leer tu prompt

Cuando envías un prompt a un modelo, sea una pregunta de una línea o una montaña de código, los tokens de entrada tienen que procesarse antes de que empiece cualquier salida. Esa fase es el prefill.

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

Decode: escribir la respuesta

Una vez procesada la entrada, el modelo produce su respuesta token a token. Eso es el decode.

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

La distinción importa porque elegir qué partes de un historial largo merecen atención es bastante fácil cuando hay una sola consulta nueva. Es otro problema cuando un fragmento entero contiene cientos o miles de consultas a la vez. Antes de llegar ahí hace falta un concepto más.

Chunked prefill

Procesar un prompt muy largo en una sola pasada de prefill es caro. Por eso los sistemas de inferencia suelen usar chunked prefill, que divide la entrada en fragmentos más pequeños procesados en orden. Para un prompt de 3.000 tokens:

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

Cada fragmento posterior al primero tiene más historial detrás:

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
...

Fragmentar hace tolerable un prefill largo, pero cada fragmento posterior puede atender a un montón cada vez mayor de historial en caché. Ese montón es lo que ataca QUOKA.

La observación detrás del método

Imagina que se procesa un fragmento tardío con 100.000 posiciones KV históricas ya guardadas. La atención densa deja que las consultas actuales miren cada una de ellas.

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

¿De verdad necesitan las consultas actuales las 100.000? Probablemente no. Algunas claves importarán mucho para lo que hace este fragmento y muchas no aportarán casi nada. Así que hay una oportunidad:

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

Esa es la intuición básica. En lugar de ejecutar una atención costosa sobre todo el historial, seleccionar un conjunto más pequeño de KV históricos que probablemente importen para el fragmento actual. Bastante simple, pero lleva a una pregunta más difícil: ¿quién decide qué KV importan?

Una consulta puede elegir por sí misma, pero un fragmento tiene miles

En la generación token a token hay una sola posición de consulta nueva cada vez, y esa consulta puede preguntar en la práctica qué claves históricas le son relevantes.

one query
    ↓
Which historical keys
are relevant to me?

Durante el chunked prefill hay muchas posiciones de consulta a la vez:

q1
q2
q3
...
q4096

Consultas distintas pueden interesarse por partes del historial completamente diferentes. Podrías reducirlas a una consulta media, pero te arriesgas a perder necesidades poco habituales pero importantes. Supón que la mayoría de las consultas apuntan en la misma dirección:

q1  ↗
q2  ↗
q3  ↗
q4  ↗
q5  →

La media quedará dominada por q1 a q4. Si q5 es la única consulta que necesita cierta clave, el promedio entierra esa interacción.

Elegir las consultas que discrepan de la multitud

QUOKA calcula primero la consulta media del fragmento actual y luego mide la similitud coseno de cada consulta con esa media. Por ejemplo:

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

En lugar de quedarse con las consultas más cercanas a la media, prioriza las menos parecidas a ella:

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

El razonamiento es directo. Las consultas muy parecidas apuntan más o menos en la misma dirección, así que quedarse con muchas de ellas da información redundante. Las de baja similitud son las que dicen que buscan otra cosa, y son justo las que probablemente descubran claves que la mayoría pasaría por alto. Así, miles de consultas se convierten en un conjunto pequeño y diverso:

thousands of queries
        ↓
a small set of diverse queries

Que esas consultas puntúen el historial

Ahora las consultas diversas recorren las claves guardadas, y cada una pregunta qué partes del pasado le parecen útiles. QUOKA vuelve a usar la similitud coseno, esta vez entre las consultas seleccionadas y las claves históricas. Digamos que se eligieron tres consultas:

qA
qB
qC

y que hay miles de claves:

K1, K2, K3, ... K10000

Cada consulta seleccionada puntúa cada clave por similitud coseno:

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

Esto es útil porque consultas distintas encuentran partes distintas del pasado. Aquí a qA le importa mucho K1, a qB K2 y a qC K3.

Pero ahora cada clave tiene varias puntuaciones de importancia, una por consulta seleccionada, mientras que QUOKA necesita un solo número para decidir si la clave sobrevive.

Por qué tomar el máximo

Mira K1:

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

Promediar sería la opción obvia:

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

Eso hace que K1 parezca moderadamente importante y oculta algo valioso: que qA lo considera importantísimo. Una clave no tiene que importar a todas las consultas para merecer un sitio. Puede contener justo lo que necesita una consulta. Por eso QUOKA toma el máximo entre las consultas seleccionadas:

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

Cada clave acaba con una puntuación final y se conservan las interacciones fuertes entre consulta y clave. Si aunque sea una consulta representativa necesita mucho una clave, no hay que dejar que las demás diluyan esa señal.

De las puntuaciones a un conjunto KV más pequeño

Con las puntuaciones en la mano, QUOKA ordena las claves y conserva solo las mejores, según su presupuesto de atención dispersa.

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  ✓

Cada vez que se elige una posición de clave, se lleva consigo su valor correspondiente. Si QUOKA conserva estas claves:

K3, K1, K2

también conserva estos valores:

V3, V1, V2

El resultado es un conjunto más pequeño de pares K,V que participa en la atención. En esencia, en lugar de obligar a la atención a considerar cada posición histórica, QUOKA construye primero un conjunto de candidatos mucho menor a partir de lo que las consultas actuales parecen necesitar de verdad.

Un matiz: grouped-query attention

Los modelos modernos no tienen necesariamente tantas cabezas de consulta como cabezas KV, y ahí entra la grouped-query attention (GQA). Supón que un modelo tiene cuatro cabezas de consulta y dos cabezas KV:

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

En lugar de que cada cabeza de consulta tenga su propia cabeza de claves y valores, varias cabezas de consulta comparten una. Eso afecta a QUOKA, porque en un modelo con GQA cada token seleccionado tiene varias representaciones por cabeza de consulta. Para un token seleccionado qA podrías tener:

qA_head1
qA_head2
qA_head3
qA_head4

Las dos primeras interactúan con la cabeza KV 1 y las otras dos con la cabeza KV 2. Puntuar cada cabeza por separado y promediar después quedaría así:

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

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

Eso repite comparaciones caras de consulta y clave contra la misma cabeza KV. QUOKA aprovecha un hecho algebraico sencillo y promedia primero las representaciones normalizadas de las consultas:

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

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

Funciona porque (Q1 K^T + Q2 K^T) / 2 es igual a ((Q1 + Q2) / 2) K^T. En lugar de dos multiplicaciones de puntuación seguidas de un promedio, promedias primero las consultas y multiplicas una vez. En el artículo de QUOKA este truco se llama preagregación (pre-aggregation).

MEAN y MAX hacen trabajos distintos

Es fácil confundir este MEAN con el MAX de antes, pero trabajan sobre dimensiones distintas.

MEAN actúa sobre las cabezas de consulta que comparten una cabeza KV y produce una consulta sustituta más barata para puntuar esa cabeza:

qA_head1 ─┐
          ├── MEAN
qA_head2 ─┘

MAX recorre distintos tokens de consulta seleccionados y protege una clave cuando aunque sea una consulta atípica la considera muy relevante:

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

Para no liarse:

MEAN across related GQA heads
          ↓
make KV selection cheaper

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

Importan dos detalles más. Las cabezas de consulta promediadas sirven estrictamente para elegir posiciones KV. Las cabezas de consulta originales del modelo siguen en su sitio y ejecutan la atención dispersa real sobre lo que se haya elegido, lo que convierte la preagregación en un atajo solo para la selección.

Fíjate además en que no se expulsa nada. La caché KV histórica completa se mantiene, los KV seleccionados son solo el subconjunto que se usa en el paso de atención actual y un fragmento nuevo puede elegir otro subconjunto según sus propias consultas. El ahorro está en cuánto historial entra en el cálculo de la atención, no en el tamaño de la caché guardada.

Dónde encontrar los números

Los resultados de benchmarks aparecen en tres tablas del artículo, que es donde conviene ver cómo se compara el método con las alternativas. Léelo en arXiv.