Optimización de Inferencia LLM: KV-Cache, PagedAttention y Eficiencia de Memoria GPU

El Cuello de Botella de Atención del que Nadie Habla

Todo el mundo sabe que la inferencia en transformers tiene dos fases: el prefill (procesar todo el prompt en paralelo) y el decode (generar tokens de a uno, autoregressivamente). Lo que la mayoría de los tutoriales pasan por arriba es que el decode es fundamentalmente memory-bound, no compute-bound.

Durante el decode, cada token nuevo necesita atender a todos los tokens previos en la secuencia. El approach naive recalcula todas las proyecciones key-value desde cero en cada paso, haciendo que la generación sea O(n²) en longitud de secuencia. El KV-cache resuelve esto almacenando los tensores K y V de todas las posiciones anteriores:

# KV-cache simplificado durante decode
# Para cada capa l, head h, y posición t:
# K_cache[l][h][t] = W_k @ hidden_states[t]
# V_cache[l][h][t] = W_v @ hidden_states[t]

# En el paso de decode T, atención para el token nuevo:
q = W_q @ hidden_states[T] # (1, d_head)
k = K_cache[l][h][:T+1] # (T+1, d_head)
v = V_cache[l][h][:T+1] # (T+1, d_head)
attn = softmax(q @ k.T / sqrt(d_head)) @ v

Esto reduce la generación a O(n) por paso. Pero el costo en memoria es brutal. Para un modelo de 70B (80 capas, 64 heads, 128 head dim, FP16):

KV-cache por token = 2 * n_layers * n_kv_heads * d_head * 2 bytes
 = 2 * 80 * 8 * 128 * 2 (GQA con 8 KV heads)
 = 327.680 bytes ≈ 320 KB por token

Para una secuencia de 4096 tokens:
320 KB * 4096 = 1,28 GB por request

Para 100 requests concurrentes:
1,28 GB * 100 = 128 GB solo de KV-cache

Los pese es del modelo son ~140 GB en FP16. O sea, el KV-cache puede competir con el modelo en consumo de memoria. Y aquí es donde se pone feo.

El Problema de Fragmentación

Los frameworks de serving tradicionales pre-asignan un bloque contiguo de memoria GPU para el KV-cache de cada request, dimensionado para la longitud máxima posible de secuencia. Si tu máximo es 4096 tokens pero el completion promedio es de 500, se está tirando ~87% de la memoria de KV-cache asignada.

Mediciones en clusters de producción muestran una utilización real promedio del KV-cache del 22%. Esto significa que el 78% del recurso más caro se está descartando.

# Cómo se ve la asignación naive:
def allocate_kv_cache(max_seq_len, batch_size, n_layers, n_heads, d_head):
 # Pre-asignar para el peor caso
 k_cache = torch.zeros(
 batch_size, n_layers, n_heads, max_seq_len, d_head,
 dtype=torch.float16, device='cuda'
 )
 v_cache = torch.zeros_like(k_cache)
 return k_cache, v_cache

# Para batch=64, max_seq=4096 en un modelo de 70B:
# Esto reclama inmediatamente ~82 GB de VRAM
# La mayoría de los cuales quedan en cero durante todo el ciclo de vida del request

PagedAttention: Conceptos de SO se Encuentran con Memoria GPU

Acá es donde PagedAttention del equipo de vLLM cambió todo. La intuición es engañosamente simple: tratar la memoria del KV-cache de la misma forma que los sistemas operativos tratan la memoria virtual.

En vez de asignaciones contiguas por request, PagedAttention divide el KV-cache en bloques de tamaño fijo (típicamente 16 tokens cada uno). Cada request mantiene una block table que mapea posiciones lógicas a ubicaciones físicas de bloques en memoria GPU, exactamente como una tabla de páginas mapea direcciones virtuales a físicas.

# Block table conceptual de PagedAttention
# Cada bloque almacena KV para 16 tokens en todas las capas y heads

class BlockTable:
 def __init__(self, block_size=16):
 self.block_size = block_size
 self.table = [] # Lista de índices de bloques físicos

 def get_physical_block(self, logical_position):
 block_idx = logical_position // self.block_size
 block_offset = logical_position % self.block_size
 physical_block = self.table[block_idx]
 return physical_block, block_offset

# Pseudocódigo del kernel para paged attention
# Para cada token de query:
# Para cada bloque en la block_table del request:
# Cargar K, V del bloque físico
# Calcular scores parciales de atención
# Acumular con online softmax (como FlashAttention)

Los beneficios se acumulan en cascada:

1. Sin fragmentación interna: Los bloques se asignan on demand. Una respuesta de 500 tokens usa exactamente ceil(500/16) = 32 bloques, no el max_seq_len completo.

