SGLang 源码完全拆解（九）：CUDA Graph 与编译路径——解码整图捕获、Prefill 分段图与静态形状约束

CPU 提交一次 kernel 要花多久

GPU 计算 kernel 的耗时由硬件决定，但提交 kernel 的耗时由 CPU 与驱动决定。一次普通的 kernel launch 需要走：Python 层调用 → framework dispatch → CUDA runtime API → 驱动命令入队 → GPU 前端取出执行。这条路径在没有任何优化时的量级是每次 5–10 微秒，其中相当一部分是 CPU 侧固定开销。

这个数字对 prefill 无关紧要：一次 prefill 前向要算几千上万个 token，矩阵乘本身耗时是毫秒级，几微秒的提交开销被完全淹没。

但对 decode 就是致命的。decode 每轮每个请求只产出 1 个 token，一次前向里的矩阵乘退化成「向量 × 矩阵」，单个 kernel 只有几微秒到十几微秒的真计算量。而一轮完整前向要走遍所有层，每层包含 RMSNorm、QKV 投影、RoPE、注意力、输出投影、MLP、残差等多个算子；一个 32 层左右的模型，一轮 decode 的 kernel 数量在数百量级。于是 CPU 侧提交开销累加起来是毫秒级，与 GPU 真计算量同阶甚至更大。

这还没算 Python 的开销。SGLang 的模型前向是纯 Python 组装的：每个 nn.Module.__call__、每次 torch.empty、每个 if forward_mode.is_decode() 分支都要走 Python 解释器。对于小 batch 的 decode，Python 解释与 framework dispatch 的时间占比可以超过 50%，GPU 在两次 kernel 之间反复空转。

CUDA Graph 解决的是什么

CUDA Graph 的思路非常直接：把一串 kernel launch 录下来，之后用一个 API 调用把整串重放出去。

上图（来源：NVIDIA CUDA Graphs 技术博客）画的就是这次替换。左侧每个 kernel 之间都有一段 CPU 提交的空隙，GPU 在空隙里闲着；右侧整个串行流被录成一张图，重放时只需一次「graph launch」，CPU 的空隙全部消失。

NVIDIA 在那篇博客里给出了一个量化的观察：在单 kernel 尺度、batch 很小时，launch 开销可以吞掉接近一半的时间。这正是 decode 的处境。

下图是捕获与重放两种状态下的时间线对照（同源）：捕获阶段要真实执行一遍并把 kernel 与依赖关系记下来；重放阶段只需发布图，所有 kernel 由 GPU 前端按记录的依赖自主调度。

PyTorch 侧的使用方式也很简洁，官方博客给出的形态大致是：

1
g = torch.cuda.CUDAGraph()
2
# 静态输入：地址固定，内容可变
3
static_input = torch.empty(batch, device="cuda")
4

5
# warmup：多次执行，让所有 kernel / 内存池进入稳定状态
6
for _ in range(3):
7
    s = model(static_input)
8

9
torch.cuda.synchronize()
10
with torch.cuda.graph(g):
11
    static_output = model(static_input)
12

13
# 之后每轮：只改内容，不改地址
14
static_input.copy_(new_input)
15
g.replay()
16
result = static_output

图源：Accelerating PyTorch with CUDA Graphs。这张图里有一个关键约定值得强调：static_input 是同一个张量对象。g.replay() 读的是捕获时刻记录下来的指针，所以「换输入」只能靠 copy_ 改内容，不能靠重新赋值。这个约定正是 SGLang 整个 InputBuffers 体系存在的理由。

三条硬约束，决定了后面所有设计

把上面的机制压成规则，只有三条：

后面讲到 capture_prepare、load_batch、PaddingPolicy、can_run_graph 时，全部可以回溯到这三条。

桶列表怎么来

get_batch_sizes_to_capture 负责把配置里的桶列表过滤成实际可用的列表，返回 (capture_bs, compile_bs)：

1
capture_bs = list(get_exec().graph.cuda_graph_config.decode.bs)
2
num_max_requests = model_runner.req_to_token_pool.size
3
...
4
mul_base = get_cuda_graph_batch_size_alignment(
5
    gathered_buffer_required=gathered_buffer_required
6
)
7
alignment_width = captured_req_width
8
if get_exec().overlap.enable_two_batch_overlap:
9
    alignment_width = 1
10

11
num_max_requests = ceil_align(num_max_requests, mul_base)
12
if max(capture_bs) > num_max_requests:
13
    capture_bs += [num_max_requests]
14

15
# Model input token count = bs * alignment_width; must be a multiple of attn_tp_size.
16
capture_bs = [bs for bs in capture_bs if bs * alignment_width % mul_base == 0]
17
capture_bs = [bs for bs in capture_bs if bs <= num_max_requests]
18
capture_bs = list(sorted(set(capture_bs)))

三处过滤都是有理由的：

1. 对齐约束。当开启 gather 缓冲（MLP-TP gather 或 attention-TP gather）时，图内会做一次跨 rank 的 token 收集，收集后的 token 总量必须是某个对齐宽度的整数倍，否则集合通信的切分会不均匀。captured_req_width 是每个请求在图内占据的 token 宽度（普通 decode 是 1，投机解码是草稿 token 数），所以约束写成 bs * alignment_width % mul_base == 0。TBO（two-batch overlap）是例外：它把每个请求的行拆到两个 micro-batch 里，对齐约束变成按请求而非按 token 行计，所以 alignment_width 被强制回 1。
2. 请求池上限。图内的 req_pool_indices 要指向 KV 池里的真实行号，行号不能超过 req_to_token_pool.size。如果配置里的最大桶超过了池容量（小 GPU 或手动限制了 max-running-requests 时会发生），代码不是报错，而是追加一个等于池容量的桶，保证「能跑满池」这种极端情形也有图可用。
3. 去重与排序。同一张图录两次没有意义，且捕获顺序影响内存池复用效果，所以去重后升序排列。

compile_bs 是从 capture_bs 里再筛一遍，只保留不超过 torch_compile_max_bs 的那些：

1
compile_bs = (
2
    [bs for bs in capture_bs if bs <= get_exec().graph.torch_compile_max_bs]
3
    if get_flags().capture.enable_torch_compile
4
    else []
5
)

