AI 大模型之 MLA

原创

ITPAPA
ITPAPA

ITPAPA

2026年10月13日 12:32
江苏

在小说阅读器读本章

去阅读

在公众号小说中沉浸阅读

Multi-head Latent Attention   MLA

多头潜在注意力     技术深化学习

范围：基本概念、技术演化、数学及计算机制、模块架构、主要配置参数、训练与推理流程、工程实现、性能评测、典型应用、风险与2026年最新公开进展。

1  技术定位与解决的问题

MLA是一种面向解码阶段KV缓存优化的注意力结构。它通过低秩下投影把每个Token的键值表征压缩为共享潜向量，再按需恢复各头所需表示，并将位置相关分量进行解耦。与减少KV头数量的GQA不同，MLA的关键是压缩保存的表示维度。

问题背景：大规模模型的结构、计算、显存、训练稳定性与生成能力之间存在不同的取舍。本技术对应的改进点，应通过与恰当的基线比较才能判断。它属于“模型结构或训练/生成机制”，不应把单篇论文的实验数字误解为通用硬件指标。

2  结构组成与数据流

图1  Multi-head Latent Attention (MLA) 核心计算/数据流示意

•  步骤1：输入隐藏状态 h\_t。

•  步骤2：下投影得到压缩KV潜向量 c\_t。

•  步骤3：分离RoPE位置分量与内容分量。

•  步骤4：各头重构/吸收投影并执行注意力。

•  步骤5：只缓存 c\_t 及必要位置分量。

3  底层原理与关键技术

3.1  机制设计与公式

常规MHA缓存与2×L×H×d\_head×bytes/token近似成正比；MLA的缓存依赖低秩潜向量及位置分量。

3.2  复杂度与状态开销

注意力输出仍来自QK相似度加权V，压缩过程并不把注意力计算复杂度直接降为线性。

3.3  硬件执行与精度

在实际实现中可吸收若干上投影矩阵，避免显式展开所有K/V头；RoPE采用解耦策略以维护位置相关计算。

架构识别原则：需要区分“数学形式”“训练目标”“GPU计算内核”“模型实例配置”“在线服务系统”。同名概念在这些层面的实现方式并不自动相同；评测必须说明层次、模型版本、数据精度、硬件与任务。

4  核心参数及工程含义

参数取值原则：除明确标注“论文示例”的值之外，本表列的是配置项与设计量，而非默认推荐值。实际有效范围受模型结构、训练方式、推理框架、显卡架构与可用显存约束。

4.1  定量公式与参数估算

KV缓存估算：传统MHA bytes≈2×L×H×d×S×b；MLA bytes≈L×(d\_c+d\_r)×S×b。L为层数，S为序列长度，b为每元素字节数。

4.2  可复核的示例计算

示例：设L=32、H=32、d=128、S=4096、b=2，则传统MHA约2048 MiB；若假定MLA潜维d\_c=512、位置分量d\_r=64，则对应压缩缓存约144 MiB。该例用于展示计算方法，并非任何发布模型的实测值。

4.3  核心伪代码与执行序列

h→W\_down→c\_KV；缓存(c\_KV,k\_RoPE)；Q与K内容/位置分量计算分数；V按重构/吸收路径聚合。

4.4  易出错的工程边界

检查潜变量与RoPE位置缓存是否重复保存；检查融合投影是否严格保持注意力输出及精度。

5  实施流程与关键技术细节

5.1  训练准备

按已发布模型配置定义Q/KV压缩矩阵、解耦RoPE维度与输出头数。

5.2  实现验证

训练阶段同时更新低秩投影与注意力权重，监测长上下文精度和重建信息损失。

5.3  测试优化

推理阶段预填充缓存与逐Token缓存分别实现；验证矩阵吸收后的数值等价。

5.4  交付评估

比较MHA/GQA/MLA相同序列长度和并发数下的缓存占用、TTFT、TPOT与精度。

6  技术验证与量化评估设计

图2  机制—参数—实现—应用的验证路径

实验控制：固定模型权重/训练Token预算、Tokenization、输入与输出长度分布、温度与采样配置、数据类型、GPU与软件版本；每组记录至少一次预热、多轮重复以及波动区间。结果按实际运行配置报告，避免跨研究条件直接比倍率。

7  典型应用场景

•  超长上下文的高并发聊天服务：先建立真实业务负载与基线，再决定是否引入该结构。

•  KV缓存预算受限的推理集群：先建立真实业务负载与基线，再决定是否引入该结构。

•  大模型推理成本优化与显存容量规划：先建立真实业务负载与基线，再决定是否引入该结构。

8  局限 代价与容易误解的地方

•  低秩压缩可能带来表达瓶颈，需要训练期共同适配。

•  优化的是KV缓存量；计算加速受kernel与硬件影响，不能将缓存节省直接等同于吞吐提升。

