DeepInfer 源码完全拆解（十二）：CUDA Graph 捕获与重放——动态批怎么冻成一张图，又在哪里放开

每一步的判定：为什么 prefill 永远不走图

真正的判定在 _prepare_input_driver 里：

1
is_prompt = seq_metadata_list[0].is_prompt
2
run_cuda_graph = (
3
    not is_prompt
4
    and self.use_cuda_graph
5
    and bool(self.cuda_graph_runner.graphs)
6
)
7

8
if run_cuda_graph:
9
    graph_size, cuda_graph_used = self.cuda_graph_runner.select(total_q)
10
    ...
11
else:
12
    graph_size, cuda_graph_used = 0, None
13
    input_metadata = self._prepare_input_tensors(seq_metadata_list, metadata_zos)
14
    input_metadata = input_metadata.recursive_fn(lambda x: x.cuda(non_blocking=True))

bool(self.cuda_graph_runner.graphs) 这一项看似多余，实际是有必要的：图的捕获发生在启动阶段的 warm_up_model 里，而在捕获之前，跑器对象已经存在但 graphs 仍是 None。如果启动流程里有任何一步在捕获完成前就进入了 execute_main_model，这项判断能保证它走 eager 而不是访问 None。

prefill 不走图有明确的技术原因，不只是”形状太长”。InputMetadata 里的 CP-mirror 字段注释写得很清楚：prepare_for_mirror_cp 会把 positions、ngram_token_ids 等字段重新绑定到新的张量地址上，理由是”prefill runs without cudagraph”。图一旦捕获，所有张量地址就固定了；CP-mirror 的重排如果发生在图内，重排后产生的新张量地址不可能被捕获进图。DeepInfer 选择的原则是：任何会重新分配张量的步骤，都放在图外。prefill 阶段要重排输入、要做跨 rank 广播、要按压缩比展开大量注意力索引，这些操作天然不具备”原地”性质，所以 prefill 整体留在 eager 路径。

解码则不同。解码的输入形状固定为每条序列若干个 token，输入内容虽然每步变化，但都可以写进预先分配好的缓冲区。这才使图成为可能，而缓冲区的设计正是下一节的全部内容。

并行与 ZOS 下的分叉

当开启上下文并行（CP）时，_prepare_input_driver 之后还有一段广播逻辑，同样对图与 eager 区别对待：

1
if PG.context_parallel_group().size() > 1:
2
    input_metadata_dict = {
3
        "is_prompt": is_prompt,
4
        "type": type(input_metadata),
5
        "run_cuda_graph": run_cuda_graph,
6
        "total_q": total_q,
7
        "num_reqs": num_reqs,
8
    }
9
    if not run_cuda_graph:
10
        input_metadata_dict.update(input_metadata.to_flatten_dict())
11
    broadcast_tensor_dict(input_metadata_dict, src=0)

走图时只广播四个标量和类型，不走图时才把整份 InputMetadata 摊平后广播。原因在于：图路径下，每个 CP rank 上的图数据对象是各自捕获时创建的，输入缓冲区在构造时就已存在；driver 只需告诉 follower “这一步是图、用哪个 total_q、真实请求数是多少”，follower 用同样的 select(total_q) 就能找到自己的同一张图。而 eager 路径没有预先存在的缓冲区，必须把字段值整个传过去，由 follower 用 from_flatten_dict 重建对象。这条分支也解释了一个容易忽略的事实：每个 rank 都要各自捕获一套图，图不能跨进程共享，因为图里写死的指针是各进程各自显存里的地址。

当打开 ZOS 时，CudaGraphRunner 的构造参数 input_on_device=True 被传入，这个标志将在捕获与回填阶段改变一串流程，后面讲 _run_model 时会看到它的作用。

档位集合的构造

build_graph_batch_sizes 只有七行：

1
def build_graph_batch_sizes(max_size: int, interval: int) -> list[int]:
2
    """Build power-of-two warmup sizes followed by fixed-width intervals."""
3
    interval = min(interval, max_size)
4
    batch_sizes: list[int] = []
5
    size = 1
6
    while size < interval:
7
        batch_sizes.append(size)
8
        size *= 2
9
    batch_sizes.extend(range(interval, max_size, interval))
10
    batch_sizes.append(max_size)
11
    return batch_sizes

假设 max_size = 64、interval = 8，得到的是 [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64]。注意第一段 while size < interval 的退出条件是严格小于：当 size 增长到等于 interval 时循环停止，所以第一段只写入 1, 2, 4，而 8 起由第二段的 range(interval, max_size, interval) 产生，两者不会重叠。如果 interval = 16，则第一段产生 1, 2, 4, 8，第二段从 16 开始。这种”先细后粗”的拼接会让档位序列在整个区间上保持严格递增。

这个两段式设计对应了两个不同的代价函数。在小 batch 段，浪费的相对比例极大：只有 3 个请求时落到 4 的档位，浪费 25%；落到 8 的档位，浪费 62%。所以 2 的幂次在 1 到 interval 之间把分辨率做细。在大 batch 段，档位间距固定为 interval，浪费最多 interval-1 个 token，相对比例随 batch 增大递减，同时档位总数被控制住。档位总数直接决定启动成本：捕获每张图都要跑一次 warmup、保存一份图对象与图内中间缓冲、把共享的输出缓冲切成视图，还要在实例化时付出驱动侧时间。档位越多，启动越慢；而共享内存池让显存开销主要随最大档位增长，不随档位数线性增长。

