训练时我们关心一个 batch 的 forward/backward/optimizer step。 推理时我们关心的是一条用户请求从进来到返回 token 的全过程: 它什么时候排队、什么时候拿到 GPU、什么时候返回第一个 token、后续 token 以多快速度流出。

LLM 推理不是一个均匀过程。prefill 阶段像“读完整段题目”, decode 阶段像“一个字一个字作答”。两者的计算形态、瓶颈和优化方向完全不同。

自回归生成时,每生成一个 token 都要看历史上下文。 如果每轮都重新计算所有历史 token 的 K/V,decode 会非常慢。 KV Cache 的做法是把每一层 attention 的 key/value 保存下来,下一轮直接复用。 它换来速度,代价是显存会随请求数、上下文长度、层数线性增长。

一个粗略估算

def kv_cache_gib(batch, seq_len, layers, kv_heads, head_dim, bytes_per_elem=2):
    bytes_total = batch * seq_len * layers * 2 * kv_heads * head_dim * bytes_per_elem
    return bytes_total / (1024 ** 3)

# 例: 32 层,32 个 KV heads,head_dim=128,bf16/fp16
print(kv_cache_gib(batch=16, seq_len=4096, layers=32, kv_heads=32, head_dim=128))
# 约 32 GiB,还没有计算模型权重和运行时 buffer

今天的代码目标不是搭服务,而是看清楚 cache 在 generation loop 中如何流动: 第一次 forward 处理完整 prompt,后续 forward 只喂最后一个 token,并把 past_key_values 传回模型。

import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "Qwen/Qwen2.5-0.5B-Instruct"
device = "cuda"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map=device,
).eval()

prompt = "用三句话解释 LLM 推理中的 KV Cache。"
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)

generated = input_ids
past = None
max_new_tokens = 64
first_token_time = None
start = time.perf_counter()

with torch.inference_mode():
    for step in range(max_new_tokens):
        # 第 0 步:喂完整 prompt。之后:只喂上一个新 token。
        cur_input = generated if past is None else generated[:, -1:]

        out = model(
            input_ids=cur_input,
            past_key_values=past,
            use_cache=True,
        )
        past = out.past_key_values

        next_token = out.logits[:, -1, :].argmax(dim=-1, keepdim=True)
        generated = torch.cat([generated, next_token], dim=-1)

        if step == 0:
            first_token_time = time.perf_counter() - start

        if next_token.item() == tokenizer.eos_token_id:
            break

elapsed = time.perf_counter() - start
new_tokens = generated.shape[-1] - input_ids.shape[-1]

print({
    "prompt_tokens": input_ids.shape[-1],
    "new_tokens": new_tokens,
    "ttft_s": first_token_time,
    "tpot_s": (elapsed - first_token_time) / max(new_tokens - 1, 1),
    "tokens_per_s": new_tokens / elapsed,
})
print(tokenizer.decode(generated[0], skip_special_tokens=True))

训练里常看 step time、MFU、loss。推理服务里更常看用户体验和系统吞吐: 首 token 等多久、后续 token 多快、单卡每秒能吐多少 token、请求排队多久、显存还剩多少 cache 空间。

1. 用自己的话画出一次 LLM 请求从 HTTP 进入到 streaming 返回的生命周期。
2. prefill 和 decode 在输入形态、attention 计算、性能瓶颈上有什么不同?
3. 写出 KV Cache 显存估算公式,解释每个变量的含义。
4. 为什么 KV Cache 会随 batch、seq_len 和 layers 线性增长?
5. TTFT、TPOT、throughput 分别衡量什么?它们之间有哪些常见权衡?
6. 如果一个服务首 token 很慢但后续 token 很快,你会优先检查哪些环节?