DeepInfer 源码拆解（二十九）：CUDA Graph 内的计时器与 decode 性能曲线——图里的耗时怎么量出来、又被谁消费

CPU 计时为什么还要参与

CPUTimer 的实现同时向两个系统报告：

1
class CPUTimer(BaseTimer):
2
    @contextmanager
3
    def record(self, name, enable=True, record_name=None):
4
        if not enable:
5
            yield
6
            return
7
        t0 = time.perf_counter()
8
        with torch.profiler.record_function(f"deepinfer:{record_name or name}"):
9
            yield
10
        elapsed = time.perf_counter() - t0
11
        self.add_time(name, elapsed)

torch.profiler.record_function 的用途是把这段 CPU 区间标成一个可被 PyTorch profiler 识别的范围（range）。也就是说，同一段 CPU 时间有两条去向：一是进入 cpu_timer.times[name] 被统计日志消费；二是在 PyTorch profiler 里出现为 deepinfer:<name> 的 range，供 trace 查看。Worker 的 _post_process_step 就用它包裹同步、flush 与后处理三段：

1
@cpu_timer.record("worker.post_process")
2
def _post_process_step(self, step_state):
3
    with cpu_timer.record("worker.post_process.sync"):
4
        step_state.sample_end_evt.synchronize()
5
    with cpu_timer.record("worker.post_process.graph_timer_flush"):
6
        cuda_graph_timer.flush_out()
7
    with cpu_timer.record("worker.post_process.do_process"):
8
        self.model_runner.post_process(...)

CPUTimer 的量纲是「秒」，而 GPU 计时器的 flush_out 返回的是「秒」（毫秒除以 1000），两者对齐后才能让 StatsLogger 把 CPU 与 GPU 两段时间放在同一条报告里比较。StatsLogger 通过 get_timers() 拿到全部三个计时器：

1
@functools.lru_cache
2
def get_timers() -> list[tuple[str, BaseTimer]]:
3
    return [("cpu", cpu_timer), ("cuda_graph", cuda_graph_timer), ("cuda", cuda_timer)]

1
# deepinfer/engine/stats_logger.py（节选）
2
def _merge_cuda_timers():
3
    ...  # 遍历 get_timers()，把 cuda 与 cuda_graph 的 times/called/max/latest 合成一份

合并在 StatsLogger 里做，意味着日志层看到的是一张统一的「组件 → GPU 时间」表，不必关心某个组件这轮用的是图内探针还是图外事件。日志这条链路已经在第二十三篇拆过，这里只说它对计时器的要求：图内计时器必须能在任意时刻被安全地读出已完成的那一轮，且不阻塞当前正在执行的图。

事件如何被录进图

CUDA 9.0 引入 CUDA Graph 时，就把「在图捕获期间记录的事件」当作图的一部分来对待。NVIDIA 文档里特别说明：在流的捕获期间调用 cudaEventRecord 会在图中创建一个事件记录节点；如果想让一个在捕获之前就已存在、或者不属于本图的 event 也能在捕获时被记录，需要使用 cudaEventRecordExternal 标志。

图源：NVIDIA CUDA Programming Guide 4.2 节 “CUDA Graphs” 中的 kernel 节点示意图。

csrc/graph_timer/graph_timer_cuda.cpp 里，EventHandle::record 就根据一个布尔值选择两种标志：

1
void EventHandle::record(int64_t stream, bool capture_external) const {
2
    unsigned int flags = capture_external ? cudaEventRecordExternal : cudaEventRecordDefault;
3
    C10_CUDA_CHECK(cudaEventRecordWithFlags(
4
        reinterpret_cast<cudaEvent_t>(raw_event_),
5
        reinterpret_cast<cudaStream_t>(stream),
6
        flags));
7
}

为什么必须是 cudaEventRecordExternal？因为这里的事件对象是在第一次预热前创建好的（下面会讲），图捕获时事件已经存在。如果不用 external 标志，捕获会报错或产生与预期不符的节点。更关键的是，后面要按「事件身份」在捕获后的图里找到对应节点——只有当记录下来的节点确实指向这个事件对象，匹配才成立：

1
EventHandle EventRecordNodeHandle::get_event() const {
2
    cudaEvent_t event;
3
    C10_CUDA_CHECK(cudaGraphEventRecordNodeGetEvent(
4
        reinterpret_cast<cudaGraphNode_t>(raw_node_), &event));
5
    return EventHandle(reinterpret_cast<int64_t>(event));
6
}
7

8
void EventRecordNodeHandle::set_event(const EventHandle &event) {
9
    C10_CUDA_CHECK(cudaGraphExecEventRecordNodeSetEvent(
10
        reinterpret_cast<cudaGraphExec_t>(raw_graph_exec_),
11
        reinterpret_cast<cudaGraphNode_t>(raw_node_),
12
        reinterpret_cast<cudaEvent_t>(raw_event_)));
13
}

cudaGraphEventRecordNodeGetEvent 从一个图节点读出它记录的 event；cudaGraphExecEventRecordNodeSetEvent 则在已实例化的图执行实例上，把某个节点改成记录另一个 event。第一招在捕获后使用，用来建立「节点 ↔ 探针」映射；第二招在每次重放前使用，用来切换本轮要用的 event。

捕获时保留原始图

一个容易被忽略的前提是：cudaGraphExecEventRecordNodeSetEvent 需要的是图节点句柄和图执行实例。PyTorch 的 torch.cuda.CUDAGraph 默认在 capture_end() 之后只暴露实例化后的 cudaGraphExec_t，原始 cudaGraph_t 会被释放。DeepInfer 用 keep_graph=True 保住原始图：

