DeepInfer 源码拆解（十七）：LLMEngine 的 step 与 Worker 握手——调度结果如何变成一次 GPU 前向

进程结构

DeepInfer 的部署形态是「每个 rank 一个进程」。LLMEngine 构造时读取 CUDA_VISIBLE_DEVICES，把第一张卡当作 Driver GPU，在上面同进程构造一个 Worker；其余用来做上下文并行（CP）的卡，各自起一个子进程，子进程里只构造 Worker 本身，不构造 LLMEngine。

1
# deepinfer/engine/llm_engine.py, _init_workers
2
for cp_rank in range(1, cp_size):
3
    worker = SpawnProcessWithCrashHeader(
4
        target=_spawn_worker,
5
        name=f"worker-cp{cp_rank}",
6
        args=(engine_type, model_config, parallel_config, scheduler_config,
7
              cache_config, cp_rank % ngpus, cp_rank, distributed_init_method, False),
8
        daemon=True,
9
    )
10
    worker.start()
11

12
driver_kwargs = {
13
    "engine_callbacks": EngineCallbacks(
14
        post_process_fn=self.post_process,
15
        pre_sampler_sync_fn=self.pre_sampler_sync,
16
    ),
17
}
18
self.driver_worker = Worker(..., local_rank=0, cp_rank=0, ..., is_driver_worker=True, **driver_kwargs)

两个细节值得注意：

• 只读的从属 Worker。子进程的 Worker(*args) 参数里没有 engine_callbacks，并且 is_driver_worker=False。在 WorkerBase.__init__ 末尾，非 driver 的 Worker 会直接进入 start_slave_worker_loop()：一个 while True 循环，持续从共享内存里接收 (method, kwargs) 并反射调用同名方法，永远不返回。
• 父进程死亡的兜底。从属进程在进入循环前调用 set_pdeathsig(SIGKILL)，通过 prctl(PR_SET_PDEATHSIG) 请求内核在父进程死亡时杀掉自己。否则当一个引擎进程崩溃，这些 daemon 子进程会残留并占住 GPU 显存。

也就是说，在这个进程模型里，只有 Driver 拥有 LLMEngine 的全部状态；从属 Worker 只持有模型权重、KV 池和 CUDA Graph，不持有调度器，也不持有请求表。它们的唯一入口是反射调用 Driver 广播过来的方法。

共享内存广播

从属 Worker 接收任务的方式是 cp_broadcast_object(None)，它读取的是一个由 torch.distributed 进程组承载的共享内存对象，而不是 pickle 序列化的 RPC。Driver 侧在 run_workers 中把 (method, kwargs) 通过同一个对象写出去：

1
# deepinfer/worker/base.py, WorkerBase.run_workers
2
def run_workers(self, method, driver_kwargs=None, **kwargs):
3
    with cpu_timer.record("broadcast_method", enable=False):
4
        assert self.is_driver_worker
5
        cp_group = PG.cpu_cp_group
6
        if cp_group and cp_group.size() > 1:
7
            cp_broadcast_object((method, kwargs))
8
    return getattr(self, method)(**(driver_kwargs or {}), **kwargs)

driver_kwargs 和 kwargs 是有区别的：driver_kwargs 只在 Driver 上生效，不会广播。step() 里 execute_model_step 调用中，seq_metadata_list 与 scheduler_outputs 是 Python 对象（C++ 扩展构造的 _C.engine 类型），里面带着 shared_ptr 和裸指针，根本不可能跨进程 pickle；所以在 CP 场景下，Driver 只广播 is_prompt / is_placeholder 这两个布尔量，从属 Worker 拿到的 seq_metadata_list 与 scheduler_outputs 都是 None。真正跨卡的输入数据走的是 broadcast_tensor_dict：_prepare_input_driver 在广播前检查 CP 组大小，把 InputMetadata 拍平成一个 dict 再发出去。这条路径在第十二篇（CUDA Graph 输入）与第十三篇（CP 拓扑）里分别讲过，这一篇只强调一个结论：

引擎到 Worker 的「控制面」广播是轻量且 CP 感知的，而「数据面」广播由 ModelRunner 在 prepare_input 中完成。两者不能混在一起。

_run_workers 的真实语义

LLMEngine 里所有”调用 Worker”的地方都写成 self._run_workers("方法名", ...)，而它的实现只是一行：

1
def _run_workers(self, method, driver_kwargs=None, **kwargs):
2
    if driver_kwargs is None:
3
        driver_kwargs = {}
4
    return self.driver_worker.run_workers(method, driver_kwargs, **kwargs)

也就是说，Python 引擎总是先在本地 Driver 上同步执行该方法，run_workers 里那条 cp_broadcast_object 只是顺手把方法名和跨卡参数广播出去，然后返回本地调用的结果。从属 Worker 在自己的循环里执行同一个方法，两边通过 torch.distributed 的集合通信（如 broadcast_tensor_dict）在方法体内部对齐。这解释了一件看起来很怪的事：load_model、init_cache_engine、execute_model_step 这些在 Driver 上”同步”调用的方法，返回值都来自 Driver 本地那次调用。至于各 rank 之间如何对齐，由方法自身决定：execute_model_step 在 prepare_input 里用 broadcast_tensor_dict 对齐输入与采样元数据；load_model 则让每个 rank 各自装载自己那份分片。从 Engine 的角度看不到 RPC 返回值的拼接，需要聚合的结果只会是 torch.distributed 的归约值。