捕获顺序与共享内存池

CudaGraphRunner.capture 里最关键的一行不是捕获本身，而是循环方向：

1
for batch_size in reversed(self.batch_sizes):

从最大的档位开始捕获。原因藏在 torch.cuda.graph 的 pool 参数里：

1
graph = torch.cuda.CUDAGraph(keep_graph=True) if PLATFORM.is_gpu else torch.cuda.CUDAGraph()
2
with torch.cuda.graph(graph, pool=self.graph_memory_pool, stream=self.graph_stream,
3
                      capture_error_mode="thread_local"):
4
    model_output = _run_model(...)
5
    output.copy_(model_output)
6
    del model_output
7
...
8
if self.graph_memory_pool is None:
9
    self.graph_memory_pool = graph.pool()

第一张被捕获的图（也就是最大的那张）分配私有内存池，后续所有图都共享同一个池。PyTorch 的 CUDA Graph 内存池在按顺序捕获时，池内内存块的地址是单调向后分配的；后捕获的图使用池里已经分配过的块，只要它用到的内存在前一张图里已经分配过且已被释放，就能复用同一地址。按从大到小的顺序捕获，意味着较大的图先占住高水位，较小的图在前一张图的地址空间内复用，不需要新的分配。反过来从小到大捕获，每换一个更大的档位就要扩池，池内地址不断变化，也就无法共享。

这也是为什么 capture 开头必须先 self.destroy()：重复捕获时旧图与旧池必须完全释放，否则池地址会叠加。

捕获在独立的 graph_stream 上进行，前后各有一个 wait_stream：捕获前 self.graph_stream.wait_stream(torch.cuda.current_stream()) 保证当前流上的已有工作（例如模型加载、权重初始化）在捕获开始前完成；捕获后 torch.cuda.current_stream().wait_stream(self.graph_stream) 让主流程等待捕获流结束。capture_error_mode="thread_local" 则把捕获状态限制在当前线程，避免其他线程（例如结构化解码的语法编译线程）的 CUDA 调用被误认为是对图的捕获调用，这会直接触发 cudaErrorStreamCaptureUnsupported。

keep_graph=True 只在 GPU 平台上生效，NPU 平台不传这个参数。它的用途在注释里写明了：让底层的 cudaGraph_t 原始对象保留下来，供后面的图级计时器按事件节点匹配。图级计时器不属于本文主线，但它反过来约束了捕获代码的写法。

预热、清空派生字段、再捕获

每张图捕获前都会先跑一次 warmup：

1
self.graph_stream.wait_stream(torch.cuda.current_stream())
2
with torch.cuda.stream(self.graph_stream):
3
    _run_model(model, kv_caches, input_metadata_source, input_metadata_device,
4
               not self.input_on_device)
5

6
for name in input_metadata_source.get_clean_list():
7
    setattr(input_metadata_source, name, None)
8
    setattr(input_metadata_device, name, None)

warmup 的作用是让惰性初始化和首次调用的开销（kernel JIT 选择、cuDNN 句柄、第一次分配）在捕获前完成，否则这些操作会进入图，而其中很多操作（分配、同步）本身不被捕获允许。

紧接着的清理循环对应模型前向的一个设计：InputMetadata.compute_attn_info 会在前向开始时把”紧凑索引字段”展开成内核能直接消费的全量索引。InputMetadataV4.get_clean_list 返回的正是这批派生字段：

1
return InputMetadata.get_clean_list() + [
2
    "kv_ca_index_len",
3
    "kv_swa_index",
4
    "kv_ca_index",
5
    "mlp_relu_schedule_metadata",
6
    "flash_mla_sched_meta_dict",
7
]

这些字段在 warmup 里被填上了，但如果带着它们进入捕获，图里就会把”用这一次 warmup 的数据算出来的索引内容”也写死。清理之后进入捕获，前向里 compute_attn_info 会重新展开这些字段。这里有一个顺序上的微妙之处：由于捕获时所有 CUDA 调用都在录制而非执行，内核消费的是这些字段的缓冲区地址，而不是内容。所以”捕获时展开出来的内容”并不重要，重要的是展开过程中 kv_ca_index、flash_mla_sched_meta_dict 这些字段被绑定到哪个缓冲区，以及这个缓冲区之后能不能被原地重写。flash_mla_sched_meta_dict 是 defaultdict(FlashMLASchedMeta)，首次访问时创建调度元数据对象；捕获时创建的对象连同其显存地址被写进图，之后每次重放前都必须用当前查询的内容原地更新它，而不是新建一个。

get_clean_list 的注释本身点出了这个机制的一半：“Fields to clear (set to None) when reusing a cudagraph; they are recomputed in forward.” 也就是说，这套字段的设计从一开始就是为”复用同一张图”服务的。

捕获产物：一个数据类，五样东西

每张图捕获完成后存成一个 CudaGraphData：

1
@dataclasses.dataclass
2
class CudaGraphData:
3
    graph: torch.cuda.CUDAGraph | None = None
4
    graph_handle: graph_timer.GraphHandle | None = None
5
    output: CausalLMForwardOutput | None = None
6
    input_metadata_source: InputMetadata | None = None
7
    input_metadata_device: InputMetadata | None = None