也就是说，同一份桶列表被两条路径共用：大桶走「纯图 + eager 前向」的捕获，小桶走「torch.compile 编译 + 捕获」的路径。为什么阈值默认是 32，后面讲 torch.compile 时会解释。

默认桶列表长什么样

桶列表的默认生成在 arg_groups/cuda_graph_hook.py 的 generate_decode_cuda_graph_batch_sizes 里：

1
if cfg.disable_cuda_graph_padding:
2
    capture_bs = list(range(1, max_bs + 1))
3
elif cfg.speculative_algorithm is None:
4
    capture_bs = (
5
        [1, 2, 4, 8, 12]
6
        + list(range(16, 257, 8))
7
        + list(range(272, 512, 16))
8
        + list(range(512, max_bs + 1, 32))
9
    )
10
else:
11
    capture_bs = (
12
        list(range(1, 9, 1))
13
        + list(range(10, 33, 2))
14
        + list(range(40, 65, 4))
15
        + list(range(72, 257, 8))
16
        + list(range(272, max_bs + 1, 16))
17
    )

这张表本身就是一段工程结论，值得逐段读：

• 小 batch 段最密（1, 2, 4, 8, 12）。低并发场景下 batch 的绝对值小，桶间空隙的相对浪费大：实际 batch=9 对上桶 12，浪费 33%；而 batch=200 对上桶 208 只浪费 4%。所以小端必须密。
• 投机解码时更密。投机解码每轮的真实 token 数是 bs × 草稿宽度，小 batch 下真实 token 数的变化步长更大，粗桶会放大浪费，所以 1–8 逐个录、10–32 每 2 录、40–64 每 4 录。
• 大端逐步放松。512 以上按 32 间隔，是因为大 batch 下相对浪费已经很小，而每多一张图就要多占一份捕获显存（记忆池 + 静态缓冲）。
• --disable-cuda-graph-padding 会把桶铺满（1..max_bs 全录）。这个选项的语义是「只有形状完全一致才用图」，等价于不允许任何 padding。它的代价是图数量爆炸、捕获时间与显存随之爆炸，收益是彻底消灭了 padding 带来的无效计算。这明显是一个给特定场景（例如对延迟极度敏感、batch 分布高度集中）用的开关。

prefill 侧的默认表完全不同：

1
capture_sizes = (
2
    list(range(4, 33, 4))
3
    + list(range(48, 257, 16))
4
    + list(range(288, 513, 32))
5
    + list(range(576, 1024 + 1, 64))
6
    + list(range(1280, 4096 + 1, 256))
7
    + list(range(4608, max_bs + 1, 512))
8
)

粒度随 token 数迅速放宽，最小桶是 4 个 token、最大间隔到 512。道理相同：prefill 的 token 数是几百到几万量级，几十个 token 的 padding 在总量里占比可以忽略。

桶选择：向上取整

实际批次进来后，选桶就是一次二分：

1
@staticmethod
2
def _pad_to_bucket(raw_size: int, buckets: Sequence[int]) -> int:
3
    assert raw_size <= buckets[-1], (
4
        f"size {raw_size} exceeds max captured bucket {buckets[-1]}; "
5
        f"can_run_graph should have rejected this batch"
6
    )
7
    index = bisect.bisect_left(buckets, raw_size)
8
    return buckets[index]

注意这里的断言顺序：先由 can_run_graph 把超限批次挡掉，再进 load_batch 选桶。如果不先挡，bisect_left 在超出所有桶时会返回 len(buckets)，随后 buckets[index] 直接 IndexError，而且报错信息毫无线索。把契约写成断言，是让这条边界显式化。

捕获编排

decode_cuda_graph_runner.capture() 的整体骨架：

1
self.warmup()
2
...
3
self.buffers.seq_lens.fill_(self.seq_len_fill_value)
4
self.buffers.seq_lens_cpu.fill_(self.seq_len_fill_value)
5
self.buffers.reset_index_buffers()
6

7
with freeze_gc(get_exec().graph.enable_cudagraph_gc):
8
    if not self.enable_pdmux:
9
        with (
10
            graph_capture(stream=get_or_create_global_graph_capture_stream())
11
            as graph_capture_context,
12
            profile_context as prof,
13
        ):
14
            self.stream = graph_capture_context.stream
15
            with self.backend.capture_session(self.stream):
16
                self._capture_one_stream()
17
    else:
18
        for i, sg in enumerate(self.stream_groups):
19
            ...

这段代码里有四个容易忽略但很关键的动作：

第一，seq_lens 先填成 seq_len_fill_value。 这个值来自后端的 get_cuda_graph_seq_len_fill_value()，FlashInfer 与 Triton 都返回 1。为什么不是 0：捕获阶段会真实执行一遍前向，kernel 会真的去读 seq_lens。填 0 会让某些注意力 kernel 进入零长度退化路径，甚至触发除零；填一个小的合法值（1）能让捕获走完正常的完整流程，而真正的内容在每轮重放前由 init_forward_metadata_out_graph 刷新。

第二，reset_index_buffers()。 捕获会真的跑前向，因此会真的写 KV。如果捕获发生在服务中途（例如弹性 EP 扩缩容后重捕），缓冲里可能残留上一个真实批次的索引值，捕获就会按那批索引去写 KV——写错位置。所以捕获前必须把索引类缓冲复位。

第三，freeze_gc。 捕获期间禁用 GC 并 gc.freeze() 现有对象，捕获结束再 unfreeze。原因是 GC 可能在捕获中途触发一次回收，而 GC 会触碰 Python 对象、可能间接释放 CUDA 内存；这既拖慢捕获，也可能改变内存池状态。--enable-cudagraph-gc 可以打开 GC，默认关掉。

第四，捕获顺序倒序。 _capture_one_stream 里写着：

1
# Reverse so cuda graphs share memory better.
2
capture_range = (
3
    tqdm.tqdm(list(reversed(self.capture_bs)))
4
    if get_parallel().tp_rank == 0
5
    else reversed(self.capture_bs)
6
)

PyTorch 的 CUDA Graph 从共享的私有内存池（graph pool）里分配图内中间张量。先捕获大形状，就先把大块内存从池里切走；之后捕获小形状时，池里已有的空闲块可以直接复用，不必再向全局 allocator 申请。反过来先捕获小形状，大形状仍然要新申请，池就被撑得更大。所以「从大到小」是显存最优的顺序。这一点在 prefill runner 里同样出现（capture_shapes = list(reversed(self.capture_num_tokens))）。