Driver 与从属 Worker 的职责差异

这张表解释了为什么 Worker.execute_model_step 里几乎每个与「输出」相关的分支前都写着 self.is_driver_worker。从属 Worker 并非完全跳过 sampler：它会走同一个 self.model.sample(...) 入口，但 prepare_sampler_inputs 在非 Driver 分支上只从广播里拿 selected_token_indices / mtp_first_token_positions，Sampler 内部的 driver_inputs 为 None，因此只参与 embedding gather 这类集合通信，产生的采样结果也只在 Driver 上有效。CP 切的是序列维，每张卡分到同一批请求的不同 token 片段，真正的采样输入只在 Driver 上完整。

DP 维度：每个 rank 一个完整引擎

EP / 大 DP 部署时，每个 rank 上都有一个完整的 LLMEngine，彼此通过 PG.data_parallel_group() 通信；DP Attention 只做注意力部分的 All-Gather / Reduce-Scatter，不交换请求表。这带来的直接后果是：

• 调度器在每个 rank 上是独立的；
• 但一次 step 的「步型」（prefill 还是 decode）必须一致，否则 All-Gather 的 seqlen_q 对不上；
• 第十六篇里的 DP 投票就是解决这个问题的：第二轮 MIN 归约把全组的步型强行拉成一致。

引擎与 Worker 的握手对 DP 层面是透明的，但握手时机的选择会直接影响 DP 投票的效果：如果引擎在调度前已经发现本 rank 空转，它会用占位步参与投票，而不是干等。

输入是 token，不是文本

签名里的 prompt_token_ids: np.ndarray[np.int32] 是硬约束。引擎内部虽然持有一个 HF tokenizer（用于 LLMInference 的便捷入口），但 add_request 不接受字符串提示。entrypoints/llm._add_request 在需要时才调用 self.llm_engine.tokenizer.encode(prompt)。这条边界把分词从引擎关键路径上移走：线上服务里分词发生在 API 网关，引擎只对 token 做调度与计算。

五道预处理

1. eos_id 注入：sampling_params.eos_id = self.eos_id。采样器需要它来判断停止，而 SamplingParams 本身不携带词表信息。
2. YoCo 截断下的 prompt_logprobs 强制关闭：源码在 yoco_prefill_truncation_mode != "off" 且请求要求 prompt_logprobs 时打印 warning 并置 False。理由是 compact 模式无法计算被裁掉行的 logprob，mask 模式对矩形之前的 token 返回退化值。与其返回错误数据，不如显式关闭。
3. 多模态提示重写：multi_modal_processor.overwrite_prompt(...) 把图像占位区间替换成模型约定的 image token 序列；如果 image_bytes 为 None，直接把序列标记为 FINISHED_VE_FAILED，让它在第一次调度时就被回收。
4. token 取值范围校验：min_token_id < 0 或 max_token_id >= vocab_size + NUM_SPECIAL（多模态时放宽到特殊 token 偏移）会直接丢弃请求并返回 None，日志写明原因。这一步防止越界 token 进入 embedding 查表。
5. 序列号与压缩提示：seq_counter 单调分配 seq_id；启用 Engram 的模型还会用 CompressedTokenizer.encode 生成 compressed_prompt_token_ids，两者形状必须一致。

结构化输出的异步编译

so_args 非空时，结构化输出的语法编译不该阻塞 add_request 的调用方。源码的处理是：

1
if so_pending:
2
    self.so_matcher_builder.submit(seq, sampling_params.so_args, _enqueue)
3
else:
4
    _enqueue(seq)

真正入队发生在编译完成的回调里。回调还检查 matcher.is_aborted()：如果编译阶段就被判定不可满足（例如语法与约束冲突），序列直接标记 FINISHED_SO_FAILED，不进入调度队列。enqueue=False 的调用方（自己管理调度）则同步 .result() 等待编译完成，因为它没有人替它驱动 pump()。

pre_sampler_sync 里的 self._so_matcher_builder.pump() 就是这条异步链的推进器：每拍检查一次已完成的编译任务，把对应序列放进 wait 队列。

abort_request 的边界

abort_request 先调 scheduler.abort_sequence(request_id) 拿回真正被取消的序列，然后：

• 通知 multi_modal_processor.abort；
• 对每条序列调用 free_multi_modal_pool_handles(seq.multi_modal_data)，归还未消费完的共享池块。

返回值是 (被取消的序列数, 被取消的 token 总数)，供上层统计。这里的关键是”真正被取消”：调度器可能已经把该序列标记为完成，abort_sequence 返回空列表时，引擎不会误释放别的序列的资源。

