Multi-head Latent Attention

与 MHA/GQA 的根本布局差异

MHA/MQA/GQA缓存标准 head-shaped K/V：每 token 约 $2H_{kv}D$ 个元素。MLA（Multi-head Latent Attention）先把 hidden state 下投影到共享低维 latent，再由每个 Query head 的上投影定义 content key/value；同时用解耦的位置编码路径保存 RoPE key 分量。

一个教学抽象：

c_t^{KV}=x_tW^{DKV}\in\mathbb{R}^{d_c}

k_{t,h}^{C}=c_t^{KV}W_h^{UK},\qquad v_{t,h}^{C}=c_t^{KV}W_h^{UV}

k_t^R=\operatorname{RoPE}(x_tW^{KR}),\qquad q_{i,h}^R=\operatorname{RoPE}(x_iW_h^{QR})

attention score 由 content 与 positional 两部分相加：

s_{i,t,h}=q_{i,h}^C(k_{t,h}^C)^\mathsf T+q_{i,h}^R(k_t^R)^\mathsf T

具体 DeepSeek 版本的符号、维度、是否还有 query compression、head dimension 和 cache layout 可能不同；本页只固定理解 latent、解耦 RoPE 与矩阵吸收所需的最小结构。

Decode 中为什么要做矩阵吸收

若每步先从 cached latent 展开全部 per-head content keys/values，会重新引入较大的中间量。利用结合律：

q_h^C(k_{t,h}^C)^\mathsf T =q_h^C(W_h^{UK})^\mathsf T(c_t^{KV})^\mathsf T =\bar q_h^C(c_t^{KV})^\mathsf T

其中 $\bar q_h^C=q_h^C(W_h^{UK})^\mathsf T\in\mathbb{R}^{d_c}$ 。于是 Query 可直接对 cached latent 打 content score。

Value 路径也可把上投影移到加权之后：

\sum_t p_t v_{t,h}^C =\sum_t p_t(c_t^{KV}W_h^{UV}) =\left(\sum_t p_tc_t^{KV}\right)W_h^{UV}

目的不是改变 attention 语义，而是避免 decode 时为所有历史位置物化 expanded per-head K/V。

INTERACTIVE EXPLAINER

MLA 的 latent content 路径与解耦 RoPE 路径

静态图标出缓存字段、Query 侧矩阵吸收、两部分 score 和 Value 上投影位置；shape 为教学抽象。

图中省略 query compression、具体 normalization/linear fusion、各模型的 head dimensions、TP 分片与生产 kernel layout。为了可读性，RoPE key 画成跨 heads 共享的一条路径；具体模型实现应以论文和 checkpoint 配置为准。

可运行的展开-vs-吸收对齐

import math,platform,torch

torch.manual_seed(107)
DEVICE=torch.device("cuda" if torch.cuda.is_available() else "cpu")
DTYPE=torch.float64

def main():
    b,t,hq=1,6,3; hidden,dc,dk,dr,dv=12,4,5,2,6
    x=torch.randn(b,t,hidden,device=DEVICE,dtype=DTYPE)
    x_new=torch.randn(b,1,hidden,device=DEVICE,dtype=DTYPE)
    w_dkv=torch.randn(hidden,dc,device=DEVICE,dtype=DTYPE)
    w_uk=torch.randn(hq,dc,dk,device=DEVICE,dtype=DTYPE)
    w_uv=torch.randn(hq,dc,dv,device=DEVICE,dtype=DTYPE)
    w_qc=torch.randn(hidden,hq*dk,device=DEVICE,dtype=DTYPE)
    w_qr=torch.randn(hidden,hq*dr,device=DEVICE,dtype=DTYPE)
    w_kr=torch.randn(hidden,dr,device=DEVICE,dtype=DTYPE)

    latent=x@w_dkv
    k_content=torch.einsum("btc,hcd->bhtd",latent,w_uk)
    v_content=torch.einsum("btc,hcv->bhtv",latent,w_uv)
    k_rope=x@w_kr
    qc=(x_new@w_qc).view(b,1,hq,dk).transpose(1,2)
    qr=(x_new@w_qr).view(b,1,hq,dr).transpose(1,2)
    content_scores=torch.einsum("bhnd,bhtd->bhnt",qc,k_content)
    rope_scores=torch.einsum("bhnr,btr->bhnt",qr,k_rope)
    probs=torch.softmax((content_scores+rope_scores)/math.sqrt(dk+dr),-1)
    expanded=torch.einsum("bhnt,bhtv->bhnv",probs,v_content)

    absorbed_q=torch.einsum("bhnd,hcd->bhnc",qc,w_uk)
    latent_scores=torch.einsum("bhnc,btc->bhnt",absorbed_q,latent)
    torch.testing.assert_close(latent_scores,content_scores,rtol=1e-10,atol=1e-10)
    weighted_latent=torch.einsum("bhnt,btc->bhnc",probs,latent)
    absorbed=torch.einsum("bhnc,hcv->bhnv",weighted_latent,w_uv)
    torch.testing.assert_close(absorbed,expanded,rtol=1e-10,atol=1e-10)
    print(f"python={platform.python_version()} torch={torch.__version__}")
    print(f"device={DEVICE} dtype={DTYPE}")
    print(f"latent={tuple(latent.shape)} k_rope={tuple(k_rope.shape)}")
    print(f"expanded_K={tuple(k_content.shape)} expanded_V={tuple(v_content.shape)}")
    print(f"query_content={tuple(qc.shape)} absorbed_query={tuple(absorbed_q.shape)}")
    print(f"score_max_abs_diff={(latent_scores-content_scores).abs().max():.3e}")
    print(f"output_max_abs_diff={(absorbed-expanded).abs().max():.3e}")
    print(f"per_token_standard_KV_elements={2*hq*dv} mla_cached_elements={dc+dr}")

if __name__=="__main__":
    with torch.inference_mode(): main()

完整文件位于 examples/mla_reference.py。实际运行中 content score 最大差 1.421e-14，输出最大差 5.329e-15；示意标准 K/V 每 token 36 元素，而 latent+RoPE cache 为 6 元素。

这个元素比较只针对教学维度，不是任何具体模型的压缩比。生产 MLA 还需要考虑 weight layout、位置分量、量化 scale、TP ownership、absorbed projection 的 kernel、prefill 与 decode 不同路径。

容量与成本

教学容量：

M_{MLA}\approx LB T(d_c+d_{rope})s

相比标准 $2LBTH_{kv}Ds$ ，收益取决于 $d_c+d_{rope}$ 与实际 KV 元素数之比。显存/带宽下降会把成本转移到 latent/query projections、absorbed score/value kernels 与位置路径。若 backend 先展开完整 K/V 再调用普通 attention，缓存可能仍省，但 decode 计算与临时内存收益会打折。

MLA 不自动解决 PagedAttention 的动态分配；latent cache 仍可按 blocks 分页。也不自动等同 FlashAttention；online softmax 仍可用于两部分 score 合并后的 attention。

上下游关系

MHA/MQA/GQA提供标准 head-shaped KV 对照；本页改变缓存字段与 decode 代数路径。
KV Cache给出通用生命周期与容量口径；MLA 页面替换每 token payload 公式。
PagedAttention仍可管理 latent/positional blocks 的物理分配。
FlashAttention负责 tile/online softmax IO，不能用来解释 latent 或矩阵吸收。

参考资料

DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model