LLM Serving CUDA Graph 生产实战：用 Shape Buckets、Piecewise Capture 与 Warmup Plan 降低 Kernel Launch 抖动

做 LLM Serving 性能优化时，人们首先会看 GPU 利用率、Attention Kernel、KV Cache 和 Batch Size。但在短 Decode Step、小 Batch 或较小模型上，另一个常见瓶颈是 CPU 到 GPU 的 Kernel Launch 开销。

一次自回归 Decode 并不是“调用一次 GPU 就结束”，而是执行大量算子。每个算子都需要 Python/C++ Runtime、CUDA Driver 和 GPU 之间完成参数准备与 Kernel 提交。单个 Kernel 很快时，提交本身就可能成为不可忽略的固定成本。

本文结合 vLLM 与 TensorRT-LLM，讲清形状分桶（Shape Buckets）、分段捕获（Piecewise Capture）、预热策略（Warmup Plan）、显存权衡与生产回退门禁，帮你把“GPU 没满，延迟却下不去”的 Host 侧瓶颈压下去。

核心原理：CUDA Graph 优化的是 Launch Overhead，不是替你减少计算量

1. Graph Replay 省掉了什么

普通执行路径大致是：

CPU prepare op A -> launch kernel A
CPU prepare op B -> launch kernel B
CPU prepare op C -> launch kernel C
...

CUDA Graph 捕获后更接近：

copy/update static inputs
        ↓
cudaGraphLaunch(graph)
        ↓
GPU replays A -> B -> C -> ...

因此它最适合 Host-bound 的推理阶段。如果某个阶段本身已经完全 Compute-bound，GPU Kernel 持续时间远大于 Launch Overhead，CUDA Graph 仍可能有收益，但通常不会像短 Kernel、低 Batch Decode 那么明显。

2. 为什么动态 Shape 是核心矛盾

CUDA Graph 依赖稳定的执行图和内存布局。一个针对 Batch=8 捕获的图，不能天然当成 Batch=37 的动态程序使用。

Serving 框架通常采取两种办法：

1. 捕获多个离散 Shape / Batch Size；
2. 运行时将实际 Shape Padding 到最近的已捕获 Bucket。

这就是 Shape Buckets。假设生产 Decode Batch 的主要分布为 1、2、4、8、12、18、30，可以设计：

运行时 Batch	处理方式
12	pad / replay 16
18	pad / replay 32
40	eager fallback / 其他路径

Bucket 太稀，Padding 浪费更多计算；Bucket 太密，则 Capture 数量、静态 Buffer 和显存占用上升。因此 Bucket Grid 本质上是延迟、额外计算、捕获时间与显存之间的成本函数。

Full CUDA Graph 与 Piecewise CUDA Graph

Full Capture：收益最高，但约束也最强

如果整段模型 Forward 都满足 CUDA Graph 的要求，可以做 Full CUDA Graph。vLLM 把 Full、Piecewise、Full Decode Only 和 Full + Piecewise 划分为不同运行模式，其中 Full Decode 对纯 Decode 工作负载尤其有意义。

但 Attention 是 LLM 中最难捕获的区域之一，其 Shape、KV 状态、Backend 能力和控制路径都可能影响 Graph Safety。因此“全图一定更快”不是一个可以脱离模型、Attention Backend 和 workload 单独成立的结论。

Piecewise Capture：把难捕获算子留在 Eager 路径

vLLM 与 TensorRT-LLM 都提供了 Piecewise CUDA Graph 思路：

CUDA Graph Segment A
        ↓
Attention / dynamic op (eager)
        ↓
CUDA Graph Segment B
        ↓
Attention / dynamic op (eager)
        ↓
CUDA Graph Segment C

这样牺牲了一部分完整图收益，但换来更好的动态兼容性。TensorRT-LLM 明确说明，Piecewise CUDA Graph 主要把不适合捕获的区域——尤其是 Attention——保留为 Eager，对其余区域捕获；并按 Token Count 配置 capture_num_tokens，运行时将实际 Token 数 Padding 到下一个已捕获值。

这说明生产调优的真正对象不是一个布尔开关，而是 Graph Partition + Capture Grid。

Shape Buckets 应该从生产流量反推，而不是照抄默认值

很多框架有默认 Capture Size，默认值适合“开箱即用”，但不意味着是你的最佳配置。更可靠的方法是先采集 24 小时或 7 天的 workload histogram：

metric: active_decode_sequences
metric: num_scheduled_tokens
metric: prefill_tokens_per_iteration
metric: mixed_batch_tokens

然后计算每个候选 Bucket 的覆盖率。例如：

区间	占比
1-4 tokens/sequences	31%
5-8	27%
9-16	21%
17-32	13%
33-64	6%
>64	2%

