AI/ML # Optimización de Inferencia LLM: KV-Cache, PagedAttention y Eficiencia de Memoria GPU 2025-12-2211 min lectura Lucio Durán Engineering Manager & AI Solutions Architect También disponible en: [English]() ## 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: ```python # 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. ```python # 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. ```python # 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 ```python # 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. ```python # 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: ```bash 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. ## Implementación reproducible: medir el serving por token [LLM Serving Bench]() registra tokens reales de SmolLM2-135M-Instruct a través de HTTP/SSE sobre CPU. Compara ocho configuraciones de hilos, concurrencia y longitud de entrada, con 48 solicitudes medidas y 1.152 tokens de salida. La ficha publica TTFT, TPOT, latencia y throughput junto con hardware, llegadas individuales, memoria muestreada y comandos de reproducción. En esa ejecución, duplicar clientes agregó principalmente cola porque había un solo slot de generación. El throughput observado fue de 10,24 a 21,50 tokens de salida por segundo. Las seis solicitudes por configuración son observaciones locales; este proyecto no mide PagedAttention, GPU ni continuous batching. El costo usa una tarifa ilustrativa declarada y la energía queda sin medir. Una cifra sin esas condiciones es decoración. ## Preguntas frecuentes ### ¿Qué es el KV-cache y por qué importa para inferencia de LLMs? El KV-cache almacena los tensores de proyección de key y value de cada token previo en la secuencia, para que el modelo no tenga que recalcular la atención sobre todo el contexto en cada paso de generación. Sin él, la inferencia sería cuadrática en longitud de secuencia. El problema es que crece linealmente con la longitud de secuencia y el batch size, muchas veces consumiendo más memoria GPU que los propios parámetros del modelo. ### ¿Cómo mejora PagedAttention respecto al manejo tradicional de KV-cache? PagedAttention toma prestados conceptos de memoria virtual de los sistemas operativos, dividiendo el KV-cache en bloques de tamaño fijo que se pueden almacenar de forma no contigua en memoria GPU. Esto elimina la fragmentación interna por sobreasignación y permite compartir memoria entre requests vía semántica copy-on-write, típicamente reduciendo el desperdicio de KV-cache un 60-70%. ### ¿Vale la pena la complejidad del speculative decoding en producción? Depende de tu tasa de aceptación. Con un draft model bien matcheado es posible sacar 2-3x de speedup en workloads pesados en generación, pero si la tasa de aceptación baja del 60% se está quemando FLOPs innecesariamente. Profileá primero con la distribución real de tu tráfico antes de comprometerte. ## Tips 1. **Tip:** Si el KV-cache consume más VRAM que los pesos del modelo, probablemente se estén sirviendo requests de contexto largo sin ` PagedAttention `. Migrar a vLLM permite observar cómo baja el consumo de memoria un 60-70%. 2. **Tip:** Continuous batching con ` --max-num-batched-tokens ` tuneado al techo de memoria de la GPU es la palanca de throughput más importante. Se ha observado equipos triplicar sus tokens/seg solo por moverse de batching estático a continuo. 3. **Tip:** Cuando se perfila inferencia, medir siempre el *time-to-first-token* (TTFT) por separado del *inter-token latency* (ITL). Responden a optimizaciones completamente distintas y a los usuarios les importan ambos. - llm - inferencia - kv-cache - vllm - paged-attention - flash-attention - speculative-decoding - gpu - optimización [Volver al Blog]()