Uniéndolo Todo: Contexto Largo en Producción

¿Cómo manejan los modelos de producción el contexto largo?

Los artículos anteriores de esta serie abordaron cada uno una pieza del rompecabezas del contexto largo: codificaciones de posición que generalizan más allá de las longitudes de entrenamiento, patrones de atención eficientes que rompen la barrera cuadrática, memoria externa mecanismos que comprimen el contexto pasado, y generación aumentada con recuperación que descarga conocimiento a un corpus externo. Cada técnica resuelve un cuello de botella real, pero ninguna es suficiente por sí sola. Los modelos de producción apilan cuatro, cinco o seis de estas ideas juntas — y la combinación es lo que realmente entrega ventanas de contexto de 128K o millones de tokens.

Este artículo examina cómo los principales modelos abiertos y propietarios logran sus longitudes de contexto. El patrón que emerge es consistente: cada sistema de producción combina un esquema de codificación de posición, una estrategia de eficiencia de atención, una técnica de gestión del KV cache, y usualmente alguna forma de receta de entrenamiento que extiende progresivamente el contexto. Los detalles difieren, pero la arquitectura en capas no.

LLaMA 3: 128K mediante escalado de RoPE

La familia LLaMA 3 de Meta (Grattafiori et al., 2024) es uno de los ejemplos mejor documentados de cómo un modelo de producción ensambla su stack de contexto largo. La idea clave es que la extensión de contexto es un paso posterior al entrenamiento, no algo incorporado desde el inicio . LLaMA 3 fue preentrenado con una ventana de contexto de 8K, luego extendido a 128K a través de una serie de etapas progresivas.

La receta es engañosamente simple: comenzar con el modelo preentrenado de 8K, luego continuar el entrenamiento a 16K, luego 32K, luego 128K — cada etapa usando una longitud de contexto más larga y ajustando la frecuencia base de RoPE para acomodar las nuevas posiciones. En cada etapa, el modelo ve documentos progresivamente más largos y aprende a atender a mayores distancias. Este enfoque progresivo es mucho más barato que entrenar desde cero a 128K, porque la mayor parte del conocimiento del modelo ya se aprendió durante la fase de preentrenamiento de 8K — las etapas de entrenamiento continuado solo necesitan enseñarle al modelo cómo usar posiciones que no ha visto antes.

¿Por qué funciona esto? Recuerda del artículo 2 que RoPE encodes position by rotating query and key vectors, and the attention score between positions $m$ and $n$ depends on the rotation angle $ heta_i$ for each dimension $i$:

where $ ext{base}$ is typically 10,000 and $d$ is the head dimension. The relative position $(m - n)$ determines the rotation difference. When we scale $ ext{base}$ to a larger value (e.g., 500,000), all the rotation frequencies decrease — the model rotates more slowly per position, so it can represent larger relative distances without the angles wrapping around chaotically. At the boundary: if $ ext{base} o \infty$, every $ heta_i o 0$, meaning no rotation at all (position information vanishes). If $ ext{base} o 0$, $ heta_i o \infty$, and the rotations become so rapid that adjacent positions look unrelated. The standard value of 10,000 sits in a sweet spot for the original training length, and scaling it upward shifts that sweet spot to accommodate longer sequences.

Pero el escalado de RoPE solo no resuelve el problema de memoria. Para el modelo de 70B, LLaMA 3 usa Grouped-Query Attention (GQA) with 8 KV heads shared across 64 query heads. This reduces the KV cache by a factor of $64/8 = 8 imes$ compared to standard multi-head attention. Without GQA, the KV cache for 128K tokens at FP16 on a 70B model would be approximately:

