KV、Prefix、Prompt 与 Context Caching

KV cache
├── layer 1
│   ├── token 1: K, V
│   ├── token 2: K, V
│   └── ...
├── layer 2
│   ├── token 1: K, V
│   ├── token 2: K, V
│   └── ...
└── layer N

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache

model_id = "HuggingFaceTB/SmolLM2-360M-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto")

inputs = tokenizer("The capital of France is", return_tensors="pt").to(model.device)
past_key_values = DynamicCache(config=model.config)

out = model.generate(**inputs, do_sample=False, max_new_tokens=20, past_key_values=past_key_values)

print(tokenizer.decode(out[0], skip_special_tokens=True))
# The capital of France is Paris. It is the largest city in France and the second-largest city in the European Union.
print("prompt tokens:", inputs["input_ids"].shape[1])    # 5
print("total tokens:", out.shape[1])                     # 25
print("cache length:", past_key_values.get_seq_length()) # 24

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache

model_id = "HuggingFaceTB/SmolLM2-360M-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto")

past_key_values = DynamicCache(config=model.config)
messages = []
questions = ["What is the capital of France?", "And its population?"]

for prompt in questions:
    messages.append({"role": "user", "content": prompt})
    inputs = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt", return_dict=True
    ).to(model.device)
    input_length = inputs["input_ids"].shape[1]
    outputs = model.generate(**inputs, do_sample=False, max_new_tokens=64,
                             past_key_values=past_key_values)
    completion = tokenizer.decode(outputs[0, input_length:], skip_special_tokens=True)
    messages.append({"role": "assistant", "content": completion})
    print(f"turn tokens in: {input_length} | cache now: {past_key_values.get_seq_length()}")

# turn tokens in: 42 | cache now: 55
# turn tokens in: 71 | cache now: 92

请求 1：A B C
请求 2：A B D

K(A)、V(A)       只依赖 A
K(B)、V(B)       依赖 A B
K(C)、V(C)       依赖 A B C

K(A)、V(A)
K(B)、V(B)

A B 的 KV：复用
D 的 KV：重新计算

输入 token
  ↓
切成 16-token block
  ↓
逐块计算链式 hash
  ↓
从第一个 block 开始查缓存
  ↓
命中：复用 KV
miss：停止查找
  ↓
miss 之后的 token 重新 prefill

BLOCK_SIZE = 16

def block_hashes(token_ids, salt=None):
    """把 token 序列切成块并链式哈希，得到每块的 key。"""
    hashes, parent = [], hash(salt)
    # 只哈希完整块，尾部不足 16 的残块跳过
    for start in range(0, len(token_ids) - BLOCK_SIZE + 1, BLOCK_SIZE):
        block = tuple(token_ids[start : start + BLOCK_SIZE])
        parent = hash((parent, block))
        hashes.append(parent)
    return hashes

def schedule(token_ids, cache):
    """返回可复用多少 token，其余重新分配。"""
    matched_blocks = 0
    for h in block_hashes(token_ids):
        if h not in cache:
            break                    # 第一个 miss 结束全部复用
        cache[h].ref_count += 1      # 钉住，防止被淘汰
        matched_blocks += 1
    reused_tokens = matched_blocks * BLOCK_SIZE
    to_prefill = token_ids[reused_tokens:]
    return reused_tokens, to_prefill

先计算共享前缀
  ↓
保存 prefix KV
  ↓
每个请求复制一份 prefix KV
  ↓
在各自副本后面继续生成

import copy
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, StaticCache

# ---------- 1. 加载模型和分词器 ----------
model_id = "HuggingFaceTB/SmolLM2-360M-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

# ---------- 2. 定义共享的前缀（所有请求共用的系统指令） ----------
SHARED_PREFIX = """You are a careful assistant.
                   Answer in one short sentence."""

# ---------- 3. 创建静态 KV 缓存对象 ----------
# StaticCache 是 transformers 提供的 KV 缓存容器，可预先分配固定长度的内存，
# 避免动态扩展带来的开销。这里指定最大缓存长度为 1024 个 token。
prompt_cache = StaticCache(
    config=model.config,   # 从模型配置中读取层数、头数、维度等参数
    max_cache_len=1024
)

