大语言模型中的注意力机制，清晰讲解

自注意力 → 交叉注意力 → 多头注意力 → MQA → GQA → FlashAttention → 稀疏注意力 → PagedAttention → RadixAttention，用图解逐一说明。

每张模型卡都会宣传其注意力机制。

多查询注意力、分组查询注意力和多头潜在注意力与参数数量和基准分数并列列出。

它们的存在都源于同一个限制。为长序列和大批次存储注意力状态会耗尽 GPU 内存。

今天，我们按它们出现的顺序逐一拆解，并说明每个机制修复了哪个缺陷。

让我们开始吧！

注意力为何存在

早期的序列模型，如循环神经网络（RNN），在标记之间传递固定大小的隐藏状态。

两个标记相距越远，连接就越弱。长距离依赖关系会逐渐消失。

注意力直接解决了这个问题。它不通过瓶颈传递状态，而是让每个标记都能查看所有其他标记，并判断每个标记的相关性。

没有任何信息隐藏在摘要步骤之后。

这种直接性正是 Transformer 强大的原因，也是它们代价高昂的原因。成本在于存储每个标记所看到的内容所需的内存。

驱动一切的约束

每次模型进行注意力计算时，它都需要记住之前每个词元的样子。

在预填充阶段，模型一次性处理你的整个提示，并为每一层的每个词元计算一个键向量和一个值向量。

这些向量被存储在所谓的 KV 缓存中，以便解码步骤能够对它们进行注意力计算，而无需从头重新计算。

缓存随着每个生成的词元而增长。对于一个 70B 模型，在 BF16 精度下，单个 128K 词元的上下文大约占用 40 GB 的 KV 缓存，这几乎相当于 4 比特量化后的模型权重大小。

约束不在于计算量，也不在于注意力公式中的数学运算，而在于存储注意力已见内容的内存。

自注意力与因果注意力

自注意力允许每个词元关注同一序列中的其他所有词元。

模型为每个词元计算查询、键和值，然后通过查询-键点积来决定每个词元应关注其他词元的程度。

这是每个 Transformer 层内部的基础操作。

因果注意力是应用了三角掩码的自注意力。每个词元只能关注它之前的词元，绝不能关注未来的词元。

这正是使仅解码器生成成为可能的原因，因为如果没有掩码，模型在生成答案之前就会看到答案。

交叉注意力在本质上有所不同。查询来自一个序列，而键和值来自另一个序列。

这就是 T5 和 Whisper 等编码器-解码器模型如何将编码器的输出连接到解码器的方式。在 Llama 和 GPT 等仅解码器模型中，交叉注意力完全不会出现。

多头注意力

多头注意力（MHA）是2017年Transformer论文中的原始设计。

每个注意力头拥有自己独立的查询、键和值权重矩阵。使用32个头时，每层就有32个独立的KV投影。

其优势在于表现力。不同的头能同时学习追踪不同类型的关系。一个头可能追踪句法结构，另一个关注语义邻近性，还有一个处理长距离指代关系。

代价则是内存。每个头维护自己的KV缓存。一个具有32层且每层32个头的模型，每个请求的每个token需存储1,024个独立的KV张量。

GPT-3每层使用96个头。在如此规模的模型上处理128K token的上下文时，仅KV缓存就足以在批次增长前占满一块GPU。

这种内存开销正是MHA之后所有设计致力于减少的目标。

多查询注意力

多查询注意力（MQA）采取了最直接的方式。所有查询头仍拥有各自的权重矩阵，但每个查询头共享一个键头和一个值头。

KV缓存缩减为原来的头数分之一。MHA存储32个独立的KV投影，而MQA仅存储一个。

解码速度得以提升，因为每一步从HBM加载的字节数大大减少，而解码受限于内存带宽。

质量上的代价是实实在在的。强制所有查询头共享一个键和一个值，会损失MHA所提供的部分表现力。

Falcon、PaLM及早期Gemini变体采用了MQA，并为了吞吐量的提升接受了这一权衡。

MQA在质量上所做的让步，在接下来的设计中得到了很大程度的弥补。

分组查询注意力

分组查询注意力（GQA）介于MHA与MQA之间。

查询头被划分为多个组，每组共享一个键头和一个值头。各组之间相互独立。

以32个查询头和8个KV组为例，你只需存储8个KV投影而非32个，相比MHA，KV缓存大小减少了4倍，同时恢复了MQA所牺牲的大部分质量。

正是这种平衡，使得GQA成为近年来几乎所有主流开源权重模型的默认选择。Llama 2 70B使用8个KV组。Llama 3、Mistral、Mixtral、Gemma和Qwen均采用GQA。

最初的GQA论文表明，在内存成本仅为MHA一小部分的情况下，其质量可与MHA相媲美，这一结论在多个模型家族中得到了验证。

GQA减少了存储的KV头数量。接下来要压缩的是头本身。

多头潜在注意力

多头潜在注意力（MLA）是DeepSeek的贡献，于2024年5月在DeepSeek-V2中首次提出。

MQA和GQA旨在减少KV头数量，而MLA则将全维度的键和值向量压缩到低秩潜在空间后再进行缓存。在注意力计算时，这些潜在向量会被解压回完整维度。

缓存的对象是潜在向量，而非完整的KV张量。这使得缓存占用比GQA更小，同时保留了更多MQA所牺牲的表达能力。

