2.5 Attention 后端与图优化：把 Decode 的 CPU 开销榨干

1. 一个反直觉的瓶颈：CPU 拖了 GPU 的后腿#

先抛一个反直觉的现象。上一章我们论证了 Decode 是 Memory Bound——瓶颈在从显存搬权重。但当模型不大、Batch 也不大时，还有一个更隐蔽的瓶颈会冒出来：CPU 发射 Kernel 的开销。

GPU 自己不会主动干活，它每做一个操作（一次矩阵乘、一次 LayerNorm、一次激活……），都要由 CPU 通过一次”启动（launch）“命令告诉它。跑一遍 Transformer，有成百上千个这样的小操作，就对应成百上千次 CPU→GPU 的启动调用。每次启动本身有固定的 CPU 开销（构造参数、驱动调度、提交队列），量级在微秒级。

平时这点开销不算什么——因为 Prefill 一步要处理成千 Token，每个 GPU Kernel 一跑就是几百微秒甚至毫秒，CPU 那点启动开销被淹没了。但 Decode 每步只处理 1 个（或很少几个）Token，每个 Kernel 的实际计算可能就几微秒，结果CPU 发射这个 Kernel 的时间比 GPU 执行它的时间还长。于是 GPU 干几微秒就停下来，眼巴巴等 CPU 发下一条命令——算力再次闲置，只不过这次不是等数据，是等指令。

下图上半部分就是这种情形：黄色是 CPU 逐条发射 Kernel，蓝色是 GPU 真正在算，灰色虚线斜线格是 GPU 空等下一条指令的时间。下半部分是第 4 节要讲的 CUDA Graph，CPU 只提交一次（绿色），8 个 Kernel 在 GPU 上首尾相接，橙色虚线标出两种方式各自完成这一步的时刻。

📌 关键点：Decode 阶段有两重”等待”——等数据（Memory Bound，靠 Batching/量化缓解）和等指令（CPU launch 开销，靠本节的图优化缓解）。小模型、小 Batch 时，后者尤其突出。这就是为什么图优化对 Decode 加速格外有效。

2. 可插拔的 Attention 后端#

Attention 是 Transformer 里最重、最讲究的算子，它的 Kernel 实现直接决定推理速度。vLLM 的设计哲学是不绑死一种实现，而是提供可插拔的 Attention 后端，根据硬件、模型结构和场景自动或手动选择最优的那一个。

常见的后端包括：

🔑 核心概念：Attention 后端是同一个数学操作的不同工程实现。它们算出的结果在数学上等价，区别在于访存模式、并行策略、对特定硬件指令和模型结构的适配程度。选对后端，同样的卡能跑出明显不同的吞吐。

💡 提示：多数情况下不用手动指定，vLLM 会根据 GPU 架构、模型类型和数据精度自动挑选合适的后端。v0.32.0 在 Ampere 和 Hopper 上默认选 FLASH_ATTN，在 Blackwell 上默认选 FLASHINFER；MLA 模型会走专门的 MLA 后端。需要显式干预时，用 --attention-backend 指定后端，用 --attention-config.flash_attn_version 指定 FlashAttention 版本（2、3 或 4）。后端支持矩阵随版本演进较快，具体某张卡、某个模型用哪个后端，以你所用版本的官方文档为准。

3. 后端如何处理 Prefill/Decode 混合批次#

这里有个上一节埋下的关键联系。vLLM V1 的统一调度器会把 Prefill 的 Chunk 和 Decode 请求塞进同一步——那么这一步的 Attention 该怎么算？毕竟 Prefill 部分要算”一整段新 Token 之间 + 对历史的注意力”，Decode 部分只算”1 个新 Token 对全部历史的注意力”，两者的计算形态不一样。

这就要求 Attention 后端原生支持混合批次（mixed prefill/decode batch）：在一次 Kernel 调用里，同时正确处理批中既有 Prefill 又有 Decode 的请求。FlashAttention 等现代后端提供的变长（varlen）接口正是为此设计的——它把批里所有请求的 Query 首尾拼成一个变长序列，配合每个请求的起止偏移，在一次 Kernel 调用里算完，不需要把 Prefill 和 Decode 拆成两次调用。

📌 关键点：“统一调度器”（软件层抹平 Prefill/Decode）和”支持混合批次的 Attention 后端”（Kernel 层抹平 Prefill/Decode）是配套的。 没有后者，V1 那套优雅的 Token 预算调度就落不了地。这也是为什么 V1 对 Attention 后端的能力有要求，早期不是所有后端都能支持。

4. CUDA Graph：把一串 Kernel 启动打包成一次#

回到第 1 节的 CPU launch 开销。CUDA Graph 是 NVIDIA 提供的解药：把一连串固定的 GPU 操作**录制（capture）成一张”图”，之后只需一条命令重放（replay）**整张图，GPU 就会依次执行录好的所有 Kernel——成百上千次 CPU 启动被压缩成一次。