1
# deepinfer/worker/cuda_graph_runner.py（节选）
2
graph = torch.cuda.CUDAGraph(keep_graph=True) if PLATFORM.is_gpu else torch.cuda.CUDAGraph()
3
with torch.cuda.graph(graph, pool=self.graph_memory_pool,
4
                      stream=self.graph_stream, capture_error_mode="thread_local"):
5
    model_output = _run_model(model, kv_caches, input_metadata_source, input_metadata_device, ...)
6
    output.copy_(model_output)
7
    del model_output
8
if PLATFORM.is_gpu:
9
    graph.instantiate()

注释写得很清楚：keep_graph=True keeps the raw cudaGraph_t alive so cuda_graph_timer can match its event-record nodes for ping-pong buffer. 而 NPU 分支只保留实例化图，没有 keep_graph 参数——这直接导致 NPU 上 graph_timer 只能退化为桩实现，后面单独说。

GraphHandle 通过 PyTorch 的 raw_cuda_graph() 与 raw_cuda_graph_exec() 拿到两个原始句柄：

1
GraphHandle::GraphHandle(py::object graph) : graph_(std::move(graph)) {
2
    raw_graph_ = graph_.attr("raw_cuda_graph")().cast<int64_t>();
3
    raw_graph_exec_ = graph_.attr("raw_cuda_graph_exec")().cast<int64_t>();
4
}

再用 cudaGraphGetNodes 枚举全部节点，按 cudaGraphGetNodeType 过滤出事件记录节点：

1
std::vector<EventRecordNodeHandle> GraphHandle::get_event_record_nodes() const {
2
    cudaGraph_t raw_graph = reinterpret_cast<cudaGraph_t>(raw_graph_);
3
    size_t num;
4
    C10_CUDA_CHECK(cudaGraphGetNodes(raw_graph, nullptr, &num));
5
    std::vector<cudaGraphNode_t> raw_nodes(num);
6
    if (num > 0) {
7
        C10_CUDA_CHECK(cudaGraphGetNodes(raw_graph, raw_nodes.data(), &num));
8
    }
9
    std::vector<EventRecordNodeHandle> nodes;
10
    for (cudaGraphNode_t raw_node : raw_nodes) {
11
        cudaGraphNodeType type;
12
        C10_CUDA_CHECK(cudaGraphNodeGetType(raw_node, &type));
13
        if (type != cudaGraphNodeTypeEventRecord) continue;
14
        nodes.push_back(EventRecordNodeHandle(reinterpret_cast<int64_t>(raw_node), raw_graph_exec_));
15
    }
16
    return nodes;
17
}

注意每个 EventRecordNodeHandle 都同时带了节点句柄与执行实例句柄。这两个句柄的生命周期不同：节点属于原始图，执行实例属于实例化结果；换绑事件要两个都用。

匹配：探针名 ↔ 起止事件 ↔ 两个节点

捕获前，模型层里的 cuda_graph_timer.record(name) 每次都会走到 record 的上下文管理逻辑。它的关键约定写在注释里：一张图里每个探针最多被调用一次。原因是每个探针名对应一堆固定的事件对象，捕获时录进图的节点数必须与之一致；若某个探针在一张图里被调用两次，就会产生四个节点（两次起、两次止），匹配逻辑无法区分。

record 的实现：

1
@contextmanager
2
def record(self, name, enabled=True, scale=None):
3
    if not enabled or not self.supported():
4
        yield
5
        return
6
    if name not in self.events[0]:  # record in warmup stage
7
        from deepinfer._C.graph_timer import EventHandle
8
        self.events[0][name] = (EventHandle(), EventHandle(), scale)
9

10
    stream = torch.cuda.current_stream().cuda_stream
11
    is_capturing = torch.cuda.graphs.is_current_stream_capturing()
12
    self.events[0][name][0].record(stream, is_capturing)
13
    yield
14
    self.events[0][name][1].record(stream, is_capturing)

两个细节值得展开：

• 事件在预热阶段创建。register_graph 前的每一次 record 如果发现 events[0] 里没有这个探针，就创建一对新 event。捕获之前必定有一次 warmup 前向（CudaGraphRunner.capture 里 _run_model 在原始 stream 上跑一遍），此时 is_capturing=False，事件以普通方式记录在流上——这些记录在流上的事件对性能没有任何影响，因为没有同步。
• is_capturing 在捕获阶段为真。此时同一对事件被以 external 标志再记录一次，成为图节点。事件对象与预热阶段是同一对，这一点保证了后续按事件身份匹配。

捕获完成后，CudaGraphRunner.capture 立刻调用 register_graph：

1
if cuda_graph_timer.supported():
2
    graph_handle = graph_timer.GraphHandle(graph)
3
    cuda_graph_timer.register_graph(graph_handle)

register_graph 做的事可以概括成「反查表」：

1
def register_graph(self, graph: GraphHandle) -> None:
2
    if not self.events[0]:
3
        return
4
    start_event_to_probe = {}   # buffer-0 start event -> probe name
5
    end_event_to_probe = {}     # buffer-0 end event -> probe name
6
    for name, (start_event, end_event, _) in self.events[0].items():
7
        start_event_to_probe[start_event] = name
8
        end_event_to_probe[end_event] = name
9

10
    nodes = graph.get_event_record_nodes()
11
    start_nodes = {}
12
    end_nodes = {}
13
    for node in nodes:
14
        event = node.get_event()
15
        if event in start_event_to_probe:
16
            start_nodes[start_event_to_probe[event]] = node
17
        elif event in end_event_to_probe:
18
            end_nodes[end_event_to_probe[event]] = node
19

20
    mismatched = start_nodes.keys() ^ end_nodes.keys()