代价是计算量。在每一步注意力计算中解压会增加FLOPs。但在推理时，内存带宽往往是比计算更常见的瓶颈，因此在大多数实际服务配置中，更小的缓存优于额外的数学运算。

DeepSeek-V2、V3和R1均采用MLA。在DeepSeek-V2的基准测试中，MLA在质量上与MHA相当或更优，同时将KV缓存削减至相同模型规模下MHA所需的大约5%至13%。

MHA、MQA、GQA和MLA都是关于存储什么的决策。接下来的技术则关乎你计算它的成本有多高。

FlashAttention

FlashAttention并不改变注意力机制计算的内容。它改变的是计算访问内存的方式。

标准注意力机制会构建完整的N x N注意力矩阵，将其写入HBM，读回以进行softmax计算，再次写入结果，然后再读回以进行加权求和。

对于一个4K令牌的序列，该矩阵包含4,096 x 4,096个值。反复将其移入和移出HBM，正是使注意力成为长上下文瓶颈的原因。

FlashAttention将计算分块处理。它按块处理注意力矩阵，这些块可容纳于芯片上的SRAM中，增量式计算softmax而无需物化完整矩阵，并将输出一次性写入HBM。

数学上是相同的。内存流量则不然。

如今，每个主要的服务引擎默认都使用FlashAttention内核。它并非一种新的注意力类型。它是执行你模型所使用的任何注意力类型的标准内核。

MHA、GQA和MLA回答了存储什么以及如何压缩的问题。FlashAttention则回答了如何高效计算的问题。

稀疏注意力则从另一个角度切入。不是缓存哪些令牌，而是究竟要关注多少个令牌。

稀疏注意力

全注意力在序列长度上是O(N⁠²)的复杂度。对于1M token的上下文，注意力矩阵有万亿个条目。即使FlashAttention减少了内存流量，对这么多token计算注意力仍然不可行。

稀疏注意力跳过了注意力矩阵的大部分区域。并非每个token都关注其他所有token，只计算选定的一小部分配对。

滑动窗口注意力（SWA）是最简单的变体。每个token只关注最近的W个token。局部上下文得以保留；远距离上下文被舍弃。

Mistral在某些层使用SWA，在其他层交替使用全注意力，以兼顾局部精度和一定的全局覆盖。

原生稀疏注意力（NSA）是DeepSeek在2025年的贡献，与MLA不同。NSA并非在推理时事后应用稀疏性，而是从一开始就用稀疏注意力训练模型。

每一层结合三个并行分支：用于全局上下文的压缩粗粒度注意力、用于重要token块的精细选择性注意力，以及用于局部上下文的滑动窗口注意力。模型在预训练期间学习哪些token是重要的。

NSA在大多数基准测试上达到全注意力的质量，同时在长序列上运行速度显著更快。

Qwen2.5-1M为其百万token上下文窗口采用了稀疏注意力方法，因为在该长度下，全注意力消耗了前向传播90%以上的时间。

训练时内置的稀疏性正被证明是长上下文扩展最根本的解决方案。

服务层

上述所有内容均存在于模型权重之中。PagedAttention与RadixAttention则驻留在服务引擎内，它们分别在不同层面应对相同的KV缓存压力。

当请求到达时，服务引擎需为其KV缓存分配GPU内存。朴素的做法是为最大可能的序列长度预留连续内存。

一个允许生成最多4,096个token的请求，无论实际使用40个还是4,000个token，都会预先保留4,096个槽位。

这种方法浪费了60-80%的GPU内存。请求间的碎片化问题使其雪上加霜。

PagedAttention，作为vLLM运行的机制，以操作系统管理虚拟内存的方式管理KV缓存。固定大小的块按需分配。

块表将每个请求的逻辑块映射到当前空闲的物理块上。无需预分配，无碎片化，内存浪费降至4%以下。

RadixAttention，作为SGLang运行的机制，则更进一步。当多个请求共享共同前缀时，例如发送给每位用户的冗长系统提示，该前缀的KV块只需计算一次。

RadixAttention将KV块存储在以token序列为索引的基数树中。新请求遍历该树，找到最长匹配前缀，复用这些块，仅计算真正的新后缀部分。

在多轮工作负载下，RadixAttention实现了75-95%的缓存命中率。服务于数千用户的系统提示仅计算一次，并在LRU策略驱逐前持续复用。

PagedAttention与RadixAttention并不改变模型所使用的注意力机制。采用GQA的Llama 3模型在任一引擎下均可正常运行。这两层相互独立。

整合归纳

贯穿所有这些方案的主线，是同一个瓶颈。KV缓存内存是制约因素，而这里的每一项设计，都是针对它的不同应对策略。

MQA、GQA和MLA减少了每个token的存储量，各自以不同程度的质量换取内存节省。

FlashAttention降低了计算访问内存的成本，却完全不改动数学本身。

稀疏注意力减少了需要关注的token数量，这是唯一能扩展到百万级token上下文的答案。

PagedAttention和RadixAttention在服务层运作，削减分配与复用中的浪费，而非在模型内部动手。

弄清楚在你的场景中，哪一个是真正的约束条件，才能决定哪一种方案能真正为你带来数字上的提升。

感谢阅读。今天就到这里。

Cheers! :)