llm inference kv-cache gpu pagedattention optimization

Pamięć KV to niewidzialna baza danych, która zużywa pamięć GPU podczas dekodowania modeli LLM.

Cache KV w cichy sposób zużywa pamięć GPU podczas inferencji modeli LLM: prefill kontra decode, matematyka KV, GQA i MQA, PagedAttention, polityki kwantyzacji i usuwania danych.

Wydawca: REACTAPP.TOPDodane na tej stronie 26 wrz 20269 min czytania2109 słów

Źródła: The KV Cache: AI’s Unseen Database Dominating GPU Memory

Prawdziwy wąskie gardło w inferencji modeli LLM

Podczas inferencji procesor GPU często pozostaje bezczynny, podczas gdy przepustowość pamięci zaopatruje coraz większą pamietć cache typu key–value. Każdy nowy token dodaje tensory K/V, które muszą być ponownie odczytywane na kolejnych etapach. Rozwiązania, które ignorują wzrost rozmiaru tej pamięci cache, napotykają ograniczenia związane z kontekstem i opóźnieniami na długo przed wyczerpaniem zasobów FLOPs.

Jak to działa: mechanizm pamięci cache typu key–value

Aby zastosować mechanizm uwagi, potrzebne są klucze i wartości każdego poprzedniego tokena. Ich przechowywanie w cache pozwala uniknąć ponownego obliczania całego prefiksu na każdym kroku. Kompromisem jest zużycie pamięci, które jest liniowe w zależności od długości sekwencji × liczby warstw × liczby głowic obliczeniowych × wymiaru.

Prefill vs dekodowanie

Prefill przetwarza treść promptu i wypełnia cache. Dekodowanie dodaje po jednym tokenie na raz, odczytując całą zawartość cache na każdym kroku. Metoda Prefill wymaga dużych zasobów obliczeniowych, natomiast dekodowanie jest uzależnione od przepustowości pamięci.

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)

Mechanizmy dodawania

W implementacjach przechowywane są buforzy K/V dla każdej warstwy, a nowe wiersze zapisywane są na każdym kroku. Błędy w tym obszarze potajemnie psują funkcjonowanie mechanizmów uwagi — lepiej używać sprawdzonych rozwiązań.

Rzeczowa matematyka pamięci

Bajty ≈ 2 (K i V) × liczba warstw × kv_heads × head_dim × seq_len × bytes_per_element. Długi kontekst w dużych modelach może wymagać więcej pamięci niż same wagi.

Techniki optymalizacji

Ponowny projekt architektury: MHA, GQA, MQA

Metody typu grouped-query i multi-query zmniejszają liczbę głów KV, dzięki czemu pamięć cache jest mniejsza przy tej samej długości sekwencji.

PagedAttention

Traktuj KV jak pamięć wirtualną z paginacją: przydzielaj bloki na żądanie, zmniejsz fragmentację i umożliwij większe grupowanie danych.

Kwantyzacja KV

Zachowuj wpisy cache o niższej precyzji, jeśli budżet jakości na to pozwala.

Szerokość okna / wyrzucanie

Pomijaj lub kompresuj dalekie tokeny, gdy zadanie toleruje ograniczony zakres uwagi.

Matryca decyzji produkcyjnych

Wybieraj modele GQA/MQA, gdy występuje ograniczenie pamięci; włącz pagowanie uwagi przy wysokiej konkurencji; kwantyzuj KV dopiero po ocenie; używaj okien do rozmów wymagających lokalnego kontekstu.

Zakończenie

Cache KV to niewidzialna baza danych używana do dekodowania autoregresywnego. Optymalizuj go wyraźnie – architekturą, pagowaniem, precyzją, mechanizmami wyrzucania – albo zapłacisz za to za pomocą GPU i wydłużonego czasu reakcji.

Mierz liczbę tokenów na sekundę oraz ilość bajtów KV na token. „Szybszy” jądro, które podwaja rozmiar cache, może przegrać w warunkach obciążenia wieloma użytkownikami.

Zdokumentuj maksymalny kontekst na SKU, aby produkt nie mógł potajemnie żądać 128k na karcie przeznaczonej dla 8k.

Podczas grupowania zadań pamiętaj, aby kroki dekodowania synchronizowały się z najdłuższą sekwencją; polityka wypełniania i ciągłe grupowanie zadań mają takie samo znaczenie jak surowa przepustowość.

Mierz tokeny na sekundę oraz bajty KV na token razem. „Szybszy” jądro, które podwaja rozmiar pamięci cache, może stracić wydajność przy obciążeniu wielu użytkowników.