21
    assert not mismatched, f"probes missing a start or end record node: {mismatched}"
22
    if not start_nodes:
23
        return
24

25
    record_nodes = {}
26
    for name, start_node in start_nodes.items():
27
        if name not in self.events[1]:
28
            from deepinfer._C.graph_timer import EventHandle
29
            self.events[1][name] = (EventHandle(), EventHandle(), self.events[0][name][2])
30
        record_nodes[name] = (start_node, end_nodes[name])
31

32
    self.graph_bindings[graph] = _CUDAGraphTimerBinding(record_nodes=record_nodes)

这段代码里有三处设计值得细读：

1. 按事件身份匹配，而不是按节点顺序。一张 decode 图里可能同时有几十个探针节点（多个同构组的注意力、MLP、索引器），cudaGraphGetNodes 返回的顺序不保证与捕获顺序一致。用事件句柄做哈希键，EventHandle 的 __eq__/__hash__ 在 C++ 侧是裸指针比较与哈希，速度 O(1)。
2. 起止节点必须成对。mismatched = start_nodes.keys() ^ end_nodes.keys() 是对称差的集合运算；任何探针只有起节点没有止节点（或反之），说明捕获时某个 record 的 enabled 状态在起止之间发生了变化，这属于逻辑错误，直接断言失败。
3. buffer 1 的事件在匹配时创建。buffer 0 的事件被「烧进」了图（捕获时以 external 方式记录），后续每次 replay 都将使用 buffer 0 的事件对象；为了让连续两轮 replay 的计时目标互不干扰，还需要一套 buffer 1 的事件。这套事件只在图实例上通过 set_event 绑定，不进入原始图，所以创建时机放在 register_graph。

为什么不在 replay 之后马上同步

如果每轮 replay 之后都 synchronize 来读时间，就会把流水线打断：GPU 必须跑完才能开始下一轮。DeepInfer 在 ZOS（零开销调度）模式下刻意保持流水线不中断。所以 flush 放在 _post_process_step 里，而且在 sample_end_evt.synchronize() 之后：

1
with cpu_timer.record("worker.post_process.sync"):
2
    step_state.sample_end_evt.synchronize()
3
with cpu_timer.record("worker.post_process.graph_timer_flush"):
4
    # flush the time of completed graph, no effect on the running graph.
5
    cuda_graph_timer.flush_out()

注释里的「no effect on the running graph」正是关键：因为采样完成事件已经同步过，说明上一轮（即被采样事件覆盖的那一轮图）已经跑完。此时读上一轮的事件不需要额外等待；而当前正在 GPU 上执行的下一轮图，用的是另一个 buffer，flush 不影响它。

如果调换这两个调用的顺序，就会产生两种问题：其一，sample_end_evt 之前 GPU 可能还在跑图，事件不同步会造成 elapsed_time 未定义或阻塞到图结束；其二，pending_buffers 弹出与当前轮的 flip 状态会错位，导致下一次 flip 的断言 len(pending_buffers) < 2 偶尔失败。

popleft 与 FIFO 语义

flush_out 取的是 pending_buffers 的队头，而不是队尾。flip_event_buffer 往队尾追加，因此这是一个先进先出队列：本轮 flush 的总是「最早被绑定的那个 buffer」，即「最早一次 replay」的那一组事件。若 ZOS 把 replay 提前了一拍，队头依然对应较早的那次 replay，时间归属不会错。

-part- 后缀与 scale

模型层里注意力探针命名是 f"attn-part-{self.layer_idx}"，MLP 是 "mlp"，索引器是 f"decode_index_attn-part-{self.layer_idx}"。flush_out 看到名字里有 -part- 时，会把它拆成前缀（attn 或 decode_index_attn）并累加：

1
if "-part-" in name:
2
    parts.append((name.split("-part-")[0], t))

这样做有两个效果：

• 日志只暴露 attn、decode_index_attn、mlp 这类逻辑组件名，不暴露具体层号。层号只在调试时才需要，日常统计不需要它。
• 名字里带 -part- 的探针已经合并了，因此同一前缀的多个探针（例如不同同构组的注意力）的时间会加在一起，得到整步所有注意力层的总耗时。

scale 是另一层：每个探针的时间在累加前会乘以 event_scale。对 attn-part-{layer} 来说，scale=len(same_layers)。同构层组里只测了一个代表层，所以这个代表层的时间乘上组内层数，就是整组层的估计。这里有一个隐含假设：同组层每层的耗时相同。DeepInfer 模型里哪些层算同组，由 layer_group_key 决定：

1
def layer_group_key(self_attn: nn.Module, mlp: nn.Module | None) -> str:
2
    """key of the timer sampling group: module structure repr + attention config"""
3
    sched_meta_key = getattr(self_attn, "sched_meta_key", "")
4
    is_yoco_consumer = getattr(self_attn, "is_yoco_hidden_consumer_layer", False)
5
    return f"{self_attn}|{mlp}|{sched_meta_key}|yoco{int(is_yoco_consumer)}"

同组的判据包括三部分：注意力模块与 MLP 模块的 repr（也就是结构一致性）、注意力调度元数据的键（由压缩比、是否 packed KV、是否 SWA、topk 数量拼成）、以及是否 YoCo 的 hidden consumer 层。这三者都相同的层，计算图在形状上一致，单层耗时趋同——用一层代表整组是合理的。若把不同形状的层（例如 SWA 层与全注意力层）混在一组，scale 推算就会失真。

mlp 这一项的 scale 在 decode 时取 self.num_mlp_layers，在 prefill 时取 self.num_prefill_mlp_layers。两者在源码里有明确来源：

