解码算法只发生在一个很小但极关键的位置: 模型 forward 产出最后一个位置的 logits, 经过一组 logits processor / warper 后,选择下一个 token, 再把 token 追加到序列并更新 KV Cache。 这个循环会重复到 EOS、max_tokens 或 stop sequence。 所以 decode 策略几乎不改变 TTFT 的 prefill 部分,但会显著影响 inter-token latency、输出长度、KV cache 占用和调度器公平性。

def decode_one_step(model, input_ids, past_kv, config):
    logits, past_kv = model(input_ids[:, -1:], past_key_values=past_kv)
    logits = logits[:, -1, :]

    logits = apply_processors(logits, config)  # repetition penalty, no-repeat ngram, bad words...
    logits = apply_warpers(logits, config)     # temperature, top-k, top-p...
    next_token = select_token(logits, config)  # argmax or multinomial sample

    input_ids = append(input_ids, next_token)
    return input_ids, past_kv

greedy 和 beam search 都属于搜索式解码。 greedy 每步选最大概率 token,成本最低、最可复现,但容易短视。 beam search 保留多个候选序列,按累计 log probability 搜索更高概率的完整序列。 它在翻译、摘要这类目标相对确定的任务中仍有价值, 但在开放式聊天里经常显得重复、保守,并且会把 decode 计算和 KV cache 近似乘以 beam 数。

Beam Search 的分数修正

采样式解码先把 logits 变成概率分布,再随机抽 token。 temperature 改变分布尖锐程度, top_k 只保留概率最高的 k 个 token, top_p 又叫 nucleus sampling,保留累计概率达到 p 的最小 token 集。 它们不是"质量按钮",而是对随机空间的约束方式。

参数语义速查

import torch

def sample_next_token(logits, temperature=0.8, top_k=50, top_p=0.95):
    logits = logits / max(temperature, 1e-6)

    if top_k is not None:
        values, _ = torch.topk(logits, top_k)
        logits = logits.masked_fill(logits < values[..., -1, None], float("-inf"))

    if top_p is not None:
        sorted_logits, sorted_idx = torch.sort(logits, descending=True)
        probs = torch.softmax(sorted_logits, dim=-1)
        cumulative = torch.cumsum(probs, dim=-1)
        remove = cumulative > top_p
        remove[..., 1:] = remove[..., :-1].clone()
        remove[..., 0] = False
        sorted_logits = sorted_logits.masked_fill(remove, float("-inf"))
        logits = torch.full_like(logits, float("-inf")).scatter(-1, sorted_idx, sorted_logits)

    probs = torch.softmax(logits, dim=-1)
    return torch.multinomial(probs, num_samples=1)

业务上看,解码策略影响"回答像不像人"。 Infra 上看,它还影响请求长度分布、KV cache 生命周期、decode batch 的形状和长尾延迟。 所以线上系统不应该让所有参数无限开放,而是按场景给出受控的参数模板。

自回归解码慢,因为 target model 通常一次只能确认一个新 token。 speculative decoding 的想法是: 用一个便宜的 draft 模型先猜多个 token, 再让大的 target 模型一次 forward 验证这串候选。 如果 draft 猜得准,一次 target forward 可以接受多个 token, inter-token latency 就下降。 严格的 speculative sampling 可以保持目标分布不变;工程实现里还会有 n-gram、EAGLE、MTP 等不同 speculator。

Spec Decode 的工程账本

今天的实验分两层。 第一层用 Transformers 在同一个 prompt 上跑不同 decoding config,观察输出和延迟。 第二层如果你本地有 vLLM 环境,尝试打开 speculative decoding,记录 token latency 与接受率相关日志。

from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
import torch, time

model_id = "gpt2"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id).cuda().eval()

prompt = "Explain why KV cache matters for LLM serving in three bullets."
inputs = tok(prompt, return_tensors="pt").to("cuda")

configs = {
    "greedy": GenerationConfig(max_new_tokens=80, do_sample=False),
    "beam4": GenerationConfig(max_new_tokens=80, num_beams=4, do_sample=False),
    "top_p": GenerationConfig(max_new_tokens=80, do_sample=True, temperature=0.8, top_p=0.95),
    "top_k": GenerationConfig(max_new_tokens=80, do_sample=True, temperature=0.8, top_k=50),
}

for name, cfg in configs.items():
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    out = model.generate(**inputs, generation_config=cfg)
    torch.cuda.synchronize()
    dt = time.perf_counter() - t0
    print("\\n===", name, f"{dt:.2f}s")
    print(tok.decode(out[0], skip_special_tokens=True))

# vLLM speculative decoding 示例:参数名会随版本变化,以当前 vLLM 文档为准
vllm serve meta-llama/Llama-2-7b-hf \
  --speculative-config '{"model":"JackFram/llama-160m","num_speculative_tokens":5}'

# 记录:
# 1. baseline inter-token latency
# 2. spec decode inter-token latency
# 3. GPU util / memory
# 4. acceptance rate 或 accepted tokens per verify step

实验记录表

1. decode loop 中 logits processor 和 logits warper 的区别是什么?
2. 为什么 beam search 会增加 KV cache 占用?
3. temperature、top-k、top-p 分别从哪个角度控制随机性?
4. top-p 为什么叫 nucleus sampling?它和 top-k 的候选集合大小有什么不同?
5. 同一个 prompt 下,greedy 和 top-p 输出差异如何解释?
6. speculative decoding 的 speedup 由哪三个因素共同决定?
7. 如果你要给 OpenAI 兼容 API 设置默认采样参数,会如何限制用户输入?