•  MLA与GQA按不同方式缩减缓存，部署对比须保证模型质量和计算负载可比。

9  与邻近架构的区别

MHA保存所有K/V头；GQA在多个Q头间共享少量KV头；MLA存储压缩后的内容潜变量及位置分量。

10  技术进展 截至2026年10月

DeepSeek-V2（2024）提出MLA；DeepSeek-V3（2024/2025技术报告）沿用MLA。DeepSeek-V2论文报告相对DeepSeek 67B的KV缓存下降93.3%，这是指定模型与基线的结果而非普适承诺。

11  推荐实验记录模板

12  原始论文与权威资料

•  [1] https://arxiv.org/abs/2405.04434

•  [2] https://arxiv.org/abs/2412.19437

说明：所有实验数值应以所引文献的相应版本、配置与对照组为准。本文的示意图为重新绘制的知识结构图，不能作为作者官方架构图。

技术深化 潜空间注意力的张量推导与缓存实现

本章在原有概览基础上展开推导与实现。文中的尺寸和计算例为教学配置，实验代码验证明确标注的简化核心；论文方法、工程建议和生产实测分别说明。

从低秩投影到可执行计算

设批量为B，历史长度为S，模型宽度为D，查询头数为H，内容头维为d，缓存潜维为r，位置维为p。以下采用列向量约定。输入h进入下投影后得到c，形状为r；每个头用独立的上投影恢复内容键和值。实际缓存为[B,S,r]的潜变量与[B,S,p]的位置键，不能再永久保存恢复后的[B,S,H,d]键值，否则缓存优势消失。

查询的内容分支q与位置分支分别计算相似度，再在同一softmax前相加。解耦RoPE并非删除位置信息，而是把位置相关的键从低秩内容恢复路径中分离。缩放分母应依模型实际拼接的查询键维度确定，不能直接把潜维r代入。

为什么矩阵吸收能避免重建全部键值

把查询内容向量左乘上投影转置，得到潜空间查询。键打分便变为潜空间点积。值聚合也可以先求加权潜变量，再一次上投影。两条等式来自矩阵乘法结合律，是对同一已训练模型的代数改写，不是对任意MHA权重的无损压缩。

这里每个头的注意力权重仍不同，因此值端不能把多个头的加权潜变量直接合并成一份。位置项含随位置变化的旋转，若强行把整套RoPE吸入固定矩阵，会破坏相对位置信息。预填充可以选择不同的计算路径；逐词解码更值得避免读取大规模展开键值。

缓存预算与性能边界

以32层、32头、d=128、r=512、p=64、4096词元和BF16为教学配置，MHA缓存为2GiB，MLA为144MiB，比例为7.03125%。并发16条时分别为32GiB与2.25GiB。只核算持久化数据，不含页表、对齐、工作区、模型权重和通信缓冲。若张量并行实现复制潜变量而不是分片，单卡缓存应按实际复制因子调整。

降低缓存读流量通常改善长上下文解码，但潜空间的点积维r可能比单头内容维d更大。低并发或短序列时，投影成本和内核启动成本可能成为主项。应分别测量缓存字节、单步带宽和端到端TPOT，不能由缓存比例直接推导加速倍率。

最小数值实验

随包脚本01\_mla.py比较显式恢复路径与矩阵吸收路径，只含一个头的内容分支，不含RoPE、残差和训练。NumPy FP64用于检查代数等价，不用于推理性能评测。

K = C @ Uk.T

V = C @ Uv.T

a = softmax(K @ q / np.sqrt(d))

explicit = a @ V

latent = Uv @ (a @ C)

absorbed\_scores = C @ (Uk.T @ q) / np.sqrt(d)

assert np.allclose(K @ q / np.sqrt(d), absorbed\_scores)

assert np.allclose(explicit, latent)

实验先固定随机种子，再分别改变S、r与d。正确实现的误差应接近浮点舍入；若误差随序列显著扩大，应检查转置、softmax维度和缩放因子。随后加入位置分支，验证位置偏移与增量缓存是否和完整前向一致。

训练与部署排障

压缩秩过低造成质量下降时，先观察长距离检索和领域任务损失，再做秩消融；仅看训练损失可能漏掉检索瓶颈。显存未下降时，检查缓存分配器是否仍按完整K/V维度预留。预填充正常而解码错位时，核对RoPE绝对位置、缓存写入索引与每层状态。FP8缓存还需要scale、反量化和长程误差测试，不能只改dtype。

验收实验设计

设置S为1024、4096、16384，并发为1、8、16；对每格记录峰值显存、缓存有效字节、P95 TPOT以及长文本精确召回率。性能比较只使用相同权重的两种MLA执行路径；比较MHA与MLA模型质量时，应另做相同训练预算的架构实验。通过条件为吸收前后输出在定义的容差内一致，缓存容量符合实测布局，且性能收益未以不可接受的质量损失换取。