五样东西各有角色。graph 是实体；graph_handle 是计时器的绑定，非 GPU 平台为 None；output 是输出缓冲的视图；两个 input_metadata 分别是主机侧与设备侧的输入元数据，它们共用同一份按最大档位分配的缓冲区。数据类的 docstring 用一句话概括了这份设计：“Source tensors are on CPU, or share device tensor storage under ZOS. Metadata objects stay separate because forward populates derived fields.” 后半句说的是 compute_attn_info 每次会往元数据对象上挂新的派生字段，所以主机与设备两侧必须是两个独立的 Python 对象——共享存储并不等于共享对象。

视图切片：同一个缓冲区，不同档位

捕获循环对每个档位只做一件事把”全量缓冲区”切出前 batch_size 行：

1
def _cut(x: torch.Tensor, _bs: int = batch_size) -> torch.Tensor:
2
    assert x.shape[0] == self.max_batch_size
3
    return x[:_bs]
4

5
input_metadata_source = input_metadata_source_full.recursive_fn(_cut)
6
input_metadata_device = input_metadata_device_full.recursive_fn(_cut)
7
output = output_full.narrow(batch_size)

input_metadata_source_full 是用 self.max_batch_size 个占位序列一次性构造出来的，设备侧再 .cuda(non_blocking=True) 整体拷一次。之后每个档位只做前缀切片，不再为每个档位单独分配。这样做有三个直接收益：所有档位的输入缓冲区首地址相同，回填逻辑只需要写”全量缓冲区”一次就同时覆盖所有档位；所有档位的输出缓冲都在同一块显存上，回放后的结果不需要额外拷贝；显存开销与档位数无关，只与最大档位有关。

recursive_fn 会递归进入 InputMetadata 的子类实例和字典字段，对每个张量调用 _cut；defaultdict 分支还有专门处理，保证 ks、ke、block_tables 这类按键分桶的字段也被正确切片。字典在捕获时是空的（派生字段刚被清空，kv_ca_index 等还是 None），所以这一步主要作用于 positions、input_ids、block_tables、context_lens 这些基础字段；block_tables 的切片尤其重要，它决定了档位图用到的块表缓冲区行数。

output = output_full.narrow(batch_size) 用的是 narrow 而非切片，返回值是 max_batch_size 缓冲上的前 batch_size 行视图。图内最后一步是 output.copy_(model_output)，所以每次重放的实际输出都落在这块共享缓冲的头部。model_output 本身在拷贝后被 del 掉，_run_model 的返回值只作为拷贝来源；replay 返回的永远是 cuda_graph_data.output，而不是图里那次前向的返回对象。

占位序列：让 padding 不花冤枉钱

补位用的是同一个占位序列对象重复 max_batch_size 次：

1
self.padding_seq_metadata_list = [padding_seq_metadata] * self.max_batch_size

占位序列由 _placeholder_seq_metadata(is_prompt=False) 构造，几个字段的设计目的是”尽量少干活、少占资源”：

1
return SequenceMetadata(
2
    request_id=PLACEHOLDER_REQUEST_ID,
3
    seq_id=0,
4
    is_prompt=is_prompt,
5
    seq_data=SequenceData([0]),
6
    compressed_seq_data=SequenceData([0]),
7
    multi_modal_data=None,
8
    sampling_params=SamplingParams(temperature=0.0),
9
    block_table=[0],
10
    seqlen_q=1,
11
    seqlen_k=1,
12
    last_output_token=0,
13
    last_accepted_token=[],
14
    last_next_n_token=[] if is_prompt else nextn_pad_token_id,
15
)

每行占位各占 1 个 query token、1 个 KV token、1 个物理块；last_next_n_token 用 nextn_pad_token_id 填充，即专门给 MTP 预留的”空草稿” token id。占位请求的采样温度是 0，且 request_id 是一个专门常量，采样与后处理环节能识别并丢弃它。mtp_info_write_slot 没有显式设置，注释说明它默认落在槽位 0，而写槽位的内核会跳过这个槽——这样占位序列不会污染真正的 MTP 信息池。

这里可以看出一个工程取向：padding 不是”算对就行”，而是尽量让它便宜。每个占位 token 仍然要读一个 KV 位置、经过全部网络层，这部分算力无法避免；但它的 KV 只指向 block 0、采样是确定性的、MTP 信息被丢弃，额外开销被压到最低。

select 的二分查找

1
def select(self, total_q: int) -> tuple[int, CudaGraphData]:
2
    index = bisect_left(self.batch_sizes, total_q)
3
    if index == len(self.batch_sizes):
4
        raise ValueError(
5
            f"No captured graph for total_q={total_q}; "
6
            f"maximum captured batch size is {self.batch_sizes[-1]}"
7
        )
8
    batch_size = self.batch_sizes[index]
9
    return batch_size, self.graphs[batch_size]

bisect_left 返回第一个 >= total_q 的档位，也就是”向上取整到最近的档位”。total_q 恰好等于某个档位时取该档位本身，不会多补一行——这是 bisect_left 与 bisect_right 的区别，选对了可以避免不必要的一个占位。

越界时直接抛异常而不是退化为 eager。原因在于前文提到的 max_batch_size = max_num_seqs * (1 + nextn)：调度器保证一步内所有序列的 query token 总数不会超过这个上界。如果越界，说明调度器或 MTP 状态出了不一致，抛异常比静默走 eager 更容易暴露问题。

回填：追加占位、原地写入、取回真实请求数

走图分支的输入准备只有三行有效操作：

