图解 KV Cache 优化：大模型推理加速的核心技术

首先实现最基础的 KV Cache 机制：

class KVCache:
    def __init__(self, num_layers, num_heads, head_dim, max_seq_len, dtype=torch.float16):
        self.num_layers = num_layers
        # K 和 V 分开缓存，shape: [num_layers, batch, num_heads, seq_len, head_dim]
        self.k_cache = torch.zeros(num_layers, 1, num_heads, max_seq_len, head_dim, dtype=dtype, device='cuda')
        self.v_cache = torch.zeros(num_layers, 1, num_heads, max_seq_len, head_dim, dtype=dtype, device='cuda')
        self.seq_len = 0

    def update(self, layer_idx, k, v):
        """更新缓存并返回完整序列的 KV"""
        bs, num_heads, _, head_dim = k.shape
        self.k_cache[layer_idx, :bs, :self.seq_len + k.shape[2], :] = k
        self.v_cache[layer_idx, :bs, :self.seq_len + v.shape[2], :] = v
        return (
            self.k_cache[layer_idx, :bs, :self.seq_len + k.shape[2], :],
            self.v_cache[layer_idx, :bs, :self.seq_len + v.shape[2], :]
        )

    def reset(self):
        """重置缓存（新请求时调用）"""
        self.k_cache.zero_()
        self.v_cache.zero_()
        self.seq_len = 0

这是最简单的全量预分配方案，适合固定长度的场景。

多轮对话场景下，系统提示词（System Prompt）是共享前缀。使用增量 Prefill 可以避免重复计算：

def incremental_prefill(model, shared_prefix, user_input):
    """
    增量 Prefill：共享前缀只计算一次
    """
    # 第一步：计算共享前缀的 KV
    with torch.no_grad():
        shared_kv = model.encode(shared_prefix)  # 只计算一次

    # 第二步：用户输入复用共享 KV
        response = model.generate(
            user_input,
            past_key_values=shared_kv,  # 复用缓存
            max_new_tokens=512
        )

    return response

# 实战示例
system_prompt = "你是一个专业的 Python 开发者助手..."
questions = [
    "如何优化 SQL 查询性能？",
    "Docker 容器网络不通怎么排查？",
    "FastAPI 如何实现 JWT 认证？"
]

# 共享 system prompt 的 KV cache
for q in questions:
    answer = incremental_prefill(model, system_prompt, q)
    print(f"Q: {q}\nA: {answer}\n")

vLLM 的核心创新 PagedAttention 解决了传统 KV Cache 的内存碎片问题。它将 KV Cache 分成固定大小的 block，动态分配给不同请求：

class PagedAttentionKVCache:
    def __init__(self, num_layers, num_heads, head_dim, block_size=16, max_blocks=1024):
        self.block_size = block_size  # 每块存 16 个 token
        self.max_blocks = max_blocks
        # 所有 block 集中存储，避免碎片
        self.k_blocks = torch.zeros(num_layers, max_blocks, num_heads, block_size, head_dim,
                                    dtype=torch.float16, device='cuda')
        self.v_blocks = torch.zeros(num_layers, max_blocks, num_heads, block_size, head_dim,
                                    dtype=torch.float16, device='cuda')
        # 记录每个序列占用了哪些 block
        self.block_tables = {}  # {seq_id: [block_id_1, block_id_2, ...]}
        self.seq_lens = {}

    def allocate(self, seq_id, num_tokens):
        """动态分配 block"""
        num_blocks = (num_tokens + self.block_size - 1) // self.block_size
        available = self._get_free_blocks()
        if len(available) < num_blocks:
            raise RuntimeError("OOM: No free blocks")

        allocated = available[:num_blocks]
        self.block_tables[seq_id] = allocated
        self.seq_lens[seq_id] = num_tokens
        return allocated

    def write(self, seq_id, token_offset, k, v):
        """写入 token 到对应 block"""
        block_idx = token_offset // self.block_size
        offset_in_block = token_offset % self.block_size
        block_id = self.block_tables[seq_id][block_idx]

        self.k_blocks[:, block_id, :, offset_in_block:offset_in_block+k.shape[-2], :] = k
        self.v_blocks[:, block_id, :, offset_in_block:offset_in_block+v.shape[-2], :] = v

传统 batching 必须等待所有请求完成才能释放 slot，而 Continuous Batching 允许请求完成后立即腾出资源：

