llm inference kv-cache gpu pagedattention optimization

La caché KV es la base de datos invisible que consume memoria en la GPU durante la decodificación de LLM.

La caché KV consume silenciosamente la memoria de la GPU durante la inferencia de LLM: prefill vs decode, cálculos KV, GQA y MQA, PagedAttention, políticas de cuantización y expulsión.

Editor: REACTAPP.TOPAñadido a este sitio 26 sept 20269 min de lectura2109 palabras

Fuentes: The KV Cache: AI’s Unseen Database Dominating GPU Memory

El verdadero cuello de botella en la inferencia de LLM

Durante la inferencia, el procesamiento en GPU suele quedar inactivo mientras el ancho de banda de la memoria alimenta un cache clave–valor que crece constantemente. Cada nuevo token agrega tensores K/V que deben leerse nuevamente en pasos posteriores. Las implementaciones que ignoran este crecimiento del cache chocan con limitaciones de contexto y latencia mucho antes de que se agoten los FLOPs.

Detrás de escena: cómo funciona el cache clave–valor

La atención necesita claves y valores para cada token anterior. Al almacenarlos en caché se evita volver a calcular todo el prefijo en cada paso. El compromiso es un consumo de memoria lineal con respecto a la longitud de la secuencia × capas × cabezas × dimensión.

Prefill vs decodificación

Prefill procesa el prompt y llena el cache. Decodificación agrega un token a la vez, leyendo todo el cache en cada paso. El Prefill requiere mucho procesamiento; la decodificación depende mucho del ancho de banda.

import torch
import torch.nn as nn

# Configuration representing a simplified LLM layer
batch_size = 1
num_heads = 8
head_dim = 64
hidden_dim = 512

# Projection layers for Keys and Values
W_k = nn.Linear(hidden_dim, num_heads * head_dim)
W_v = nn.Linear(hidden_dim, num_heads * head_dim)

# --- 1. PREFILL PHASE ---
# The prompt has 5 tokens. We process them simultaneously.
prompt_len = 5
prompt_embeddings = torch.randn(batch_size, prompt_len, hidden_dim)

# Project to K and V, then reshape to: [Batch, Heads, SeqLen, HeadDim]
kv_cache_k = W_k(prompt_embeddings).view(batch_size, prompt_len, num_heads, head_dim).transpose(1, 2)
kv_cache_v = W_v(prompt_embeddings).view(batch_size, prompt_len, num_heads, head_dim).transpose(1, 2)

print(f"Prefill complete. Cache K shape: {kv_cache_k.shape}")
# Output: [1, 8, 5, 64] (5 tokens cached)

# --- 2. DECODING PHASE ---
# We generate 3 subsequent tokens, one by one.
for step in range(3):
    # Input is only the single, newly generated token (SeqLen = 1)
    new_token_embedding = torch.randn(batch_size, 1, hidden_dim)

    # Project only the new token's K and V
    new_k = W_k(new_token_embedding).view(batch_size, 1, num_heads, head_dim).transpose(1, 2)
    new_v = W_v(new_token_embedding).view(batch_size, 1, num_heads, head_dim).transpose(1, 2)

    # Append the new token's representations to our historical cache
    # This prevents recalculating the past prompt_len + step tokens
    kv_cache_k = torch.cat([kv_cache_k, new_k], dim=2)
    kv_cache_v = torch.cat([kv_cache_v, new_v], dim=2)

    print(f"Decoding Step {step + 1} | Cache K shape: {kv_cache_k.shape}")
    # Step 1 Output: [1, 8, 6, 64]
    # Step 2 Output: [1, 8, 7, 64]
    # Step 3 Output: [1, 8, 8, 64]

def calculate_kv_cache_gb(batch_size: int, seq_len: int, layers: int, hidden_dim: int, precision_bytes: int = 2) -> float:
    # 2 represents the key and value states
    total_bytes = batch_size * seq_len * layers * hidden_dim * 2 * precision_bytes
    return total_bytes / (1024 ** 3)  # Convert to Gigabytes

# Llama 2 70B architecture parameters (assuming standard MHA)
llama_70b_gb = calculate_kv_cache_gb(
    batch_size=32,
    seq_len=8192,
    layers=80,
    hidden_dim=8192
)

print(f"KV Cache Size: {llama_70b_gb:.2f} GB")
# Output: KV Cache Size: 320.00 GB

+---------------------------------------------------------------+
|                      GPU VRAM (80 GB)                         |
+------------------------------+--------------------------------+
| Static Model Weights (Fixed) |  KV Cache (Dynamic / Variable) |
| [==========================] | [=======================>    ] |
| FP16 Weights (e.g., ~35GB/GPU| Grows linearly with:           |
| across split tensor setups)  | Batch Size * Sequence Length   |
+------------------------------+--------------------------------+