1
graph_size, cuda_graph_used = self.cuda_graph_runner.select(total_q)
2
self._prepare_input_tensors(
3
    seq_metadata_list + self.cuda_graph_runner.padding_seq_metadata_list[: graph_size - total_q],
4
    out_dict=cuda_graph_used.input_metadata_source.to_dict(),
5
    metadata_zos=metadata_zos,
6
)
7
input_metadata = self._graph_input_metadata(cuda_graph_used, num_reqs)

第一步，把占位序列追加在真实序列之后，凑成 graph_size 个序列。第二步是关键：out_dict 传入的是那张档位图 input_metadata_source 的 __dict__，_prepare_input_tensors 检测到 out_dict 非空时，会走 prepare_inputmetadata_dict 的”就地填充”分支，把结果写回这个字典并返回 None，而不是新建一个 InputMetadata。第三步，把设备侧元数据对象上的 num_reqs 改成真实请求数：

1
@staticmethod
2
def _graph_input_metadata(cuda_graph_used: CudaGraphData, num_reqs: int) -> InputMetadata:
3
    input_metadata = cuda_graph_used.input_metadata_device
4
    # 实际的请求个数，不包含 pad 的。要和 sampler inputs 对齐，sampler inputs 不包括 pad 的部分。
5
    # 写在 input_metadata 里是因为 cp 的时候只有 driver 有 sampling inputs。
6
    input_metadata.num_reqs = num_reqs
7
    return input_metadata

为什么是改 num_reqs 而不是重新切片？因为图里的内核读的是缓冲区地址，形状由捕获时的 graph_size 固定；num_reqs 只作为元数据字段供前向之后的采样阶段使用。如果这里把它切成 num_reqs，采样阶段拿到的就是一份形状不一致的对象，而采样输入是按真实请求数构造的。注释里括号里的原因更具体：在 CP 场景下，只有 driver 才持有采样输入，follower 的前向阶段需要用 input_metadata.num_reqs 来对齐自己的采样阶段，所以这个值必须写进被各 rank 共享语义所依赖的字段。

“就地填充”是这个设计能成立的核心。_prepare_input_tensors 里：

1
input_metadata_dict = self.model.model_metadata.prepare_inputmetadata_dict(
2
    self.block_size,
3
    seq_metadata_list,
4
    out_dict,
5
    max_num_blocks=self.max_num_blocks,
6
    window_size=self.window_size,
7
    swa_block_size=self.swa_block_size,
8
    metadata_zos=metadata_zos,
9
)
10
if out_dict is not None:
11
    return None

out_dict 就是那张档位图的缓冲区字典，它的值都是从最大档位缓冲上切出的、长度为 graph_size 的前缀视图；prepare_inputmetadata_dict 在 out_dict 非空时会把每个字段写进这个字典里对应的已有张量。序列列表里真实请求在前、占位序列在后，所以真实请求的 total_q 行与补齐的 graph_size - total_q 行都落在档位图访问的地址范围内。这意味着无论调度器这一步给出多少真实 token，回填都是一次写入完成，既不需要重建缓冲区，也不需要额外的拼接张量。

这也解释了 CudaGraphData 为什么同时保存 input_metadata_source 与 input_metadata_device：前者是主机侧（或 ZOS 下的设备共享视图），回填时改写它的 __dict__ 里的张量；后者是设备侧元数据对象，图内内核访问的是它的字段。两者在捕获时由同一份缓冲区切出，回填时只改主机侧，设备侧要读取的主机侧数据通过一次拷贝或共享存储传过去（见下一节的 _run_model）。

input_on_device 与两种拷贝方式

_run_model 在捕获与重放时都被调用，它处理了主机侧到设备侧的输入同步：

1
def _run_model(model, kv_caches, input_metadata_source, input_metadata_device, copy_inputs):
2
    cp_group = PG.context_parallel_group()
3
    if copy_inputs and cp_group.rank() == 0:
4
        input_metadata_device.recursive_copy_(input_metadata_source)
5
    if cp_group.size() > 1:
6
        input_metadata_device.broadcast()
7
    return model(kv_caches=kv_caches, input_metadata=input_metadata_device)

copy_inputs = not self.input_on_device，而 input_on_device 等于 ZOS 开关。ZOS 关闭时，input_metadata_source 在 CPU 上，每一步需要一次 H2D 拷贝到设备侧元数据；ZOS 开启时，两者共享设备存储，不需要这一步拷贝。这正是 CudaGraphData 的 docstring 里”Source tensors are on CPU, or share device tensor storage under ZOS”的含义。

recursive_copy_ 只拷贝张量字段与字典里的张量字段，非张量字段（如 num_reqs、is_prompt）不拷贝——因为 num_reqs 是直接写在设备侧对象上的，而 is_prompt 在图路径下恒为 False。这种”只拷张量”的策略保证了拷贝量最小。

一次重放之后，主干的控制流

replay 本身只有三行，但它的调用处揭示了与内核计时器的耦合：

1
def replay(self, cuda_graph_data: CudaGraphData) -> CausalLMForwardOutput:
2
    if cuda_graph_data.graph_handle is not None:
3
        cuda_graph_timer.flip_event_buffer(cuda_graph_data.graph_handle)
4
    cuda_graph_data.graph.replay()
5
    return cuda_graph_data.output

flip_event_buffer 是图级计时器的乒乓缓冲切换，在每次重放前把图内的事件记录节点指向下一组事件缓冲。对非 GPU 平台（NPU）或未启用计时器时，graph_handle 为 None，跳过。

execute_main_model 包住这段调用的上下文是另一处设计：

