从零理解
KV Cache 与推理加速

安装 Python 依赖

本教程需要 Python 3.10+、PyTorch 2.x。推荐使用 conda 或 virtualenv 隔离环境。

# 创建虚拟环境
python -m venv kvcache-env
source kvcache-env/bin/activate  # Windows: kvcache-env\Scripts\activate

# 安装依赖
pip install torch transformers numpy

# 验证 PyTorch 版本
python -c "import torch; print(torch.__version__)"  # >= 2.0.0

验证 GPU 可用性（可选）

本教程的代码在 CPU 上也可运行，但 GPU 能更直观地体现加速效果。

import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"使用设备: {device}")

if device == "cuda":
    print(f"GPU 型号: {torch.cuda.get_device_name(0)}")
    print(f"GPU 显存: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")

实现无缓存版 Self-Attention

import torch
import torch.nn as nn
import torch.nn.functional as F
import time

class SelfAttentionNaive(nn.Module):
    """无 KV Cache 的朴素 Self-Attention"""
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads

        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x):
        """
        x: (batch, seq_len, d_model)
        每次调用都对整个序列重新计算 K 和 V！
        """
        B, T, C = x.shape

        # 计算所有 token 的 Q, K, V
        Q = self.W_q(x)  # (B, T, d_model)
        K = self.W_k(x)  # (B, T, d_model)  ← 重复计算！
        V = self.W_v(x)  # (B, T, d_model)  ← 重复计算！

        # 多头分割
        Q = Q.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        K = K.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        V = V.view(B, T, self.n_heads, self.d_head).transpose(1, 2)

        # Attention 得分 (因果 mask)
        scale = self.d_head ** -0.5
        scores = (Q @ K.transpose(-2, -1)) * scale  # (B, heads, T, T)
        mask = torch.tril(torch.ones(T, T, device=x.device)).bool()
        scores = scores.masked_fill(~mask, float('-inf'))
        attn = F.softmax(scores, dim=-1)

        out = attn @ V  # (B, heads, T, d_head)
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.W_o(out)


def simulate_naive_decode(model, prompt_emb, n_new_tokens):
    """模拟无缓存自回归解码，每步都把整个序列传入"""
    B, T, C = prompt_emb.shape
    tokens = prompt_emb.clone()

    times = []
    for i in range(n_new_tokens):
        t0 = time.perf_counter()
        # 每步传入全部 token，包含所有历史
        out = model(tokens)               # (B, T+i, C)
        new_token = out[:, -1:, :]       # 只取最后一个 token 的输出
        tokens = torch.cat([tokens, new_token], dim=1)
        times.append(time.perf_counter() - t0)

    return tokens, times


# 测试
d_model, n_heads, seq_len = 512, 8, 50
model_naive = SelfAttentionNaive(d_model, n_heads)
prompt = torch.randn(1, seq_len, d_model)

with torch.no_grad():
    _, naive_times = simulate_naive_decode(model_naive, prompt, n_new_tokens=20)

print(f"无缓存 - 首 token: {naive_times[0]*1000:.1f}ms")
print(f"无缓存 - 末 token: {naive_times[-1]*1000:.1f}ms")
print(f"无缓存 - 总耗时: {sum(naive_times)*1000:.1f}ms")

实现带 KV Cache 的 Self-Attention

class KVCache:
    """单层 Attention 的 KV 缓存"""
    def __init__(self):
        self.k_cache = None  # (B, heads, T_cached, d_head)
        self.v_cache = None  # (B, heads, T_cached, d_head)

    def update(self, k_new, v_new):
        """追加新 token 的 K/V 到缓存"""
        if self.k_cache is None:
            self.k_cache = k_new
            self.v_cache = v_new
        else:
            self.k_cache = torch.cat([self.k_cache, k_new], dim=2)
            self.v_cache = torch.cat([self.v_cache, v_new], dim=2)
        return self.k_cache, self.v_cache

    def get_seq_len(self):
        return 0 if self.k_cache is None else self.k_cache.shape[2]

    def clear(self):
        self.k_cache = None
        self.v_cache = None