每个形状的捕获

capture_one_shape 的骨架是「建 dummy 批次 → 进 ForwardContext → 建元数据 → 录制」：

1
forward_batch, attn_backend, pp_proxy_tensors = self.capture_prepare(
2
    bs, stream_idx=stream_idx, num_tokens=num_tokens
3
)
4

5
with forward_context(ForwardContext(attn_backend=attn_backend)):
6
    self.tbo_plugin.capture_one_batch_size(forward_batch, num_tokens=num_tokens)
7
    if forward_batch.lora_ids is not None:
8
        self.model_runner.lora_manager.prepare_lora_batch(forward_batch)
9
    attn_backend.init_forward_metadata_out_graph(forward_batch, in_capture=True)
10

11
    def run_once():
12
        attn_backend.init_forward_metadata_in_graph(forward_batch)
13
        self._record_in_graph_metadata_prep_done()
14
        ...
15
        out = forward(
16
            forward_batch.input_ids,
17
            forward_batch.positions,
18
            forward_batch,
19
            **kwargs,
20
        )
21
        for capture_hook in self.model_runner.capture_tail_hooks:
22
            capture_hook(self, out, forward_batch, num_tokens)
23
        return out
24
    ...
25
    self.backend.capture_one(
26
        shape_key, run_once, capture_inputs=None, post_warmup_hook=post_warmup_hook,
27
    )

注意 init_forward_metadata_out_graph 在 run_once 外面调用，而 init_forward_metadata_in_graph 在里面。这正是第八篇讲的「元数据两段式」：主机侧准备（涉及 Python 分支、host 端计算、可能的 .item() 同步）必须在捕获区间之外做完；只有纯设备算子才放进 run_once，才会被录进图。

如果反过来，把一个含 host 读取的元数据准备录进图，捕获会直接失败（CUDA 不允许在捕获中做同步）。第八篇里 MetadataGlueGraph 的注释把更隐蔽的失败讲了出来：某些后端的元数据准备每个重放步都要在 host 上重算值，这种后端绝不能录进图——录了不会报错，捕获成功、输出「正确」，但 host 写入的计划输入被冻结在捕获时刻的值，只有接受长度会悄悄塌掉。 这类「静默错误」是图机制里最难查的一类。

重放

每轮 decode 的入口是 execute：

1
with timer_ctx, self.backend.replay_session():
2
    self.load_batch(forward_batch, pp_proxy_tensors)
3
    ...
4
    if shared_read_ends is SharedReadEnds.PRE_REPLAY:
5
        self._publish_read_done(in_graph=False)
6

7
    output = self.backend.replay(self._replay_graph_key, forward_batch)
8

9
    if shared_read_ends is SharedReadEnds.IN_REPLAY:
10
        self._publish_read_done(in_graph=True)
11

12
    if shared_read_ends is SharedReadEnds.POST_REPLAY:
13
        self._publish_read_done(in_graph=False)
14

15
    output = self._process_output_after_replay(output, forward_batch)

load_batch 负责把真实批次拷进静态缓冲并算出图键，backend.replay 负责执行图，_process_output_after_replay 负责切出真实的 token 行数。SharedReadEnds 的三个时机（预重放前 / 重放中 / 重放后）就是第八篇讲的「后端自证什么时候读完调度器共享数据」，它决定了重叠调度器能提前多久准备下一轮。

输出切片是必要的：load_batch 把 batch 向上取整到桶，图算出来的是桶大小行数的结果，而真实 batch 只有 raw_num_token 行：

1
next_token_logits = (
2
    output.next_token_logits[: self.raw_num_token]
3
    if output.next_token_logits is not None
4
    else None
5
)

注册表的设计

CudaGraphBufferRegistry 的模块注释说得很清楚：它是 ForwardBatch 到图内缓冲的镜像，取代了原先手写的逐 runner 缓冲填充逻辑；采用 DecodeInputBuffers / PrefillInputBuffers 已经分配好的存储（同一个物理张量、data_ptr 在捕获与重放间稳定），只负责填充。后端私有的缓冲（kernel workspace、派生页表）留在 AttentionBackend.cuda_graph_* 上，注册表只管 FB 共享槽位——槽位名与 FB 属性名一一对应。

每个槽位是一个 GraphSlot：

1
@dataclass
2
class GraphSlot:
3
    name: str
4
    shape_fn: Callable[[int, int], Tuple[int, ...]]
5
    dtype: torch.dtype
6
    axis: str = "tokens"          # "bs" | "tokens" | "none"
7
    device: Optional[torch.device] = None
8
    padding_policy: PaddingPolicy = PaddingPolicy.FOREACH_COPY
9
    pad_value: Optional[Any] = None
10
    enabled: bool = True
11
    copy_from_fb: bool = True
12
    post_fill: Optional[Callable[[torch.Tensor, ForwardBatch, FillContext], None]] = None
13
    slice_fn: Optional[Callable[[torch.Tensor, int], Tensor]] = None
14
    source_fn: Optional[Callable[[ForwardBatch, FillContext], Optional[torch.Tensor]]] = None

几个字段值得展开：

• axis：槽位在 batch 轴（[:bs]）还是 token 轴（[:num_tokens]）切片，或者不切片（全局计数类）。批量填充时按轴分组，才能一次性 foreach_copy_。
• slice_fn：非平凡布局的切片覆盖，例如 mrope_positions 形状是 [3, T]，要沿 1 轴切而不是 0 轴。
• copy_from_fb：默认 True 表示「把 FB 上同名张量拷进缓冲头部」。设成 False 用于算出来的槽位（如 global_num_tokens_*），它们的值不来自 FB 拷贝，由 post_fill 算出。
• source_fn：拷贝源不是同名 FB 属性时的覆盖。例如流水线并行的 proxy 张量根本不在 FB 上，而是作为 load_batch 的独立参数传入、挂在 FillContext 上；再如 ngram embedding 的信息嵌在 FB 的嵌套 dataclass 里。这类槽位用点号命名，并被 extract_buffer 跳过。

五种 PaddingPolicy

这是全篇最值得记住的一张表。当真实长度 raw_n 小于桶长度 padded_n 时，多出来的尾巴怎么处理：

逐个解释「为什么」：