1
# deepinfer/model_executor/models/decoder_layer.py（节选）
2
# MLP timer 按实际执行的层数缩放；关闭 MTP 或执行 DFlash prefill 时只有主模型 MLP。
3
self.num_mlp_layers = cfg.num_hidden_layers + (
4
    getattr(cfg, "num_nextn_predict_layers", 0) if envs.DEEPINFER_USE_NEXTN_DECODE else 0
5
)
6
self.num_prefill_mlp_layers = cfg.num_hidden_layers

decode 时若启用了 MTP（DEEPINFER_USE_NEXTN_DECODE），实际执行 MLP 的层数包含 num_nextn_predict_layers 个 MTP 模块；prefill 只跑主模型，因此是 num_hidden_layers。scale 与「这一步实际跑了多少层 MLP」对齐，mlp 探针测的是 1 号层那一次，乘上真实层数即为整步总量。这也解释了一个容易观察到的现象：开启 MTP 后 mlp 的均值显著上升，不是单层变慢，而是分母里的层数变多了。

组里选第几层

get_enable_timer_layer：

1
def get_enable_timer_layer(layer_idx: int) -> tuple[int, int, list[int]]:
2
    same_layers = DeepseekDecoderLayerType[DeepseekDecoderLayerIdx[layer_idx]]
3
    second_layer = same_layers[1] if len(same_layers) > 1 else same_layers[0]
4
    second_last_layer = same_layers[-2] if len(same_layers) > 1 else same_layers[0]
5
    return second_layer, second_last_layer, same_layers

decode 选组内第二层，prefill 选组内倒数第二层。源码没有为这两个选择写注释，从代码结构可以读出两条合理动机：

• 避开组内首层。同构组的第一层往往要建立跨层运行时状态（例如 YOCO 的 shared hidden、跨层复用的缓冲区），这些一次性初始化在后续层不再出现。用首层测量会把一次性开销乘到整组头上。
• prefill 取倒数第二层：prefill 路径的最后一层之后马上接 lm_head 与采样，容易混入结尾段的特殊处理；倒数第二层避开了组尾，又比第二层更靠近稳态执行的后段。这里属于工程上的启发式选择，没有严格的性能证明。

换句话说，这两个选择的目标一致：选一个「干净且稳态」的代表层，让 scale 推算尽量贴近同组层的真实平均。

图内探针与 eager 探针是否会双计

模型层代码中两个 record 是并排放置的，enabled 条件互斥（一个 is_prompt，一个 not is_prompt），因此同一时刻只会有一个探针生效。但 record 的实现中，decode 探针即使 enabled=False（prefill 阶段）也不会创建事件，因为事件创建发生在 if name not in self.events[0] 的分支里，而这个分支只有在 enabled 为真且 supported() 为真时才会走到。这个设计避免了 prefill 阶段白白创建一堆永远不会匹配上的事件。

测量流程

profile_decode_perf 在启动完成后跑一次：

1
def profile_decode_perf(self) -> list[tuple[int, float]]:
2
    if not self.use_cuda_graph:
3
        return []
4
    graph_runner = self.cuda_graph_runner
5
    group = PG.expert_parallel_group()
6
    decode_perf_cache = DecodePerfCache(
7
        self.model_config.model, self.parallel_config,
8
        graph_runner.batch_sizes, group,
9
    )
10
    cached_results = decode_perf_cache.load_decode_perf_cache()
11
    if cached_results is not None:
12
        ...
13
        return cached_results
14

15
    input_metadata_source = graph_runner.graphs[graph_runner.batch_sizes[-1]].input_metadata_source
16
    input_metadata_source.valid_token_mask.fill_(True)
17
    input_metadata_source.gate_force_random.fill_(True)
18

19
    NUM_WARMUP, NUM_MEASURE, NUM_TRIM = 1, 5, 1
20
    results = []
21
    for batch_size in reversed(graph_runner.batch_sizes):
22
        graph = graph_runner.graphs[batch_size].graph
23
        for _ in range(NUM_WARMUP):
24
            graph.replay()
25
            dist.all_reduce(torch.zeros(1, device="cuda"), group=group)
26
        events = []
27
        for _ in range(NUM_MEASURE):
28
            start_event = torch.cuda.Event(enable_timing=True)
29
            end_event = torch.cuda.Event(enable_timing=True)
30
            start_event.record()
31
            graph.replay()
32
            end_event.record()
33
            dist.all_reduce(torch.zeros(1, device="cuda"), group=group)
34
            events.append((start_event, end_event))
35
        times_ms = []
36
        for start_event, end_event in events:
37
            start_event.synchronize()
38
            end_event.synchronize()
39
            times_ms.append(start_event.elapsed_time(end_event))
40
        times_ms.sort()
41
        trimmed = times_ms[NUM_TRIM : NUM_MEASURE - NUM_TRIM]
42
        avg_time_ms = sum(trimmed) / len(trimmed)
43
        ...
44
        results.append((batch_size, avg_time_ms))
45
    results.reverse()
46
    ...
47
    decode_perf_cache.write_decode_perf_cache(results)
48
    return results

逐段拆解：