class SelfAttentionWithCache(nn.Module):
    """带 KV Cache 的 Self-Attention"""
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads

        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)

        self.kv_cache = KVCache()

    def forward(self, x, use_cache=True):
        """
        x:
          - Prefill 阶段: (B, T_prompt, d_model)
          - Decode  阶段: (B, 1, d_model) ← 只传入 1 个新 token！
        """
        B, T, C = x.shape

        # 只为当前输入的 token 计算 Q、K、V
        Q = self.W_q(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        K = self.W_k(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        V = self.W_v(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2)

        if use_cache:
            # 将新的 K/V 追加到缓存，返回完整的历史 K/V
            K_full, V_full = self.kv_cache.update(K, V)
        else:
            K_full, V_full = K, V

        # Q 只是当前 token 的，但 K_full/V_full 包含所有历史
        T_kv = K_full.shape[2]  # 完整的 KV 长度
        scale = self.d_head ** -0.5
        scores = (Q @ K_full.transpose(-2, -1)) * scale  # (B, heads, T, T_kv)

        # 因果 mask：Q 的位置 i 只能 attend 到 KV 中 <= i 的位置
        T_q = Q.shape[2]
        offset = T_kv - T_q  # 解码时 offset = 已缓存的长度
        mask = torch.ones(T_q, T_kv, device=x.device).tril(diagonal=offset).bool()
        scores = scores.masked_fill(~mask, float('-inf'))

        attn = F.softmax(scores, dim=-1)
        out = attn @ V_full  # (B, heads, T_q, d_head)
        out = out.transpose(1, 2).contiguous().view(B, T_q, C)
        return self.W_o(out)


def simulate_cached_decode(model, prompt_emb, n_new_tokens):
    """带 KV Cache 的自回归解码"""
    # Step 1: Prefill — 一次性处理完整 prompt，填充缓存
    model.kv_cache.clear()
    t_prefill = time.perf_counter()
    with torch.no_grad():
        out = model(prompt_emb, use_cache=True)
    last_emb = out[:, -1:, :]  # 取最后一个 token 的输出作为第一个新 token 输入
    prefill_time = time.perf_counter() - t_prefill

    # Step 2: Decode — 每步只传入 1 个 token
    times = []
    current_emb = last_emb
    for i in range(n_new_tokens):
        t0 = time.perf_counter()
        with torch.no_grad():
            out = model(current_emb, use_cache=True)  # (B, 1, C)
        current_emb = out  # 下一步的输入
        times.append(time.perf_counter() - t0)

    return times, prefill_time


# 对比测试
d_model, n_heads, seq_len = 512, 8, 50
model_cached = SelfAttentionWithCache(d_model, n_heads)
# 共享权重（保证公平对比）
model_cached.W_q.weight.data = model_naive.W_q.weight.data.clone()
model_cached.W_k.weight.data = model_naive.W_k.weight.data.clone()
model_cached.W_v.weight.data = model_naive.W_v.weight.data.clone()
model_cached.W_o.weight.data = model_naive.W_o.weight.data.clone()

prompt = torch.randn(1, seq_len, d_model)
cached_times, prefill_time = simulate_cached_decode(model_cached, prompt, n_new_tokens=20)

print(f"\n===== 对比结果 =====")
print(f"无缓存 - 总耗时: {sum(naive_times)*1000:.1f}ms")
print(f"有缓存 - Prefill: {prefill_time*1000:.1f}ms | Decode 总计: {sum(cached_times)*1000:.1f}ms")
print(f"解码加速比: {sum(naive_times)/sum(cached_times):.1f}x")

用 HuggingFace Transformers 对比有/无缓存的实际速度

HuggingFace 的 transformers 库默认开启 KV Cache（use_cache=True），你也可以手动关闭来做对比测试。

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import time

model_name = "gpt2"  # 小模型，无需 GPU 也能快速测试
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32)
model.eval()

prompt = "The future of artificial intelligence is"
inputs = tokenizer(prompt, return_tensors="pt")
input_ids = inputs["input_ids"]

n_new_tokens = 50

# ─── 方案 A: 禁用 KV Cache ───
t0 = time.perf_counter()
with torch.no_grad():
    out_no_cache = model.generate(
        input_ids,
        max_new_tokens=n_new_tokens,
        use_cache=False,        # ← 关键：禁用 KV Cache
        do_sample=False,
    )
time_no_cache = time.perf_counter() - t0

# ─── 方案 B: 启用 KV Cache（默认）───
t0 = time.perf_counter()
with torch.no_grad():
    out_with_cache = model.generate(
        input_ids,
        max_new_tokens=n_new_tokens,
        use_cache=True,         # ← 默认开启
        do_sample=False,
        cache_implementation="static",  # 可选: "dynamic"(默认), "static", "quantized_static"
    )
time_with_cache = time.perf_counter() - t0

print(f"\n===== 实际模型对比 =====")
print(f"无 KV Cache — 生成 {n_new_tokens} tokens 耗时: {time_no_cache:.2f}s")
print(f"有 KV Cache — 生成 {n_new_tokens} tokens 耗时: {time_with_cache:.2f}s")
print(f"加速比: {time_no_cache / time_with_cache:.1f}x")
print(f"\n生成文本: {tokenizer.decode(out_with_cache[0], skip_special_tokens=True)}")