本册配套实验验证记录

配套文件为 实验代码/01\_mla.py。本次使用NumPy CPU固定随机输入运行，执行通过。结果验证本章所标注的公式、索引、容量或事务边界，不代表完整模型训练效果或GPU性能实测。

max\_abs\_error 8.881784197001252e-16

工程实操 矩阵吸收的逐步计算与解码接口

固定输入与检查目标

使用单头二维内容分支，潜变量C的三行为[1,0]、[0,2]、[1,1]；Uk两行为[2,1]、[0,1]，Uv两行为[1,-1]、[2,0]，查询q=[1,2]。本例不加位置项，也不使用头维缩放，专门核对吸收的代数顺序。

逐步计算与预期结果

1  先恢复K=C×Uk转置，得到三行[2,0]、[2,2]、[3,1]；与q点积，分数依次为2、6、5。

2  另一条路径先求潜空间查询Uk转置×q=[2,3]，再由C点积得到同一分数。稳定softmax权重约为[0.013213,0.721399,0.265388]。

3  恢复后的V为[1,2]、[-2,0]、[0,2]。权重聚合输出约为[-1.429585,0.557202]；先聚合C再左乘Uv也应得到此值。

4  验收对象是两条执行路径的输出与分数。没有进行模型训练，不能由本例证明给任意MHA模型加低秩投影后质量不变。

接口与实现约束

增量接口至少传入当前查询、每层潜变量缓存、位置键缓存、有效历史长度与绝对位置。缓存写入后才能把相应位置纳入可见范围；异步执行还需确保消费者读取前写入已完成。

把显式恢复路径保留为小尺寸参考实现。定位错误时分别比对内容分数、位置分数、softmax权重和最终输出，避免只比较最后logits。长上下文部署另核算分片或复制的潜变量、对齐字节和临时展开区。

故障核查表

扩展验证组合

在S=3、17、129及不同r/d比值下对照两条路径；再加入位置分支、非零起始位置和单步缓存追加。通过代数检查后才进入设备计时。

配套代码与结果

进阶实验/boundary\_cases.py 中的 c01 函数对应本节。可在进阶实验目录运行 python boundary\_cases.py，或执行 python -c \"from boundary\_cases import c01; print(c01())\" 单独检查。本节固定样例执行通过，结果保存在边界实验结果.json。

验证范围为正文明确标注的简化机制与错误反例。工程接口、扩展测试和故障注入属于本资料提出的实施建议；完整模型训练、GPU数值与吞吐、跨进程恢复仍须在目标环境验证。

训练与部署 潜变量的双路梯度与吸收路径训练

目标与验证范围

在单头内容分支中，潜变量C同时参与分数和输出值。检查吸收前后前向一致，还不够证明训练路径一致；反向必须把这两条贡献相加。以下沿用列向量与行存储混合约定，C每行是一条缓存潜向量。

解析关系与梯度路径

令J=u转置o，把u作为上游输出梯度；a=softmax(C×潜查询)。值路径给出对C的梯度a与Uv转置u的外积，分数路径则先算dP，再经softmax雅可比得到δs。

把C作为唯一缓存，并不会让键恢复投影Uk停止训练。Uk通过潜查询影响分数，Uv通过聚合值影响输出；随包分别扰动C、Uk、Uv和q，检查四组梯度。测试使用FP64和中心差分，不依赖自动微分框架。

训练配置与误用边界

前向吸收只是一种执行重排。若训练实现预先缓存由权重算出的吸收矩阵，权重更新后必须重建，不能复用上一优化步的矩阵。冻结投影与训练投影应在参数组中明确区分。

梯度核查先使用小尺寸、无dropout、无量化输入。随后加入位置分支，验证内容梯度与位置梯度分别流向正确投影；不要把位置键错误当成不可训练的常量。

部署验收与恢复

部署路径可选择不保留反向图，但在线解码状态与训练缓存的语义不同。训练中把C随意detach会阻断前缀隐藏表示的梯度，需要明确这是截断训练方案。

验收同时保存分数误差、输出误差和参数梯度误差。比较执行路径使用同一权重；改变潜维后的质量与训练预算另行评价。

配套检查与本次数值结果

训练部署实验/training\_checks.py 中的 t01 函数已执行通过。本例解析梯度与中心差分的最大绝对误差为 1.375e-10。完整输入、解析实现和实际输出见代码及训练与状态验证结果.json。

可在训练部署实验目录执行 python training\_checks.py，或执行 python -c \"from training\_checks import t01; print(t01())\"。所示结果来自NumPy CPU教学验证；目标模型的自动微分、低精度和分布式执行需另做对应验收。