为什么 seq_lens 要 FILL_SENTINEL。 补齐出来的假请求，其真实 KV 长度是 0，但图内 kernel 不能读 0（第八篇讲过：零长度会让某些 kernel 退化、甚至除零）。填 1 让 kernel 走完整的合法路径，算出来的结果落在被丢弃的尾部。这个值是后端声明的，不是 runner 硬编码——因为不同后端对「最小合法长度」的要求不同。

为什么 out_cache_loc 和 req_pool_indices 要 ZERO。 这是最关键的一条。out_cache_loc 是写 KV 的目标槽位：图内注意力 kernel 会按它把当前 token 的 K/V 写进 KV 池。补齐出来的假行如果不处理，其值就是上一轮真实批次留下的槽位号——图会照着写一遍，把别人已经算好的 KV 覆盖掉。清零的语义是「指向槽位 0」，一个安全、无害的目标。注释原话是：padded rows must point at slot 0 so dummy attention reads land harmlessly。

req_pool_indices 同理：它指向 req_to_token_pool 里的请求行，kernel 用它查页表。假行指向第 0 行，读到一份合法但无意义的映射，不会越界。

为什么有 KEEP_PAD 和 FOREACH_COPY 的分工。 FOREACH_COPY 是默认：只拷真实部分，尾巴不管。这依赖调用者证明「尾巴不会被读」或「被读也无害」。KEEP_PAD 更显式地把这个责任写在策略名上，用于那些已经被证明尾部安全的槽位——写成 KEEP_PAD 而不是默认值，是为了让审阅者一眼看到「这里有人做过论证」。

为什么 encoder_lens 要 FILL_ONCE。 encoder-decoder 模型的 cross-attention 需要每个请求的 encoder 长度，图内用 encoder_len_fill_value（来自 max_source_positions）非零填充，保证 cross-attention kernel 被真实录进图。这个值在整个生命周期里不需要变，所以只在分配时填一次。

fill_from 的执行分两阶段，注释里写明：阶段一按每个槽位的策略重置 padding 尾巴；阶段二按 (dst_dtype, src_dtype) 分组做 torch._foreach_copy_——一次 foreach 调用要求 dtype 对一致，所以先分组。分批的原因是：一轮重放要填几十个槽位，逐个 copy_ 就是几十次 kernel launch，正好抵消了图带来的收益；分成几组 foreach 可以把内核数压到一个很小的常数。

两个绕不开的特征

特征一：形状连续变化且范围极大。 prefill 的 token 数从 1 到几万，不可能像 decode 那样「用几十个桶覆盖」。SGLang 的做法是分桶 + 向上取整，并给 padding 浪费设一个上限：

1
# No exact-shape check: load_batch bucket-pads; only reject disproportionate padding waste.
2
padded_num_tokens = self._pad_to_bucket(num_tokens, self.capture_num_tokens)
3
if padded_num_tokens > num_tokens * _MAX_PREFILL_CUDA_GRAPH_PADDING_FACTOR:
4
    return False

特征二，也是决定性的特征：prefill 的注意力是「变长」的。 decode 时每个请求的 query 长度固定为 1（或投机时的固定宽度），不同请求只在 KV 长度上不同，而 KV 长度的差异可以被分页注意力 + kv_indptr 自然吸收，图内的 kernel 形状依然统一。prefill 则不同：每个请求的 query 长度各不相同（从几十到几万），且因果掩码的形状随之变化。FlashAttention 这类 kernel 是按 query 分块在 grid 上展开的，query 长度直接决定 grid 与执行路径。

结论：把 prefill 的注意力层录进 CUDA Graph，要求所有请求的 query 长度完全一致。 但真实服务里的请求长度是连续分布，这就意味着捕获只能覆盖「所有人长度恰好相同」这个测度为零的情形。于是要么放弃图，要么不录注意力层。

分段捕获 BCG 把图切在哪里

BreakableCudaGraphBackend 的类注释直接给出了答案：

Segmented capture: graphs break at attention / mamba boundaries; attention metadata is recomputed at replay outside captured segments.

机制本体在 runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py。核心是 @eager_on_graph 装饰器：

1
def eager_on_graph(fn=None, capture_stub=None, *, enable=None):
2
    def decorator(inner: Callable):
3
        @functools.wraps(inner)
4
        def wrapper(*args, **kwargs):
5
            capture = _current_capture_var.get()
6
            if capture is None:
7
                return inner(*args, **kwargs)
8

9
            # End the segment that captured up to this break point.
10
            capture._end_current_segment()
11

12
            # Re-sync ranks after segment teardown before break fns with
13
            # rank-coupled collectives and hard timeouts。
14
            if capture._barrier_fn is not None:
15
                capture._barrier_fn()
16

17
            if capture_stub is not None:
18
                output = capture_stub(*args, **kwargs)
19
            else:
20
                output = inner(*args, **kwargs)
21

22
            captured_inner = inner
23
            captured_args = tuple(_weak_ref_if_tensor(a) for a in args)
24
            captured_kwargs = {k: _weak_ref_if_tensor(v) for k, v in kwargs.items()}
25
            captured_output = output
26

27
            def replay_fn():
28
                new_out = captured_inner(*captured_args, **captured_kwargs)
29
                return _copy_output(captured_output, new_out)
30

31
            capture.cuda_graph._break_fns.append(replay_fn)
32

33
            # Start a fresh CUDAGraph segment for the remainder of the forward.
34
            capture._begin_new_segment()
35
            return output
36

37
        return functools.update_wrapper(_EagerOnGraphCallable(wrapper, inner), inner)
38

39
    return decorator(fn) if fn is not None else decorator

逐步拆开这段逻辑：

第一步，结束当前段。 _end_current_segment() 对当前 torch.cuda.CUDAGraph 调 capture_end()，把它追加进 BreakableCUDAGraph._segments。注意注释里那句：A segment that enqueued no kernels captures an empty graph, which replays as a no-op——两个断点紧挨着会录出一张空图，重放是空操作，无害；torch 每次都会对空图发警告，所以这里显式过滤掉了那个警告。

第二步，跨 rank 屏障。 断点处各 rank 要分开录制，段拆除是个「慢且耗时可变」的步骤。如果断点函数里有与 rank 耦合的集合通信且带硬超时（注释点名 DeepEP NORMAL 的 100 秒超时），一段录得快、另一段录得慢就会把对端拖到超时。所以在段拆除后插一次 tp_group.barrier()。这个屏障只在捕获期需要，重放期不走这个 wrapper。