阶段一：调度与状态克隆

进入 step() 后先更新时间戳与运行时估计器，然后进入调度：

1
with cpu_timer.record("schedule"):
2
    seq_metadata_list, scheduler_outputs = self.scheduler.schedule(self.sequence_metadata_pool)
3
    # NOTE 开启 yoco 截断时，允许抢占导致的结果变化
4
    with cpu_timer.record("schedule:clone"):
5
        clone_states, clone_blocks, zero_states = (
6
            scheduler_outputs.clone_states,
7
            scheduler_outputs.clone_blocks,
8
            scheduler_outputs.zero_states,
9
        )
10
        if clone_states.size > 0 or clone_blocks.size > 0 or zero_states.size > 0:
11
            to_release, to_release2 = self._run_workers(
12
                "clone_and_zero_states_and_blocks_on_default_stream",
13
                clone_states=clone_states,
14
                clone_blocks=clone_blocks,
15
                zero_states=zero_states,
16
            )
17
            self.deferred_release.collect_clone_unpins(to_release, to_release2)
18
    with cpu_timer.record("schedule:post"):
19
        self.post_schedule(scheduler_outputs)

这里有几个互相咬合的设计：

• scheduler.schedule() 是 C++ 实现，它在同一个循环里按下标构造 seq_metadata_list 与 scheduled_seqs。源码注释明确说明两者「严格同序」，后面所有 zip(..., strict=True) 都依赖这一点。
• clone_states / clone_blocks / zero_states 是序列派生（fork）所需的状态克隆。C++ 侧的 clone_states 是一个 vector<int> 的 (src, dst) 对，clone_blocks 同样，dst=0 表示「只归还 src、不拷贝」，zero_states 则是需要清零的状态槽编号。
• 三个数组都是零拷贝 numpy 视图：pybind 用 py::array_t<int>(..., data(), py::capsule([](){ })) 直接指向调度器内部 vector，属性文档写明「下一次 schedule() 会清空复用：须在此之前消费完，禁止跨 step 持有」。
• 克隆被放到 GPU 的默认流上执行，返回的 states_to_unpin / blocks_to_unpin 由 Worker 给出，但并不是立刻归还的，而是交给 DeferredRelease.collect_clone_unpins()。原因是 clone kernel 是异步的，GPU 还在读源页，CPU 侧归还池编号会让调度器在本步内把同一页再分出去，产生数据竞争。
• 克隆完成后才进入 post_schedule()。post_schedule 是虚函数，基类是空实现，P/D 分离的引擎在子类里会利用这个钩子做预取或 KV 保存的准备工作。

阶段二：前置副作用

调度之后、模型执行之前，step() 还处理了几类与「GPU 前向无关但与请求生命周期有关」的事情：

1
if not scheduler_outputs.prompt_run and self.copy_state is not None:  # 处理 kvcache 落盘相关
2
    to_copy_state_seqs = [seq for seq in
3
        scheduler_outputs.scheduled_seqs + scheduler_outputs.finished_seqs
4
        if self.should_copy_state(seq)]
5
    if to_copy_state_seqs:
6
        self.do_copy_state(to_copy_state_seqs)
7
        self.copy_state_event.record()
8

9
with cpu_timer.record("multi_modal_fetch"):
10
    if self.multi_modal_processor is not None:
11
        self.multi_modal_processor.fetch()
12
        abort_requests = self.multi_modal_processor.fetch_to_abort_requests()
13
        if abort_requests:
14
            logger.error(...)
15
            self.abort_request(abort_requests)

• do_copy_state 把某个序列当前的状态槽拷进双缓冲 copy_state，只下发任务、不等待，消费方自行同步。源码的实现直接遍历 to_copy_state_seqs，对每个 state_id != -1 的序列做 copy_state[state_id].copy_(gpu_states[state_id])，并记录 last_copy_state_token_idx = seq.num_cached_tokens。这个机制服务于 KV 落盘与跨机迁移：源状态槽即将被改写，而落盘线程还在读它，所以必须先拷一份。copy_state_event 会在下一次 pre_sampler_sync 里被同步，同步之后才调用 do_save_kvcache。
• 多模态的 fetch() 是预取。图片预处理放在 CPU 上会很慢（解码 JPEG、缩放），提前一步 fetch 可以让预处理和 GPU 前向重叠。如果某张图损坏，fetch_to_abort_requests() 会返回要中止的请求列表，引擎当场 abort 这些请求，而不是等到 embedding 阶段再暴露异常——这保证了后续 _process_model_outputs 不会见到一个缺少 embedding 的序列。
• 源码里还保留了引擎自己的 tokenizer，但 add_request 要求调用方传入 prompt_token_ids（np.ndarray[np.int32]），entrypoints/llm._run_engine 只在缺失时才用引擎自带 tokenizer 编码。这说明分词位于引擎之外的设计意图：在线上服务里，分词由 API 网关层的分词器集群完成，引擎只接收 token。

阶段三：Worker 前向 + 采样

模型执行被封装成一行调用：