class ContinuousBatchScheduler:
    def __init__(self, max_batch_size, kv_cache):
        self.max_batch_size = max_batch_size
        self.kv_cache = kv_cache
        self.active_requests = {}  # {req_id: {"tokens": [...], "done": bool}}
        self.pending_queue = deque()

    def schedule_step(self):
        """单步调度：处理已完成的请求，加入新请求"""
        # 1. 移除已完成的请求
        done_ids = [rid for rid, req in self.active_requests.items() if req["done"]]
        for rid in done_ids:
            del self.active_requests[rid]
            self.kv_cache.free(rid)  # 立即释放 KV 缓存

        # 2. 从队列填充新请求（有空位时）
        while len(self.active_requests) < self.max_batch_size and self.pending_queue:
            new_req = self.pending_queue.popleft()
            self.active_requests[new_req["id"]] = new_req

        # 3. 对当前 batch 执行推理
        if self.active_requests:
            self._run_inference_step()

    def _run_inference_step(self):
        """执行单步推理"""
        # 收集当前 batch 的所有 token
        batch_tokens = [req["tokens"][-1:] for req in self.active_requests.values()]

        # 模型 forward（关键：一次 forward 服务多个请求）
        logits, new_kv = self.model.forward(batch_tokens, use_cache=True)

        # 更新 KV cache
        for idx, (rid, req) in enumerate(self.active_requests.items()):
            self.kv_cache.update(rid, new_kv[idx])

            # 检查是否生成结束符
            if EOS_TOKEN in logits[idx]:
                req["done"] = True

2026 年 vLLM 和 SGLang 都支持 Prefix Caching，自动识别并缓存共享前缀：

from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Llama-3-8B-Instruct",
    enable_prefix_caching=True,  # 开启前缀缓存
    gpu_memory_utilization=0.95,
    max_num_seqs=256,
)

# 模拟多轮对话场景（共享 system prompt）
prompts = [
    "System: 你是代码助手。\nUser: 写个快速排序",
    "System: 你是代码助手。\nUser: 解释一下上面的时间复杂度",
    "System: 你是代码助手。\nUser: 改成原地排序版本",
]

sampling_params = SamplingParams(temperature=0.7, max_tokens=256)
outputs = llm.generate(prompts, sampling_params)

# 第二次请求时，"System: 你是代码助手。" 的 KV 会被复用

2026 年新趋势是使用量化减少 KV Cache 显存占用。INT8 量化可将 KV Cache 压缩 50%：

class QuantizedKVCache:
    def __init__(self, shape, dtype=torch.int8):
        self.scale = None
        self.k_cache_int8 = torch.zeros(shape, dtype=torch.int8, device='cuda')
        self.v_cache_int8 = torch.zeros(shape, dtype=torch.int8, device='cuda')

    def quantize(self, tensor):
        """FP16 -> INT8 量化"""
        scale = tensor.abs().max() / 127  # 对称量化
        return (tensor / scale).to(torch.int8), scale

    def dequantize(self, tensor_int8, scale):
        """INT8 -> FP16 反量化"""
        return tensor_int8.float() * scale

    def update(self, k, v):
        # 量化后存储
        k_int8, k_scale = self.quantize(k)
        v_int8, v_scale = self.quantize(v)

        self.k_cache_int8[...] = k_int8
        self.v_cache_int8[...] = v_int8
        self.scale = (k_scale, v_scale)

    def get(self):
        # 使用时反量化
        k_scale, v_scale = self.scale
        k = self.dequantize(self.k_cache_int8, k_scale)
        v = self.dequantize(self.v_cache_int8, v_scale)
        return k, v

# 显存对比
# FP16: 2GB -> INT8: 1GB (节省 50%)

更激进的方案是 KV Cache 卸载 (Offload)：将旧 token 的 KV 移到 CPU 内存，需要时再加载回 GPU。

使用 vLLM 的内置指标监控 KV Cache 使用率：

import requests

# 查询 vLLM 指标
response = requests.get("http://localhost:8000/metrics")
metrics = response.text

# 关键指标
# vllm:gpu_cache_usage_percent: GPU KV 缓存使用率
# vllm:cpu_cache_usage_percent: CPU KV 缓存使用率
# vllm:num_requests_waiting: 等待调度的请求数
# vllm:time_to_first_token_seconds: 首 token 延迟

# 调优建议
# - gpu_cache_usage_percent > 90%：增加 max_num_blocks 或减小 max_seq_len
# - num_requests_waiting > 10：增加 max_batch_size
# - time_to_first_token_seconds > 1：开启 Prefix Caching