llm inference kv-cache gpu pagedattention optimization

Der KV-Cache ist die unsichtbare Datenbank, die im Decodieren von LLMs GPU-Speicher verbraucht.

Der KV-Cache verbraucht während der LLM-Infersion leise den GPU-Speicher: Vorausfüllen versus Dekodieren, KV-Mathematik, GQA und MQA, PagedAttention, Quantisierungs- und Ausschließungsrichtlinien.

Verlag: REACTAPP.TOPAuf dieser Website hinzugefügt 26. Sept. 20269 Min. Lesezeit2109 Wörter

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

Das eigentliche Engpassproblem bei der Inferenz von LLMs

Während der Inferenz bleibt die GPU-Berechnung oft untätig, während das Speicherbandbreite einen stetig wachsenden Key–Value-Cache versorgt. Jedes neue Token fügt K/V-Tensoren hinzu, die in späteren Schritten erneut gelesen werden müssen. Implementierungen, die das Wachstum des KV-Caches ignorieren, stoßen bereits lange vor dem Auslaufen der FLOPs an Grenzen hinsichtlich Kontextmenge und Latenz.

Hinter den Kulissen: Wie der KV-Cache funktioniert

Die Attention-Verarbeitung benötigt für jedes vorherige Token Schlüssel und Werte. Durch das Caching wird verhindert, dass der gesamte Präfix in jedem Schritt neu berechnet werden muss. Der Kompromiss besteht darin, dass der Speicherverbrauch linear mit der Länge der Sequenz, der Anzahl der Schichten, der Anzahl der Heads sowie der Dimension wächst.

Prefill gegen Decoding

Prefill verarbeitet den Prompt und füllt den Cache. Decoding fügt einen Token nach dem anderen hinzu und liest in jedem Schritt den gesamten Cache. Prefill ist rechenintensiv, während Decoding bandbreitenintensiv ist.

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)

Mechaniken des Hinzufügens

In Implementierungen werden pro Schicht K/V-Puffer gespeichert und jede Schritt wird eine neue Zeile geschrieben. Fehler hier führen heimlich zu Störungen in der Aufmerksamkeitsverarbeitung – es sind lieber erprobte Kernel zu verwenden.

Die ernüchternde Rechenlogik bezüglich des Speichers

Bytes ≈ 2 (K und V) × Schichten × kv_heads × head_dim × seq_len × bytes_per_element. Lange Kontexte in großen Modellen können mehr Speicher verbrauchen als die Gewichte selbst.

Optimierungstechniken

Architektonische Neugestaltung: MHA, GQA, MQA

Gruppierte- und mehrfache Abfragen zur Aufmerksamkeitsverarbeitung verringern die Anzahl der KV-Heads, wodurch für dieselbe Sequenzlänge weniger Cache benötigt wird.

PagedAttention

KV-Daten wie paginierte virtuellen Speicher behandeln: Blöcke nach Bedarf zuweisen, Fragmentierung verringern und höheres Batching ermöglichen.

KV-Quantisierung

Speichern Sie Cache-Einträge bei geringerer Präzision, wenn die Qualitätsvorgaben es zulassen.

Schiebefenster-/Eviction-Verfahren

Löschen oder komprimieren Sie entfernte Token, wenn die Aufgabe eine begrenzte Aufmerksamkeitsspanne zulässt.

Entscheidungsmatrix für die Produktion

Wählen Sie GQA/MQA-Modelle, wenn der Speicher begrenzt ist; aktivieren Sie paginierte Aufmerksamkeit für hohe Konkurrenzgrade; quantisieren Sie KV nur nach der Bewertung; verwenden Sie Schiebefenster für Chats mit lokalen Kontextanforderungen.

Fazit

Der KV-Cache ist die unsichtbare Datenbank der autoregressiven Dekodierung. Optimieren Sie ihn gezielt – in Bezug auf Architektur, Paging, Präzision und Eviction – oder zahlen Sie dafür mit GPUs und Verzögerungen.

Messen Sie sowohl Tokens pro Sekunde als auch KV-Bytes pro Token. Ein „schnellerer“ Kernel, der die Cache-Größe verdoppelt, kann unter mehrfach genutzten Lasten an Leistung verlieren.

Dokumentieren Sie den maximal zulässigen Kontext pro SKU, damit Produkte nicht heimlich 128k anfordern können, obwohl nur 8k bereitgestellt sind.

Beim Batchen sollten Sie beachten, dass die Dekodierungs Schritte mit der längsten Sequenz synchronisiert werden müssen; die Pufferverwaltungspolitik sowie das kontinuierliche Batchen sind genauso wichtig wie die reine Bandbreite.

Messen Sie sowohl Tokens pro Sekunde als auch KV-Bytes pro Token. Ein „schnellerer“ Kernel, der den Cachebedarf verdoppelt, kann unter Mehrbenutzerlast Nachteile aufweisen.