第三步，执行断点函数（真身）。 断点函数在这里被真实执行一次，目的是让它的输出被分配出来、地址被记录下来。同时，被捕获段产出的中间张量用弱引用记录（_weak_ref_if_tensor）——因为它们的存储被段的 CUDAGraph mempool 的 use-count 钉住了，Python 侧不必强引用。但断点函数的输出必须强引用（captured_output = output），因为它是「下一个被捕获段的静态输入地址」，需要跨重放存活。

第四步，登记重放函数。

1
def replay_fn():
2
    new_out = captured_inner(*captured_args, **captured_kwargs)
3
    return _copy_output(captured_output, new_out)

重放时真实调用断点函数（用弱引用拿回的捕获期参数），把结果 _copy_output 拷进当初记录的那块桥接缓冲。也就是说：断点在捕获期与重放期都会真跑，且都写到同一块地址。

第五步，开新段。 _begin_new_segment() 新建一个 torch.cuda.CUDAGraph，capture_begin(pool=..., capture_error_mode="global")。新段与旧段共用同一个 memory pool（构造函数里的 self._pool），所以段间分配的中间张量可以复用池里的空闲块——注释说明：任何段存活时，它的 beginAllocateToPool 都让 mempool 的 use_count 保持大于 0，因此对段内分配张量做 weak_ref_tensor 在后续重放中是安全的。

BreakableCUDAGraph.replay() 就是把二者交替跑一遍：

1
def replay(self) -> None:
2
    stream = get_device_module().current_stream()
3
    token = _current_stream_var.set(stream)
4
    try:
5
        for i, seg in enumerate(self._segments):
6
            seg.replay()
7
            if i < len(self._break_fns):
8
                self._break_fns[i]()
9
    finally:
10
        _current_stream_var.reset(token)

断点插在哪里：注意力层

断点函数由 @eager_on_graph 标注在注意力层与 mamba 层上（backend 注释原话：eager break markers (eager_on_graph decorators on attention / mamba layers)）。于是 prefill 的一次前向被切成：

1
[段1: embedding + layer0 的 norm/qkv/rope]  → eager: layer0 注意力（读真实元数据）
2
[段2: layer0 o_proj + MLP + layer1 前半]     → eager: layer1 注意力
3
...
4
[eager tail: lm_head + logits 处理]

每个被捕获段内的张量形状是静态的：hidden states 的 token 数就是捕获桶大小，权重形状不变。唯一随请求变化的注意力计算被赶到 eager 段里，它可以自由读真实的 extend_seq_lens / extend_start_loc / block table。

这个结构的收益是：Transformer 主体的所有线性层、norm、激活，占了一次前向上绝大部分的 kernel 数量与计算量，全部进了图；只有注意力层（数量上只有每层一个）留在图外。

为什么这仍然值得

有人会问：既然注意力层还是 eager，收益是不是就打折了？

打折，但仍然显著。理由有三：

1. 注意力层是变长的，且通常被高度优化过。 FlashAttention / FlashInfer 的 prefill kernel 一次 launch 就覆盖整个变长 batch，kernel 数量本来就少（每个注意力层 1–2 个 launch）。其余部分——每一层的多个线性投影、RMSNorm、激活、残差——才是「碎 kernel」的大头，数十倍于注意力层数。
2. prefill 的 chunked prefill 让形状集中在少数几个值。 SGLang 默认按 chunked_prefill_size 切分长提示，真实 prefill 桶常常是「一整块」或「一小块」，命中率不低。
3. BCG 与 prefill 是默认开启的（在 CUDA 上）：

1
def default_prefill_backend() -> str:
2
    """Enable BCG by default on CUDA; other platforms remain opt-in."""
3
    from sglang.srt.utils import is_cuda
4

5
    return Backend.BREAKABLE if is_cuda() else Backend.DISABLED

而 decode 默认是 Backend.FULL（整图）。这个默认值的分工本身就是结论：decode 用整图（形状规整、收益最大），prefill 用分段图（必须绕开变长注意力）。

BCG 的额外约束与配套机制

分段捕获带来了几个整图没有的约束：

一、与 memory saver 不兼容。 BreakableCudaGraphBackend.__init__ 直接抛 NotImplementedError：

1
if (
2
    self._memory_saver_adapter is not None
3
    and self._memory_saver_adapter.enabled
4
):
5
    raise NotImplementedError(
6
        "Breakable CUDA graph is not compatible with memory saver mode"
7
    )

原因是 memory saver 的暂停/恢复会对图池做 CPU 备份，而分段捕获的跨段引用语义（弱引用 + 桥接缓冲）让「恢复后地址仍有效」这个隐含前提不再可靠。

二、需要 post-warmup hook。 BCG 的 capture_one 会跑两次 warmup，并在每次之后调用 post_warmup_hook：

1
warmup_out = None
2
for _ in range(2):
3
    self._device_module.synchronize()
4
    self._tp_group.barrier()
5
    with self._precarve.measure():
6
        warmup_out = forward_fn()
7
    if post_warmup_hook is not None:
8
        post_warmup_hook()

这个 hook 的作用是「重置被 warmup 改动的后端状态」。prefill runner 里有一段具体的例子：mamba 池在 warmup 中被写入，如果不清理，第二次 warmup 和真正的捕获就会读到脏状态：

1
if isinstance(self.backend, BreakableCudaGraphBackend):
2
    post_warmup_hook = None
3
    req_pool = self.model_runner.req_to_token_pool
4
    mamba_pool = getattr(req_pool, "mamba_pool", None)
5
    if mamba_pool is not None and not prefix_num_chunks:
6
        capture_state_indices = req_pool.translate_mamba_indices(
7
            req_pool.get_mamba_indices(forward_batch.req_pool_indices)
8
        ).unique()
9

10
        def post_warmup_hook():
11
            mamba_pool.clear_slots(capture_state_indices)
12

13
        post_warmup_hook()

注意这里的写法：hook 先被手动调用一次，再作为参数传给 capture_one（后者在两次 warmup 后各调一次）。因为 capture_one 的第一段 warmup 前没有 hook 机会，脏状态必须在进入前先清一遍。

三、BCG 只捕获 Transformer 主体，不捕获 lm_head。 prefill runner 的 _run_forward 注释：