1. 复用最大批的输入元数据。所有档位的图共享同一份最大批缓冲的切片视图（capture 时 _cut 出来的），因此只要在最大批的源对象上填 valid_token_mask=True 和 gate_force_random=True，每张小图都能拿到「所有行都有效、门控随机」的输入。gate_force_random 让 MoE 路由每次落到随机专家，避免固定路由让专家访存命中缓存导致测量偏低——这是一步很细的工程处理，如果忘了它，测出来的图耗时会比真实 decode 低一截。
2. 每个档位先预热一次，再测 5 次。各档位的图共享同一个显存池，第一次 replay 会触发一些内部状态初始化（例如 NCCL 相关缓冲区）。预热一次能把这些一次性开销排除在外。
3. 每次测量后做一次 all-reduce。专家并行（EP）下，一次前向包含 dispatch/combine 通信；dist.all_reduce(torch.zeros(1)) 在这里的作用是确保前向中的集合通信真正完成，防止测出只覆盖本地计算的虚假图时间。它与 ZOS 模式下「不阻塞」的意图相反——这是启动阶段的测量，可以接受同步。
4. 裁剪均值。NUM_MEASURE=5，NUM_TRIM=1，排序后去掉最小与最大各一个，剩 3 个取平均。这么小的样本量下裁剪均值比算术均值抗一次抖动，比中位数更平滑。
5. 跨 rank 平均。avg_time_ms 被 all_reduce(SUM) 后除以 group.size()。不同 rank 上同一批大小的图耗时可能因专家分布而不一致，跨 rank 平均给出一个「集群视角」的统一曲线。
6. 从大到小测量，最后反转。大图先跑能让显存池的「暖」状态一致；结果反转后按批大小升序，与 batch_sizes 的顺序对齐。

DecodePerfCache：重启不重复测

测量需要几十秒，如果每次重启都跑一遍，服务冷启动时间会被拉长。DecodePerfCache 把结果写到一个文件里，用 SHA256 键区分不同的模型与并行配置：

1
hash_inputs = {
2
    "time_unit": "ms",
3
    "model": model,
4
    "batch_sizes": self.batch_sizes,
5
    "parallel_config": {
6
        "data_parallel_size": parallel_config.data_parallel_size,
7
        "context_parallel_size": parallel_config.context_parallel_size,
8
        "expert_parallel_size": parallel_config.expert_parallel_size,
9
    },
10
}
11
cache_key = json.dumps(hash_inputs, sort_keys=True, separators=(",", ":"))
12
return hashlib.sha256(cache_key.encode("utf-8")).hexdigest()

键里包含的字段各自对应一个「曲线会变」的原因：

• time_unit：旧版本缓存是 SPS 口径，新版本是毫秒口径。加一个 "ms" 标记让旧文件自动失效，不用手工清理。
• model：模型标识（名字或路径），不同模型图的耗时不同。
• batch_sizes：图档位列表，档位变了曲线点数就变了。
• parallel_config：DP、CP、EP 三种并行度都会改变一次前向里的通信量与专家分布，进而改变图耗时。注意 TP 不在其中——DeepInfer 的拓扑里不做张量并行（见并行拓扑）。

文件的读写在 rank 0 上进行，再通过 broadcast_object_list 广播给整个 EP 组：

1
cached_results_list = [cached_results]
2
dist.broadcast_object_list(cached_results_list, src=dist.get_global_rank(self.group, 0), group=self.group)
3
return cached_results_list[0]

写入流程是「临时文件 + fsync + os.replace」：

1
tmp_path = os.path.join(self.cache_dir, f"{filename}.tmp")
2
with open(tmp_path, "w") as f:
3
    json.dump(results, f, separators=(",", ":"))
4
    f.flush()
5
    os.fsync(f.fileno())
6
os.replace(tmp_path, final_path)

这样即使进程在写入中途被杀，半写文件也不会被当成有效缓存——只会留下一个 .tmp 文件，在下次清理时被当成脏文件删除。

两条清理阈值：读 TTL 与删 TTL

1
def _cleanup(self, now_s: int, ttl_s: int) -> None:
2
    ...
3
    for filename in filenames:
4
        parsed = self._parse_filename(filename)
5
        if parsed is not None and now_s - parsed[1] <= ttl_s:
6
            continue
7
        os.remove(os.path.join(self.cache_dir, filename))

调用处传的清理阈值比读阈值宽两小时：

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

为什么需要这个「宽窗口」？设想多个服务实例（不同 DP 组或不同 EP 组）共享同一个缓存目录。rank 0 的清理可能在另一个实例正读到一半时把文件删掉。读侧本身有容错（except (OSError, TypeError, ValueError, json.JSONDecodeError) 会跳过读坏的文件，继续尝试下一个候选），但「先删后读」会浪费一次无效尝试。让文件在被清理前有 2 小时的「已过期但仍可读」窗口，所有实例都有机会读到同一份过期曲线，然后在各自的下一次启动时重新测量并写入新文件。

读取侧还有一个细节：候选文件按时间戳降序排序，取最新的一个；同时校验文件里的批大小列表与当前 batch_sizes 一致，不一致则跳过。这个二次校验比只看 hash 更严格——hash 已经覆盖了 batch_sizes，但万一有人手动改了文件内容，这里能兜住。

为什么需要曲线加运行时修正

启动曲线只反映了「固定批大小、固定上下文长度」的图耗时。真实 decode 里，context length 随生成增长而增长——注意力 kernel 读取的 KV 长度变长，单次前向变慢。曲线的点数有限（每个档位一个点），无法覆盖上下文长度这个维度。

同时 sampler 部分不在图里，是单独跑的 kernel 链（惩罚、采样、logits 处理），其耗时随批大小与采样参数变化，也不在曲线里。

MTPVarlen 用两个运行时估计器分别补这两块：

1
void update_runtime(int num_seqs, int graph_size, double forward_ms, double sampler_ms) {
2
    const auto curve_it = std::lower_bound(_decode_perf_curve.begin(), _decode_perf_curve.end(), graph_size,
3
        [](const auto &point, int size) { return point.first < size; });
4
    const double baseline_ms = curve_it->second;
5
    _forward_time_slope.update(graph_size, (forward_ms - baseline_ms) / graph_size);
6

7
    _sampler_time_by_num_seqs[num_seqs].update(sampler_ms);
8
    _sampler_time_fallback.update(sampler_ms);
9
}

