llm inference kv-cache gpu pagedattention optimization

KV缓存是隐藏在背后的数据库，正在消耗LLM解码过程中的GPU内存。

在大语言模型推理过程中，KV缓存会悄悄消耗GPU内存：预填充与解码、KV运算、GQA与MQA、分页注意力机制、量化及驱逐策略。

发布者：REACTAPP.TOP加入本站 2026年9月26日阅读约 9 分钟2109 词

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

大语言模型推理中的真正瓶颈

在推理过程中，GPU的计算能力往往处于闲置状态，而内存带宽则持续为不断扩大的键值缓存提供数据。每个新生成的标记都会增加K/V张量，这些张量在后续步骤中需要被重新读取。那些忽视键值缓存增长问题的系统，在浮点运算次数耗尽之前早已面临上下文限制和延迟问题。

底层原理：键值缓存的工作机制

注意力机制需要每个先前标记的键和值，将它们缓存起来可以避免在每一步都重新计算整个前缀内容。但相应的代价是内存占用与序列长度×层数×头数×维度成正比。

预填充与解码

预填充模式会先处理提示语并填满缓存；而解码模式则是逐个添加标记，每一步都需要读取整个缓存。预填充模式计算量大，解码模式则对带宽要求较高。

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)

追加机制

实现方式会为每一层存储K/V缓冲区，并在每一步写入新行。此处的错误会悄无声息地破坏注意力机制——建议使用经过实际验证的核函数。

令人警醒的内存计算

字节数 ≈ 2（K和V）× 层数 × kv_heads × head_dim × 序列长度 × 每个元素的字节数。大型模型处理长上下文时所需的内存可能超过权重本身占用的内存。

优化技术

架构重构：MHA、GQA、MQA

通过分组查询和多查询注意力机制缩小KV头数量，从而在相同序列长度下减少缓存需求。

PagedAttention

将KV视为分页虚拟内存：按需分配块，降低碎片化程度，支持更高的批量处理。

KV量化

在质量预算允许的情况下，以较低的精度存储缓存条目。

滑动窗口/淘汰机制

当任务允许的注意力范围有限时，丢弃或压缩较远的标记。

生产环境决策矩阵

在内存受限时选择GQA/MQA模型；在高并发场景下启用分页注意力机制；仅在评估完成后对KV进行量化处理；在需要本地上下文的聊天场景中使用窗口机制。