用类比来说：没有 CUDA Graph，就像老师每讲一句话都要等秘书传一次话（每个 Kernel 都要 CPU 单独发射）；有了 CUDA Graph，相当于老师提前把整段讲稿录好，上课时按一下播放键，一气呵成放完，中间不用秘书再一句句传。第 1 节那张时间线的下半部分就是重放时的样子：CPU 提交一次后就空出来，可以去准备下一步的调度。

但 CUDA Graph 有个硬性前提：录制的操作序列和张量形状必须固定。这对 Decode 是天作之合——每步都是”处理固定 Batch 个 Token”，形状稳定，可以针对不同的 Batch Size 分别录制好图，运行时按当前 Batch Size 选对应的图重放。

⚠️ 注意：正因为要求形状固定，vLLM 只为一组预设的 Batch Size 捕获 CUDA Graph。v0.32.0 的默认捕获尺寸是 [1, 2, 4]，加上 8 到 248 之间每隔 8 取一个，再加上 256 到 max_cudagraph_capture_size 之间每隔 16 取一个；max_cudagraph_capture_size 默认 512（数据中心级 Blackwell GPU 上为 1024），两者都可以通过 compilation_config 调整。运行时实际 Batch 会被向上 Padding 到最近的已捕获尺寸再重放，比如 13 个 Token 的批次按 16 的图跑，多出的 3 个位置是空算；只有超过最大捕获尺寸的批次才回退到逐 Kernel 启动。这也意味着 CUDA Graph 的收益主要体现在 Decode 这类小而规整的批次，大块的 Prefill 较难直接套用。

5. torch.compile 与 Piecewise CUDA Graph#

CUDA Graph 解决了”启动开销”，但还有一层优化空间：能不能把很多小 Kernel 本身合并成更少、更大的 Kernel？ 这就是 torch.compile 干的事——它把模型的前向计算编译成优化后的图，通过算子融合（fusion）等手段减少 Kernel 数量。在 vLLM V1 中，torch.compile 默认开启，且所有编译在开始服务前完成，避免请求过程中触发编译导致延迟尖刺。

但这里有个矛盾：CUDA Graph 要求整段操作形状固定，而 Attention 恰恰是最难被 CUDA Graph 兼容的部分——它涉及变长的 KV、复杂的 Mask，形状不那么规整。如果因为 Attention 不好处理就放弃整张图的 CUDA Graph，就太可惜了。

vLLM V1 的巧解是 Piecewise CUDA Graph（分段 CUDA Graph）：

• 把整个前向图在 Attention 算子处切开（Attention 被包装成一个”不透明”的自定义算子，编译器不去动它）。
• Attention 之外的部分（LayerNorm、各种 GEMM、激活、残差……这些都是规整的 token-wise 操作）用 CUDA Graph 捕获、重放。
• Attention 本身保持 eager（即时）执行，保留它处理变长 KV 的灵活性。

🔑 核心概念：Piecewise CUDA Graph = 对规整的部分用图优化榨干 CPU 开销，对灵活的 Attention 网开一面保留 eager。 二者各取所长：CPU launch 开销大头（那成百上千个小算子）被 CUDA Graph 消灭，而 Attention 的灵活性（支持混合批次、变长 KV）不受损。这正是 V1 相比 V0 在 Decode 上提速的关键工程之一。

vLLM 还支持 Full CUDA Graph，把 Attention 也录进图里，前提是后端在这类批次上支持图捕获。v0.32.0 的 V1 默认模式是两者结合的 FULL_AND_PIECEWISE：纯 Decode 批次形状最规整，用一张 FULL 图整步重放；含 Prefill 的混合批次退回 PIECEWISE，Attention 在图外执行。早期版本的默认是 PIECEWISE，Pooling 类模型现在仍默认 PIECEWISE。完整的模式还有 NONE、PIECEWISE、FULL、FULL_DECODE_ONLY，可通过 compilation_config 的 cudagraph_mode 设置。

下图上半部分是切图后的前向：蓝色是 torch.compile 编译出的子图，黄色是留在图外的 Attention，中间各层共用同一种子图。下半部分是 FULL_AND_PIECEWISE 在运行时的选择：绿框表示整步一张 FULL 图，黑框表示每段子图各是一张 CUDA Graph。

💡 提示：FULL 图能不能用，取决于 Attention 后端对 CUDA Graph 的支持程度。官方文档把它分成几档：FA3 和 Triton 在任意批次上都支持（ALWAYS）；FA2 和 FlashMLA 只支持形状统一的批次（UNIFORM_BATCH）；FlashInfer 只支持每个请求 1 个 Token 的纯 Decode 批次（UNIFORM_SINGLE_TOKEN_DECODE）。不满足条件的批次会自动退回 PIECEWISE。文档也指出，FULL 图对小模型和 MoE 模型的 Decode 提速最明显。

6. 源码走读：vLLM 怎么切图与捕获#