2. Compartir con copy-on-write: Para beam search o sampling paralelo, múltiples secuencias que comparten un prefijo común pueden apuntar a los mismos bloques físicos. Cuando una secuencia diverge, solo se copia el bloque que diverge.

3. Memoria dinámica: Los bloques se liberan inmediatamente cuando un request termina, vuelven al pool para el siguiente. Sin memoria reservada-pero-idle.

En un cluster de producción representativo, pasar de HuggingFace TGI a vLLM con PagedAttention aumentó la capacidad de requests concurrentes de ~40 a ~110 en el mismo hardware. El throughput pasó de 480 a 1.340 tokens/segundo agregados.

Continuous Batching: La Otra Mitad de la Ecuación

El batching estático espera un batch completo de requests, los procesa juntos, y no arranca el siguiente batch hasta que todas las secuencias del batch actual terminaron de generar. Esto significa que los completions cortos se quedan innecesariamente mientras los largos siguen laburando.

El continuous batching (también llamado iteration-level scheduling) es distinto. En cada paso de decode, el scheduler puede:

• Sacar secuencias completadas
• Admitir secuencias nuevas
• Desalojar secuencias de menor prioridad si la memoria aprieta

# Scheduler simplificado de continuous batching
class ContinuousBatchScheduler:
 def __init__(self, max_batch_tokens, block_manager):
 self.waiting = deque() # Requests esperando para arrancar
 self.running = [] # Generando actualmente
 self.max_batch_tokens = max_batch_tokens
 self.block_manager = block_manager

 def schedule(self):
 # Sacar secuencias completadas
 self.running = [s for s in self.running if not s.is_finished()]

 # Intentar admitir secuencias nuevas
 while self.waiting:
 next_req = self.waiting[0]
 total_tokens = sum(s.get_len() for s in self.running)
 total_tokens += next_req.prompt_len

 if total_tokens > self.max_batch_tokens:
 break
 if not self.block_manager.can_allocate(next_req):
 break # No hay suficientes bloques de KV-cache

 self.running.append(self.waiting.popleft())
 self.block_manager.allocate(self.running[-1])

 return self.running

El impacto práctico: en un workload con completions mixtos cortos (50 tokens) y largos (2000 tokens), continuous batching da 2.8x más throughput comparado con batching estático. Los requests cortos dejaron de ser rehenes de los largos.

FlashAttention-3: Porque el Ancho de Banda de Memoria Es el Verdadero Enemigo

FlashAttention reescribe el cálculo de atención para ser IO-aware, minimizando lecturas/escrituras a HBM (High Bandwidth Memory) de la GPU manteniendo datos en SRAM (la memoria rápida on-chip) el mayor tiempo posible. La versión 3 agrega varios trucos específicos para GPUs con arquitectura Hopper (H100):

Atención estándar:
 1. Calcular S = Q @ K.T → escribir S a HBM
 2. Calcular P = softmax(S) → escribir P a HBM
 3. Calcular O = P @ V → escribir O a HBM
 Accese es a HBM totales: O(N² * d)

FlashAttention-3:
 1. Dividir Q, K, V en bloques del tamaño de SRAM
 2. Para cada tile: calcular S, P, O locales en SRAM
 3. Acumular con reescalado de online softmax
 4. Escribir O final a HBM una sola vez
 Accese es a HBM totales: O(N * d), lineal, no cuadrático

FA3 en H100s explota específicamente:

• Warp specialization: Warps separados para MMA (matrix multiply-accumulate) y softmax, solapando cómputo con movimiento de datos
• Soporte FP8: Tensor cores FP8 nativos para el matmul QK^T, con acumulación en FP32
• Softmax asincrónico por bloques: Usando TMA (Tensor Memory Accelerator) para copias asincrónicas de memoria global a compartida

Para resultados reproducibles, consultá las configuraciones y mediciones del artículo de FlashAttention 3 y medí tu carga concreta. Esta guía no aporta una medición reproducible para la configuración H100 SXM que antes figuraba aquí.

La integración con PagedAttention de vLLM es donde resulta interesante, el kernel flash_attn_with_kvcache maneja la indirección de la block table directamente, así que teners tanto gestión de memoria paginada como atención IO-aware en un solo kernel fusionado.

Speculative Decoding: Intercambiar Cómputo por Latencia

La idea: usar un modelo "draft" chiquito para generar N tokens candidatos rápidamente, y después verificarlos todos de una con el modelo grande "target" en un solo forward pass. Si el modelo target está de acuerdo con los tokens del draft (chequeado vía un esquema modificado de rejection sampling), efectivamente generaste N tokens en el tiempo que le toma al target hacer un forward pass.

