Inicio / Artículos / Mover la caché KV entre modelos Qwen3 en lugar de repetir el prefill

Mover la caché KV entre modelos Qwen3 en lugar de repetir el prefill

Una investigación de NVIDIA traslada el estado de atención de Qwen3-14B a Qwen3-32B con un puente de regresión ridge. Qué significan las cifras de latencia, dónde falla y por qué importa a los agentes.

1817 palabras

Imagina un agente de código que lleva veinte minutos funcionando. Ha leído un repositorio, ha abierto la documentación, ha llamado a un puñado de herramientas y ha acumulado cerca de 30K tokens de contexto en Qwen3-14B. Entonces llega al único paso en el que un modelo más potente ayudaría de verdad.

Enviar el siguiente turno a Qwen3-32B es la parte fácil. Conservar el trabajo ya hecho es la difícil. El modelo de 32B tiene que masticar de nuevo todo el contexto acumulado antes de poder emitir un solo token. Esa pasada se llama prefill, y lo que produce es la caché KV del modelo.

Un preprint reciente de investigación de NVIDIA plantea una pregunta incisiva. El modelo pequeño ya construyó una caché KV para esta conversación; ¿se podría convertir ese estado en la caché del modelo grande en vez de obligar a este a reconstruirla desde el texto? En el experimento descrito con Qwen3 14B → 32B, el mapper tarda 277,6 ms con 32.768 tokens, frente a 6.975,3 ms de repetir el prefill en el modelo destino. Son unas 25x para esa fase concreta. Es un resultado prometedor, pero mucho más acotado que «inferencia de LLM 25 veces más rápida».

Por qué cambiar de modelo desperdicia trabajo

La caché KV es el estado de atención que un transformer autorregresivo construye mientras lee un prompt. Una vez que existe, el modelo no necesita recalcular la atención sobre todo el prefijo para cada token que genera.

Dentro de un mismo modelo, esta reutilización funciona con naturalidad. Los sistemas de serving ya cachean prefijos y mueven estado KV entre workers compatibles; NVIDIA TensorRT-LLM, por ejemplo, admite reutilización de caché KV e intercambio de KV para serving desagregado.

Al cambiar de modelo es donde todo se desmorona.

Qwen3-14B
   │
   │  30K-token conversation
   ▼
14B KV cache
   │
   ✕ cannot simply hand this to 32B
   │
Qwen3-32B normally re-prefills the text

Los modelos de distinto tamaño tienen distinto número de capas y representaciones internas diferentes. Compartir tokenizador o familia de modelos no hace que sus cachés sean intercambiables. NVIDIA lo trata como un problema de mapeo de representaciones.

14B KV cache
     ↓
model-pair mapper
     ↓
32B-compatible KV cache
     ↓
32B starts decoding

El estudio se limita a propósito a pares dentro de una misma familia con KV coincidente: origen y destino deben coincidir en número de cabezas KV y en dimensión de cabeza. Los pares probados son de Qwen3, Llama 3.1 y Ministral 3. Pasar de una familia a otra, por ejemplo de Qwen a Llama, queda como pregunta abierta.

Un puente lineal basta, sorprendentemente

Lo interesante no es que se pueda entrenar un mapper. Es que un mapper de forma cerrada bastante sencillo funciona bien en varios pares de modelos.

El análisis de estructura lineal del artículo muestra que, para Qwen3 14B → 32B, la mejor capa de origen por sí sola explica cerca del 56 % de la varianza de las keys del destino y el 32 % de la de los values. Si se combina una selección de capas de origen, esas cifras suben a cerca del 79 % y el 65 %.

El mapper de producción se apoya en tres ideas.

Elegir capas de origen útiles. Una capa destino no copia sin más la capa de origen a la misma profundidad. El método busca las capas de origen cuyos estados KV mejor predicen esa capa destino y concatena sus características.