With $L = 80$ layers, $H = 64$ heads, $d_k = 128$, and $n = 128{,}000$, that's $2 imes 80 imes 64 imes 128 imes 128{,}000 imes 2 \approx 335$ GB — far more than a single GPU can hold. GQA with 8 KV heads reduces this to $2 imes 80 imes 8 imes 128 imes 128{,}000 imes 2 \approx 42$ GB. At the boundaries: if we had just 1 KV head (multi-query attention), the cache drops to $\sim$5 GB but quality may suffer. If we use the full 64 KV heads (standard MHA), we get the best quality but cannot fit the cache in memory. 8 KV heads is the production compromise.

Finalmente, LLaMA 3 usa FlashAttention during both training and inference. FlashAttention doesn't reduce the $O(n^2)$ FLOPs — it reduces the memory footprint from $O(n^2)$ to $O(n)$ by never materialising the full attention matrix, instead computing attention in tiles that fit in SRAM. This is what makes 128K contexts feasible on hardware that couldn't store a $128{,}000 imes 128{,}000$ attention matrix.

Mistral y Mixtral: Ventana deslizante + RoPE

Donde LLaMA 3 usa atención completa en cada capa, Mistral 7B (Jiang et al., 2023) toma un enfoque fundamentalmente diferente: atención de ventana deslizante (SWA) with a window size of $w = 4{,}096$. Each token only attends to the $w$ tokens immediately before it, not the entire sequence. This means the attention cost per layer is $O(n \cdot w)$ instead of $O(n^2)$ — linear in $n$ when $w$ is a fixed constant.

¿Pero esto no significa que un token en la posición 50,000 no tiene acceso directo a un token en la posición 1,000? En una sola capa, sí. Pero aquí está la idea clave: en un transformer multicapa, la información se propaga a través de las capas. Después de la capa 1, el token 50,000 ha atendido a los tokens 45,904–50,000. Después de la capa 2, esos tokens han atendido a sus propias ventanas, así que el token 50,000 tiene acceso indirecto a los tokens 41,808–50,000. Después de $L$ capas, el campo receptivo efectivo es:

Para Mistral 7B con $L = 32$ capas y $w = 4{,}096$:

So a 7B model with only 4K local attention effectively "sees" 131K tokens through layer stacking. At the boundary: if $L = 1$ (single layer), the receptive field is just $w$ — the model is truly local with no long-range access. As $L$ grows, long-range information can propagate, though with increasing attenuation. The trade-off is that information from distant positions passes through many attention steps and may be lossy compared to direct full attention.

La implementación del KV cache es elegantemente simple: un buffer circular (circular buffer) of size $w$. When the cache fills up, new entries overwrite the oldest ones. Position $t$ maps to cache index $t \mod w$. This means the KV cache uses a fixed $O(w \cdot d_k \cdot H \cdot L)$ memory regardless of sequence length — a token at position 1,000,000 uses the same cache memory as a token at position 5,000. Compare this to full attention, where the KV cache grows linearly with $n$:

import json, js

# Rolling buffer KV cache vs full KV cache
w = 4096       # sliding window size
d_k = 128      # head dimension
H = 32         # number of KV heads
L = 32         # layers
bytes_per_param = 2  # FP16

# Fixed rolling buffer size
rolling_bytes = w * d_k * H * L * 2 * bytes_per_param  # 2 for K+V
rolling_gb = rolling_bytes / 1e9

seq_lengths = [4096, 32768, 131072, 524288, 1048576]
labels = ["4K", "32K", "128K", "512K", "1M"]

rows = []
for n, label in zip(seq_lengths, labels):
    full_bytes = n * d_k * H * L * 2 * bytes_per_param
    full_gb = full_bytes / 1e9
    saving = (1 - rolling_bytes / full_bytes) * 100
    rows.append([
        label,
        f"{full_gb:.1f} GB",
        f"{rolling_gb:.1f} GB",
        f"{saving:.0f}%" if saving > 0 else "0%"
    ])

js.window.py_table_data = json.dumps({
    "headers": ["Seq Length", "Full KV Cache", "Rolling KV Cache", "Memory Saved"],
    "rows": rows
})