BCG path: captures only the inner layer_model.forward (transformer stack), excluding the outer model.forward tail (logits_processor / pooler). The captured output is bs=1 hidden states; replay then runs the outer tail eagerly with live multi-req metadata.

原因是显存：如果把 lm_head 与 logits 处理也录进图，图内就要绑定一块形状为 (req_slots, vocab_size) 的输出缓冲，几十万词表乘以请求槽位数是几百 MB 到 GB 级，而且必须按最大请求槽位分配。把尾巴留在图外 eager 跑，图的输出就只是 hidden states，形状只跟 token 数有关，req_slots 不再约束缓冲大小。

这段注释还提到同一函数里 @torch.no_grad 是刻意补上的：外层 *ForCausalLM.forward 本身带这个装饰器，但 BCG 直接调 layer_model.forward 就绕过了它，导致某些 MoE 的 @torch.compile kernel（torch.sum(out=...)）在 dynamo 里报「out= doesn’t support autograd」，mamba 状态算子也会莫名开始追踪梯度。

Prefill 的重放准入

prefill 的重放判定被集中到一个函数 can_replay_locally，它的 docstring 讲清了为什么要集中：

Rank-local replay eligibility: the single source of truth for can_run_graph (ForwardBatch, forward time) and the dp mlp-sync vote (ScheduleBatch, schedule time) — all dp ranks must reach the same replay-vs-eager decision or their collectives mismatch.

这是 prefill 图与 decode 图最大的结构差异：prefill 的判定必须提前到调度时刻。 DP attention 下所有 rank 要一起进入集合通信，如果 rank A 重放图、rank B 走 eager，通信就分叉了。而重放与否取决于 num_tokens、prefix_lens、return_logprob 等必须等批次组好才知道的信息，所以判定函数被抽出来，让调度期的投票与执行期的判断用同一个逻辑（注释里叫 “dp mlp-sync vote”，以 ScheduleBatch 为输入）。

判定条件里几个值得注意的：

1
if self._is_full_backend and self._fa4_prefill and contains_mm_inputs:
2
    return False
3
...
4
if lora_ineligible:
5
    return False
6
if is_mixed and getattr(self, "prefer_eager_mixed_prefill", False):
7
    return False
8
if input_embeds is not None:
9
    return False
10
if replace_embeds is not None:
11
    return False
12
...
13
if is_target_verify:
14
    return False
15
if (
16
    capture_hidden_mode is not None
17
    and self.capture_hidden_mode < capture_hidden_mode
18
):
19
    return False
20
if return_logprob and not self._uses_eager_prefill_tail():
21
    return False
22
for context_limit in (self.max_context_size, self.max_seq_len):
23
    if context_limit is not None and (
24
        batch_max_context_len is None or batch_max_context_len > context_limit
25
    ):
26
        return False

• 多模态输入在某些后端下禁：Full 图重放会绕过 forward_extend 里的 Python 图像掩码守卫，FA4 的完整图路径尚未支持，所以保持 eager。
• LoRA 不合格则禁：LoRA 重放需要 prepare_lora_batch 建好的静态元数据；在 DP attention 下 LoRA prefill 在所有 rank 上都保持 eager，保证投票一致。
• mixed batch 在特定组合下禁：BCG + DP attention 时倾向 eager mixed prefill。
• return_logprob 有讲究：它只在图外有 eager 尾巴时才能重放——因为 logprob 要在图外对完整 logits 做处理。
• 上下文长度双上限：max_context_size 与 max_seq_len 两个都要判，注释原话是 “Vote on both bounds so DP ranks agree on graph versus eager”。原因是 DeepSeek-V4 的上下文形状元数据与 indexer logits 是按固定 max_context_size 分配的，而 max_seq_len 是另一维约束；两个上限不同时满足会导致 rank 间决策不一致。

分块前缀的额外拓扑

当启用 full_prefill_prefix_chunk_tokens（注释里叫 chunked-prefix），Full prefill 图会额外捕获不同 Cached Prefix 拓扑的变体：

Only meaningful for Full prefill CUDA graphs that capture a distinct cached-prefix topology: aggregate cached-prefix tokens represented by one fixed-capacity chunk across all request slots. FullCG captures 1/2/4/8/16 chunk variants and chooses the smallest one covering a batch.

也就是说，请求命中 RadixCache 的前缀长度也会影响图内 kernel 的形状，所以按「一个固定容量的 chunk」为单位录 1/2/4/8/16 五种变体，重放时选最小的能覆盖的那一个。这解释了 ShapeKey.variant_label 为什么要存在。

自动降级：什么时候必须关掉图

arg_groups/cuda_graph_hook.py 里有一族 apply_* / disable_* 函数，代表性的是：

• disable_breakable_cudagraph_if_incompatible：BCG 有一张不兼容规则表，命中就降级。从源码里的分支看，涉及 prefill context parallelism + pipeline parallelism 的组合、非标准 GQA 层、LoRA 后端能力不足等。
• disable_full_prefill_cudagraph_if_incompatible：注释写得很直白——Full prefill CG: empty rule list today，目前是空规则表，即 Full prefill 还处在「实验性、无自动降级规则」的阶段。
• disable_prefill_cuda_graph_for_deepseek_trtllm_mla：在使用 trtllm_mla 后端时默认关掉 prefill 图。
• apply_deepep_adjustments：DeepEP 需要对齐的桶，会把 max_bs or 2048 转成 generate_prefill_cuda_graph_batch_sizes(max_bs)。
• apply_cuda_graph_disaggregation_roles：PD 分离下按角色关掉对应相位——prefill 实例关 decode 图，decode 实例关 prefill 图：

1
if cfg.disaggregation_mode == "prefill":
2
    if (Phase.DECODE, "backend") not in server_args._cuda_graph_config_locked:
3
        declare_resolution(..., cuda_graph_config=with_phase(cfg, Phase.DECODE, backend=Backend.DISABLED))

理由直接：PD 分离后，prefill 实例永远只跑 prefill，decode 实例永远只跑 decode。给 prefill 实例录 decode 图纯粹是浪费显存、浪费启动时间。

在 runner 构造层面还有一层兜底。capture_prefill_graph 里有一长串 return-to-eager 的分支：

1
if check_cuda_graph_backend(Phase.PREFILL, Backend.DISABLED):
2
    ...