前面讲了 Piecewise CUDA Graph 的思路，这一节看 vLLM V1 在代码层面怎么把它做出来。核心机制在 torch.compile 集成里（官方设计文档 design/torch_compile）。代码持续演进，以下描述基于 main 分支文档，具体 API 以你所用版本为准。

1
from vllm import LLM
2
from vllm.config import CompilationConfig
3

4
llm = LLM(
5
    model="meta-llama/Llama-3.1-8B-Instruct",
6
    compilation_config=CompilationConfig(
7
        # 只为这些 batch size 捕获 CUDA Graph；批次向上 Padding 到最近的尺寸重放，
8
        # 超过最大值（这里是 32）的批次不走 CUDA Graph
9
        cudagraph_capture_sizes=[1, 2, 4, 8, 16, 32],
10
        # 可选：为特定静态形状做 Inductor 自动调优（首次较慢，默认关闭）
11
        compile_sizes=[1, 2, 4, 8],
12
    ),
13
)

7. 实践中的开关与权衡#

对使用者来说，这些优化大多是默认开启、自动生效的，日常无需操心。但了解它们的权衡有助于排查问题和调优：

• ✅ 收益：Decode 阶段 CPU 开销大幅下降，小模型/小 Batch 下的吞吐和延迟明显改善；torch.compile 的算子融合进一步减少 Kernel 数量。
• ❌ 代价：
  • 启动变慢：服务启动时要做编译和 CUDA Graph 捕获，首次拉起耗时增加（vLLM 用编译缓存缓解，缓存目录可复制以加速后续启动）。
  • 显存占用：捕获多个 Batch Size 的 CUDA Graph 需要额外显存。
  • 调试复杂度：编译后的图不如 eager 模式好调试，排查数值问题时可能需要临时关掉。

⚠️ 注意：遇到疑难的数值异常或想确认某个 Kernel 行为时，可以临时加 --enforce-eager 启动，它会同时关闭 torch.compile 和 CUDA Graph，用纯 eager 模式对比。定位完再去掉——生产环境务必保持开启，否则会白白损失 Decode 性能。更细的开关（如只关 CUDA Graph 可设 cudagraph_mode 为 NONE）以对应 vLLM 版本的官方文档为准。

📝 总结#

• Decode 阶段除了”等数据”（Memory Bound），还有”等指令”——CPU 逐 Kernel 启动的开销，小模型/小 Batch 下尤为突出。
• vLLM 提供可插拔 Attention 后端（FlashAttention/FlashInfer/FlashMLA/Triton），数学等价而工程各异，通常按 GPU 架构自动选择（v0.32.0 在 Ampere/Hopper 默认 FlashAttention，Blackwell 默认 FlashInfer）。
• 现代后端需原生支持 Prefill/Decode 混合批次，与 V1 的统一调度器配套。
• CUDA Graph 把成百上千次 Kernel 启动压成一次重放，消除 CPU 开销，但要求形状固定：批次被 Padding 到最近的捕获尺寸，超过最大捕获尺寸的批次不走图。
• Piecewise CUDA Graph 在 Attention 处切图：规整部分走 CUDA Graph，Attention 保持 eager，兼得性能与灵活性。V1 默认的 FULL_AND_PIECEWISE 再进一步，纯 Decode 批次整步用一张 FULL 图。
• 源码上，Attention 被包成不透明算子 unified_attention_with_output，图按 splitting_ops 切成三种可复用子图，再按 cudagraph_capture_sizes 分尺寸捕获、编译产物写入可拷贝的缓存。
• 这些优化默认开启，代价是启动变慢、额外显存和调试复杂度；排障时可用 --enforce-eager 临时关闭，生产环境应保持开启。

🎯 自我检验清单#

• 能解释为什么 Decode 阶段 CPU 的 Kernel 启动开销会成为瓶颈
• 能说出至少三种 vLLM Attention 后端及各自的适用场景
• 能解释”数学等价、工程不同”的 Attention 后端意味着什么
• 能说明为什么统一调度器需要后端支持 Prefill/Decode 混合批次
• 能描述 CUDA Graph 通过”录制—重放”消除 CPU 开销的原理及其形状固定的前提
• 能解释 Piecewise CUDA Graph 为什么在 Attention 处切图，以及 FULL_AND_PIECEWISE 如何按批次类型选图
• 能说清为什么切图后只产生三种子图，以及 cudagraph_capture_sizes 与 Padding 的关系
• 能说清这些图优化的代价（启动、显存、调试）与生产环境的取舍

📚 参考资料#

• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• vLLM Design — torch.compile Integration
• vLLM Design — CUDA Graphs（cudagraph_mode 与后端支持级别）
• vLLM Design — Attention Backends（默认后端选择）
• vLLM V1 Alpha Release（Piecewise CUDA Graph、FlashAttention 3）
• NVIDIA Developer Blog — Getting Started with CUDA Graphs
• vLLM — Engine Arguments（—attention-backend、—enforce-eager）