1
request_outputs = self._execute_model_async_step(seq_metadata_list, scheduler_outputs)

它本身只是 _run_workers("execute_model_step", ...) 的外壳。真正的四段流水在 Worker.execute_model_step 里。这个方法的名字里带 async，但它是同步返回的——「异步」指的是它内部的 CPU 与 GPU 重叠，而不是调用方式。

阶段四：延后收账

step() 返回前，把本步的 request_outputs 和 seq_metadata_list 交给 DeferredRelease：

1
self.deferred_release.collect_request_outputs(request_outputs)
2
return request_outputs

seq_metadata_list 在 Worker 侧的 _post_process_step 之后、经 self.engine_callbacks.post_process_fn(step_state) 回到引擎的 post_process() 时被收集（下文详述）。

第 1 段：主模型前向

1
if not is_placeholder:
2
    main_model_output = self.model_runner.execute_main_model(
3
        seq_metadata_list, self.gpu_cache, metadata_zos=metadata_zos, total_q=total_q
4
    )
5
else:
6
    seq_metadata_list = [self.model_runner._placeholder_seq_metadata(is_prompt)]
7
    main_model_output = self.model_runner.execute_main_model(seq_metadata_list, self.gpu_cache, metadata_zos=metadata_zos)

execute_main_model 内部做两件事：_prepare_input 构造 InputMetadata，然后要么 cuda_graph_runner.replay() 要么直接调用 self.model(...)。源码注释用全大写的语气写着：

IMPORTANT: This function must NOT contain any cuda synchronize.

这不是风格要求，而是正确性要求：ZOS 之所以能把 CPU 调度挪到 GPU 前向旁边，依赖的就是前向调用本身是纯异步 launch。如果 execute_main_model 里出现一次隐式同步（比如某个 Tensor 没有 non_blocking=True 的 H2D 拷贝，或某个 .item()），ZOS 的收益会被一次同步抹平。源码里为此专门提供了 forbid_cuda_sync() 上下文管理器，包住 launch_forward 与 launch_sample 两段。

第 2 段：上一拍的 post-process

1
request_outputs = self._post_process_step(self.prev_step_state)

_post_process_step 会先 sample_end_evt.synchronize()，再调用 model_runner.post_process(...) 和引擎回调 post_process_fn。这个同步点是整个 step 里唯一一次 CPU 等待 GPU 的地方。它等待的是上一拍的采样结束事件，而此刻 GPU 正在执行本拍的前向，所以这次等待通常已经完成，真实的停顿接近零。

1
with cpu_timer.record("worker.post_process.sync"):
2
    step_state.sample_end_evt.synchronize()

prompt 步或关闭 ZOS 时，prev_step_state 为 None，_post_process_step(None) 直接返回空数组；本拍的 post-process 会在第 4 段立即执行，此时 CPU 必须等待本拍采样结束——这就是非 ZOS 路径的同步开销。

第 3 段：采样本拍并回调引擎

1
if correction is not None:
2
    if correction.d2h_evt is not None:
3
        correction.d2h_evt.synchronize()
4
    scheduler_outputs.update_metadata_for_sampling_zos(
5
        seq_metadata_list, correction.seqlens_q_cpu, correction.metadata_zos.remap_cpu, self.nextn
6
    )
7

8
sampler_output = self.model_runner.execute_sample(
9
    seq_metadata_list, input_metadata, model_output, self.gpu_cache,
10
)
11

12
sample_events = cuda_timer.latest_event_pair("device_sample") if record_runtime else None
13
with cpu_timer.record("pre_sampler_sync"):
14
    if self.is_driver_worker:
15
        self.engine_callbacks.pre_sampler_sync_fn(scheduler_outputs)
16
    self.maybe_gc_collect()

两处关键设计：

• correction 的 D2H 同步。第二篇详细讲过 ZOS 的 seqlens 修正，下文的「跨拍的 seqlens 修正」会把这三个 kernel 的算术完整展开；这里要注意的是：D2H 同步只发生在 varlen 模式（DEEPINFER_ENABLE_MTP_VARLEN），而且同步的是修正后的 seqlens_q，不是采样结果。非 varlen 路径下 CPU 拿到的 seqlens_q 直接沿用调度器里的值。
• pre_sampler_sync 是一个「引擎钩子」。名字说的是「采样器同步前」，但实际做的远不止同步：释放上一拍的 DeferredRelease 资源、KV 落盘、结构化输出的 matcher pump、完成序列 finalize_seq。它被放在采样 kernel launch 之后、sample_end_evt.record() 之前，是因为此刻 GPU 在跑采样，而 CPU 做这些纯 CPU 工作完全不需要碰 GPU——它们被「藏」在采样 kernel 后面。

采样之后立即记录 sample_end_evt，并把它打进 WorkerStepState。这个事件对象是下一拍的同步依据。

第 4 段：本拍 post-process 或延后

1
if is_prompt or not envs.DEEPINFER_ZERO_OVERHEAD_SCHEDULING:
2
    request_outputs = np.concatenate([request_outputs, self._post_process_step(curr_step_state)])
