feat: execute DeepSeek absorbed MLA cache

2026-07-29 14:57:06 +08:00
parent ab19381046
commit 9ba26da545
10 changed files with 2290 additions and 32 deletions
@@ -0,0 +1,420 @@
+# DeepSeek MLA 吸收式缓存：真实权重执行与 FlashMLA 架构边界
+
+> 状态：DeepSeek 第三轮第二里程碑  
+> 证据等级：`O` 官方源码 / `X` 本机真实执行 / `D` 确定推导 / `U` 未覆盖  
+> 结论边界：正确性与缓存格式实验，不是服务吞吐 benchmark
+
+## 0. 先给结论
+
+这轮补上了上一次实验最重要的空白：
+
+- 上一次已经观察到 V2-Lite 的真实压缩投影是 `[4, 27, 576]`；
+- 但 Hugging Face 官方 eager 路径实际缓存展开后的
+  `K [4,16,27,192]` 与 `V [4,16,27,128]`；
+- 这一次把同一组真实 V2-Lite layer-1 权重映射到 DeepSeek-V3 官方
+  `inference/model.py`，真正执行了 `attn_impl="absorb"`；
+- 25-token prefill + 1-token decode 后，官方 buffer 中实际存在：
+  - `kv_cache [1,32,512]`，其中 26-token active slice 为 `26,624 B`；
+  - `pe_cache [1,32,64]`，其中 26-token active slice 为 `3,328 B`；
+  - active 总计 `29,952 B`。
+- 同一官方参考类的 naive 路径实际 active cache 为：
+  - `k_cache`：`159,744 B`；
+  - `v_cache`：`106,496 B`；
+  - active 总计 `266,240 B`。
+- 因而在这个 V2-Lite 配置、BF16、同一 26-token 输入下：
+
+```text
+naive / absorb = 266,240 / 29,952 = 8.888888…
+reduction       = 1 - 29,952 / 266,240 = 88.75%
+```
+
+这不再只是 shape 算术：两组官方 cache buffers 都已在 RTX 5090 上实际分配、写入并用于
+incremental decode。
+
+但是另一条边界同样重要：
+
+- DeepSeek-V3 的官方纯 PyTorch `absorb` 参考路径：**已执行**；
+- DeepSeek FlashMLA 优化 kernel：**未执行**；
+- 原因不是把环境问题猜成性能结论，而是 pinned FlashMLA 官方支持矩阵与源码只覆盖
+  SM90 / SM100，本机 RTX 5090 是 SM120。
+
+## 1. 可复现身份
+
+### 1.1 模型与源码
+
+| 工件 | 固定身份 |
+|---|---|
+| Hugging Face 模型 | `deepseek-ai/DeepSeek-V2-Lite` |
+| HF revision | `604d5664dddd88a0433dbae533b7fe9472482de0` |
+| DeepSeek-V3 仓库 | `deepseek-ai/DeepSeek-V3` |
+| V3 revision | `9b4e9788e4a3a731f7567338ed15d3ec549ce03b` |
+| FlashMLA 仓库 | `deepseek-ai/FlashMLA` |
+| FlashMLA revision | `15f13e5030374295491c5ce31b02d7e63a7772c6` |
+| FlashMLA CUTLASS submodule | `147f5673d0c1c3dcf66f78d677fd647e4a020219` |
+
+官方入口：
+
+- [DeepSeek-V2-Lite 模型](https://huggingface.co/deepseek-ai/DeepSeek-V2-Lite)
+- [DeepSeek-V2 仓库](https://github.com/deepseek-ai/DeepSeek-V2)
+- [DeepSeek-V2 技术报告](https://arxiv.org/abs/2405.04434)
+- [DeepSeek-V3 inference/model.py](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/model.py)
+- [FlashMLA](https://github.com/deepseek-ai/FlashMLA)
+
+### 1.2 本机环境
+
+| 字段 | 值 |
+|---|---|
+| GPU | NVIDIA GeForce RTX 5090 |
+| compute capability | `12.0` / SM120 |
+| driver | `595.84` |
+| torch | `2.11.0+cu128` |
+| CUDA runtime | `12.8` |
+| Transformers | 隔离的 `4.41.2` |
+| checkpoint dtype | BF16 |
+
+V2-Lite 的 2024 官方 remote code 不能直接在当前 Transformers 5.5 中导入，因此延续上一轮合同：
+官方模型源码不打补丁，只把依赖固定到隔离目录中的 Transformers 4.41.2。
+
+### 1.3 输出
+
+- 主结果：`src/data/deepseek-v2-lite-absorb.json`
+- 独立复跑：`src/data/deepseek-v2-lite-absorb-repro.json`
+- 两份完整 JSON SHA-256：
+  `6b4c714a470baa6d6b04982cd289e693ec863398b5def4e6df541cb10525e63d`
+- `cmp`：逐字节一致
+
+实验没有计时字段，因此不需要用“排除 timing”掩盖不同；两次输出完整 byte-exact。
+
+## 2. 为什么权重可以“吸收”
+
+### 2.1 普通展开路径
+
+令每个历史 token 的 512 维压缩状态为 `cₜ`。忽略位置分支时，每个 attention head 的
+key/value 可以写成：
+
+```text
+kₜ = Wᴷ cₜ
+vₜ = Wⱽ cₜ
+```
+
+naive 实现先把所有历史 `cₜ` 展开成每个 head 的 `kₜ` 和 `vₜ`，再缓存展开结果。V2-Lite
+有 16 个 head：
+
+```text
+key   = 16 × (128 no-RoPE + 64 RoPE) = 3,072 elements/token
+value = 16 × 128                     = 2,048 elements/token
+total = 5,120 elements/token/layer
+```
+
+问题不是这些张量算错了，而是同一个 512 维内容状态被投影并复制到各个 head 后才进入缓存。
+
+### 2.2 把 key 投影搬到 query
+
+注意力分数里的 no-RoPE 部分：
+
+```text
+qᵀkₜ
+= qᵀ(Wᴷcₜ)
+= (Wᴷᵀq)ᵀcₜ
+```
+
+矩阵乘法结合律允许先把 query 变成 `Wᴷᵀq`，再直接与历史 latent `cₜ` 点积。于是历史状态不必
+缓存每个 head 展开后的 `Wᴷcₜ`。
+
+这里的“吸收”不是删掉参数：
+
+- `Wᴷ` 仍然存在；
+- 它从“写 cache 前作用在每个历史 token 上”，搬到“读 cache 时作用在当前 query 上”；
+- 权重的作用位置改变，数学函数在精确算术中不变。
+
+### 2.3 把 value 投影搬到 attention 后
+
+attention 的 value 聚合：
+
+```text
+Σₜ pₜvₜ
+= Σₜ pₜ(Wⱽcₜ)
+= Wⱽ(Σₜ pₜcₜ)
+```
+
+因此可以先用注意力概率 `pₜ` 对 512 维 latent 做加权和，再把聚合结果通过 `Wⱽ` 展开到各个
+head 的 value 空间。
+
+naive 顺序：
+
+```text
+每个历史 cₜ → Wⱽ → 展开 vₜ → cache → 按 pₜ 加权
+```
+
+absorb 顺序：
+
+```text
+每个历史 cₜ → cache → 按 pₜ 加权 → Wⱽ → 展开输出
+```
+
+### 2.4 为什么 64 维 RoPE key 仍要缓存
+
+RoPE 是位置相关变换。历史 token `t` 的 key 旋转角度依赖位置 `t`，不能像固定线性投影那样无条件
+搬到当前 query 一侧。因此 V2-Lite 的可缓存状态不是只有 512：
+
+```text
+512 KV latent + 64 decoupled RoPE key = 576 elements/token/layer
+```
+
+这也是网站从第一轮起始终把“512 latent”与“576 完整缓存合同”分开的原因。
+
+## 3. 实验怎样避免偷换输入
+
+### 3.1 输入不是随机 hidden state
+
+固定中文 prompt：
+
+```text
+用通俗的语言解释，为什么稀疏专家模型可以拥有很多参数，但每个 token 只使用其中一小部分。
+```
+
+官方 tokenizer 得到 26 个 token。执行链是：
+
+```text
+官方 tokenizer
+  → 官方 embedding 权重
+  → 官方 V2-Lite decoder layer 0 完整 forward
+  → 官方 layer-1 input RMSNorm
+  → 真实 layer-1 attention input
+```
+
+因此 attention 的输入来自 checkpoint 执行轨迹，不是为了让误差更小而构造的随机数。
+
+### 3.2 为什么只测 layer 1 attention
+
+第一 checkpoint 分片完整包含：
+
+- embedding；
+- layer 0 全部参数；
+- layer 1 全部 attention 参数与 input RMSNorm。
+
+本里程碑问的是“吸收路径能否用真实权重兑现 576 元素缓存”，不需要再次运行 MoE。只隔离 layer-1
+attention，可以同时做到：
+
+- 输入由真实前层产生；
+- attention 权重全部来自官方 checkpoint；
+- naive 与 absorb 使用完全相同的输入和参数；
+- 不把未下载的完整模型写成已复现。
+
+### 3.3 权重怎样映射
+
+V2-Lite HF remote code 与 V3 官方 inference reference 的张量语义一致：
+
+| V2-Lite HF | V3 reference |
+|---|---|
+| `q_proj.weight` | `wq.weight` |
+| `kv_a_proj_with_mqa.weight` | `wkv_a.weight` |
+| `kv_a_layernorm.weight` | `kv_norm.weight` |
+| `kv_b_proj.weight` | `wkv_b.weight` |
+| `o_proj.weight` | `wo.weight` |
+
+没有转置 checkpoint、没有重新训练、没有拟合中间 adapter。
+
+### 3.4 三条 BF16 路径
+
+同一 26-token 输入切成：
+
+- prefill：前 25 token；
+- incremental decode：最后 1 token。
+
+执行：
+
+1. 官方 V2-Lite Hugging Face eager attention；
+2. pinned V3 官方 `attn_impl="naive"`；
+3. pinned V3 官方 `attn_impl="absorb"`。
+
+HF eager 与 V3 naive 的对照先验证“权重映射 + YaRN 配置 + causal/incremental 路径”没有发生大幅漂移；
+V3 naive 与 V3 absorb 的对照再隔离权重吸收本身。
+
+## 4. 真实 cache 观测
+
+### 4.1 BF16 active slice
+
+| 路径 | 实际 shape | bytes |
+|---|---|---:|
+| V3 naive key | `[1,26,16,192]` | 159,744 |
+| V3 naive value | `[1,26,16,128]` | 106,496 |
+| **V3 naive total** | 26 × 5,120 BF16 | **266,240** |
+| V3 absorb latent | `[1,26,512]` | 26,624 |
+| V3 absorb RoPE | `[1,26,64]` | 3,328 |
+| **V3 absorb total** | 26 × 576 BF16 | **29,952** |
+
+### 4.2 32-slot 实际 buffer
+
+| 路径 | 分配 shape | allocated bytes |
+|---|---|---:|
+| V3 naive key | `[1,32,16,192]` | 196,608 |
+| V3 naive value | `[1,32,16,128]` | 131,072 |
+| **V3 naive total** |  | **327,680** |
+| V3 absorb latent | `[1,32,512]` | 32,768 |
+| V3 absorb RoPE | `[1,32,64]` | 4,096 |
+| **V3 absorb total** |  | **36,864** |
+
+active 与 allocated 分开记录，是为了不把 32-slot 容器的尾部空位算成 26-token 状态。
+
+## 5. 数值正确性
+
+所有输出 shape 都是 `[1,1,2048]`，全部 finite。
+
+### 5.1 HF eager 与 V3 naive
+
+| 指标 | BF16 |
+|---|---:|
+| max absolute error | 0.001220703125 |
+| mean absolute error | 0.0000888268 |
+| p99 absolute error | 0.0003823861 |
+| exact fraction | 0.284668 |
+
+两份官方实现的张量布局、einsum/matmul 顺序不同，BF16 不能要求逐元素完全相等。这个对照的作用是检查
+映射与配置；它不是“证明两个代码库 bit-exact”。
+
+### 5.2 V3 naive 与 V3 absorb
+
+| 指标 | BF16 |
+|---|---:|
+| max absolute error | 0.00390625 |
+| mean absolute error | 0.0002703113 |
+| p99 absolute error | 0.0009236149 |
+| exact fraction | 0.100098 |
+
+absorb 改变了矩阵乘法结合顺序；有限精度下，先乘什么会改变舍入位置。因此 BF16 输出接近但不
+bit-exact，符合预期。
+
+### 5.3 FP32 代数审计
+
+为了区分“公式不等价”和“BF16 舍入顺序”，实验把同一组 BF16 checkpoint 权重与真实输入提升为 FP32，
+重新执行 V3 naive/absorb：
+
+| 指标 | FP32 |
+|---|---:|
+| max absolute error | `1.1920929e-7` |
+| mean absolute error | `9.4644861e-9` |
+| p99 absolute error | `3.1446120e-8` |
+
+这不是声称存在 FP32 checkpoint，也不是模型精度评测。它是一个代数隔离实验：精度提高后，两种结合
+顺序的误差下降到 FP32 舍入量级。
+
+## 6. FlashMLA：为什么没有在 RTX 5090 上宣称成功
+
+### 6.1 官方支持矩阵
+
+pinned FlashMLA README：
+
+| kernel | GPU architecture | mode | KV cache |
+|---|---|---|---|
+| dense decoding | SM90 | MQA | BF16 |
+| sparse decoding | SM90 & SM100 | MQA | FP8 |
+| dense prefill | SM100 | MHA | — |
+| sparse prefill | SM90 & SM100 | MQA | — |
+
+README 对 MQA mode 的定义是：
+
+```text
+head_dim_k = 576
+head_dim_v = 512
+```
+
+这里的 `head_dim_v=512` 是吸收后的 latent 聚合空间，不应与 naive 路径每个 query head 的
+`v_head_dim=128` 混为一谈。
+
+### 6.2 编译目标与运行时 guard
+
+pinned `setup.py` 只加入：
+
+```text
+arch=compute_100f,code=sm_100f
+arch=compute_90a,code=sm_90a
+```
+
+`csrc/api/dense_decode.h` 还显式检查：
+
+```text
+Dense decode MLA is only supported on SM90a architecture
+```
+
+本机是 SM120。因此即使 CUDA 工具链能把扩展编译完成，也不能从这份官方源码推导出 SM120 dense
+decode 支持。
+
+### 6.3 两次构建尝试
+
+| 环境 | 结果 | 停止位置 |
+|---|---|---|
+| host CUDA 12.8 + g++ 15.2 | 失败 | CUDA 12.8 在正式编译前拒绝过新的 host compiler；本机没有 g++-13 |
+| 隔离 CUDA 13.0.2 + torch 2.11.0+cu130 + g++ 13.3 | 失败 | 官方 extension 编译 `csrc/api.cpp` 时找不到 `cuda/std/utility` |
+
+第二次构建中的其它 CUDA translation units 继续输出日志，是 Ninja 并行任务已经启动的结果；最终
+`ninja: build stopped: subcommand failed`，没有 wheel。
+
+正确的实验状态是：
+
+```text
+official V3 pure-PyTorch absorb reference: EXECUTED
+official FlashMLA optimized kernel on SM120: NOT EXECUTED / NOT LISTED AS SUPPORTED
+```
+
+不能把前者的成功写成后者的成功，也不能把后者的未支持写成“5090 做不了 MLA”。
+
+## 7. 与 V2 报告 93.3% 的关系
+
+本实验得到的 `88.75%` 是：
+
+```text
+V2-Lite 576-element latent contract
+vs
+V2-Lite 5,120-element HF/V3-naive expanded cache
+```
+
+V2 技术报告的 `93.3%` 属于论文中的完整 V2 模型、论文基线与论文配置。两者回答的口径不同。
+
+可以说：
+
+- V2-Lite 本地吸收缓存相对本地展开格式减少 88.75%；
+- 这轮已经从 shape 推导升级为真实 buffer 与 decode 执行；
+- 结果支持 MLA 权重吸收机制。
+
+不能说：
+
+- 本实验“复现了论文 93.3%”；
+- 88.75% 推翻了 93.3%；
+- reference 实现的缓存降幅就是生产吞吐提升；
+- SM120 已运行 FlashMLA。
+
+## 8. 复现命令
+
+```bash
+PYTHONPATH=/tmp/deepseek-v2-lite-pydeps.6ZJAH3 \
+/home/wuyang/.pyenv/versions/navi-router-cu128/bin/python \
+  experiments/deepseek/v2_lite_absorb_probe.py \
+  --artifact-dir /tmp/deepseek-v2-lite-artifacts.OEjfce \
+  --v3-repo /tmp/llm-atlas-deepseek-v3-official \
+  --flashmla-repo /tmp/llm-atlas-flashmla-official \
+  --output src/data/deepseek-v2-lite-absorb.json \
+  --captured-at 2026-07-29T06:35:00Z
+```
+
+第二次改写 `--output` 到
+`src/data/deepseek-v2-lite-absorb-repro.json`，然后：
+
+```bash
+sha256sum \
+  src/data/deepseek-v2-lite-absorb.json \
+  src/data/deepseek-v2-lite-absorb-repro.json
+
+cmp -s \
+  src/data/deepseek-v2-lite-absorb.json \
+  src/data/deepseek-v2-lite-absorb-repro.json
+```
+
+## 9. 下一轮应该做什么
+
+1. 在官方明确支持的 SM90/Hopper 环境执行 FlashMLA dense decode correctness test；
+2. 把同一组 latent cache 接入 serving runtime，分开测 prefill、decode、batch 与 context；
+3. 同时记录吞吐、TTFT、TPOT、峰值显存和数值误差，避免只报一个“快多少”；
+4. 扩展到多个 layer 与多个 prompt，检查误差如何累计；
+5. 若 FlashMLA 官方未来加入 SM120，固定新 revision 后重开支持矩阵审计，不能把未来支持倒写到本次记录。
+