一个请求走完 Kimi K3 的 prefill 与 decode
这次不再逐个孤立地看模块。我们让提示词 token [A,B,C,D] 一次穿过 93 层，采样得到 E；再把 E 作为下一轮输入，更新 KDA 与 MLA cache，得到用于采样 F 的 logits。
f4b161d7（2026-08-18）。为了看清时间线，示例只有一个请求、4 个 prompt token、非 speculative decode；A...F 是教学符号，不代表真实词表 id。真实模型仍是 93 层、hidden 7168、词表 160K。0. 最先澄清：模型输出的是 logits，不是“下一 token”
因此 prefill 已经计算出了第一个生成 token E 的分布；第一次 decode 消费的是 E，并产生 F 的分布。把“decode 输入 E”和“decode 输出 E”混为一谈，会让 cache 的写入时点整体错一格。
1. 先记住三类东西的生命周期
| 类别 | prefill 后是否留下 | 例子 |
|---|---|---|
| 请求持久状态 | 是，decode 下一步读取并更新 | 69 层的 KDA short-conv history 与 recurrent state；24 层的 MLA latent pages；长度、slot mapping、page table |
| 当前 forward 临时量 | 否，使用完成即可复用/释放 | Q/K/V、AttnRes depth buffer、router logits、expert activation、hidden states、LM logits |
| 模型常驻权重 | 一直在模型实例中，不属于某个请求 | 投影矩阵、896 个 expert 权重、AttnRes score vectors、RMSNorm 参数、decode 吸收后的 W_UK_T/W_UV |
2. Prefill：一次输入四个 token
2.1 embedding 与本轮 packed token 维
vLLM 常把同一调度轮的多个请求 pack 到 token 维，所以生产代码里的第一维通常是本轮总 token 数 T，而不是整齐的 [batch,seq]。本例只有一个长度 4 的请求，所以 T=4。
2.2 第 0 层是 KDA：四个位置并行准备，状态按顺序组合
逻辑上仍按 A→B→C→D 更新每个 head 的 S:[128,128]；prefill kernel 把 chunk 内可改写成 GEMM 的 token 交互并行计算，只在 chunk 边界传 state。该层完成后，请求保留最终 S_D:[96,128,128]，而不是四份 state。
short conv 也只保留继续计算所需的最后 3 个 raw Q/K/V 输入。全局逻辑 shape 为 [36864,3]；TP=8 时每 rank 是 [4608,3]，recurrent state 是 [12,128,128]。
2.3 经过第三个 KDA 后，进入第一层 Gated MLA
第 D 个位置只能看 A...D；第 A 个位置只能看自己。prefill 为了高吞吐会临时展开完整 K/V，但请求持久 cache 只留下四个 576 维 latent：[c_A,c_B,c_C,c_D]。相同动作在 24 个 MLA 层分别发生。
2.4 每层的 AttnRes 与 Stable LatentMoE
R 是该层可见的深度来源数，随深度 block 推进而增加，最多是 embedding、已完成的 block 和当前 block 部分和。它不是序列长度 4。四个 token 的 AttnRes buffer 在这次 93 层 forward 内流转，最后即失去用途；decode E 时会为 E 新建一份。
2.5 第 92 层、LM head 与第一次采样
服务引擎通常无需保留 prompt 每个位置的全词表 logits；生成只需要最后一个有效位置。至此 prefill 结束，cache 已经包含 A...D，但还不包含刚采样出的 E。
3. Prefill 结束后的请求状态快照
| 持久项 | 每层/每请求逻辑内容 | 共几组 | 此刻代表的历史 |
|---|---|---|---|
| KDA recurrent state | [96,128,128]；TP=8 每 rank [12,128,128] | 69 层 | A...D 经各层输入形成的最终 state |
| KDA conv state | Q/K/V 各通道最近 3 项；全局 [36864,3] | 69 层 | 能继续对 E 做 causal Conv4 |
| MLA latent pages | 每 token [576] | 24 层 | 每层各有 [c_A,c_B,c_C,c_D] |
| 调度/cache metadata | request length=4、block table、slot mapping 等 | 请求级 | 下一步把 E 写到逻辑位置 4 |
4. Decode：把 E 作为一个新输入 token
4.1 新 forward 从头穿过 93 层
4.2 在每个 KDA 层原地前进一步
纯 decode 正是 vLLM kda.py 的 fused recurrent 路径：按 metadata 找到该请求的 state slot，更新 conv state 与 recurrent state。计算量不会随 A...D 的长度线性增长，但必须为 69 个 KDA 层分别维护状态。
4.3 在每个 MLA 层追加 E，并查询 A...E
decode 路径没有为五个历史位置物化 K:[5,96,192] 和 V:[5,96,128]。vLLM 在模型加载后的 weight processing 阶段准备 W_UK_T/W_UV，运行时走 BMM1 → latent MQA → BMM2。这是 MLA 真正降低 decode cache 读取与中间展开的地方。
4.4 MoE 与 AttnRes 没有历史 cache
对 E，router 重新生成 [1,896] logits，并可能选择与 D 完全不同的 16 个专家；expert activation 用完就释放。AttnRes 也只混合 E 在深度方向的表示，不会读取 A...D 的 AttnRes buffer。跨 token 信息来自 KDA state 与 MLA latent history。
4.5 得到 F 的分布
5. vLLM 为什么需要一份混合 metadata
一次 scheduler batch 可以同时放入长 prompt、chunked prefill 的续块、普通 decode token，甚至 speculative token。token 张量会被 pack 在一起，但 KDA kernel 要区分“可 chunk 并行的 prefill”和“原地前进一步的 decode”，MLA backend 又要知道每个 query 对应哪些 latent pages。
| metadata 作用 | KDA 使用方式 | MLA 使用方式 |
|---|---|---|
| 请求边界与 query 长度 | query_start_loc 切分各请求递推序列 | 构造 causal mask 与每请求 context 长度 |
| state/cache 索引 | state_indices 定位 conv/recurrent slot | slot_mapping/block_table 定位 latent page |
| prefill/decode 分类 | 选择 FlashKDA/chunk 或 fused recurrent | 选择 full prefill attention 或 absorbed latent MQA |
| 重新排序 | spec 与 non-spec token 可先拆开再写回 packed 顺序 | decode query 与 cache insert 在 fused kernel 中对齐 |
所以 Python forward 看起来是一条统一数据流，底层却按 token 类型分派不同 kernel；这属于 vLLM 的批处理实现，不是模型公式额外增加了一种 attention。
6. Speculative decode 为什么对 KDA 更难
普通 MLA 可以给 draft token 分配临时 cache slots，拒绝后丢弃。KDA 的问题是 state 更新具有顺序性：若一次验证 [E,F,G]，执行完后 S 已前进三步；若只接受 E，不能通过删除两条 KV 记录恢复 S_E。
vLLM 当前 K3 支持把 speculative token 与 non-spec token 分开。ReplaySSM/RecoverSSM 路径保留 checkpoint 与更小的校正/投影记录，在验证结果已知后重放或恢复正确边界，而不是为每个 draft 位置复制完整 [96,128,128] state。这是后端优化；不开 speculative decode 时，本课前面的 plain decode 时间线不需要这些附加记录。
7. 用一张表复盘整轮生成
| 时刻 | 本轮模型输入 | cache 提交后包含 | 模型给出的分布 | 随后采样 |
|---|---|---|---|---|
| prefill | A,B,C,D | A...D | p(next | A...D) | E |
| decode #1 | E | A...E | p(next | A...E) | F |
| decode #2 | F | A...F | p(next | A...F) | G |
检索练习
问题：prefill 刚采样出 E、但还没执行第一次 decode 时，哪些请求状态已经包含 E？
先读：第一课：四个模块的公式与 shape。实现对照：vLLM Kimi K3 NVIDIA backend at f4b161d7。速查：推理状态台账。