# Loop de speculative decoding
def speculative_decode(target_model, draft_model, prompt, gamma=5):
 tokens = prompt
 while not is_done(tokens):
 # Fase draft: generar gamma candidatos barato
 draft_tokens = []
 draft_probs = []
 for _ in range(gamma):
 logits = draft_model(tokens + draft_tokens)
 p = softmax(logits[-1])
 t = sample(p)
 draft_tokens.append(t)
 draft_probs.append(p)

 # Fase verificación: un solo forward pass en target
 all_logits = target_model(tokens + draft_tokens)

 # Aceptación: rejection sampling modificado
 accepted = 0
 for i in range(gamma):
 q = softmax(all_logits[len(tokens) + i])
 p = draft_probs[i]
 r = random()
 if r < min(1, q[draft_tokens[i]] / p[draft_tokens[i]]):
 accepted += 1
 else:
 # Re-samplear de distribución ajustada
 adjusted = normalize(max(0, q - p))
 tokens.append(sample(adjusted))
 break
 else:
 # Todos aceptados, samplear uno más del target
 tokens.extend(draft_tokens)
 q = softmax(all_logits[-1])
 tokens.append(sample(q))
 continue
 tokens.extend(draft_tokens[:accepted])

La matemática cierra de forma que la distribución de salida es probadamente idéntica a samplear del modelo target solo. No se está sacrificando calidad.

En la práctica, la tasa de aceptación depende fuertemente de qué tan bien el draft model matchea al target. Resultados observados:

• Llama-3-8B drafteando para Llama-3-70B: ~75% tasa de aceptación, ~2.3x speedup
• Drafter custom chiquito de 100M fine-tuneado en salidas del target: ~65%, ~1.8x
• Tareas de code completion (alta predictibilidad): ~85%, ~3.1x

El punto óptimo para gamma (longitud de especulación) suele ser 4-6. Si lo subís más, la probabilidad de al menos un rechazo por ronda se acerca a 1, desperdiciando el cómputo extra del draft.

Juntando Todo: Configuración de Producción

Esta es una configuración de vLLM que puede llevar un cluster de 12 tok/s/usuario a 45 tok/s/usuario:

python -m vllm.entrypoints.openai.api_server \
 --model meta-llama/Llama-3-70B-Instruct \
 --tensor-parallel-size 8 \
 --max-num-batched-tokens 32768 \
 --max-num-seqs 256 \
 --block-size 16 \
 --swap-space 4 \
 --gpu-memory-utilization 0.92 \
 --enable-prefix-caching \
 --use-v2-block-manager \
 --speculative-model meta-llama/Llama-3-8B-Instruct \
 --num-speculative-tokens 5 \
 --quantization awq

Decisiones clave:

• --max-num-batched-tokens 32768: Permite hasta 32K tokens en vuelo por iteración. Tuneado justo por debajo del umbral de OOM.
• --enable-prefix-caching: Reutiliza bloques de KV-cache para system prompts compartidos. Con nuestro system prompt de 1.500 tokens, ahorra ~480 KB por request.
• --gpu-memory-utilization 0.92: Dejamos 8% de headroom para picos de memoria de activaciones durante prefill de prompts largos.
• --swap-space 4: 4 GB de swap en CPU por GPU para requests desalojados, mejor que rechazarlos directamente.

Los Números Que Realmente Importan

En 8xA100 80GB, sirviendo Llama-3-70B-Instruct (AWQ 4-bit):

Métrica	Naive (TGI estático)	Optimizado (vLLM)	Mejora
Throughput (tok/s)	480	1.850	3,85x
Requests concurrentes	40	180	4,5x
TTFT p50	820 ms	210 ms	3,9x
TTFT p99	4.200 ms	890 ms	4,7x
ITL p50	38 ms	22 ms	1,7x
Costo por 1M tokens	$2,40	$0,62	3,87x

La reducción de costos sola puede pagar el esfuerzo de ingeniería en unos cuatro días.

Qué Viene Después

Varias áreas merecen seguimiento para la próxima ronda de optimizaciones:

• Prefill/decode desagregado: Correr prefill en GPUs distintas que decode, ya que tienen perfiles de cómputo completamente diferentes. Los papers de Splitwise y DistServe pintan prometedores.
• KV-cache cuantizado: Almacenar KV-cache en FP8 o incluso INT4, separado de la cuantización de pesos. Resultados tempranos muestran pérdida mínima de calidad con 2x de reducción de memoria.
• Ring attention: Para contextos ultra-largos (100K+), distribuir el cómputo de atención entre múltiples GPUs a lo largo de la dimensión de secuencia.

El campo se mueve increíblemente rápido. Hace seis meses PagedAttention era novedoso; ahora es lo mínimo indispensable. La ventaja de infraestructura la tiene el que integre estas técnicas más rápido en su stack de serving.

Para cualquiera que corra inferencia en producción, examinar la utilización del KV-cache es una actividad de alto valor. Probablemente ahí es donde se concentra la mayoría de los costos.