这类流量应在小 Shape 区间配置更密的 Bucket，而不是平均铺开。

一个实用的 Bucket 选择原则

可以把候选 Bucket 设计为：

1, 2, 4, 8, 12, 16, 24, 32, 48, 64, 96, 128

然后用真实 Trace 比较：

• Graph Coverage；
• Padding Ratio；
• Graph Memory；
• p95/p99 TTFT；
• p95/p99 TPOT；
• 最大并发数。

不要只比较 Tokens/s。vLLM 当前配置逻辑本身也采用离散 Capture Size，运行时针对可覆盖的 Batch 使用相应 CUDA Graph；超过最大 Capture Size 的请求则不会走对应图路径。框架默认 Bucket 可作为 Baseline，但生产配置仍应由真实分布校准。

Warmup Plan：避免把 Compile/Capture 延迟留给第一个真实用户

CUDA Graph 最容易在测试环境被忽略的问题，是 首次请求成本。首次进入某个新 Shape 时，系统可能发生：

1. Torch/Inductor 编译；
2. Lazy CUDA 初始化；
3. Attention Backend 初始化；
4. Workspace 分配；
5. CUDA Graph Warmup；
6. Capture 与 Instantiate；
7. 静态 Buffer 建立。

PyTorch 官方 CUDA Graph 文档明确建议在 Capture 前先做 Warmup，避免把懒初始化行为捕获进图或导致 Capture 失败。生产部署不应该只做一个“健康检查请求”，而应该有 Shape-aware Warmup Plan：

warmup_plan:
  decode_buckets: [1, 2, 4, 8, 16, 32]
  mixed_token_buckets: [64, 128, 256, 512]
  repeat_per_bucket: 3
  block_ready_until_complete: true

Ready 不等于进程启动成功

Kubernetes 或服务注册层建议区分：

Process Started
   ↓
Model Loaded
   ↓
Compile Complete
   ↓
Required Graph Buckets Captured
   ↓
READY

否则滚动发布时，新实例刚加入负载均衡，就可能把真实用户流量当成 Warmup 请求。

捕获越多，显存并不越省

CUDA Graph 常被描述为“减少 CPU 开销”，但生产配置中必须同时考虑 Graph Memory Footprint。捕获更多 Shape 往往意味着更多 Graph Executable、静态输入输出 Buffer 或相关 Runtime 状态。TensorRT-LLM 也明确指出，更广泛的 Capture Token Count 会增加 GPU Memory，并可能降低可达到的并发度。

因此需要把 CUDA Graph 与 KV Cache 放在同一个显存预算里看：

GPU Memory
├── Model Weights
├── KV Cache
├── Runtime Workspace
├── CUDA Graph / Static Buffers
└── Safety Margin

如果为了把 Graph Coverage 从 97% 提升到 99.8%，却减少了 15% 的可并发序列数，这个优化未必划算。更合理的指标是：在满足 Tail SLO 的前提下，每 GB 显存能承载多少有效并发和多少有效 Tokens/s。

vLLM：先控制 Capture Grid，再评估 Full/Piecewise

下面是一个简化示意，具体字段请以部署版本文档为准：

vllm serve /models/your-model \
  --compilation-config '{"cudagraph_capture_sizes":[1,2,4,8,16,32,64]}'

不要一开始就追求“最大 Capture Grid”。更稳妥的步骤是：

1. Eager / 默认模式跑 Baseline；
2. 采集真实 Batch/Token Shape；
3. 先覆盖 90%～95% 高频 Shape；
4. 测 Graph Hit 与 Padding；
5. 观察显存和并发变化；
6. 再决定是否增加更大的 Capture Size。

vLLM 当前文档还提供 Full、Piecewise、Full Decode Only 和 Full + Piecewise 等 CUDA Graph Mode。不同版本支持能力仍在快速变化，因此 Runtime Version 必须作为 Benchmark Artifact 的一部分固定记录。

TensorRT-LLM：Token Count Bucket 比“一个 max_batch_size”更重要

TensorRT-LLM 的 Piecewise CUDA Graph 配置示意：

torch_compile_config:
  capture_num_tokens: [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]
  enable_userbuffers: false
  enable_piecewise_cuda_graph: true

官方文档特别提醒：

• Piecewise Capture 的 Token Bucket 应根据 硬件、模型与并行策略 调优；
• Bucket 越多，Padding 越少，但图内存越高；
• Context Phase 中较大的 Token Count 本身 Host Overhead 占比下降，继续增加 Capture Point 的边际收益可能变小。

因此 Bucket 的合理密度通常是：小 Shape 密、Shape 越大越稀。

生产观测：至少增加六类 CUDA Graph 指标

如果只看 TTFT 和 GPU Utilization，很难解释启用 Graph 后为什么偶尔还是抖。建议增加：