Ajustar keys y values por separado, por cabeza destino. El puente en sí es una regresión ridge que se resuelve de forma independiente para cada capa y cabeza destino, con un ajuste para las keys y otro para los values. Una versión simplificada de las matemáticas es esta:

import numpy as np

def fit_ridge(x, y, lam=0.01):
    """Fit Y ≈ XW + b. Conceptual demo, not NVIDIA's implementation."""
    x_mean = x.mean(axis=0, keepdims=True)
    y_mean = y.mean(axis=0, keepdims=True)
    xc = x - x_mean
    yc = y - y_mean
    eye = np.eye(xc.shape[1], dtype=xc.dtype)
    w = np.linalg.solve(xc.T @ xc + lam * eye, xc.T @ yc)
    b = (y_mean - x_mean @ w).squeeze(0)
    return w, b

def map_cache(x, w, b):
    return x @ w + b

En el mapper real, x es mucho mayor, porque concatena características KV de todas las capas de origen seleccionadas, y cada cabeza destino tiene su propio ajuste. La calibración usa 500 secuencias de FineWeb-Edu de 1.024 tokens cada una, submuestreadas a unas 128K observaciones a nivel de token por cabeza destino.

Quitar la posición antes de mapear las keys. Es el detalle que más vale la pena recordar. RoPE rota las keys según la posición del token, así que mapear directamente keys rotadas corre el riesgo de aprender la distribución de posiciones de los datos de calibración de 1.024 tokens en lugar del contenido. El método deshace primero la rotación RoPE del modelo de origen, mapea la representación del contenido sin posición y después aplica el RoPE del modelo destino.

source key
    ↓
strip source RoPE
    ↓
map content representation
    ↓
apply target RoPE
    ↓
target key

Gracias a esto, el mapper sigue funcionando en posiciones que nunca vio durante la calibración, incluida la prueba de latencia de 32K.

Qué cubre la cifra del titular

Estas son las latencias publicadas para Qwen3 14B → 32B en un nodo 8xH100:

  • 64 tokens: 14,0 ms de mapeo frente a 61,7 ms de repetir el prefill, unas 4 veces más rápido.
  • 8K tokens: 67,8 ms frente a 1.154,8 ms, unas 17 veces más rápido.
  • 32K tokens: 277,6 ms frente a 6.975,3 ms, unas 25 veces más rápido.

Hay que leerlas con cuidado. Comparan el tiempo de ejecutar el mapper con el que el modelo receptor necesita para repetir el prefill. No son una aceleración de 25x de la generación en conjunto, porque la decodificación tiene que hacerse después igualmente. El artículo señala además una salvedad de despliegue: en un serving real con varios procesos quizá haya que enviar la caché mapeada al proceso destino, y el benchmark del mapper no lo convierte en una cifra que incluya la red.

La calidad es buena en algunos pares y mala en otros. Qwen3 14B → 32B conserva el 97,6 % de la precisión media del 32B destino por sí solo en los cinco benchmarks publicados. Qwen3 8B → 32B conserva el 87,5 %, Llama 3.1 8B → 70B el 72,8 % y Ministral 3B → 8B el 76,2 %. En cambio, Ministral 3B → 14B y 8B → 14B caen a solo el 44,2 % y el 41,6 % de retención media con el mapper lineal. En este estudio, que coincidan las dimensiones KV es un requisito, no una promesa de traspaso limpio.

Dónde ayuda a los agentes

Los agentes de larga duración son justo el sitio donde repetir el prefill duele.

Así podría ser una política de enrutamiento basada en esto:

"""Illustrative serving API. The paper does not provide this interface."""
state = agent.run(model="qwen3-14b", task=task)

if state.needs_stronger_reasoning:
    mapped_kv = kv_bridge.convert(
        source_model="qwen3-14b",
        target_model="qwen3-32b",
        kv_cache=state.kv_cache,
    )
    state = agent.resume(
        model="qwen3-32b",
        kv_cache=mapped_kv,
    )

