wuyang
421 lines
13 KiB
Markdown
421 lines
13 KiB
Markdown
# 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 后重开支持矩阵审计，不能把未来支持倒写到本次记录。
|
||
|