1. Graph Hit Ratio：cuda_graph_hit_requests / eligible_requests，最好按 Bucket 分层。
2. Eager Fallback Ratio：统计因为 Shape 超界、Backend 不兼容或 Graph Invalid 而回退到 Eager 的比例。
3. Padding Overhead：padding_ratio = padded_tokens / actual_tokens。
4. Capture / Compile Latency：区分 model_load_ms / compile_ms / warmup_ms / capture_ms / ready_ms，对滚动扩容和故障重建尤其重要。
5. Graph Memory Footprint：记录启用前后的 GPU used memory、KV cache capacity、max concurrent sequences。
6. Tail Latency：至少比较 TTFT p50/p95/p99、TPOT/ITL p50/p95/p99、E2E p95/p99。

CUDA Graph 的价值通常更容易反映在小 Batch Decode 与 Tail Latency，而不只是平均吞吐。

推荐的上线流程

阶段一：建立 Eager Baseline。 固定 Model Revision、Runtime Version、CUDA/Driver、GPU 型号、Tensor Parallel/Expert Parallel、workload trace，否则不同 Benchmark 之间无法比较。

阶段二：只打开默认 CUDA Graph。 确认框架默认配置的真实收益和显存成本。

阶段三：用生产 Trace 调 Capture Grid。 对每组 Bucket 记录 coverage / padding_ratio / graph_memory / TTFT p99 / TPOT p99 / max_concurrency。

阶段四：加入 Warmup Gate。 新 Pod / Worker 必须完成关键 Bucket Capture 后才进入 Ready。

阶段五：Canary。 不要用平均吞吐作为唯一放量门禁，至少设置：

• TTFT p99 regression <= threshold；
• TPOT p99 regression <= threshold；
• OOM = 0；
• graph fallback ratio <= threshold；
• max concurrency regression <= threshold。

阶段六：保留 Eager / Previous Runtime 回退。 CUDA Graph、Torch Compile 和 Attention Backend 的兼容矩阵会随版本变化，发布系统必须能回滚到上一个 Runtime Artifact。

适用场景

这套方法尤其适合：

• Decode 占比较高的在线聊天；
• 小模型或低 Batch 低延迟服务；
• GPU 未持续满载、但 CPU Launch 明显占据 Step Gap 的场景；
• 请求 Shape 有明显稳定分布的线上系统；
• 已完成 KV Cache、Batching 和 Kernel 基础优化，准备继续压低尾延迟的 Serving 平台。

如果服务主要是超大 Prefill、单次 Kernel 已经长时间占满 GPU，CUDA Graph 可能不是优先级最高的优化项。

常见误区

误区一：Capture Size 越多越好。 更多 Graph 会提高覆盖率，但也会增加捕获成本和显存占用。目标应该是 最小 Graph Set 覆盖大部分生产流量。

误区二：Full CUDA Graph 一定优于 Piecewise。 只有在 Attention Backend、动态 Shape 和控制流都满足条件时，Full Capture 才值得使用。生产系统更看重稳定性和兼容性。

误区三：启动后跑一个 Warmup 请求就够了。 如果有多个 Shape Bucket，只 Warmup 一个 Shape 并不能避免其他 Bucket 第一次命中时的 Capture Spike。

误区四：只看 Tokens/s。 如果 Tokens/s 上升 5%，但 TTFT p99、OOM 或最大并发恶化，这未必是生产优化。

误区五：忽略 Runtime 版本。 vLLM、TensorRT-LLM 与 PyTorch 对 CUDA Graph 的实现仍持续演进。Capture Mode、兼容 Attention Backend、默认 Bucket 与内存策略都可能变化，Runtime Version 应与 Model Revision 一样被版本化。

上线检查

上线前至少确认：

• 已有真实 workload shape histogram；
• Capture Bucket 来自流量分布，而不是直接照抄示例；
• 已测量 Graph Hit Ratio 与 Eager Fallback；
• 已测量 Padding Overhead；
• 已记录 Capture 前后 GPU Memory 与 KV Cache Capacity；
• Warmup 覆盖所有关键 Bucket；
• Readiness 在 Warmup/Capture 完成后才放流量；
• 已比较 TTFT/TPOT/ITL p95 与 p99；
• 已保留 Eager 或上一 Runtime 版本回退路径；
• Runtime、Driver、CUDA、模型版本和配置均进入发布 Artifact。

参考资料

1. vLLM Compilation Configuration
2. vLLM Compilation API
3. vLLM Runtime / CUDA Graph Capture Size Configuration
4. TensorRT-LLM — Torch Compile & Piecewise CUDA Graph
5. PyTorch — CUDA Semantics / CUDA Graphs
6. Torch-TensorRT — CUDAGraphs