llm-atlas/research/DEEPSEEK_MLA_ABSORB_AUDIT.md

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