• forward 斜率：实际一帧的 forward 时间减去该档图大小的启动曲线值，除以图大小，得到「每 query 的额外耗时」。这个差分值主要体现了上下文长度带来的注意力开销随 query 数线性增长的部分。
• sampler 估计：没有可用的启动数据，完全靠历史。按 num_seqs 分桶记录滑动中位数；若该桶样本不足，退回全局中位数。

两层中位数

ForwardRuntimeEstimator 的注释写得很细：维护最近 _window_size = 5000 个 forward step，先按 graph size 聚合取中位数，再对这些中位数取中位数；只考虑至少 _min_samples = 10 个样本的 graph size，各 graph 等权。

实现上用了两个 multiset 维持动态中位数（MedianTracker），外层 _graph_medians 是所有「已达样本门槛的 graph size 的中位数」的中位数。为什么不用简单平均？

• 斜率分布对上下文长度敏感：长上下文阶段的斜率显著高于短上下文。平均会被长上下文样本拉高，导致下一步决策偏保守（选中更大的图，期望吞吐变低）。
• 按 graph size 分层再加权：不同图大小的采样频率不同（小图档位多、被选中的次数多），直接全局取中位数会让小图主导。先按 graph size 分层、再对层中位数取中位数，相当于让每个档位等权。
• 复杂度的控制：每步只更新新增 / 淘汰样本所属的 graph，以及它们在第二层中位数里的贡献；两层中位数与 graph 查找都用平衡树，更新 O(log 5000)，总空间 O(5000)。

吞吐估计与 best_bsz 选择

1
const double sampler_ms = estimate.value_or(0.0);
2
const double slope_ms = std::max(0.0, _forward_time_slope.get());
3

4
double best_throughput = 0.0;
5
int best_bsz = 0;
6

7
float cur_expected = static_cast<float>(num_seqs);
8
int cur_bsz = num_seqs;
9
int cur_bsz_idx = 0;
10
auto update_best = [&]() {
11
    while (cur_bsz_idx < num_bsz_selections && cur_bsz > _decode_perf_curve[cur_bsz_idx].first) {
12
        ++cur_bsz_idx;
13
    }
14
    if (cur_bsz_idx == num_bsz_selections) return;
15
    const auto &[graph_size, baseline_ms] = _decode_perf_curve[cur_bsz_idx];
16
    const double time_ms = baseline_ms + slope_ms * graph_size + sampler_ms;
17
    const double cur_throughput = cur_expected / time_ms;
18
    if (cur_throughput > best_throughput) {
19
        best_throughput = cur_throughput;
20
        best_bsz = cur_bsz;
21
    }
22
};
23
update_best();  // All the golden tokens
24
for (auto &acc : acc_queue_for_bsz) {
25
    cur_expected += acc.first;
26
    cur_bsz += 1;
27
    update_best();
28
}

这段是整个变长验证策略的核心。先解释变量：

• num_seqs：这一轮被调度的请求个数。每个请求至少有 1 个 query（必然生成的那个 token）。
• cur_expected：当前假设下期望被接受的 token 数。从 num_seqs 开始（每个请求至少接受 1 个），每多分配一个草稿 token，就按这个 token 的 accept_rate 增加期望值。
• acc_queue_for_bsz：把每个请求的每个草稿 token 的接受概率拆成单独一行，按概率降序排列。每行是 (accept_rate, seq_idx)。
• cur_bsz：以 num_seqs 为起点，每加入一个候选草稿 token 就加一——这就是「如果这个 token 也分配了 query，总 query 数是多少」。
• time_ms：用曲线中满足 graph_size >= cur_bsz 的最小档位的图时间，加上 slope_ms * graph_size 与 sampler_ms。
• cur_throughput = cur_expected / time_ms：每毫秒期望接受多少 token。

这是一个标准的边际收益贪心：按接受概率从高到低逐个尝试「再加一个草稿 token」，每加一个都重新评估吞吐；如果吞吐不再上升就停止。最终得到 best_bsz——这一轮总共应该验证的 query 数。

值得强调的几个约束：

• 档位选择的 while 循环：cur_bsz 增加时，cur_bsz_idx 单调推进找到第一个 graph_size >= cur_bsz 的档位。因为曲线按档位升序，这个指针只会前进不回退，复杂度 O(档位数 + 候选数)。
• 时间估计随档位跳变：当 cur_bsz 跨过某个档位边界，time_ms 会跳到下一档的图时间——多分配一个 query 可能换来一个更贵但容量更大的图。贪心在这里隐式处理了「档位跳变」的成本，但因为贪心假设每次增量只影响「当前档位」，实际上在跨档时可能不是全局最优。这是工程上的近似：真正的组合优化是二维背包（每个请求可选 1 到 k 个 query），逐 token 贪心只保证在单档内部最优。
• 以启动曲线为下界：slope_ms = std::max(0.0, _forward_time_slope.get())。注释解释得很明确——避免负斜率外推到更大的图后产生非正耗时。启动曲线本身是「短上下文」的乐观估计，运行时斜率只会往上加；若某段时间上下文极短导致观测到负斜率，不能让它把时间估成负数。
• 没有 graph 样本时用启动曲线：estimate.value_or(0.0)；slope 在没达到样本门槛时返回 0。因此冷启动初期的决策完全依赖启动曲线，这是可以接受的保守行为。

每请求的 seqlen_q 分配