Enrutar a un modelo más grande no es nuevo; eso ya funciona hoy. El cambio es que el estado de cómputo podría viajar con la conversación. Un agente de investigación podría explorar barato en un miembro pequeño de la familia, pasar a uno mayor para un tramo de razonamiento difícil y volver abajo para el trabajo rutinario. Los mapeos son direccionales, así que 14B → 32B y 32B → 14B necesitan mappers distintos, y el artículo justifica tanto la mejora de calidad como la reducción de costes.

El experimento multiturno alternó Qwen3 14B y 32B en 100 conversaciones de CoQA. En el turno 10, la brecha de pequeño a grande se había ampliado en unos 1,7 puntos porcentuales, y la deriva de grande a pequeño creció unos 0,33 puntos por turno. Es una prueba razonable a favor de cambiar de modelo dentro de una conversación, aunque los propios investigadores advierten de que la deriva lineal podría acumularse en sesiones mucho más largas.

Límites que importan en producción

El artículo es lo bastante bueno como para cambiar cómo se piensa el enrutamiento de modelos. Pero todavía no es algo que se pueda activar en un stack de serving.

  • El margen de compatibilidad es estrecho. Todos los experimentos publicados se quedan dentro de familias de modelos, usan modelos densos con atención completa y prueban seis pares con el mismo número de cabezas KV y la misma dimensión de cabeza. El método no se ha probado entre familias ni con formas KV distintas.
  • Cada par necesita su propia validación. Dos mapeos de Ministral fallan estrepitosamente con regresión ridge aunque su arquitectura KV coincida. Un MLP no lineal recupera entre 24 y 37 puntos porcentuales de retención en HellaSwag en esos pares, lo que es alentador, pero demuestra que la historia lineal simple no se cumple siempre.
  • El mapper no es pequeño. Los mappers evaluados tienen entre 1.010 y 3.360 millones de parámetros y ocupan de 4 a 12 GB por dirección. Qwen3 14B → 32B ronda los 1.070 millones de parámetros y 4 GB. Ajustarlo lleva unos 47 a 87 minutos por par en un solo nodo 8xH100.

Como los mapeos son direccionales, un router con muchos modelos acumula estado específico de cada par. Con el tamaño medio de mapper del artículo, salen 39 GB con tres modelos, 79 GB con cuatro y 131 GB con cinco. Los mappers pueden estar en la memoria del host o en almacenamiento en vez de en VRAM, pero siguen siendo infraestructura que hay que mantener.

Una cosa más: hoy no hay implementación oficial asociada al artículo. El registro de arXiv tiene el artículo y su código fuente, pero no un repositorio de NVIDIA para este mapper, así que los fragmentos de arriba son explicaciones y no código listo para usar. Para el contexto del lado del serving, consulta la documentación de TensorRT-LLM sobre reutilización de caché KV y el registro del artículo en arXiv.

Qué tener en cuenta

Hasta ahora, la caché KV se ha tratado casi siempre como un estado que pertenece a un único modelo. Si la transferencia entre modelos llega a ser lo bastante fiable, esa suposición desaparece.

Para la infraestructura de agentes, el futuro interesante no es solo un router que elige modelo por petición. Es un router que mueve el estado de una conversación ya avanzada entre modelos compatibles sin pagar por segunda vez el coste completo del contexto.

Nada de esto hace que la caché KV sea portable en general. No elimina el coste de decodificación, no resuelve el enrutamiento entre familias y los pares más flojos muestran por qué hay que probar cada par. Pero el resultado de Qwen3 14B → 32B demuestra que el trabajo guardado en el estado de atención de un modelo puede reutilizarlo otro modelo por menos de lo que cuesta recalcularlo desde el texto.

Si eso se traslada a sistemas de serving reales, encadenar modelos se vuelve mucho más interesante para los agentes de larga duración. La pregunta deja de ser solo qué modelo debe responder al siguiente turno. También pasa a ser cuánto del cómputo del modelo anterior puede venir con él.