3
    self.prev_step_state = None
4
else:
5
    self.prev_step_state = curr_step_state

prompt 步不走延后路径，因为 prefill 是「一次性」的：它的 token 不需要立即进入下一次调度，延后没有收益，反而会拖慢 TTFT 之外的其他路径。decode 步在 ZOS 打开时把 curr_step_state 挂到 self.prev_step_state，留给下一拍。

WorkerStepState 的字段必须活到下一次 _post_process_step，源码注释写得很直白：

Must remain valid until the next step’s _post_process_step, where engine post_process releases them back to the pool.

这是理解 DeferredRelease 的钥匙：seq_metadata_list 来自 SequenceMetadataPool，request_outputs 来自 RequestOutputPool，两者的所有权都在引擎，因为池在 C++ 侧。Worker 只是借用，必须在用完之后通过回调归还。

零拷贝的 SchedulerOutputs

SchedulerOutputs 是 C++ 对象（csrc/engine/scheduler.hpp 第 46 行起），Python 只能看到 pybind 暴露的字段：

1
.def_readonly("prompt_run", &SchedulerOutputs::prompt_run)
2
.def("is_empty", &SchedulerOutputs::is_empty)

clone_states、clone_blocks、zero_states 是 torch::Tensor，它们之所以能被 np.int32 数组接收，是因为 C++ 侧用零拷贝视图把它们包成 numpy 数组交给 Python。step() 里直接写 clone_states.size > 0，而不是 len(clone_states)，说明 Python 看到的是 numpy 风格的数组接口。

生命周期约束：这些视图指向调度器内部 vector 的存储，pybind 注释要求「下一次 schedule() 之前消费完」。第十六篇结尾已经点过这个约束——如果 Python 侧把它们存进队列跨步使用，就会在下一次调度时被静默覆写。

SequenceMetadata 与对象池

seq_metadata_list 与 scheduled_seqs 同序。SequenceMetadata 由 SequenceMetadataPool 分配，C++ 侧持有 shared_ptr，Python 侧拿到的是一组引用。schedule() 在构造输出时顺手把已完成序列的状态写回，并释放对应的序列组。此时序列的 KV 块和状态槽可能还在被 GPU 使用（本拍前向尚未执行），所以释放是「立即」的，但块的归还通过 zero_states / clone_unpins 的延后机制处理——这是调度器和引擎之间的另一层默契。

outputs.hpp 里的 RequestOutputPool 更简单：调用方拿到 uint64 指针数组，解码为 Python 对象时才知道具体字段。build_request_outputs 在 C++ 中直接构造，Python 侧只做薄包装。这种「指针数组」设计的好处是：

• 输出构造的循环完全在 C++ 里跑，Python 不需要为每个请求做一次对象分配；
• 返回值跨 _run_workers 时只是一个小 numpy 数组，没有 pickle 开销；
• 释放与否由 DeferredRelease 统一决定，生命周期是集中的。

代价是 RequestOutput.clone(raw_output) 这一步——entrypoints/llm._run_engine 必须显式克隆，否则下游拿到的对象会在下一拍被池回收并覆写。源码注释写得很明确：

RequestOutput will be automatically released in the next step() call. If you need to retain them beyond the current step, you must call pool.clone(...).

InputMetadata 的另一条路

SchedulerOutputs 与 SequenceMetadata 只到 Worker 的 execute_model_step；从 Worker 到模型，还要再构造一次 InputMetadata。这一步由 ModelRunner._prepare_input 完成，走的是完全不同的路径：

• 非 CUDA Graph 路径：_prepare_input_tensors 调用模型元数据的 prepare_inputmetadata_dict，得到一个 dict，构造为 InputMetadata，再 recursive_fn(lambda x: x.cuda(non_blocking=True)) 全部非阻塞搬到 GPU。
• CUDA Graph 路径：_prepare_input_tensors 直接往 cuda_graph_used.input_metadata_source.to_dict() 里原地填，不返回新对象；GPU 侧看到的是同一块 input_metadata_device 缓冲区。第十二篇讲过 ZOS 下 input_on_device=True 的捕获行为，这里补一句：这正是「输入在设备上构造」的真正含义——CPU 写的是 pinned source，Device 端持有的是另一份只读副本。

graph_size 把这两条路径统一起来：它不是 total_q，而是选中的那张图对应的 batch 大小。padding_seq_metadata_list[: graph_size - total_q] 就是补齐用的虚拟序列。graph_size 被写进 WorkerStepState，再喂给 update_mtp_runtime——估计器按图的规模而不是真实请求数来分桶，这与第十二篇讲的「按档位缓存」口径一致。

问题：调度器在第 N+1 拍时并不知道第 N 拍采了几个