1
with (
2
    cpu_timer.record("launch_forward", record_name=f"launch_forward_N{...}_B{total_q}"),
3
    cuda_timer.record("device_forward", scale=1),
4
    forbid_cuda_sync(),
5
):
6
    if run_cuda_graph:
7
        model_output = self.cuda_graph_runner.replay(cuda_graph_used)
8
    else:
9
        model_output = self.model(kv_caches=kv_caches, input_metadata=input_metadata)

forbid_cuda_sync 打开 torch 的同步检测模式，任何 CPU-GPU 同步都会抛异常。图路径尤其需要这条保护：回填、拷贝、重放这一整条链路都要保持异步，任何一次意外的同步都会把”每步只有几毫秒”的解码变成”每步几十毫秒”。

图路径的宏观收益可以用仓库给出的 decode 曲线看一个侧面：下面这张 pareto 图横轴是每用户解码速度，纵轴是输出吞吐，曲线是在 64 张 SM100 上、按同一 scale-up 域内测得的（README 注明 prefill 用 4 卡、decode 用 64 卡）。本文不深入这张图的具体数字，只借用它说明一个事实：解码吞吐与单用户速度之间存在明确的前沿，而这条前沿是”每步执行开销”与”每步有效计算”之比的直接体现。图存在的意义，就是把每一步的提交开销压下去，让同一硬件在同样的前沿上整体右移。

原地更新的对象清单

DeepInfer 在 ZOS 解码路径上有大量”原地写”的内核，tilelang_ops 里有几个从名字就能看出用途的模块：

这套设计的核心原则可以从 compute_attn_info 在 attn.py 里的调用位置看出来。注意力层的 forward 开头有一行：

1
block_size = kv_cache.ca_cache.shape[1] * self.kv_compress_ratio if kv_cache.ca_cache is not None else None
2
input_metadata.compute_attn_info(self.max_model_len, self.window_size, block_size)

compute_attn_info 是幂等的：派生字段已有值时直接跳过，只有为 None 时才展开。这解释了为什么捕获前要清空这些字段：清空后捕获时展开一次，展开结果被绑定到图里的缓冲区；之后每次重放，字段已有值，方法不做任何事，图内内核读到的就是”上一次展开写入缓冲区的内容”——而这个内容在每次重放前已经被 ZOS 的原地更新内核改成了这一步的正确值。

decode 分支就是这套机制的落点：

1
if self.kv_swa_index is None:
2
    swa_block_size = get_kvstore_swa_block_size(window_size)
3
    self.kv_swa_index = get_kv_swa_index_decode(
4
        self.kv_swa_index_offset_len, window_size, swa_block_size
5
    )
6
if self.kv_ca_index is None and block_size is not None:
7
    self.kv_ca_index = {}
8
    for ratio, offset_len in self.kv_ca_index_offset_len.items():
9
        n = int(ratio)
