13 KiB
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 输入下:
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 |
官方入口:
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 可以写成:
kₜ = Wᴷ cₜ
vₜ = Wⱽ cₜ
naive 实现先把所有历史 cₜ 展开成每个 head 的 kₜ 和 vₜ,再缓存展开结果。V2-Lite
有 16 个 head:
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 部分:
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 聚合:
Σₜ pₜvₜ
= Σₜ pₜ(Wⱽcₜ)
= Wⱽ(Σₜ pₜcₜ)
因此可以先用注意力概率 pₜ 对 512 维 latent 做加权和,再把聚合结果通过 Wⱽ 展开到各个
head 的 value 空间。
naive 顺序:
每个历史 cₜ → Wⱽ → 展开 vₜ → cache → 按 pₜ 加权
absorb 顺序:
每个历史 cₜ → cache → 按 pₜ 加权 → Wⱽ → 展开输出
2.4 为什么 64 维 RoPE key 仍要缓存
RoPE 是位置相关变换。历史 token t 的 key 旋转角度依赖位置 t,不能像固定线性投影那样无条件
搬到当前 query 一侧。因此 V2-Lite 的可缓存状态不是只有 512:
512 KV latent + 64 decoupled RoPE key = 576 elements/token/layer
这也是网站从第一轮起始终把“512 latent”与“576 完整缓存合同”分开的原因。
3. 实验怎样避免偷换输入
3.1 输入不是随机 hidden state
固定中文 prompt:
用通俗的语言解释,为什么稀疏专家模型可以拥有很多参数,但每个 token 只使用其中一小部分。
官方 tokenizer 得到 26 个 token。执行链是:
官方 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。
执行:
- 官方 V2-Lite Hugging Face eager attention;
- pinned V3 官方
attn_impl="naive"; - 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 的定义是:
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 只加入:
arch=compute_100f,code=sm_100f
arch=compute_90a,code=sm_90a
csrc/api/dense_decode.h 还显式检查:
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。
正确的实验状态是:
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% 是:
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. 复现命令
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,然后:
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. 下一轮应该做什么
- 在官方明确支持的 SM90/Hopper 环境执行 FlashMLA dense decode correctness test;
- 把同一组 latent cache 接入 serving runtime,分开测 prefill、decode、batch 与 context;
- 同时记录吞吐、TTFT、TPOT、峰值显存和数值误差,避免只报一个“快多少”;
- 扩展到多个 layer 与多个 prompt,检查误差如何累计;
- 若 FlashMLA 官方未来加入 SM120,固定新 revision 后重开支持矩阵审计,不能把未来支持倒写到本次记录。