回顾一次 decode 步的正常口径：一个序列在第 N 拍输入 1 个 token（或 nextn+1 个 token），采样后得到 1 个（或若干被接受的）token，于是第 N+1 拍的 seqlen_k 应该等于”历史长度 + 第 N 拍实际写入 KV 的 token 数”。在非 ZOS 路径里，post_process 在采样后立即把接受数写回 seq.num_tokens，第 N+1 拍的调度器读到的是准确值。ZOS 路径把 post_process 挪到第 N+1 拍，调度器在 schedule(N+1) 的那一刻拿到的是乐观估计：它按”本拍最多能接受 nextn+1 个 token”来记账。于是：

• seqlens_k 会被高估，最多高出 nextn - accepted；
• 在 varlen 置信度调度下，seqlens_q 是”按置信度把总 token 预算分给各序列”的结果，也必须用上一拍真实的接受数与置信度来重算。

如果不修正，注意力 kernel 会在第 N+1 拍读到一个比真实 KV 更长的 seqlen_k，读到尚未写入的槽位；seqlen_q 对不上采样时的输入位置则会让 KV 写入的 mtp_info_write_slot 错位。

解法：在 GPU 上用两个融合 kernel 改回来

Worker 在 execute_model_step 的第 1 段之前调用 _correct_seqlens_qk，分三步：

1
remap_cpu, seqlens_q_cpu, seqlens_k_cpu, has_overlap = prepare_zos_inputs(
2
    seq_metadata_list, prev.seq_metadata_list, n_seq_with_padding, self.nextn
3
)
4
if not has_overlap:
5
    return None
6
remap = remap_cpu.cuda(non_blocking=True)
7
seqlens_q = seqlens_q_cpu.cuda(non_blocking=True)
8
seqlens_k = seqlens_k_cpu.cuda(non_blocking=True)
9

10
if envs.DEEPINFER_ENABLE_MTP_VARLEN:
11
    compute_seqlens_q_by_confidence_rates(total_q, remap[:n_seq], prev.sampler_output.confidence_rates_gpu, seqlens_q[:n_seq])
12
if self.nextn > 0:
13
    compute_seqlens_k_zos(seqlens_k[:n_seq], seqlens_q[:n_seq], remap[:n_seq],
14
                          prev.sampler_output.num_accepted_tokens_gpu, nextn=self.nextn)

第一步：prepare_zos_inputs（C++）建 remap。 它在上一拍与这一拍的批次之间建立映射 remap[i] = j，要求 remap[i] 指的上一批序列与本批第 i 条是同一个 seq_id，否则记 -1。它先尝试一条快路径：如果两批长度相同、且逐下标 seq_id 都相等（并且上一拍不是占位步），直接 remap[i] = i；否则退化成哈希表查找。全部为 -1 时返回 has_overlap=false，Python 侧直接跳过修正——这一拍的序列全部是新加入的，没有”上一拍”可参考。

这个函数同时给出两个初值：

其中 seqlens_k 初值先把”本拍将要写入的 q 个 token”从 k 里扣掉，留作下一步再加回实际接受的部分。padding 槽位（i >= len(seq_metadata_list)）填 remap=-1, seqlens_q=1, seqlens_k=1，与 CUDA Graph 的图规模对齐。

第二步：compute_seqlens_q_by_confidence_rates（TileLang kernel）。 它只处理上一拍留下的”尚未真实消费”的 token：先把上一拍的置信度按 remap 散射到本拍序列上（新序列视为 0），沿 nextn 维做累乘得到 cum_rates[i, p]——含义是”第 p 个草稿 token 及它之前全部被接受的概率”。然后把 total_q - num_seqs 个额外 token 预算按 cum_rates 从大到小逐个分配，平局按下标优先。结果是整批 seqlens_q 之和恰好等于 total_q，预算没有浪费。

第三步：compute_seqlens_k_zos（TileLang kernel）。 它把实际接受数补回 k：

源码 compute_seqlens_k_zos.py 的参考实现把这段写成三行 torch 运算：先 seqlens_k += seqlens_q，再用 torch.where(mask, nextn - accepted, 0) 减去”上一拍没被接受、被乐观记进去的那部分”。融合 kernel 把这三次遍历合成一次。

为什么修正必须发生在 GPU 上、并且在第 1 段之前

修正依赖的输入（confidence_rates_gpu、num_accepted_tokens_gpu）都在 GPU 上，调度时还没有 D2H 回来。如果让 CPU 先同步取值再算，ZOS 的整条重叠链就断了。所以 DeepInfer 的选择是：在 GPU 上把 seqlens_q/k 原地改好，作为 MetadataZOS 的一部分直接喂给 execute_main_model，CPU 侧只需要拿到 remap（小数组，由 C++ 在 pinned 内存里生成）。

MetadataZOS 这个 dataclass 把修正结果打包得非常完整：

注意 remap 的文档写法：remap[i] = j 若本批第 i 条与上一批第 j 条匹配，否则 -1。这个方向（本→上）决定了后面两个 kernel 的散射语义，方向写反会让整批序列的接受数错配，症状是输出出现重复或跳字。

D2H 与 update_metadata_for_sampling_zos

第三步之后，varlen 模式还会把修正后的 seqlens_q 从 pinned 缓冲拷回 CPU，并记录一个 d2h_evt：