# ---------- 4. 对共享前缀进行分词，并移到模型设备 ----------
prefix_inputs = tokenizer(SHARED_PREFIX, return_tensors="pt").to(model.device)

# ---------- 5. 第一次前向传播：只处理前缀（Prefill），不生成新 token ----------
# 目的：将前缀的所有 token 对应的 Key/Value 计算出来并填入 StaticCache 中。
# 之后所有请求都可以复用这份已计算好的 KV 缓存，避免重复计算前缀部分的注意力。
with torch.no_grad():  # 禁用梯度计算，节省显存和计算
    # 调用模型时传入 past_key_values=prompt_cache，模型会将当前输入的 KV
    # 追加到缓存中（而不是覆盖）。返回的 past_key_values 就是更新后的缓存。
    prompt_cache = model(
        **prefix_inputs,
        past_key_values=prompt_cache
    ).past_key_values

# ---------- 6. 准备多个不同的用户问题 ----------
questions = ["What is the capital of France?", "Name one ocean."]

# ---------- 7. 对每个问题分别生成回答 ----------
for question in questions:
    # 7.1 将共享前缀和当前问题拼接后分词
    # 注意：此时 tokenizer 会把整个字符串（前缀+问题）一起编码，
    # 但由于前缀部分已经在缓存中，后续生成时会自动跳过重复计算。
    inputs = tokenizer(
        SHARED_PREFIX + question,
        return_tensors="pt"
    ).to(model.device)

    # 7.2 **关键操作：深拷贝（deepcopy）共享缓存**
    # 原因：每个请求的生成过程会不断向缓存中追加新的 KV（生成的新 token），
    # 如果直接使用同一个缓存对象，多个请求会互相污染（串行时也会保留其他请求的生成历史）。
    # 深拷贝确保每个请求拥有自己独立的缓存副本，初始状态完全一致（仅含前缀内容）。
    past_key_values = copy.deepcopy(prompt_cache)

    # 7.3 生成答案
    # 传入 past_key_values 后，模型知道前导 token 的 KV 已经存在，
    # 因此只需处理“问题部分”（即 inputs 中除前缀以外的 token）并逐步生成新 token。
    # do_sample=False 表示使用贪心解码（确定性）。
    outputs = model.generate(
        **inputs,
        past_key_values=past_key_values,
        do_sample=False
    )

    # 7.4 解码并打印完整回复（包含用户问题）
    print(tokenizer.decode(outputs[0], skip_special_tokens=True))

共享物理 KV block
  + 每个请求维护自己的逻辑路径
  + 引用计数管理生命周期

Transformers demo：复制缓存，容易理解
生产推理引擎：共享物理块，节省内存

缓存占得越多
→ 留给新请求的 block 越少
→ 并发能力可能下降

命中收益
vs
显存占用和并发损失

import anthropic

client = anthropic.Anthropic()   # 从环境变量读 ANTHROPIC_API_KEY

# 必须超过模型的最小可缓存长度，否则什么都不会被缓存
LONG_INSTRUCTIONS = "You are a precise technical editor. " * 400

def ask(question: str):
    return client.messages.create(
        model="claude-sonnet-4-6",
        max_tokens=512,
        system=[
            {
                "type": "text",
                "text": LONG_INSTRUCTIONS,
                "cache_control": {"type": "ephemeral"},   # 标记以上全部可缓存
            }
        ],
        messages=[{"role": "user", "content": question}],
    )

for question in ["Summarize section 3.", "Now rewrite it for a beginner."]:
    resp = ask(question)
    u = resp.usage
    print(f"write={u.cache_creation_input_tokens} read={u.cache_read_input_tokens} uncached={u.input_tokens}")

# write=2823  read=0     uncached=14
# write=0     read=2823  uncached=17

系统指令 + 固定上下文 + query

系统指令 + 检索块 A/B/C + query

请求 1：A + B + C
请求 2：C + A + B

A → B → C
C → A → B

普通 prefix caching：只能复用连续前缀
CacheBlend：可以复用乱序、分散的检索块