10
        width = (max_model_len // n + 127) // 128 * 128
11
        self.kv_ca_index[ratio] = get_kv_ca_index_decode(
12
            offset_len, width, block_size // n, self.block_tables
13
        )

get_kv_swa_index_decode 与 get_kv_ca_index_decode 是 TileLang 内核，它们把 [start, len] 形式的紧凑描述展开成稠密索引矩阵。offset_len 形状是 [total_q, 2]，内容每步不同；内核原地写入 kv_swa_index 与 kv_ca_index 的前 total_q 行。捕获时创建的矩阵是 [max_batch, width]，档位图看到的是它的前缀；内核按本步 offset_len 的行数原地写入这张矩阵，写入范围与档位图访问的行范围一致。

width 的表达式 (max_model_len // n + 127) // 128 * 128 值得单独解释。注释写着”the sparse-attention kernel requires padding to 128”。压缩比越大，可见的压缩 token 越少，max_model_len / N 是可见压缩键数的上界；补齐到 128 的倍数是因为稀疏注意力内核按 128 列分块处理索引矩阵，列数不齐会让内核多做一次不满的块。

为什么必须有 ZOS

一个自然的问题：如果不开 ZOS，这套原地更新不行吗？答案是可以，但会回到”从 CPU 拷”的老路。非 ZOS 路径下，input_metadata 在 CPU 上构造，每一步用 recursive_fn(lambda x: x.cuda(non_blocking=True)) 整体拷到设备。对于图路径，recursive_copy_ 把整份元数据里的张量从主机拷贝到图缓冲区；kv_ca_index_len、kv_swa_index 等派生字段在 CPU 侧展开后再拷，拷完再被图内核使用。

问题在于 CPU 侧展开这些字段是一串 Python 循环，batch 大、压缩比种类多时，展开本身就要几毫秒；展开完还要再拷一次 H2D。解码步的 GPU 时间是毫秒量级，CPU 侧的展开与拷贝会把整个流水线拖住。ZOS 把这些展开操作挪到 GPU 上的小内核里，与其他的准备步骤共享同一个设备侧缓冲区，让每一步的 CPU 侧工作只剩下发指令。

forbid_cuda_sync 与 ZOS 的组合还有一条正确性边界：ZOS 解码时 seqlens_q 依赖上一步采样结果，CPU 侧拿到的可能是滞后一步的旧值，所以 worker._correct_seqlens_qk 用 compute_seqlens_q_by_confidence_rates 与 compute_seqlens_k_zos 在 GPU 上修正，并在 MTP 变长模式下把修正后的 seqlens_q 拷回主机一份（d2h_evt），供采样阶段使用。这条时序在调度篇里已经展开，这里只需要记住：图内所有地址固定，凡是依赖”上一步结果”的内容，必须先用小的设备内核原地更新，再重放。

计时代价：图内不能放同步

capture 开头有一段降级：

1
if not cuda_graph_timer.supported():
2
    logger.info("cuda_graph_timer probes not supported on this platform.")

图级计时器用成对的事件记录节点包住图内各段（注意力、MoE 等），用乒乓事件缓冲在多次重放间轮换，以便在不带同步的情况下读出每一段的耗时。它在捕获时把事件记录节点嵌进图，重放前通过 set_event 把节点指向另一组事件，重放后从事件里读时间。这套机制要求捕获时带 keep_graph=True，并且要求 register_graph 时按事件身份匹配节点：

1
mismatched = start_nodes.keys() ^ end_nodes.keys()
2
assert not mismatched, f"probes missing a start or end record node: {mismatched}"

这个断言看起来是在检查计时器，实际上是在为图捕获本身做健康检查：每个计时探针都必须同时有起始与结束记录节点，否则说明某个探针的结束事件落在了图外，进而说明图被意外截断。对生产系统而言，“图被截断”往往意味着某段代码走出了当前的捕获流。

输入形状：按请求数，而不是按 token 数

DFlashCudaGraphRunner.__init__ 的档位集合是：

1
self.batch_sizes = build_graph_batch_sizes(
2
    max_num_seqs, envs.DEEPINFER_DFLASH_CUDAGRAPH_BATCH_SIZE_INTERVAL
3
)

max_num_seqs 而不是 max_num_seqs * (1 + nextn)。草稿阶段每步处理的单位是”请求”：每个请求一段草稿。主图处理的单位是”token”，因为验证阶段一条序列可能一次验证多个 token。两种图对 batch 的度量口径不同，档位含义也不同，不能共用。

从主图借用元数据、借用内存池与流

capture_dflash_model 里有一处值得细看：

1
dflash_model: DFlash = self.model.MTP_module
2
free_mem_before, _ = torch.cuda.mem_get_info()
3
self.dflash_cuda_graph_runner.capture(
4
    dflash_model,
5
    kv_caches,
6
    main_input_metadata=main_graph_runner.graphs[
7
        max(main_graph_runner.graphs)
8
    ].input_metadata_device,
9
    graph_memory_pool=main_graph_runner.graph_memory_pool,
10
    graph_stream=main_graph_runner.graph_stream,
11
)
12
dflash_model.decode_graph_runner = self.dflash_cuda_graph_runner
13
free_mem_after, _ = torch.cuda.mem_get_info()
14
graph_mem_delta = (free_mem_before - free_mem_after) / 1024**2

三样东西被主跑器借给了 DFlash：最大主图的设备侧元数据、主图的内存池、主图的捕获流。

借用最大主图的元数据，是因为草稿阶段的输入要基于”主图这一轮的输入”来构造。DFlash.run_drafts 接收 main_input_metadata，并用 build_dflash_input 从它派生草稿阶段的输入元数据：

1
dflash_input_metadata = main_input_metadata.build_dflash_input(
2
    decode_input.first_token_positions,
3
    last_valid_offset,
4
    self.dflash_block_size,
5
    self.window_size,
6
)

这份派生需要主图元数据里的位置、块表等字段；借最大的那份是为了覆盖主图所有档位的字段范围，因为草稿阶段在任何一个主图档位都可能被触发。

借用主图的内存池，是为了让两张图共享同一块捕获内存，而不是各自单独撑起一个池。捕获内存池的大小由第一张捕获的图决定，共享池意味着总显存开销更可控。借用同一个捕获流，则是为了让捕获顺序对齐：主图必须先捕获（capture_dflash_model 里断言了 main_graph_runner.graphs 非空），DFlash 图在同一个流上捕获，池地址的分配顺序与前文讨论的”从大到小”原则在同一个流上得到延续。

graph_mem_delta 用 torch.cuda.mem_get_info 的前后差值来计算 DFlash 图额外的显存占用，并写进日志。这是一个诚实的度量方式：图捕获时图内分配的中间张量来自 CUDA Graph 的专用池，不进入 PyTorch 默认分配器的统计视野，所以只能用设备剩下的空闲内存来推。日志里同时记下了 rank，方便在 EP、CP 场景下对多卡上的差异做归因。

草稿图的捕获形状与选择

DFlashCudaGraphRunner.capture 与主图跑器同构，但有几处不同：

1
self.decode_input_padding = MTPDecodeInput.build_padding(
2
    self.max_num_seqs, module.noise_token_id, self.do_sample
3
)
4
decode_input_full = MTPDecodeInput.build_padding(
5
    self.max_num_seqs, module.noise_token_id, self.do_sample
6
)
7
output_full = module.allocate_output(self.max_num_seqs) if self.do_sample else None
8

9
for num_seqs in reversed(self.batch_sizes):
10
    decode_input = decode_input_full.cut(num_seqs)
11
    output = output_full.narrow(num_seqs) if self.do_sample else None
12

13
    graph_stream.wait_stream(torch.cuda.current_stream())
14
    with torch.cuda.stream(graph_stream):
15
        module.run_drafts(main_input_metadata, decode_input, nextn_kv_caches)
16

17
    graph = torch.cuda.CUDAGraph()
18
    with torch.cuda.graph(graph, pool=graph_memory_pool, stream=graph_stream,
19
                          capture_error_mode="thread_local"):
20
        dflash_output = module.run_drafts(main_input_metadata, decode_input, nextn_kv_caches)
21
        if self.do_sample:
22
            output.copy_(dflash_output)

三个不同点。第一，占位输入是 MTPDecodeInput.build_padding，它的 first_token_positions 全填 -1，注释写着”-1 marks CUDA graph padding”；driver 侧用 golden_token_id 构造占位，golden_token_id 就是模块的 noise_token_id（扩散草稿的噪声起始 token）。第二，只有 driver worker 做采样（do_sample = self.is_driver_worker），follower 只跑前向，所以 follower 的图没有输出缓冲，output_full 为 None。第三，草稿图没有预热后清空派生字段的步骤。前向 run_drafts 内部从主图元数据派生草稿元数据，每次派生生成的元数据结构可能不同，但图里内核访问的是这些对象的地址；占位输入在捕获时就把结构确定下来，运行时用 copy_dflash_decode_input 原地更新即可。

replay 的设计是这个跑器里最有代表性的一段：

1
def replay(self, decode_input: MTPDecodeInput) -> MTPPredictInfo | None:
2
    req_cnt = decode_input.req_cnt
3
    _, data = self.select(req_cnt)
4
    copy_dflash_decode_input(data.decode_input, decode_input, self.decode_input_padding)
5

6
    data.graph.replay()
7
    if self.do_sample:
8
        return data.output.narrow(req_cnt)  # remove CUDA graph padding

入口先按请求数选档位，再原地拷贝输入到捕获时用到的 decode_input 缓冲区，decode_input_padding 提供占位内容——copy_dflash_decode_input 只把真实 req_cnt 个请求写进去，剩余行保持占位。重放后，输出用 narrow(req_cnt) 切回真实请求数，注释直白地写着这是”去除 CUDA Graph 的 padding”。填充在输入侧补进来、在输出侧切出去，这是 padding 策略最干净的形态：图内的计算是固定的，但对外界暴露的数据接口始终按真实请求数。

decode_forward 里还有一层封装：

1
@cpu_timer.record("launch_dflash_decode_forward")
2
def decode_forward(self, main_input_metadata, nextn_kv_caches, decode_input):
3
    if self.decode_graph_runner is not None:
4
        return self.decode_graph_runner.replay(decode_input)
5
    return self.run_drafts(main_input_metadata, decode_input, nextn_kv_caches)

CPU timer 用装饰器在每次重放时打点，名称为 launch_dflash_decode_forward，它与主图的 launch_forward 配合，让调度器能分别观察两段 GPU 工作之间的 CPU 间隙。这个位置暗示了一个调度事实：草稿阶段与验证阶段的启动是串行的，草稿重放完成之后目标模型的重放才会被发出；两段之间若有 CPU 间隔，GPU 会空转。

草稿图的业务位置

草稿图的业务位置可以用下图说明，它是 DFlash 论文中的推理设计图：草稿头在目标模型的一步 forward 之后、下一轮验证之前运行，依赖目标模型产出的 hidden state 来生成多 token 草稿，然后由目标模型在下一步并行验证。

DeepInfer 的 DFlash 图把草稿头这一整段（从投影、MTP 层前向、采样到写信息池）都捕获进了图。DFlash.run_drafts 的 docstring 写得很明确：“With the DFlash CUDA graph this is captured as a unit: the driver’s broadcast (followers receive into buffers allocated here, whose addresses the graph fixes), the fixed-stride metadata, the draft layers and the sampling (which includes the pool write) are all replayed.”

这意味着在 CP 大于 1 的部署里，图里包含通信。driver 在草稿阶段广播 last_valid_offset 与 golden_tokens 给 follower，follower 接收进图内缓冲区。这又回到了前文提到的问题：图内的 NCCL 通信在销毁进程组前必须先释放图，DFlash 图同样纳入 destroy_cuda_graphs 的范围。

“fixed-stride metadata” 这个词则指向 MTPDecodeInput 与 main_input_metadata.build_dflash_input 的构造策略：草稿阶段每个请求产出一段定长的草稿，元数据的 stride 由 dflash_block_size 与 max_MTP_value 决定，不随本步接受了几张草稿而改变。这也是为什么草稿图可以存在而”变长验证”部分不行——变长的部分在主图上，在那里通过主图档位和占位补位处理。

性能曲线测量：用重放本身去测量

profile_decode_perf 的测量方式值得拆开：

1
graph_runner = self.cuda_graph_runner
2
group = PG.expert_parallel_group()
3
decode_perf_cache = DecodePerfCache(
4
    self.model_config.model,
5
    self.parallel_config,
6
    graph_runner.batch_sizes,
7
    group,
8
)
9
cached_results = decode_perf_cache.load_decode_perf_cache()
10
if cached_results is not None:
11
    logger.info("Loaded cached decode performance curve: "
12
                + ", ".join(f"(bs={bs} ms={ms:.3f})" for bs, ms in cached_results))
13
    return cached_results

缓存命中就直接返回。缓存键 _get_hash_key 用 json.dumps 以 sort_keys=True 序列化 {time_unit, model, batch_sizes, parallel_config}，再取 SHA-256。显式把 time_unit: "ms" 放进键里，注释写着是为了与旧版以 SPS（每秒）为单位记录的缓存区分，避免读到旧口径的数据后误用。

未命中时：

1
input_metadata_source = graph_runner.graphs[graph_runner.batch_sizes[-1]].input_metadata_source
2
input_metadata_source.valid_token_mask.fill_(True)
3
input_metadata_source.gate_force_random.fill_(True)
4

5
NUM_WARMUP = 1
6
NUM_MEASURE = 5
7
NUM_TRIM = 1
8

9
results = []
10
for batch_size in reversed(graph_runner.batch_sizes):
11
    graph = graph_runner.graphs[batch_size].graph
12

13
    for _ in range(NUM_WARMUP):
14
        graph.replay()
15
        dist.all_reduce(torch.zeros(1, device="cuda"), group=group)
16

17
    events = []
18
    for _ in range(NUM_MEASURE):
19
        start_event = torch.cuda.Event(enable_timing=True)
20
        end_event = torch.cuda.Event(enable_timing=True)
21
        start_event.record()
22
        graph.replay()
23
        end_event.record()
24
        dist.all_reduce(torch.zeros(1, device="cuda"), group=group)
25
        events.append((start_event, end_event))

几个细节值得讲。

valid_token_mask.fill_(True) 和 gate_force_random.fill_(True) 只写在最大档位的缓冲区上，注释解释得很清楚：“各 batch size 共享最大 batch buffer 的切片视图，填一次即可作用于所有 graph.”这与捕获时按前缀切片的设计是同一个逻辑的两面：捕获时共享缓冲区，所以这里只需要填一次。

gate_force_random 是一个有趣的字段：它的注释写着”True: replace the token’s gate logits with random values to balance MoE load.” 即强制把 MoE 门控 logits 替换为随机值以平衡负载。在测量性能曲线时把它设为真，意味着测量的是”最平衡 MoE 负载”下的耗时，而不是”当前实际路由”下的耗时。这与前面 profile_run 里的 DEEPINFER_FORCE_IMBALANCE_RATIO = 1.8 正好成对：前者测显存占用时故意制造不平衡，后者测耗时曲线时故意做平衡。两个测量都在构造”最不利”或”最典型”的输入，而不是随手用当前状态测一遍。

每次测量后紧跟一个 dist.all_reduce(torch.zeros(1), group=expert_parallel_group())。这个 all_reduce 的用处是在 EP 组上插入一道逻辑屏障：图重放本身是异步的，紧随其后的空 all_reduce 必须等本 rank 的重放把图内的集合通信都执行完才能完成，而它又要求 EP 组内所有 rank 都达到同一位置，于是下一次重放不会在任一 rank 上一张图还没跑完时发起。测量循环因此始终按”整组同步”的节奏推进，各档位的耗时读数也在同一批 rank 状态上取得。

每次测量使用独立的 start_event/end_event 对，而不是复用一对。这样每次测量的时间事件彼此独立，避免后一次记录的结束时间覆盖前一次。事件全部 synchronize 之后得到一组毫秒数，排序、去掉最小和最大各一个（NUM_TRIM = 1），取均值。这个截尾均值策略在系统测量里很常见：去掉最快的样本（可能因为某些部件还在热状态）与最慢的样本（可能被系统噪声干扰）。

最后各 rank 的均值再做一次 SUM all_reduce 并除以组大小：

1
sum_time = torch.full((1,), avg_time_ms, dtype=torch.float32, device="cuda")
2
dist.all_reduce(sum_time, op=dist.ReduceOp.SUM, group=group)
3
avg_time_ms = sum_time.item() / group.size()

这一步得到的是 EP 组内各 rank 的平均耗时。为什么不是最大值？因为图内的 EP 通信是同步的：所有 rank 在同一张图里互相等待，图重放的完成时间对每一个 rank 而言都等于”最慢的那个 rank 完成的时间”。因此每个 rank 自己测出来的时间其实已经包含了其他 rank 的等待；再取平均反而会稀释掉最慢 rank 的影响。从记录的角度看，这个平均更像是一个”群体视角的典型值”。这条实现选择本身有讨论空间，但它的意图很清楚——调度器需要的是一个能反映群体上单步耗时的代表值，而不是某一 rank 的局部值。

DecodePerfCache 的写盘也体现了生产化的一面：

1
timestamp = int(time.time())
2
file_uuid = uuid.uuid4().hex
3
filename = f"{self.hash_key}.{file_uuid}.{timestamp}"
4
tmp_path = os.path.join(self.cache_dir, f"{filename}.tmp")
5
final_path = os.path.join(self.cache_dir, filename)
6
os.makedirs(self.cache_dir, exist_ok=True)
7
with open(tmp_path, "w") as f:
8
    json.dump(results, f, separators=(",", ":"))
9
    f.flush()
10
    os.fsync(f.fileno())
11
os.replace(tmp_path, final_path)

文件名由哈希键、UUID、时间戳三段构成。哈希键保证不同模型/档位/并行配置互不覆盖；UUID 保证同一哈希键的并发写入互不覆盖；时间戳支持按时间选择最新的有效文件。写入用临时文件加 os.replace，保证读取方永远看不到半截文件。清理阈值比读取 TTL 宽 2 小时：

1
# 删除阈值比读取 TTL 宽 2h：文件先对读取过期，再被清理，避免并发进程读到一半被删
2
self._cleanup(now_s, self.ttl_s + 2 * 60 * 60)

这句注释是一个典型的分布式系统细节：读取过期与删除过期分离，避免”清理”与”读取”并发时读到一半文件被删。只有 expert_parallel 组的 rank 0 做读写，结果用 broadcast_object_list 广播给组内其他 rank，这也对应了”曲线是群体属性，不必各 rank 各测一遍”的设定。