图例说明
矩阵图中,宽度代表列数(dim),高度代表行数。两个矩阵相乘时,左矩阵的列数必须等于右矩阵的行数(即相邻边匹配)。
第一步:回顾标准 Attention 的 GEMM
标准MHA score = Q * K^T,其中 Q 和 K 的 head_dim 都是 192(128 nope + 64 rope):
第二步:Non-Absorbed 完整 GEMM 链
Non-Absorbed 先解压 K = c_t * W_UK,再 score = Q * K^T。两步 GEMM:
第三步:矩阵吸收的关键洞察 — 结合律换括号
核心 把上面两步 GEMM 写成一个表达式,然后利用矩阵乘法结合律重新分组:
关键区别一目了然:左图需要先做一个大 GEMM (c_t x W_UK = [L,512] x [512,192]) 生成完整 K,右图把 W_UK 挪到 Q 侧 (Q x W_UK^T = [1,192] x [192,512]),只是 1 行的小 GEMM,然后直接用压缩后的 c_t 做 attention。核心差异在于与历史长度 L 相关的那步 GEMM 被消除了。
对比:Non-Absorbed vs Absorbed 实际 GEMM 形状
第四步:V 侧的矩阵吸收(同理)
V侧 原始 O = Attn x V = Attn x (c_t x W_UV),吸收后 O = (Attn x c_t) x W_UV:
第五步:矩阵吸收分离 — 为什么 Key 要「拆」成两部分
核心问题 上面我们假设整个 K 都可以吸收,但实际上 K 由两部分组成,其中一部分不可吸收!
5.1 Key 的组成结构
在 DeepSeek-V2/V3 的 MLA 中,每个 head 的 Key 向量 [192维] = K_nope [128维] + K_rope [64维],它们的来源完全不同:
5.2 为什么 RoPE 部分不能吸收?— GEMM 图解
吸收的本质是利用结合律把 W 挪到 Q 侧。但 RoPE 不是固定矩阵,每个位置都不同:
直觉理解:W_UK 就像一个固定的「翻译模板」,所有 token 用同一个模板,所以可以提前合并到 Q 侧。但 RoPE 是「随位置旋转的角度」,每个 token 位置角度不同,无法合并成一个固定模板——就像翻译每个人的话都要用不同的翻译标准,无法统一处理。
5.3 分离后的完整计算流程 — 两路合并
「矩阵吸收分离」总结:把 Key 的 192 维拆成 nope [128维] + rope [64维]。nope 部分是纯线性投影 (K_nope = c_t x W_UK),可以通过结合律将 W_UK 吸收到 Q 侧;rope 部分经过 RoPE 位置旋转 (K_rope = RoPE(k_pe, pos)),每个位置不同无法吸收,必须单独缓存 k_rope 做标准 attention。最后两路 score 相加。这就是「分离」的含义——把能吸收的分离出来吸收,不能吸收的单独保留计算。
第六步:一图总结 — Absorbed Decode 完整 GEMM 路径
总结 Absorbed Decode 的三步 GEMM:
GEMM-a [1,128] x [128,512] -> Q'_nope [1,512](吸收 W_UK,极小计算)
GEMM-b [1,512] x [512,L] + [1,64] x [64,L] -> score [1,L](在潜空间计算,读 c_t 和 k_rope)
GEMM-c [1,L] x [L,512] -> O'[1,512],再 [1,512] x [512,128] -> O[1,128](V 侧吸收,W_UV 最后做)
核心收益:所有与历史长度 L 相关的 GEMM 都在 512 维潜空间完成(而非展开的 192 维),KV Cache 只需存 576 维/token(vs 传统 MHA 的 32768 维)。