mla-analysis

MLA 成本分析与运行状态指南

运行状态选择

运行状态	s 范围	最佳内核	原因
解码	s=1	FlashMLA	相较 FlashAttention，延迟降低 16 倍（压缩 KV）
推测解码	s=2-32	MLAvar6+ 或 FlashMLA	MLAvar6+ 应能胜过 FlashMLA 和 FlashAttention
预填充	s>128	FlashAttention	避免潜空间计算带来的 4 倍 FLOP 开销

性能交叉点：对于 DeepSeek-V3 参数，s 约为 16-32 时，FlashAttention 的速度会超过 FlashMLA。

成本模型（类似 DeepSeek-V3：h=128，d=128，k=512，p=64）

FlashAttention

• FLOP：2bhst(2d + p) = 2bhst * 320
• 字节数：w * bh(s+t)(2d + p) = w * bh(s+t) * 320
• s=1 时：AI ≈ 1 FLOP/字节（严重受内存限制）
• s=1024 时：AI ≈ 819 FLOP/字节（严重受计算限制）

FlashMLA（通过吸收实现潜空间注意力）

• FLOP：2bhst(2k + p) = 2bhst * 1088
• 字节数：w * (bhs(2k+p) + bt(k+p))
• s=1 时：AI ≈ 228 FLOP/字节（即使在解码时也受计算限制）
• 问题：k=512 时，FLOP 随 s 线性增长，而非使用 d=128 ——每个 token 的 FLOP 多 3.4 倍

MLAvar6+（拆分潜空间与解压数据）

• FLOP：2bhstp + 4bhsnd + 4bhsok
• 字节数：w * (bhsp + bhsd + bhsk + bok + btp + bhnd + bhnd + bhsd + bhsk)
• 关键点：选择 n 可将运算强度调至接近屋顶线拐点

吸收技巧（FlashMLA）

Score:  Q @ K^T = Q @ (Z @ Wk)^T = (Q @ Wk^T) @ Z^T = Qz @ Z^T
Value:  softmax(A) @ V = softmax(A) @ (Z @ Wv) = (softmax(A) @ Z) @ Wv

完全在潜空间（k 维）中运算，输出解压 O_latent @ W_kvb2^T 则由单独的内核完成。

MLAvar6+ 设计

将 KV 缓存拆分为：

• o 个旧 token：潜变量 Z（b, o, k）——计算密集、带宽需求低
• n 个新 token：解压后的 K,V（b, n, h, d）——带宽密集、计算需求低
• t = o + n

在两者之间插值：n=0 → 类似 FlashMLA，o=0 → 类似 FlashAttention。

历史基准（RTX 5090，b=32，t=4096，bfloat16）

注意：这些结果来自 RTX 5090 开发阶段。将其用作优化目标前，请在当前设备上重新进行性能分析。不同设备的拐点、带宽和计算上限各不相同——参见 src/mla_var3/conf/devices.json。

配置	FlashMLA	FlashAttention	MLAvar6+ V3（最佳）
s=1	419 μs	6,781 μs	829 μs（V2）
s=16	5,161 μs	6,727 μs	4,444 μs

MLAvar6+ V3 目前在推测解码（s=16）方面表现最佳，比 FlashMLA 快 14%。

张量形状（DeepSeek-V3 默认值）

符号	描述	默认值
b	批量大小	64
h	注意力头数量	128
s	查询序列长度	可变
t	KV 上下文长度	4096
d	头维度	128
p	位置嵌入维度	64
k	潜空间（压缩后）维度	512

张量布局：

• Q, Qc：[B, H, S, D] 或 [B, H, S, K]
• K, V：[B, T, H, D]（已解压）
• Z：[B, T, K]（潜空间 KV 缓存）
• QPE, KPE：[B, H, S, P] 或 [B, T, P]

详细分析

完整的成本模型推导和屋顶线分析：docs/cost-analysis.md