1
seqlens_q.assign(num_seqs, 1);
2
for (int i = 0; i < best_bsz - num_seqs; ++i) {
3
    seqlens_q[acc_queue_for_seqs_q[i].second] += 1;
4
}

确定 best_bsz 后，按 acc_queue_for_seqs_q（同样按接受概率降序，但用的是另一组 accept_rate）把多余的 best_bsz - num_seqs 个 query 依次分配给请求：排在最前面的草稿 token 所属的请求获得一个额外的 query。这等价于「优先验证最可能被接受的草稿」。

为什么用两套 accept rate

schedule 的注释把这一点讲透了：计算 best_bsz 时永远用 N-2 轮的接受率，计算每个请求的 seqlen_q 时用 N-1 轮的接受率。原因是：

1
bool post_process_paired = mtp.process_cnt == mtp.schedule_cnt;
2
// 永远用 N-2 轮的 accept_rate 去计算 best_bsz
3
push_rates(acc_queue_for_bsz, post_process_paired ? mtp.prev_accept_rates : mtp.accept_rates, nd, i, has_draft);
4
// 永远用 accept_rates 去计算每个请求的 seqlen_q
5
push_rates(acc_queue_for_seqs_q, mtp.accept_rates, nd, i, has_draft);

• best_bsz 用 N-2：决定总数会改变 CUDA Graph 档位。而「决定档位」的决策发生在 GPU 执行之前，此时 N-1 轮的 forward 可能还挂在 GPU 上（ZOS 模式下延后一拍）。为了保证 best_bsz 的决策可复现、并且不会因为「看到最新一轮的信息又改变了决定」而破坏图重放的形状稳定性，它刻意用更旧的一轮反馈。
• seqlen_q 用 N-1：每个请求分几个 query 是一个可以图内修正的量——在 ZOS 下，seqlens_q 可以在 GPU 上被改写（见第二篇里讲的 seqlens 修正）。因此它可以用最新一轮的接受率，获得更准确的个体分配。
• has_draft 的兜底：若某请求没有 accept_rate（例如预先调度、还没有结果），push_rates 用 has_draft ? 1.0 : 0.0 填充——有草稿就当作必然接受（乐观），没有草稿就当作必然拒绝（不分配额外 query）。这个选择让「预调度」的请求不被无端多分配 query。

这个设计里还有一个细节：断言 TORCH_CHECK(mtp.process_cnt <= mtp.schedule_cnt && mtp.schedule_cnt <= mtp.process_cnt + 1)。schedule_cnt 与 process_cnt 之间的差不超过 1，正是 ZOS「最多领先一拍」的不变量。曲线加运行时估计的整套决策，是建立在这个不变量上的。

探针埋在哪里：三个层内组件

模型层在三个位置埋了 decode 探针，各自覆盖一段计算：

• decoder_layer 的 attn-part-*：包住整个注意力子层的入口（含 mHC 残差融合、norm、注意力调用、FFN norm）。
• v4_modules/attn 的 decode_core_attn-part-*：只包住核心的稀疏 MLA 调用 decode_mla_mqa_sparse，不含前后的 norm 与残差。
• v4_modules/index_attn 的 decode_index_attn-part-*：包住压缩 KV 上的索引器计算与 Top-K 候选块选择。

这两层嵌套的注意力探针不是重复：attn 是子层总量，decode_core_attn 是核心注意力本身，两者之差就是 norm、残差融合、RoPE 等外围开销。分析「注意力慢是因为核心 kernel 慢，还是因为外围变换慢」时，这个差值是第一个要看的东西。

decode_index_attn 是在索引器模块单独埋的探针：

1
# deepinfer/model_executor/models/v4_modules/index_attn.py（节选）
2
second_layer, _, same_layers = get_enable_timer_layer(self.layer_idx)
3
with cuda_graph_timer.record(
4
    f"decode_index_attn-part-{self.layer_idx}", enabled=self.layer_idx == second_layer, scale=len(same_layers)
5
):
6
    p_sum = self._compute_decode_p_sum(...)
7
    ...

routed_experts 则只在 MLP 里：

1
with cuda_graph_timer.record(
2
    "routed_experts",
3
    enabled=(self.layer_idx == 1 and not input_metadata.is_prompt),
4
    scale=eplb.current().num_moe_layers,
5
):
6
    moe_out_hidden = self.routed_experts_forward(...)

注意这里 scale 取的是 eplb.current().num_moe_layers，不是 num_mlp_layers。原因是路由专家计算的层数与 MoE 层的层数关系由 EPLB 的当前视图决定（例如某些层可能是稠密层或共享专家层），用 EPLB 给出的层数更准确。这与 mlp 探针的 scale 不同，是有意为之的两套口径。

routed_experts 探针只包住专家计算，不包 dispatch / combine。所以「通信」在 decode 路径里并不出现在 routed_experts 里；想看通信占比，需要用 PyTorch profiler 的 trace 或 limit_num_sms 相关日志配合。这种细分方式有一个明显的取舍：探针尽量贴在最核心的计算上，通信和共享专家计算另算。若把通信包进去，routed_experts 在 EP 下就会混入网络时间，无法反映 GPU 计算真实占用。

与 StatsLogger 报告的衔接

StatsLogger._merge_cuda_timers 会遍历 get_timers() 并合并所有标签以 cuda 开头的计时器（即 cuda 与 cuda_graph 两个）。这一步在 flush_out 之后、日志周期结束之前发生，合并规则按字段不同：

1
def _merge_cuda_timers():
2
    times, called, max_time, latest = {}, {}, {}, {}
3
    for lbl, tmr in get_timers():
4
        if lbl.startswith("cuda"):
5
            for n, t in tmr.times.items():
