Module 3：自回归生成、KV Cache 与采样

理解一个 LLM 如何从 prompt 开始逐 token 生成文本，以及为什么 KV cache、采样策略和吞吐指标决定推理系统的体验。

3.1自回归生成：一次只预测一个 token

decoder-only LLM 采用 autoregressive generation（自回归生成）：给定 prompt tokens，模型输出最后一个位置的 logits；采样器选出下一个 token；该 token 追加到上下文末尾后，模型继续预测下一步。循环在生成 EOS 或达到 max_new_tokens 时停止。xinfer-runtime::generate 按这一控制流实现 prefill 与逐 token decode。

如果使用 greedy decoding，符号 $\sim$ 可以改写成取最大值：

其中 $z_i$ 是最后一行 logits 中第 $i$ 个 token 的分数。Qwen2.5-0.5B 的词表大小为 151936，因此这一行 logits 含 151936 个分量。

3.2Prefill 与 Decode：同一个模型，两种工作负载

推理过程通常分为 prefill 与 decode 两个阶段。二者调用同一组 decoder layer，但输入长度和系统瓶颈不同：

prefill 一次处理整个 prompt，并把每层新产生的 K/V 写入 cache。decode 每步只输入刚生成的 1 个 token；第 $t+2$ 个 token 的输入依赖第 $t+1$ 个 token 的采样结果，因此不能直接并行生成多个未来 token。

性能报告需要分开给出两个吞吐量：

3.3为什么需要 KV cache？

每个 decoder layer 都会由 hidden states 计算 key 和 value。设已有 $t$ 个历史 token，下一步只新增 1 个 token。若没有 KV cache，每一步都必须把完整上下文重新送入所有层，历史 token 的 Q/K/V projection 和 attention 输入会被反复计算。

使用 KV cache 后，模型只为新 token 计算一次 $K,V$，并把结果追加到每层缓存的下一个槽位。xinfer-model::run_forward 在 RoPE 后通过 copy_region 将新的 K/V 写入 cache.k[li] 与 cache.v[li]：

计算 attention 时，新 query 读取缓存中已经写入的 keys/values。decode 单步的典型形状为 $Q:[1,14,64]$，$K,V:[T,2,64]$，输出再展开为 $[1,896]$：

KV cache 的显存公式

对本项目的 f32 KV cache，容量按下式计算：

其中 2 表示 K 和 V，$L$ 是层数，$T$ 是 cache capacity，$n_{\text{kv}}$ 是 KV head 数。Qwen2.5-0.5B 取 $L=24$、$n_{\text{kv}}=2$、$d_{\text{head}}=64$、dtype=f32。若 prompt 为 26 tokens、max_new_tokens=64，CLI 分配的容量为 $T=90$，KV cache 约 $2\cdot24\cdot90\cdot2\cdot64\cdot4=2{,}211{,}840$ bytes，即约 2.1 MiB。若没有 GQA 而使用 14 个 KV head，该项会扩大为 7 倍。

3.4采样：从 logits 到下一个 token

模型输出 logits $z$ 后，采样器把它转换为下一个 token id。xinfer-runtime::sampling 支持 greedy、temperature、top-k 与 top-p。默认配置为 temperature=0.0、top_k=0、top_p=1.0，即 greedy decoding。

Greedy / Argmax

greedy 每步选择最大 logit 对应的 token，输出确定且可复现。本项目对 greedy 走 decode_argmax 快路径：shaders/argmax.hlsl 在 GPU 上归约最后一行 logits，只读回一个 u32 token id。若读回完整 f32 logits，Qwen2.5-0.5B 每步需要传输约 $151936\times4\approx608$ KB。

Temperature

当 temperature > 0 时，采样器先按温度缩放 logits：

$\tau$ 越小，概率分布越集中；$\tau$ 越大，低 logit token 获得更高相对概率。代码中 temperature <= 0.0 被视为 greedy，不执行随机采样。

Top-k 与 Top-p

Top-k 在排序后只保留 logit 最高的 $k$ 个候选；top_k=0 表示不启用该截断。Top-p（nucleus sampling）在 softmax 后按概率从高到低累加，只保留累积概率首次达到 $p$ 的前缀；top_p>=1.0 表示不启用该截断：

3.5Tokenizer、chat template 与 special tokens

模型输入输出均为 token id，而不是 Unicode 字符串。xinfer-tokenizer 封装 HuggingFace tokenizers，从模型目录加载 tokenizer.json，提供 encode、decode 与 token_to_id：

对于 instruct 模型，裸 user prompt 通常需要先套用 chat template。qwen_chat_prompt(system, user) 在 CLI 默认启用，生成的格式为：

<|im_start|>system
You are a helpful assistant.<|im_end|>
<|im_start|>user
What is the capital of France?<|im_end|>
<|im_start|>assistant

EOS token 表示模型认为当前回答已经结束。read_eos_token 优先读取 generation_config.json；Qwen2.5-0.5B-Instruct 的 generation config 中 EOS 列表为 151645 与 151643，当前实现取列表中的第一个 id。生成循环在以下条件之一成立时停止：

• 生成了 EOS token；
• 达到 max_new_tokens；
• 外部系统要求中断（例如 UI 停止按钮）。

3.6吞吐、延迟与用户体验

推理性能不能只用单个平均值描述。至少需要区分首 token 前的等待、持续生成速度、显存占用与服务端吞吐：

本项目的示例运行报告为 prefill: 26 tok in 69 ms (378 tok/s) | decode: 7 tok in 0.23 s (30.5 tok/s)。decode 的关键优化是 greedy 模式下避免每 token 将完整 vocab logits 读回 CPU；temperature、top-k、top-p 路径仍调用 model.decode 读回 logits，并在 CPU 上执行采样。

小结

自回归生成把一次回答拆成 prefill 和 decode。prefill 处理完整 prompt 并建立 KV cache；decode 每步处理 1 个新 token，从 cache 中读取历史 K/V，并将新 K/V 追加到每层缓存。Qwen2.5-0.5B 的 logits 长度为 151936；greedy 模式可用 GPU argmax 只读回 token id，而 temperature、top-k、top-p 仍需把 logits 读回 CPU 后采样。tokenizer 与 chat template 决定模型实际看到的 token 序列，EOS 与 max_new_tokens 决定生成何时结束。

Lab 3纯 CPU 生成循环：KV cache + samplers

本实验不需要 GPU。用 tiny 模型或假 logits 实现完整生成循环，重点检查 prefill、cache 追加、EOS 判断与采样器调用顺序。

def generate(prompt_tokens, max_new_tokens):
    cache = empty_kv_cache()
    logits = model_prefill(prompt_tokens, cache)
    next_token = sample(logits[-1])

    output = []
    for _ in range(max_new_tokens):
        if next_token == EOS:
            break
        output.append(next_token)
        logits = model_decode_one(next_token, cache)
        next_token = sample(logits[-1])
    return output