llm inference kv-cache gpu pagedattention optimization

La mémoire KV est la base de données invisible qui consomme la mémoire GPU lors du décodage des LLM.

Le cache KV consomme discrètement la mémoire GPU pendant l’inférence des LLM : préremplissage contre décodage, calculs KV, GQA et MQA, PagedAttention, politiques de quantification et d’éviction.

Éditeur : REACTAPP.TOPAjouté sur ce site 26 sept. 20269 min de lecture2109 mots

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

Le véritable goulot d’étranglement dans l’inférence des LLM

Pendant l’inférence, les processeurs GPU restent souvent inactifs tandis que la bande passante mémoire alimente un cache clé-valeur en constante expansion. Chaque nouveau token ajoute des tenseurs K/V qui doivent être relus lors des étapes suivantes. Les déploiements qui ignorent cette croissance du cache atteignent rapidement des limites en termes de contexte et de latence, bien avant que les FLOPs ne soient épuisés.

En détail : comment fonctionne le cache clé-valeur

L’attention a besoin des clés et des valeurs de chaque token précédent. Leur mise en cache permet d’éviter de recalculer l’intégralité du préfixe à chaque étape. Le compromis réside dans une consommation mémoire proportionnelle à la longueur de la séquence × le nombre de couches × le nombre de têtes × la dimension.

Prefill contre décodage

Prefill traite la requête et remplit le cache. Décodage ajoute un token à la fois, en lisant l’intégralité du cache à chaque étape. Le Prefill est gourmand en calculs ; le décodage l’est en bande passante.

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)

Mécanismes d’ajout

Les implémentations stockent des buffers K/V par couche et écrivent la nouvelle ligne à chaque étape. Les erreurs dans ce domaine corrompent silencieusement le mécanisme d’attention — privilégiez des noyaux éprouvés en pratique.

Les calculs de mémoire complexes

Octets ≈ 2 (K et V) × nombre de couches × kv_heads × head_dim × longueur_de_la_séquence × octets_par_élément. Un contexte long sur de grands modèles peut consommer plus de mémoire que les poids eux-mêmes.

Techniques d’optimisation

Ré conception architecturale : MHA, GQA, MQA

L’attention à requêtes groupées et multi-requêtes réduit le nombre de KV heads, ce qui diminue la taille du cache pour une même longueur de séquence.

PagedAttention

Traitez les données KV comme de la mémoire virtuelle paginée : allouez des blocs sur demande, réduisez la fragmentation et permettez un lotage plus important.

Quantification des données KV

Stockez les entrées du cache avec une précision moindre lorsque le budget de qualité le permet.

Fenêtre glissante / éviction

Éliminez ou compressez les tokens éloignés lorsque la tâche tolère une portée d’attention limitée.

Matrice de décision pour la production

Sélectionnez les modèles GQA/MQA lorsque la mémoire est un facteur limitant ; activez l’attention paginée pour une forte concurrence ; quantifiez uniquement les données KV après évaluation ; utilisez des fenêtres pour les conversations nécessitant un contexte local.

Conclusion

Le cache KV est la base de données invisible utilisée pour le décodage autoregressif. Optimisez-le explicitement — architecture, pagination, précision, éviction — ou payez le prix en termes de GPU et de latence.

Mesurez à la fois le nombre de tokens par seconde et les octets KV par token. Un noyau « plus rapide » qui double l’espace occupé par le cache peut se révéler moins performant sous une charge multi-locataire.

Dokumentez la taille maximale du contexte par SKU afin que le produit ne puisse pas demander silencieusement 128 k de mémoire sur une carte prévue pour 8 k.

Lors du batchage, n’oubliez pas que les étapes de décodage doivent être synchronisées sur la séquence la plus longue ; la politique de remplissage et le batchage continu sont tout aussi importants que la bande passante brute.