3
    return result(eager_runner)
4
try:
5
    language_model = resolve_language_model(model_runner.model)
6
except AttributeError:
7
    logger.warning("Disable prefill CUDA graph because the model is not a language model")
8
    return result(None)
9
if not get_exec().graph.cuda_graph_config.prefill.bs:
10
    logger.warning("Disable prefill CUDA graph because the capture size is not set")
11
    return result(None)
12
...
13
if not has_standard_gqa_for_all_local_layers(...):
14
    log_info_on_rank0(logger, "Disable prefill CUDA graph because some layers do not apply Standard GQA")
15
    return result(None)

其中「非标准 GQA」这条尤其能说明图的边界：注意力层如果用了非标准的 query/key/value 头排布，图内 kernel 的 reshape 布局就不通用，图机制直接放弃这个模型。

decode 路径：小桶直接编译

_capture_one_stream 里对每个桶包了一层 patch_model：

1
with torch_compile_decoration.patch_model(
2
    self.model_runner.model,
3
    bs in self.compile_bs,
4
    num_tokens=bs * self.captured_req_width,
5
    tp_group=get_parallel().tp_group,
6
) as forward:
7
    self.capture_one_shape(bs, forward, stream_idx, variant_label, attention_variant)

patch_model 的行为是「落在编译桶里就返回一个被 torch.compile 包过的 forward，否则返回原始 forward」：

1
@contextmanager
2
def patch_model(model, enable_compile, num_tokens, tp_group):
3
    backup_ca_comm = None
4
    try:
5
        if enable_compile:
6
            _to_torch(model, reverse=False, num_tokens=num_tokens)
7
            backup_ca_comm = tp_group.ca_comm
8
            yield torch.compile(
9
                torch.no_grad()(model.forward),
10
                mode=os.environ.get(
11
                    "SGLANG_TORCH_COMPILE_MODE", "max-autotune-no-cudagraphs"
12
                ),
13
                dynamic=_is_hip and get_bool_env_var("SGLANG_TORCH_DYNAMIC_SHAPE"),
14
            )
15
        else:
16
            yield model.forward
17
    finally:
18
        if enable_compile:
19
            _to_torch(model, reverse=True, num_tokens=num_tokens)
20
            tp_group.ca_comm = backup_ca_comm

三处细节：

• max-autotune-no-cudagraphs 这个 mode 名字很关键。max-autotune 会做自动调优（对候选 kernel 配置做 benchmark），而 -no-cudagraphs 后缀明确禁止 Inductor 自己再做一层 CUDA Graph 包装——因为外层 SGLang 已经在录图了，两层图嵌套是无效也是危险的。
• 编译前后切换融合算子：_to_torch 递归遍历模型，把 BaseFusedOp 子模块在「编译态」与「普通态」之间切换（enter_torch_compile(num_tokens=...) / leave_torch_compile()）。原因是某些融合算子在编译态下需要暴露可被 Dynamo 追踪的实现，而运行时用高度手写 kernel。
• tp_group.ca_comm 被备份并临时置空：编译期间的追踪不应触发自定义 all-reduce 的图缓冲注册（那个注册必须在捕获上下文之外完成，否则与 custom_all_reduce.register_graph_buffers 死锁——BaseRunner.warmup 的注释也专门警告过这一点）。

torch_compile_max_bs 默认 32 的理由由此清楚：Inductor 对每个新的形状都要重新编译一份 kernel，编译桶里每多一个 batch size，就多一份编译产物与编译时间（启动耗时）。小 batch 的收益最大而形状种类最少，所以只在 32 以下编译；大 batch 靠 CUDA Graph 本身的收益就够了。

prefill 路径：编译是另一条腿

prefill 侧的结构更微妙。_run_forward 的注释把两条路径摆在一起：

1
BCG path: captures only the inner layer_model.forward (transformer stack),
2
excluding the outer model.forward tail ...
3
TC_PIECEWISE path: captures the outer model.forward; torch.compile
4
FX-traces produce bs-invariant kernels.

tc_piecewise 是第三种捕获后端，它的思路是「用编译把形状敏感性消除掉」：让 torch.compile 把模型追踪成 FX 图，Inductor 生成的主 kernel 用 bs-invariant 的循环结构（用 tl.range 之类表达动态行程），于是同一份编译产物能服务多个 batch size，图内形状不再随桶变化。这就是 “piecewise” 与 “bs-invariant” 的含义。

不过当前代码里 tc_piecewise 已经退役。cuda_graph_config.py 的解析器会直接拒绝：

1
if phase_settings.get("backend") == "tc_piecewise":
2
    raise ValueError(
3
        "tc_piecewise was removed; select breakable, full, or disabled"
4
    )
5
if "tc_compiler" in phase_settings:
6
    raise ValueError("tc_compiler was removed with tc_piecewise")

Backend.TC_PIECEWISE 字段本身还留着，注释是 Temporary internal name while unreachable legacy branches are removed。prefill 的默认现在是 BCG，Full 是可选的实验路径。

这个演变的信号值得注意：「编译 + 整图」这条路线在 prefill 上被「分段 + 绕开变长算子」取代了。原因不难推断——编译路线要求把变长注意力也表达成图内可重放的 kernel，这需要后端 kernel 本身支持「形状由张量内容决定」；而分段路线只需要在 Python 层面把注意力层摘出去，工程复杂度低得多，对后端的改动也小。

编译与捕获的顺序：warmup 先行

BaseRunner.warmup 是所有 runner 共用的前置动作：

1
def warmup(self) -> None:
2
    """Run kernel warmup + autotune once, gated by mr._kernel_warmed_up."""
3
    mr = self.model_runner
4
    if getattr(mr, "_kernel_warmed_up", False):
5
        return
6
    mr._kernel_warmed_up = True
7
    if mr.device != "cuda":
8
        return
9
    self._pre_initialize_flashinfer_allreduce_workspace()
10
    self._pre_initialize_fi_a2a_workspace()
11
    self._pre_initialize_pcie_ipc_workspace()
12
    prepare_model_resources = getattr(
13
        mr.model, "prepare_before_cuda_graph_capture", None
14
    )
15
    if prepare_model_resources is not None:
16
        prepare_model_resources(mr)
17
    ...

_kernel_warmed_up 这个标志意味着 warmup 在整个进程里只跑一次（decode runner 与 prefill runner 共享），因为两件事都必须在任何捕获之前完成：