print(f"Config: w={w}, d_k={d_k}, H={H}, L={L}, FP16")
print(f"Rolling buffer is always {rolling_gb:.1f} GB regardless of sequence length")
print(f"At 1M tokens, this saves {rows[-1][3]} of KV cache memory")

Mixtral 8x7B (Jiang et al., 2024) extiende este enfoque combinando el mismo patrón de atención de ventana deslizante con una arquitectura de Mezcla de Expertos (MoE) . Cada capa tiene 8 redes feed-forward expertas, y un router selecciona las 2 principales para cada token. El patrón de atención permanece local (misma ventana deslizante), pero el conteo total de parámetros del modelo (47B) y su capacidad son mucho mayores que su conteo de parámetros activos por token (13B). El mecanismo de atención maneja la pregunta "dónde mirar" con ventanas deslizantes, mientras que la MoE maneja la pregunta "cuánta capacidad" con expertos activados de forma dispersa.

Gemini: Millones de tokens

Gemini 1.5 Pro de Google (Gemini Team, 2024) representa el extremo más radical del espectro de longitud de contexto: hasta 10 millones de tokens de contexto en entornos de investigación, con 1–2 millones de tokens disponibles en producción. Esto es aproximadamente 7,500 páginas de texto — una base de código completa, una serie de libros completa, u horas de video procesadas como secuencias de tokens.

Los detalles arquitectónicos son propietarios, pero del artículo y análisis externo, varias técnicas probablemente están en juego. Primero, Ring Attention o un esquema de atención distribuida similar para distribuir el cómputo a través de los pods de TPU de Google. Ring Attention particiona la secuencia entre dispositivos, con cada dispositivo calculando la atención para su fragmento y pasando bloques KV en un patrón de anillo. Para una secuencia de $n$ tokens distribuida entre $P$ dispositivos, cada dispositivo maneja $n/P$ tokens y la memoria por dispositivo es $O(n^2/P)$ para las puntuaciones de atención (o $O(n/P)$ con mosaico estilo FlashAttention). Con $P = 256$ TPUs, una secuencia de 10M de tokens se convierte en 39K tokens por dispositivo — bien dentro del rango operativo normal.

Segundo, Gemini 1.5 es un modelo de Mezcla de Expertos , lo que significa que aunque puede tener una capacidad total enorme, solo una fracción de los parámetros están activos por token. Esto reduce el costo de cómputo por token, liberando más del presupuesto de hardware para el cómputo de atención que escala con la longitud de la secuencia.

Tercero, el artículo demuestra un rendimiento notable en la evaluación de aguja en un pajar . Gemini 1.5 Pro logra una recuperación casi perfecta (>99.7% de precisión) de hechos colocados en posiciones arbitrarias dentro de contextos de hasta 10M de tokens. Esto sugiere que cualquier patrón de atención y receta de entrenamiento que use Google, ha resuelto sustancialmente el problema de "perdido en el medio" descrito en el artículo 1 — al menos para tareas simples de recuperación factual.

Pero esta capacidad tiene un costo enorme. Procesar 10 millones de tokens a través de incluso un transformer eficiente requiere un cómputo asombroso. Si estimamos los FLOPs de atención para una sola capa con $d_k = 128$ y $H = 16$ (una configuración conservadora):

That's $4.1 imes 10^{20}$ FLOPs for una sola capa of attention alone. A model with 64 layers would need $\sim 2.6 imes 10^{22}$ FLOPs just for attention — roughly equivalent to several hours of compute on a large TPU pod. At the boundary: for $n = 1{,}000$ (a normal prompt), this same formula gives $\sim 4.1 imes 10^{12}$ FLOPs per layer — eight orders of magnitude less. The quadratic wall doesn't disappear; it just gets pushed very far with enough hardware and engineering.

El stack completo

Ahora podemos ver cómo todas las técnicas de esta serie se componen en un sistema moderno de contexto largo. Ninguna técnica individual resuelve el problema. En cambio, los modelos de producción apilan múltiples optimizaciones, cada una dirigida a un cuello de botella diferente. Aquí está el stack completo:

• 1. Codificación de posición: RoPE con escalado de frecuencia base o interpolación NTK-aware estilo YaRN. Esto es lo que permite al modelo representar posiciones más allá de su longitud de entrenamiento sin colapsar. Casi universal en modelos abiertos modernos.
• 2. Patrón de atención: atención completa con FlashAttention (LLaMA, Gemini) or sliding window attention (Mistral). FlashAttention is an IO optimisation that keeps the FLOPs at $O(n^2)$ but avoids the $O(n^2)$ memory cost. Sliding window reduces FLOPs to $O(n \cdot w)$ but trades off direct long-range attention.
• 3. Gestión del KV cache: GQA (fewer KV heads) + quantised cache (INT8 or INT4 values) + PagedAttention (virtual memory for KV blocks). These three together can reduce KV cache memory by 16–32$ imes$ compared to naive full-precision multi-head attention.
• 4. Memoria (opcional): mecanismos de memoria compresiva como Infini-Attention o Titans para conversaciones que exceden incluso la ventana de contexto extendida. Estos almacenan un estado comprimido de tamaño fijo que resume un historial arbitrariamente largo. Aún experimental a escala de producción.
• 5. Recuperación (opcional): RAG para conocimiento que vive completamente fuera del contexto — documentos empresariales, contenido web actualizado, bases de datos privadas. La ventana de contexto maneja la memoria de trabajo; la recuperación maneja el almacén de conocimiento a largo plazo.
• 6. Distribución: paralelismo tensorial entre GPUs para los pesos del modelo, y Ring Attention o paralelismo de secuencia para distribuir la secuencia misma entre dispositivos. Esencial para secuencias muy largas (1M+ tokens) donde incluso el KV cache de una sola capa excede la memoria de un dispositivo.

La tabla a continuación resume qué usa cada modelo principal a través de estas seis dimensiones:

import json, js

rows = [
    ["LLaMA 3 70B", "128K", "RoPE (scaled base)", "Full + FlashAttn", "GQA (8 KV heads)", "Progressive (8K->128K)"],
    ["LLaMA 3 405B", "128K", "RoPE (scaled base)", "Full + FlashAttn", "GQA (8 KV heads)", "Progressive (8K->128K)"],
    ["Mistral 7B", "32K*", "RoPE", "Sliding window (w=4K)", "GQA + Rolling buffer", "Standard training"],
    ["Mixtral 8x7B", "32K*", "RoPE", "Sliding window (w=4K)", "GQA + Rolling buffer", "Standard training"],
    ["Gemini 1.5 Pro", "1-10M", "Likely RoPE variant", "Full + Ring Attention", "Likely GQA + distributed", "Proprietary recipe"],
    ["GPT-4 Turbo", "128K", "Unknown", "Full + FlashAttn (likely)", "Unknown", "Unknown"],
    ["Claude 3.5", "200K", "Unknown", "Full (likely)", "Unknown", "Unknown"],
]

js.window.py_table_data = json.dumps({
    "headers": ["Model", "Context", "Position Enc.", "Attention", "KV Cache", "Training Recipe"],
    "rows": rows
})

print("* Mistral/Mixtral: 32K sequence length, but effective receptive field of ~131K via layer stacking")
print()
print("Key pattern: every model uses at least 3 of the 6 stack components.")
print("Open models (LLaMA, Mistral) document their choices; proprietary models (GPT-4, Claude, Gemini) are inferred.")

La idea clave es que ninguna técnica individual resuelve el contexto largo . Los sistemas de producción apilan 4–6 optimizaciones, cada una dirigida a un cuello de botella diferente. El escalado de RoPE resuelve el problema de codificación de posición pero no reduce el cómputo. FlashAttention resuelve el problema de memoria pero no reduce los FLOPs. GQA reduce el KV cache pero no ayuda con el cómputo de atención. La atención de ventana deslizante reduce el cómputo pero limita el acceso directo de largo alcance. Ring Attention distribuye la carga de trabajo pero requiere configuraciones multi-dispositivo. Cada técnica cubre un vacío; la combinación los cubre todos.