MHA (1:1 Ratio)       GQA (Grouped Ratio)     MQA (Many:1 Ratio)
[Q1][Q2][Q3][Q4]      [Q1][Q2]  [Q3][Q4]      [Q1][Q2][Q3][Q4]
  |   |   |   |         \   /     \   /         \   |   /   /
[K1][K2][K3][K4]        [ K1 ]   [ K2 ]          [   K1    ]
[V1][V2][V3][V4]        [ V1 ]   [ V2 ]          [   V1    ]

# A simple simulation of PagedAttention block mapping
class VirtualPageTable:
    def __init__(self, block_size: int):
        self.block_size = block_size
        self.logical_to_physical = {} # Maps (request_id, block_index) -> physical_block_id
        self.free_blocks = list(range(1000)) # Pool of physical memory blocks

    def allocate_block_for_token(self, request_id: int, token_index: int) -> int:
        # Determine which logical block the token falls into
        logical_block_idx = token_index // self.block_size
        key = (request_id, logical_block_idx)

        # Allocate a new physical block only when the previous one is full
        if key not in self.logical_to_physical:
            physical_block = self.free_blocks.pop(0)
            self.logical_to_physical[key] = physical_block

        return self.logical_to_physical[key]

# Instantiating a table with a block size of 4 tokens
table = VirtualPageTable(block_size=4)
# Token 5 belongs to logical block 1 (tokens 4-7). We allocate physical space on demand.
physical_id = table.allocate_block_for_token(request_id=42, token_index=5)
print(f"Logical token mapped to physical block: {physical_id}")

FP16/BF16 Value: [0.12345678]  (16 bits of precision)
     ↓ Quantization via Scaling Factor (e.g., scale = 127.0)
INT8 Value:      [15]          (8 bits of precision, 50% memory saved)

Sequence: [T1] [T2] [T3] [T4] [T5] [T6] [T7]
          └───── Window Size (4) ─────┘
          (T1 and T2 are evicted from active memory)

# Example: Initializing a vLLM engine with memory optimizations
from vllm import LLMEngine, EngineArgs

# Define engine arguments optimized for memory management
engine_args = EngineArgs(
    model="meta-llama/Llama-3-8B-Instruct",
    # 1. Quantization: Reduces KV cache footprint
    quantization="fp8",
    # 2. Paged KV Cache: Eliminates physical fragmentation
    gpu_memory_utilization=0.90,
    # 3. Dynamic Batching: Prevents static padding overhead
    max_num_seqs=256
)

# Initialize the engine with built-in memory guardrails
engine = LLMEngine.from_engine_args(engine_args)

Mecánicas de anexión

Las implementaciones almacenan buffers K/V por capa y escriben la nueva fila en cada paso. Los errores aquí corrompen silenciosamente el mecanismo de atención; es mejor utilizar kernels probados en la práctica.

La matemática de la memoria preocupante

Bytes ≈ 2 (K y V) × capas × kv_heads × head_dim × seq_len × bytes_per_element. Un contexto largo en modelos grandes puede consumir más memoria que los pesos propiamente dichos.

Técnicas de optimización

Rediseño arquitectónico: MHA, GQA, MQA

La atención de consulta agrupada y la de múltiples consultas reducen el número de heads KV, lo que hace que la caché sea más pequeña para la misma longitud de secuencia.

PagedAttention

Trate los datos KV como memoria virtual paginada: asigne bloques según sea necesario, reduzca la fragmentación y permita un mayor agrupamiento de datos.

Cuantización de KV

Almacene las entradas de caché con menor precisión cuando lo permitan los presupuestos de calidad.

Ventana deslizante / expulsión

Descarte o comprima los tokens alejados cuando la tarea tolere un rango de atención limitado.

Matriz de decisiones para producción

Elija modelos GQA/MQA cuando haya limitaciones de memoria; active la atención paginada para alta concurrencia; cuantifique solo KV después de la evaluación; utilice ventanas para chats que requieran contexto local.

Cierre

El caché KV es la base de datos invisible del decodificación autoregresiva. Óptimícelo explícitamente: arquitectura, paginación, precisión, expulsión; de lo contrario, se pagarán precios más altos en términos de GPUs y latencia.

Mida tanto los tokens por segundo como los bytes KV por token. Un núcleo “más rápido” que duplique el tamaño del caché puede perder rendimiento bajo cargas multiinquilino.

Documente el máximo contexto por SKU para que el producto no solicite silenciosamente 128k en una tarjeta provista para 8k.

Al trabajar con lotes, recuerde que los pasos de decodificación deben sincronizarse con la secuencia más larga; la política de relleno y el procesamiento continuo en lotes son tan importantes como el ancho de banda bruto.

Mida tanto los tokens por segundo como los bytes KV por token. Un núcleo “más rápido” que duplica el espacio de caché puede perder rendimiento bajo cargas multiinquilino.