1. 所有一次性初始化必须在捕获之外。FlashInfer 的 allreduce workspace、FI all2all workspace、PCIe IPC workspace 的分配都会做同步与集合通信；注释原话是 “must run before CG capture to keep broadcasts/barriers outside the capture context (else deadlock with custom_all_reduce.register_graph_buffers)”。
2. JIT 编译与 autotune 必须在捕获之外。Triton kernel 第一次执行要编译，FlashInfer 的 plan 要做 autotune，两者都会打断捕获。

_dummy_run 则是 warmup 的执行体：构造一个形状合适的 dummy 批次，跑一遍前向。它有一个约束写在断言里：

1
assert (
2
    buffers is not None
3
    and num_tokens <= buffers.input_ids.shape[0]
4
    and batch_size <= buffers.seq_lens.shape[0]
5
), (...)

_dummy_run 从不分配、从不重新 padding——调用者拥有形状，必须传入一份不小于 dummy 形状的静态缓冲。注释解释了原因：autotune 必须在「最终会用的形状」上跑，否则调优结果不适用；而且复用已有缓冲能让 dummy 前向分配出的临时张量与后续捕获共用池，避免池被撑大。

捕获成本的测量

capture_prefill_graph 与 capture_decode_graph 都用同一个模式测成本：

1
tic = time.perf_counter()
2
before_mem = get_available_gpu_memory(model_runner.device, model_runner.gpu_id)
3
...
4
prefill_runner = PrefillCudaGraphRunner(model_runner)
5
after_mem = get_available_gpu_memory(model_runner.device, model_runner.gpu_id)
6
mem_usage = before_mem - after_mem
7
capture_time = time.perf_counter() - tic
8
logger.info(
9
    f"Capture {capture_name} CUDA graph end. "
10
    f"elapsed={capture_time:.2f} s, "
11
    f"mem usage={mem_usage:.2f} GB, avail mem={after_mem:.2f} GB."
12
)

测量结果按相位归到 graph_memory_usage 字典（decode / prefill / target_verify / draft_decode / draft_prefill），GraphCapture.memory_usage 与 time_usage 是它的读写接口，最终由调度器上报：

1
graph_memory_usage_gb=combine_graph_memory_usage(
2
    self.tp_worker.graph_memory_usage,
3
    (None if self.draft_worker is None else self.draft_worker.graph_memory_usage),
4
),

启动日志里的 elapsed=... mem usage=... GB 与 /get_server_info 暴露出来的 graph 内存用量，都是从这里来的。

谁来为图留出显存

关键的一步在 KV 池的容量计算。KV cache configurator 的 _profile_available_bytes 决定了 KV 池能拿多少：

1
gc.collect()
2
available_gpu_memory = get_available_gpu_memory(...)
3
slack_gb = pre_model_load_memory * (1 - get_schedule().mem_fraction_static)
4
...
5
if (
6
    self.mambaish_config is not None
7
    and self.post_capture_kv_active
8
    and not capped_full_prefill_role
9
):
10
    slack_gb = max(
11
        slack_gb,
12
        pre_capture_activation_reserve_mb(
13
            get_device_memory_capacity(self.device)
14
        ) / 1024,
15
    )
16
...
17
rest_memory = available_gpu_memory - slack_gb - mm_reservation_gb

这里有两个概念：

• mem_fraction_static：静态内存（权重 + KV 池 + 图）占 GPU 总显存的比例。pre_model_load_memory * (1 - mem_fraction_static) 就是留给「非静态运行时」的余量 slack_gb。
• pre_capture_activation_reserve_mb：捕获之前必须预留的激活内存估算：

1
def pre_capture_activation_reserve_mb(gpu_mem: float | None) -> float:
2
    """Return the activation reserve in MB before CUDA graph capture, using current config."""
3
    schedule = get_schedule()
4
    if get_disagg().disaggregation_mode == "decode":
5
        running_requests = (
6
            schedule.max_running_requests
7
            or get_exec().graph.cuda_graph_config.decode.max_bs
8
            or 1
9
        )
10
        activation_tokens = max(
11
            running_requests * (get_spec().speculative_num_draft_tokens or 1), 2048
12
        )
13
    elif schedule.chunked_prefill_size > 0:
14
        activation_tokens = max(schedule.chunked_prefill_size, 2048)
15
    else:
16
        activation_tokens = max(schedule.max_prefill_tokens, 2048)
17
    parallel = get_parallel()
18
    reserved_mem = (
19
        512 + activation_tokens * 1.5 + parallel.tp_size * parallel.pp_size / 8 * 1024
20
    )
21
    if gpu_mem is not None and gpu_mem > 60 * 1024:
22
        reserved_mem = max(reserved_mem, 10 * 1024)
23
    return reserved_mem

公式把三块加起来：512 MB 固定开销、activation_tokens × 1.5（激活的经验系数）、以及并行度相关的通信缓冲（tp_size × pp_size / 8 × 1024 MB，即每 8 个 rank 预留 1 GB）。大显存卡（> 60 GB）再兜一个 10 GB 的下限。

为什么这块余量必须留。 因为 CUDA Graph 捕获时真实跑前向，图外还有 eager 路径（大 batch 回退、prefill 尾巴、logits 处理）要跑，Mamba 状态是固定预捕获分配。如果按「理论上需要」精确抠出来，KV 池会把显存占满，运行时的激活分配立刻 OOM。_profile_available_bytes 的注释专门留了一句实战教训：如果不先 gc.collect()，权重加载的临时张量仍被引用，empty_cache() 收不回，KV 池会被算成小几个数量级，而 GPU 显存还在闲置。

与调度器的闭环

调度器那边能看到两组数字：一是 max_total_num_tokens 与 max_running_requests（KV 池容量，已被图预留扣减过），二是 decode_cuda_graph_capture_bs（实际捕获的桶列表，通过 /get_server_info 暴露）。批次尺寸超过 max_bs 的请求会走 eager，这在功能上正确，在性能上退化成「无图 decode」。

这个闭环的形态是：配置桶列表 → 捕获占用显存 → KV 池容量按剩余显存计算 → max_running_requests 受池容量约束 → 实际 batch 上限受资源约束。任何一个环节设得太满都会互相挤压，这就是为什么桶列表默认值那么克制。