1
if envs.DEEPINFER_ENABLE_MTP_VARLEN:
2
    seqlens_q_cpu[:n_seq].copy_(seqlens_q[:n_seq], non_blocking=True)
3
    d2h_evt = torch.cuda.Event()
4
    d2h_evt.record()

采样前，Worker 同步这个事件，然后调用 C++ 的 scheduler_outputs.update_metadata_for_sampling_zos(...)，把 CPU 侧的采样元数据（token 位置、投机验证的起点）刷新到修正后的长度上。如果少了这一步，采样 kernel 会按”乐观记的 seqlen_q”去读 logits 位置，接受率统计就会和实际前向对不上。

这条 D2H 只在 varlen 模式发生，而且它同步的是 seqlens_q 的拷贝而不是整条采样链——这也是为什么 pre_sampler_sync 仍然能藏在采样 kernel 后面：真正挡住 CPU 的只有这个小数组的往返，不是整个 batch 的 logits。

工程含义

• 修正只在有交集时发生。大批量 decode 时，相邻两拍的批次交集通常很高，修正几乎每拍都跑；小批量时容易整批换人，has_overlap=false 直接跳过。这等价于”修正成本随批次稳定性自动增减”，而不是固定开销。
• CUDA Graph 与 padding 槽。_correct_seqlens_qk 里的 n_seq_with_padding 用的是图选中的 cuda_graph_batch_size，而不是真实请求数；padding 槽的 remap=-1 保证它们不会被散射到上一拍的任意序列上。
• 它与投机解码的置信度调度是同一个账本。compute_seqlens_q_by_confidence_rates 决定”下拍每条序列带几个 token 进前向”，compute_seqlens_k_zos 决定”注意力读到多长”。前者是预算，后者是边界；两边必须基于同一份 prev.sampler_output，否则会出现”预算算对了但 KV 读错了”的隐性 bug。

为什么没有未完成请求时还要 flush

ZOS 把上一拍 decode 的 post_process 挂在 prev_step_state 里。当调度器认为”所有请求已完成”时，上一拍的采样结果可能还没被消费。如果直接退出循环，最后 1–nextn 个 token 永远不会出现在 RequestOutput 里，用户看到的输出会短一截。

flush_inflight_reqs 正是为此存在：

1
def flush_inflight_reqs(self):
2
    request_outputs = self._run_workers("flush_inflight_reqs")
3
    self.deferred_release.collect_request_outputs(request_outputs)
4
    self.process_output(request_outputs)
5
    return request_outputs

Worker 侧只做一件事：_post_process_step(self.prev_step_state)，然后清空 prev_step_state。引擎侧把它收集进 DeferredRelease，再调用虚函数 process_output（ZOS 下处理的是”上一拍输出的 token”）。注释也承认这一步”may cause CPU-GPU synchronization”，因为等待的是最后一拍的采样结束事件，而不是被下一拍前向遮住的事件。

has_unfinished_requests 为什么要跨 DP 归约

1
unfinished = self.scheduler.has_unfinished_seqs()
2
unfinished = unfinished or self.so_builder_has_pending()
3
if dp_group.size() > 1:
4
    unfinished = torch.full((1,), unfinished, dtype=torch.int64, device=device)
5
    dist.all_reduce(unfinished, op=dist.ReduceOp.MAX, group=dp_group)
6
    unfinished = unfinished.item() > 0

如果每个 rank 独立判断”我这边空了就退出”，会出现两种灾难：

1. 还有请求的 rank 继续调用 step()，空了的 rank 已退出循环，下一次 DP Attention 的 All-Gather 永远等不到对端；
2. 空了的 rank 提前销毁进程组，未完成的 rank 在 collective 上挂起。

MAX 归约保证”只要全组有一个 rank 还有活，所有 rank 都继续走”，这就是”一起停”的语义。ZOS 打开时，代码把通信组换成 PG.cpu_data_parallel_group()、设备换成 CPU：这样这次归约不触碰 CUDA 流，不会在重叠窗口里插入一次设备同步。非 ZOS 时用 CUDA 组，unfinished 直接留在 device 上，省一次 D2H，但 item() 本身已经是一次同步——因为非 ZOS 路径本来就要在采样后同步，多这一次不改变关键路径。

三种引擎类型

EngineType 只有三个值，C++ 侧用同名枚举与 Python 保持一致（engine_type.py 的注释明确要求这一点）：

1
class EngineType(enum.Enum):
2
    NORMAL = 0       # process prompts & completions
3
    PROMPT = auto()  # only process prompts
4
    COMPLETION = auto()  # only process completions

EngineArgs.__post_init__ 用它来设置两个环境变量的默认值：

1
os.environ.setdefault("DEEPINFER_ZERO_OVERHEAD_SCHEDULING",
2
                      "1" if self.engine_type.default_zero_overhead_scheduling else "0")
3
os.environ.setdefault("DEEPINFER_USE_MULTIMODAL_POOL",
4
                      "1" if self.engine_type != EngineType.COMPLETION else "0")

