Kimi K3 四个核心模块的公式与真实张量流
目标不是记住“KDA 省缓存、MoE 省计算”这类结论，而是能从输入 X 开始，写出每个投影的 shape、状态更新式、输出 shape，以及为什么 vLLM 必须为 prefill 和 decode 选择不同执行路径。
f4b161d7（2026-08-18）。以下 T 是本轮实际 token 数，B 是 decode 请求数，d=7168。TP=8 的局部 shape 会单独标注。0. 先看一层到底做什么
K3 有 93 个 Transformer 层：23 次重复 KDA → KDA → KDA → Gated MLA，最后再加一个 Gated MLA。每个 attention 后配一个 channel mixer：第一层是 Dense MLP，其余约 92 层是 Stable LatentMoE。Block AttnRes 不替代 token mixer；它先在深度轴构造 attention 或 MLP 的输入。
embedding / 当前 block 部分和 / 已完成 block 表示
↓ Block AttnRes，得到 [T, 7168]
KDA 或 Gated MLA，得到 [T, 7168]
↓ 再做一次 Block AttnRes，得到 [T, 7168]
Dense MLP 或 Stable LatentMoE，得到 [T, 7168]
↓ 输出进入当前深度 block 的部分和
| 模块 | 核心真实维度 | 跨 decode step 保留什么 | 替代/修改了什么 |
|---|---|---|---|
| KDA | 96 heads，head dim 128 | 每层卷积历史 + [96,128,128] 递归 state（按 TP 分片） | 大部分全局 softmax attention |
| Gated MLA | q latent 1536；KV latent 576；96 heads；qk 192；v 128 | 每层每 token 576 个 latent 元素 | 周期性全局 MHA，并增加输出门 |
| Block AttnRes | 8 个深度 block；每个最多 12 层；hidden 7168 | 不跨 token 持久；只在当前 forward 内维护 block 表示 | 普通逐层相加的 residual stream |
| Stable LatentMoE | 7168→3584；896 选 16；expert hidden 3072；2 shared | 不产生请求 cache | 全宽 Dense FFN / 全宽 routed MoE |
1. KDA：完整计算链，而不只是一个递推式
1.1 从 X:[T,7168] 生成六组量
对第 t 个 token、head h，KDA 需要 q,k,v、写强度 β、逐 key-channel 保留率 α 和输出门 g²。vLLM 把第一批投影打包成一个 GEMM：
TP=8 时每 rank 有 12 heads：Q/K/V/G² 各为 [T,12,128]，F_a:[T,128] 被复制，β:[T,12]。生产 decode 的合并投影输出被对齐到 [T,6288]。
1.2 短卷积不是装饰，它也产生 cache
q/k/v 各自先经过宽度 4 的 causal depthwise short convolution，再经过 Swish；q/k 还做 L2 normalization：
所以 decode 下一步除了 S_t，还必须知道每个 Q/K/V channel 的最近 3 个输入。vLLM 的全局逻辑 shape 是 [3×96×128, 3] = [36864,3]；TP=8 是 [4608,3]。这正是“只说 KDA 有一个固定矩阵 state”会漏掉的实现事实。
1.3 decay、delta write 与读取
对单 head，q_t,k_t,v_t∈R^128，S_t∈R^(128×128)。K3 把 log-decay 限制在 (-5,0)：
v_t-S̄^Tk_t 是“当前 state 对 key k_t 的预测误差”。外积写入的是这份误差，而不是盲目叠加 k_tv_t^T；因此重复 key 会修正旧关联。Diag(α) 又允许不同 key channel 以不同速度遗忘。
1.4 输出门与投影
这叫 full-rank gate，因为 G² 直接从 7168 维输入投到所有 96×128 通道；它能按 token、head、channel 决定递归读出的哪些部分通过。
1.5 prefill 与 decode 为什么不是同一个 kernel
数学上 KDA 沿 token 递推；prefill 若真的逐 token 启动 kernel 会失去 GPU 并行性。K3 把序列切 chunk，chunk 内用因果下三角矩阵乘并行计算，chunk 间只传 S。vLLM 优先走 FlashKDA，fallback 是 chunk_kda_with_fused_gate。decode 每请求只有一个新 token，反而适合把 packed projection、conv update、KDA recurrence、head RMSNorm 和输出 gate 融合。
| 传统 MHA/GQA | KDA | |
|---|---|---|
| 请求状态 | 每层保存历史 token K/V，随长度线性增长 | 每层固定卷积历史和递归矩阵 |
| decode 历史读取 | 读取历史 KV，随上下文增长 | 读取并原地更新固定 state |
| 精确回看位置 | 强：query 可直接选择某个历史 token | 弱：历史已压缩进 state |
| 新增代价 | KV 容量和带宽 | 递归依赖、state 恢复、spec rollback 更复杂 |
2. Gated MLA：prefill 展开 K/V，decode 吸收投影
2.1 前端投影与真实 shape
K3 的 MLA 层是 NoPE，所以 qk_rope_head_dim=0，cache entry 就是 576 维 c_t。普通 96-head MHA 若直接保存每头 192 维 K 和 128 维 V，每 token 每层是 96×320=30720 个元素；MLA 是 576 个，元素数约缩小 53.3 倍。与 GQA 比倍率会更小，所以文章不能把 53.3 倍写成对所有传统架构都成立。
2.2 prefill：为了并行 attention，临时展开
prefill 中 K:[T,96,192] 和 V:[T,96,128] 可以是临时张量，最终请求 cache 仍只保留 c:[T,576]。vLLM 的 fused epilogue 同时拼 key、插入 paged latent cache，并调用 prefill backend；chunked prefill 读取旧 latent page 后才临时做 kv_b_proj。
2.3 decode：把 W_K 和 W_V 吸收到两侧
令每头的 KV up-projection 拆为 W_UK,h:[576,192] 与 W_UV,h:[576,128]。原本分数是 q_h(W_UK,h c_j)^T，结合律允许先变换 query：
这样 decode 不必为所有历史 token 重建 96 份 K/V。vLLM 中这两块预处理权重叫 W_UK_T 和 W_UV，decode 路径是 BMM1 → latent MQA → BMM2。
2.4 full-rank output gate
相比 K2/K2.5 的 MLA，K3 取消所有 MLA RoPE，并增加这个输入相关的 channel gate。NoPE 避免 1M 扩展时重调 RoPE，但位置敏感性必须由相邻 KDA 的递推历史和训练数据承担；这是职责转移，不是“位置消失”。
3. Block AttnRes：在深度来源上做 softmax
普通 residual 只有 h_(l+1)=h_l+f_l(h_l)。早期层信息必须穿过单一累计向量。K3 把 93 层划成 8 个至多 12 层的 block，embedding 作为 b_0，已完成 block 的和为 b_1...b_(n-1)，当前 block 的部分和为 p。
vLLM 在每层 attention 前和 MLP 前各调用一次 attn_res，两次有独立的 7168 维打分权重；kernel 同时完成当前部分和的更新、必要的 block 写入、online softmax 混合与下一层 RMSNorm。buffer 是 [T,num_blocks,7168]。
[T,R,7168] 是当前模型 forward 的深度 activation。prefill 完整跑完后，它不会像 MLA latent 或 KDA state 那样留给下一个 decode token；新 token 会在穿过 93 层时重新建立自己的深度来源。优势是每层能直接选择早期深度表示，深度信息不再只能等权累加。代价是每个 token 要读取多个 7168 维 block source，并引入额外 softmax、buffer 和 TP 通信；Block 版只是把 full AttnRes 的 O(Ld) 活跃来源降为 O(Nd)，不是零成本。
4. Stable LatentMoE：router 看全宽，专家吃半宽
4.1 两条并行路径
关键细节：Quantile Balancing 训练出的 bias:[896] 只参与 Top-16 选择，不进入 mixture weight。vLLM 以 FP32 生成 router logits，并冻结使用 checkpoint 中的 e_score_correction_bias。
4.2 单个 routed expert 内部
输出逐元素绝对值上界是 4×25=100。它解决的是 W_down → expert 的两次输入投影 → W_up 这条长矩阵链在 2.8T 规模下的 activation outlier。聚合后的 routed latent 在升回 7168 之前再做 RMSNorm。
shared path 等价于把两个始终开启的全宽 expert 合并成一个中间宽度 2×3072=6144 的 MLP。它不经过 3584 latent bottleneck，用来承载所有 token 都需要的通用变换。
4.3 相比普通 MoE 到底省在哪里
| 普通全宽 MoE | K3 Stable LatentMoE | |
|---|---|---|
| router 输入 | 通常全宽 | 仍是 7168，全局语义不丢 |
| dispatch 给专家 | 每次 7168 维 | 每次 3584 维，16 路通信约减半 |
| 专家输入/输出宽度 | 7168 | 3584 |
| 额外成本 | 无 latent 投影 | 每 token 多一次 7168→3584 和 3584→7168 |
| 稳定机制 | 常见 SwiGLU | RMSNorm + 有界 SiTU-GLU + QB bias |
| 请求 cache | 无 | 无；expert 权重属于模型常驻状态 |
5. 把四个模块放回同一层
AttnRes 从 embedding、已完成深度 block 和当前部分和混合出 [T,7168]。S；MLA 层追加 [T,576] latent page。两者都输出 [T,7168]。AttnRes 将 attention 输出纳入当前 block 部分和，再混合出 MoE 输入。[T,7168]。检索练习
问题：MLA decode 为什么可以不为全部历史 token 重建 96 份 K/V？
主读物：Kimi K3 Technical Report §2。实现对照：vLLM Kimi K3 at f4b161d7。下一课会让同一个请求逐步走过这些张量。速查：核心 shape 与公式表。任何一步不清楚，直接指出输入和输出，我会沿那条边继续展开。