12. 注意力算子、算子融合与 CUDA Graph
前几篇讨论的都是调度层面的事：谁来算、算多少、KV 放在哪里。这一篇深入到一次前向计算的内部，看 GPU 上的计算本身怎样变快。三条路线分别是：FlashAttention 减少注意力计算对显存的反复读写，算子融合把相邻的小运算合并，CUDA Graph 消除 CPU 一次次下达指令的开销。读完你应该能说清它们各自优化的是哪一段时间，以及为什么调试时常常要先把它们关掉。
GPU 上的计算由一个个“内核”组成
模型的一次前向计算，在 GPU 上被拆成一连串内核（kernel）：一次矩阵乘法是一个内核，一次归一化是一个内核，一次激活函数又是一个内核。每个内核的工作方式大致相同：从显存（HBM）读取输入，在计算单元里算完，把结果写回显存，下一个内核再从显存读出来接着算。
第 6 篇说过，HBM 的带宽是有限的，而 GPU 片上的 SRAM（共享内存和寄存器）快得多，只是容量很小，H100 每个 SM 只有两百多 KB。所以让计算变快的一个通用思路是：数据一旦搬上片，就尽量在片上把能做的事做完，少往 HBM 写中间结果。
标准注意力的问题：一个巨大的中间矩阵
按第 3 篇的公式，注意力分三步：算分数矩阵 S = Q·Kᵀ，对每一行做 softmax 得到权重 P，再算 P·V。最直接的实现是三个内核，每一步把结果写回 HBM。
问题出在中间的 S 和 P。输入有 N 个 token 时，它们是 N × N 的矩阵。以 Prefill 8192 个 token 为例，一个头的分数矩阵就有 8192 × 8192 ≈ 6700 万个元素，按 2 字节算约 134 MB。一层有 32 个头，S 和 P 各要写一遍、读一遍，一层就要在 HBM 上来回搬十几 GB。而 Q、K、V 本身加起来只有约 100 MB。
真正的计算量并不大，时间都花在搬运这个中间矩阵上了。
FlashAttention：分块计算，不写出中间矩阵
FlashAttention 的做法是把 Q、K、V 切成小块（tile），每次把一块 Q 和一块 K、V 搬到片上 SRAM，在片上算出这一小块的分数、softmax 和加权求和，再处理下一块 K、V，最后只把最终结果写回 HBM。完整的 N × N 矩阵从头到尾没有出现在显存里。
难点在 softmax：它需要知道一整行的最大值和总和，才能算出每个元素的权重，而分块计算时一次只能看到一行的一部分。FlashAttention 用的是在线 softmax：每处理一块，就更新这一行目前为止的最大值和总和，并用一个缩放系数修正之前已经累加的结果。全部块处理完，得到的结果与一次性计算完全相同，不是近似。
几点需要分清：
- 计算量没有减少。 注意力仍然要算 N × N 个分数，计算量仍随长度平方增长。FlashAttention 减少的是 HBM 读写，以及中间矩阵占用的显存。
- 显存占用从平方级降到线性。 不再需要存放 N × N 的矩阵，长上下文的 Prefill 才变得可行。
- 它是一个算子，不是一个框架。 它和 PagedAttention、前缀缓存解决不同的问题，可以同时使用。
FlashAttention 之后又有第二、第三代，分别改进了 GPU 上的并行划分方式，并针对 H100 这一代 GPU 的新硬件特性做了优化。
Decode 的注意力是另一种形状
Prefill 时，Q 有几千行，计算密集；Decode 时，每个请求的 Q 只有 1 行，却要和几千个历史 K、V 做注意力。这是一个典型的访存受限计算，主要工作是把 KV 从显存读出来。
Decode 注意力还有一个特殊的麻烦：批很小、上下文很长时，可并行的工作太少。只有一个请求、只有 8 个 KV 头时，如果每个头只交给一个 SM 处理，H100 的 132 个 SM 大部分都闲着。解决办法是沿着 KV 的长度方向切分：把 3 万个历史 token 分成若干段，交给不同的 SM 同时处理，最后再把各段的结果合并（合并方法和在线 softmax 一样）。这类做法常被称为 Flash-Decoding 或 split-KV。
Decode 注意力还要读分页存放的 KV（第 8 篇）、处理批里长短不一的请求。推理引擎通常提供多个注意力后端可选，比如 FlashAttention、FlashInfer 等库的实现，它们对硬件、数据类型、块大小和功能的支持各有不同。引擎会按硬件和模型自动选择，一般不需要手动指定。
算子融合：把相邻的小运算合在一起
注意力之外，每一层还有很多小运算：残差相加、RMSNorm 归一化、旋转位置编码、SiLU 激活后与另一路相乘……它们各自的数学细节在这里不重要，只需要知道它们都是逐元素的简单计算，每个都是一个单独的内核，都要从 HBM 读一遍、写一遍。它们的计算量很小，几乎全部时间都花在搬运上。
算子融合（kernel fusion）把相邻的几个运算合进一个内核：数据读进来一次，在片上依次完成相加、归一化、缩放，只写出最后的结果。比如 nano-vllm 把“残差相加 + RMSNorm”融合成一个函数，把 SiLU 和乘法融合成另一个，都用 torch.compile 生成融合后的内核。同理，Q、K、V 三个投影矩阵也常被拼成一个大矩阵，一次矩阵乘法算出三份结果。
torch.compile 会分析模型的计算图，自动找出可以融合的运算，生成对应的 GPU 内核。代价是启动时要花时间编译，有时还会因为动态形状而重新编译。vLLM 会把编译结果缓存下来，第二次启动就快多了。
CPU 也会拖慢 GPU
GPU 上的每个内核，都要由 CPU 发出一次“启动”指令。每次启动大约要几微秒。一次 Decode 前向，每层要启动十几个内核，32 层就是几百次启动，加起来可能有一两毫秒。
第 6 篇算过，小批量 Decode 一步的 GPU 时间下界只有约 5 毫秒。如果 CPU 发指令的速度跟不上，GPU 算完一个内核就要空等下一个内核的指令，这一两毫秒就全部成了额外开销。批越小、模型越小，这个问题越明显。
CUDA Graph 的办法是“录制和回放”：先把一次 Decode 前向要启动的全部内核，连同它们的参数和依赖关系，录制成一张图；之后每一步只需要让 GPU“回放”这张图，CPU 只发一次指令，GPU 上的内核一个接一个紧密执行。
录制有一个限制：图里的每个内核，输入输出的形状和显存地址都是固定的。所以：
- 只适合形状固定的计算。 Decode 每个请求每步固定 1 个 token，批大小有限几种，适合录制；Prefill 的长度千变万化，很难整段录制。
- 要为多种批大小分别录制。 nano-vllm 为批大小 1、2、4、8，再从 16 到 512 每隔 16 各录一张图，运行时选不小于实际批大小的那张，多出的位置用空数据填充。
- 录制需要时间和显存。 引擎启动时要先预热、再逐个录制，几十张图要占用一些额外显存。
nano-vllm 的录制代码（engine/model_runner.py:234-248，略去了设置上下文和同步的几行）：
python
self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16))
for bs in reversed(self.graph_bs):
graph = torch.cuda.CUDAGraph()
outputs[:bs] = self.model(input_ids[:bs], positions[:bs]) # 预热
with torch.cuda.graph(graph, self.graph_pool):
outputs[:bs] = self.model(input_ids[:bs], positions[:bs]) # 录制vLLM 把 torch.compile 和 CUDA Graph 结合使用。注意力的输入形状随每个请求的长度变化，最难录制，vLLM 的做法是分情况处理：纯 Decode 的批录制完整的图，注意力也在图内；含 Prefill 的混合批，只把注意力以外的部分分段录制，注意力即时执行。从 v0.11 起这是默认模式（FULL_AND_PIECEWISE）。批大小超过最大录制尺寸时，才退回完全即时执行。
调试时为什么先关掉它们
融合内核、编译和 CUDA Graph 都在改变计算的执行方式。出现奇怪的错误或者结果异常时，一个常用的排查方法是先切回即时执行模式（eager mode），比如 vLLM 的 --enforce-eager 参数，关掉图录制和部分编译：
- eager 模式下问题消失，说明问题可能出在图录制或编译上。
- eager 模式下问题依旧，就可以排除这一层，继续往模型、模板或数据上查。
要注意，eager 模式只是缩小排查范围的工具，它的性能通常明显更差，不能作为性能测试的基线。
一次前向计算由许多内核组成，每个内核都要读写显存、由 CPU 启动。FlashAttention 分块计算注意力，用在线 softmax 保证结果精确，不再把 N × N 的中间矩阵写进显存；Decode 注意力沿 KV 长度切分来提高并行度；算子融合减少小运算的往返搬运；CUDA Graph 把几百次内核启动变成一次回放，消除小批量 Decode 中 CPU 的开销。
第三部分讲到这里，已经覆盖了推理引擎的主要机制。下一篇换一种学法：读一个一千多行的最小推理引擎，把分页、调度、前缀缓存、CUDA Graph 和代码一一对上：nano-vllm 源码导读。