[Tech] 为什么 MLA + MTP 容易吃亏?从 Arithmetic Intensity 看 Attention
前几天在看苏剑林老师的这篇文章:https://kexue.fm/archives/11848,里面有一句话让我停了一下: 但现在除 KV Cache 外,Decoding 还有一个新的变数——MTP,或者说推测解码,其思想是计算换速度。然而 MLA 在 Decoding 时表现为 head_dims=512+ 的 MQA,已经提前消耗了大部分算力,所以“MLA+MTP”很容易吃亏。 我第一反应是:这是什么意思?为什么 MLA 在 decode 时会“提前消耗算力”?为什么它在 decode 时的 FLOPs,相当于一个 head dim 512+ 的 MHA,反而和 MTP 有冲突? 顺着这个问题往下看,最后发现背后其实是一个非常经典、但放在 Attention 上又很有意思的概念:Arithmetic Intensity(AI,计算密度)。 而且最后得到的结论意外地简单(下面都按 BF16 的 KV cache 算): 把 MHA 一路约下来,AI 居然刚好就是 1; GQA、MQA 也一样干净,只和 head 的数量有关 —— context length、head dim 全部约掉了,一个都不剩; MLA 还是同一个形状,同样和 latent dim 无关,只是多了一个不到 2 的常数。 把四种结构排在一起,single-token decode 下 attention core 的 AI 是这样的: Attention 历史 cache 里存什么 这时的 AI 大约是 MHA 每个 Query head 各有 K、V 1 GQA 一组 Query head 共用 K、V Query head 数 / KV head 数 MQA 所有 Query head 共用 K、V Query head 数 MLA 只存一份 latent,K、V 都从它展开 约 2 × Query head 数 第 2 节和第 4 节会把这四行分别推一遍。前三行其实是同一个公式的三个 special case,只是 KV head 数取值不同;最后一行多出来的那个常数,来源完全不同。 ...