6
                times[n] = times.get(n, 0.0) + t
7
                called[n] = called.get(n, 0) + tmr.called[n]
8
                max_time[n] = max(max_time.get(n, 0.0), tmr.max_time[n])
9
            for n, t in tmr.latest_times.items():
10
                latest[n] = latest.get(n, 0.0) + t
11
    return times, called, max_time, latest

• times 与 called 相加：同一名字来自 eager 与图内两套测量时，总量是两者之和，调用次数也是两者之和。
• max_time 取大：长尾只关心最大值，不关心来自哪个计时器。
• latest 相加：这一步中同名组件可能既走了 prefill 也走了 decode，latest 相加后给出的「当前一步总耗时」才是正确的。

这个「周期合并」是必要的：一个日志周期内可能既有 prefill 又有 decode，两段路径的 mlp 探针都有效，只是测量尺度不同（prefill 的 mlp scale 是 num_prefill_mlp_layers，decode 是 num_mlp_layers）。合并后得到的 mlp 总量是正确加总的；但要注意，scale 不一致意味着同名探针的「单次平均」不具可比性，比较时要看总量的变化趋势，而不是单次值。

日志行里的四个数字

_build_log_messages 为每个 GPU 计时项输出四个数字：

1
f"{n} cur{1000 * lat.get(n, 0):.1f}ms/mean{1000 * t / called[n]:.1f}ms"
2
f"/max{1000 * max_t[n]:.1f}ms/{called[n]}calls/{100 * t / engine_t:.1f}%"

以 attn cur4.2ms/mean3.9ms/max6.1ms/32calls/18.7% 为例：

• cur：这一步各 latest 之和，即「本步 attn 实际花了多少毫秒」。因为 attn 已经经过 -part- 合并与 scale 放大，这里的 4.2ms 是整步所有注意力层的估计值，而不是某一层。
• mean：周期内累计时间除以调用次数。由于 scale 已经乘在单次时间上，mean 就是「每步注意力层均摊时间」。
• max：周期内单次最大值，用来发现偶发的长尾（例如某一步上下文突然变长、或某一步回退了图档位）。
• %：除以 engine_step 的 CPU 时间，即 GPU 计算相对引擎步长的占比。这个百分比的参照系是 CPU 侧一步的墙钟时间，因此在 ZOS 模式下可以超过 100%——GPU 的多个步被重叠执行，这是正常的。

NPU 上 graph_timer 只是一个桩

csrc/graph_timer/graph_timer_asc.cpp 的全部实现都是抛异常：

1
EventHandle::EventHandle() { throw std::runtime_error("NPU graph timer primitives are not implemented yet"); }
2
void EventHandle::record(int64_t, bool) const {
3
    throw std::runtime_error("NPU graph timer primitives are not implemented yet");
4
}

Python 侧 supported() 返回 PLATFORM.is_gpu，因此 NPU 上所有 cuda_graph_timer.record 与 flip_event_buffer 都会直接 yield 或返回。计时器本身不崩溃，但图内探针在 NPU 上完全失效。同时 CudaGraphRunner.capture 在 NPU 分支也不传 keep_graph=True，与 register_graph 所依赖的原始图句柄一致地缺失。这意味着：

• NPU 上解码路径没有 attn、mlp、router 等图内细分耗时，性能分析只能看 CPU 计时与整体曲线。
• 曲线测量本身只依赖 graph.replay() 与设备计时事件，不依赖图内探针；只要平台提供可用的计时事件，曲线仍能建立，但无法细分到组件。

这是一个明确的平台差异，不是设计缺陷。graph_timer 的三个类里，EventHandle、GraphHandle 的接口在 CUDA 与 NPU 上一致，未来 NPU 后端补齐时只需实现 cudaGraphExecEventRecordNodeSetEvent 对应的接口即可。

图内探针的一次性限制

注释里写得很清楚：「Called at most once in one CUDA graph.」若某个探针在一张图里被调用两次，register_graph 的断言会失败，或者更糟——匹配到错误的节点。这个约束限制了它的表达能力：无法在同一个 decode 图内区分「第 1 次调用」与「第 2 次调用」的同名探针。DeepInfer 的应对方式是让每个探针名带层号（attn-part-{layer}），通过命名空间区分。

同构假设

scale = len(same_layers) 假设同组层耗时一致。在 MoE 模型里，同组的 MLP 层走的是同一份专家路由逻辑，各层专家权重分布不同、访存命中不同，实际耗时会略有差异。layer_group_key 把 sched_meta_key、yoco 都纳入判据，已经把最显著的结构差异分开了，但无法区分「同一结构不同权重」导致的访存差异。这个误差在日志层是可见的（max_time 与 times/called 的分布），但在决策层会被 scale 直接放大。

启动曲线与运行时估计的耦合

ForwardRuntimeEstimator 把「实际 forward 时间减去启动曲线值」当作斜率来学习。一旦启动曲线本身因为硬件、NCCL 版本、拓扑变化而与实际不符，斜率就把这个偏差吸收进去并持续沿用——它不会自己纠正，因为启动曲线只测一次。DecodePerfCache 的 hash 没有包含驱动版本、NCCL 版本与硬件型号，所以在这些环境变化后，旧曲线仍会被复用。这是一个可以被外部观察到的薄弱点：如果服务在升级驱动后出现吞吐下降但曲线缓存未失效，先清缓存重测。

pending_buffers 容量 2

Ping-pong 只有两套 buffer，因此「一步里最多两张图可以同时被计时」。若未来出现「同一 step 内三类图交错」（主图 + DFlash + 索引器前向各自独立捕获）的情况，容量会不够；当前的模型结构刚好没有这么深的交错。