Problemas abiertos

A pesar del impresionante progreso — de 1K tokens en 2019 a millones en 2024 — varios problemas difíciles permanecen sin resolver.

Perdido en el medio. Como se discutió en el artículo 1, Liu et al. (2023) mostró que los modelos tienen dificultades para recuperar información colocada en el medio de contextos largos. Mientras que algunos modelos han mejorado en benchmarks de aguja en pajar (Gemini 1.5 reporta >99.7% de recuperación en todas las posiciones), el rendimiento en tareas más complejas — razonamiento multi-salto, síntesis entre pasajes distantes, resolución de contradicciones entre contexto temprano y tardío — sigue siendo desigual. El problema de "perdido en el medio" puede estar mayormente resuelto para recuperación simple pero sigue abierto para tareas que requieren razonamiento intensivo e integración de información a través de todo el contexto.

Eficiencia a escala. Procesar 1 millón de tokens es técnicamente posible, pero es caro. Incluso con FlashAttention y GQA, un paso hacia adelante sobre 1M de tokens en un modelo de 70B requiere cientos de TeraFLOPs de cómputo de atención por capa y decenas de gigabytes de KV cache. Para servicio por lotes (muchos usuarios simultáneamente), estos costos se multiplican. Hacer que la inferencia de un millón de tokens sea lo suficientemente barata para uso rutinario — no solo demostraciones impresionantes — sigue siendo un desafío central de ingeniería. Técnicas como la decodificación especulativa, la compresión de KV cache y la poda dinámica de contexto son direcciones de investigación activas.

Calidad de la memoria. Compressive memory systems like Infini-Attention (Google, 2024) and Titans (Google, 2025) offer a theoretically elegant path: store a fixed-size compressed state that grows sublinearly with context length. But these systems face a quality-compression trade-off. A compressed memory that stores $c$ dimensions to represent $n$ tokens necessarily discards information — the compression ratio is $n/c$, and as this ratio grows, retrieval of specific facts from the compressed state degrades. At the boundary: if $c = n \cdot d$ (no compression), we recover full attention. If $c = d$ (maximum compression), we're cramming the entire history into a single vector. Finding the right operating point on this curve, and training models to use compressed memories effectively, remains an open problem. No compressive memory system has been demonstrated at production scale as of early 2025.

Evaluación. ¿Cómo medimos siquiera la calidad de contexto largo? La prueba de aguja en un pajar es el estándar actual, pero es demasiado simple — prueba la recuperación de un solo hecho, no el razonamiento complejo multi-documento que los casos de uso reales demandan. Considera las tareas que realmente se benefician del contexto largo: entender una base de código completa para corregir un error, sintetizar información a través de cientos de documentos legales, o seguir una conversación de múltiples turnos que abarca semanas. Estas requieren que el modelo integre información a través de múltiples pasajes, rastree entidades a largas distancias y priorice detalles relevantes de un mar de ruido. Carecemos de benchmarks estandarizados que midan estas capacidades de manera confiable.

El campo está evolucionando rápido. Técnicas que parecen experimentales hoy — módulos de memoria neuronal, arquitecturas de mezcla de profundidades, atención dispersa consciente del hardware — pueden ser estándar dentro de un año. Lo que permanece constante es el enfoque en capas: los sistemas de contexto largo de producción continuarán combinando múltiples técnicas ortogonales, porque el problema tiene múltiples cuellos de botella ortogonales. La codificación de posición, el patrón de atención, la gestión de memoria, la receta de entrenamiento y la estrategia de distribución abordan cada una una restricción diferente, y el progreso en cualquiera de ellas eleva el techo para todas las demás.

Quiz

Pon a prueba tu comprensión de cómo los modelos de producción combinan técnicas de contexto largo.