而 default_zero_overhead_scheduling 的属性是：

1
return self in (EngineType.NORMAL, EngineType.COMPLETION)

这行代码给出了 P/D 分离的第一个非对称性：PROMPT 引擎默认不启用 ZOS，COMPLETION 引擎默认启用。原因不难从 step 的语义推出来：prefill 的「后处理」没有跨拍依赖，启用 ZOS 只会把 TTFT 拉长；decode 的每一步都要重复「调度—前向—采样—收账」，才有把收账藏到下一拍的价值。这也解释了 step() 里那条断言：

1
if not scheduler_outputs.is_empty() and self.engine_type == EngineType.PROMPT:
2
    assert scheduler_outputs.prompt_run, "Prompt engines should not perform decoding"
3
    assert seq_metadata_list[0].seqlen_q > 0, "seqlen_q should be > 0"

PROMPT 引擎永远不应该出现 decode 步。反方向上，COMPLETION 引擎在构造 C++ 调度器时被 TORCH_CHECK 禁止「允许抢占后重算」（见第十六篇），因为它的显存预算按「只跑 decode、短序列、大并发」配置，重算一条长 prompt 的成本不可接受。

钩子分布在握手的哪几个位置

P/D 分离需要跨进程交接 KV。交接点不在 Worker 里，而在引擎这一侧的两个回调里——这一点从虚函数的调用位置可以直接读出来：

1
def pre_sampler_sync(self, scheduler_outputs):
2
    self.deferred_release.release()
3
    if not scheduler_outputs.prompt_run:
4
        if self.copy_state_event is not None:
5
            self.copy_state_event.synchronize()
6
        self.do_save_kvcache(scheduler_outputs.scheduled_seqs, scheduler_outputs.finished_seqs)
7
    else:
8
        self.pe_prepare_save_kvcache(scheduler_outputs.scheduled_seqs, scheduler_outputs.finished_seqs)
9
    if self._so_matcher_builder is not None:
10
        self._so_matcher_builder.pump()
11
    if self.engine_type in [EngineType.COMPLETION, EngineType.NORMAL]:
12
        for seq in scheduler_outputs.finished_seqs:
13
            self.finalize_seq(seq)

四个钩子各有分工：

pe_prepare_save_kvcache 与 pe_do_save_kvcache 的两段式位置很有讲究：prepare 位于采样 kernel 已经 launch、但采样结果还没被 CPU 消费的时刻；do 位于 post_process 拿到采样结果之后。前者可以做与采样重叠的准备（比如把 KV 块句柄、状态槽编号整理进待发送队列），后者才能确定这一条 prompt 最终产生了多少 token、块表长什么样。prompt 步通常没有投机解码，所以两段之间的差异不是「接受数未知」，而是「采样输出（含 logprobs、SO status）未知」——pe_do_save_kvcache 放在 _process_model_outputs 之前，说明它只需要序列长度和块表，不需要采样结果；两段式更像是把「可在采样期间做的准备」和「必须在采样后做的提交」分开。

send_processed_prompts 与 finalize_seq 是互斥的两条收尾路径：PROMPT 引擎把完成的 prompt 序列送走，所以不能在 pre_sampler_sync 里 scheduler.free_seq；COMPLETION / NORMAL 引擎没有下游，才在 pre_sampler_sync 里真正释放。

开源仓库能读到哪，读不到哪

必须把证据边界说清楚：上面这些钩子在开源仓库里全部是空实现，注释统一写着「虚函数」。具体的 PROMPT / COMPLETION 子类不在 deepseek-ai/DeepInfer 里。官方部署问答也明确了原因：

The KV cache management and load balancing components of the inference system are tightly coupled to our internal infrastructure, so that code is not yet open-sourced.

所以对 P/D 分离的源码级分析只能停在「引擎在什么时刻调用哪个钩子、传什么参数」这一层，不能进一步断言子类内部如何编码 KV、如何选路。能确认的是：

• 钩子挂在引擎而不是 Worker 上，说明 KV 的保存/发送以 scheduler_outputs 的序列账本为准，而不是以 GPU 上的张量为准；
• pre_sampler_sync 先 deferred_release.release() 再落盘，说明「释放上一拍资源」必须早于「为本拍保存 KV」，否则被释放的块可能在落盘之前被重新分配；
• decode 侧的 copy_state_event.synchronize() 在 do_save_kvcache 之前，说明落盘读的是拷贝缓冲而不是活跃状态槽，活跃槽在本拍之后还会被继续写。

官方问答给出的部署口径补上了外部约束：prefill 用 EP2–8 / DP2–8、decode 用 EP64–128 / DP64–128，P 的 GPU 比例大致 1:1；KV 主缓存放在 3FS，SWA 缓存留在本地主机内存，跨机命中走 400G VLAN。这些数字属于部署层，和本篇的握手协议是正交的，但它解释了为什么 pe_* 钩子必须存在：P 侧跑完的序列不是「结束」，而是「移交给另一个进程继续 decode」，引擎的账本因此要跨进程转移。