平衡调度重构¶

摘要现有的 patch_balance_schedule.py 逐字复制了三个大型上游单元，以注入大约5行实际逻辑：Scheduler.schedule()（约520行）、DPEngineCoreProc.run_busy_loop()（约40行）和 EngineCoreProc.run_engine_core()（约55行）。本次重构在**严格保留 balance_flag 语义**的前提下，首先删除了两个已经过时的副本 run_busy_loop() / run_engine_core()（将其替换为基于 _has_global_unfinished_reqs 的引擎核心钩子，以及用于条件激活的模块级名称交换）。schedule() 的副本**暂时保留**——上游没有提供更细粒度的钩子供借用，删除它依赖于向上游贡献一个覆盖接缝，计划在后续的 Phase 2B 中完成。因此，该文件并未缩减到几十行：schedule() 的主体仍然是上游的逐字副本（与发布标签 v0.23.0 逐字对齐，仅包含3处平衡增量），文件约830行。与稳定的发布标签（而非变动的主分支验证提交哈希）对齐，使得“与上游逐字比较”成为可复现的漂移检查——固定标签在每次CI运行中都指向相同的源代码。本轮实际移除的是 run_busy_loop / run_engine_core 副本的陈旧漂移风险，并修复了一个启用平衡时的死锁（最初是“在 schedule() 内部收集”，后来在一次迭代中改为“在 _process_engine_step 内部收集”）；收集操作现在紧跟在 _has_global_unfinished_reqs 跨秩全归约之后。

背景¶

平衡调度的作用¶

当 data-parallel-size 较大且并发度 ≈ DP × max-num-seqs 时，请求往往会堆积在部分DP秩上：饱和的秩同时处理预填充和解码，速度变慢，而其他秩则不断接纳新请求，导致差距扩大。平衡调度**不会**主动重新平衡每个秩的运行计数。相反，它提供了一个**全局准入门控**：一旦**任何一个秩**的运行计数达到上限，所有秩**都停止从WAITING队列接纳新请求，让饱和的秩有机会排空其正在处理的请求，从而阻止差距继续扩大。它**不是“让落后的秩追上领先的秩”（该语义在此处**被明确拒绝**——请参见行为保持契约第1条）。

该功能通过 additional_config.enable_balance_scheduling = true 启用（环境变量 VLLM_ASCEND_BALANCE_SCHEDULING 已弃用）。它仅支持PD混合模式；验证逻辑位于 vllm_ascend/platform.py 和 vllm_ascend/ascend_config.py 中。

实际逻辑（仅两处）¶

整个功能归结为两个操作：

运行计数的跨秩同步 — 每个引擎步骤执行一次 all_gather，收集每个秩的 len(self.running)：

def balance_gather(self):  # dp_group is injected into self.dp_group by the engine core
    running_tensor = torch.tensor([len(self.running)], dtype=torch.int, device="cpu")
    dist.all_gather(self.balance_queue, running_tensor, group=self.dp_group)

WAITING调度循环内部的准入门控 — 由于每个秩都持有相同的收集向量，此检查在每个秩上产生相同的结果。只要在上一步骤结束时**任何一个秩的**运行计数达到上限，**所有秩**在此步骤中都停止接纳新的WAITING请求：

balance_flag = max(t.item() for t in self.balance_queue) == self.max_num_running_reqs
if balance_flag:
    break

语义（必须逐位保留）： "领先者达到上限 ⇒ 全局冻结准入"。它**不是** "让落后的秩追上领先的秩"。请参见行为保持契约。

问题陈述¶

为了注入上述两个片段，现有的 patch_balance_schedule.py 逐字复制了三个大型上游单元：

复制的单元	行数	复制原因
`Scheduler.schedule()`	~520	在循环中间插入3行的 `balance_flag` 门控
`DPEngineCoreProc.run_busy_loop()`	~40	在每一步之后调用 `balance_gather`
`EngineCoreProc.run_engine_core()`	~55	当DP>1时替换为 `BalanceDPEngineCoreProc`

这种“复制整个单元”的方法有三个具体危害：

schedule() 的副本现已与发布标签（生产固定版本，当前为 v0.23.0）逐字对齐。 发布标签的**唯一真实来源**是 .github/vllm-release-tag.commit（CI通过 tr -d '[:space:]' < .github/vllm-release-tag.commit 读取同一文件），当前为 v0.23.0；开发/CI实际安装的是 .github/vllm-main-verified.commit 指向的主分支验证提交（该提交已包含 throttle_prefills 和其他 v0.23.1+ 的演进）。旧补丁复制了一个比 v0.23.0 更旧的 vLLM 中的 schedule()，因此整体已经过时。本轮将 schedule() 副本**逐字对齐到发布标签的 Scheduler.schedule()，仅保留3处平衡增量（禁用路径的提前返回、balance_flag 门控、if request_queue is None: break）；run_busy_loop() / run_engine_core() 的副本已在 Phase 1 中删除。**注意：本文档中任何具体的 v0.23.0 只是固定版本文件当前值的快照——随着固定版本推进它会过时，绝不能用作版本权威依据；任何需要此标签的代码/测试必须在运行时读取该文件。

为什么对齐到 v0.23.0 标签而不是已安装的 main-verified 提交？ 两个原因：(a) 生产环境实际运行的是 v0.23.0 版本，因此将副本对齐到该版本可保持生产行为与运行时一致；(b) 固定的 git 标签在每次 CI 运行时都指向**相同**的源码，因此“对副本与上游进行逐字比较（仅允许 3 处差异）”成为**可重现**的漂移检查——而移动的 main-verified 哈希会随着每次提交向前漂移，无法作为稳定的护栏。

代价与边界： 副本（v0.23.0 逻辑）与 main-verified 运行时的行为略有差异，但 balance 的实际调度路径仅在 NPU + DP + MoE 下触发，CPU 单元测试不会涉及（参见测试计划）；这些差异不影响门控本身的语义。关键点：vllm-ascend CI 同时运行两个 vllm 版本（发布标签 v0.23.0 + main-verified 提交 1f486d96），它们的引擎调用 schedule() 的方式不同——v0.23.0 调用 schedule()，1f486d96 调用 schedule(throttle_prefills)。因此，重写签名采用**两个版本的并集**（schedule(self, throttle_prefills=False)，一个带默认值的超集），使得两个引擎都能调用它；禁用路径委托给 super() 并使用**签名自省**（_SUPER_SCHEDULE_HAS_THROTTLE，在导入时决定一次）来决定是否传递 throttle_prefills，而不是使用版本字符串——这在两条路径上都是正确的，并且不受开发检出分支的非标准 PEP 440 __version__ 影响（后者会导致 vllm_version_is 抛出异常）。即：主体对齐发布标签；签名采用两版本并集；禁用路径使用签名自省。

它违反了 AGENTS.md 补丁策略。 该策略要求补丁“最小且聚焦”，并“有长期计划向上游贡献”。一个 500 行的逐字副本是无法审查的（不对比上游差异就无法看到真正的改动），并且每次 vLLM 升级都必须手动重新同步。
静默且未记录的偏差会累积。 例如，重写悄悄地将上游的 assert request_queue is not None 改成了 if request_queue is None: break。此类偏差会使未来的差异变得不可信。

本轮经验教训（我们遇到的一个坑——schedule() 签名在不同版本间漂移 + 双版本 CI）。vLLM 的 Scheduler.schedule() 签名在不同版本间会变化：发布标签 v0.23.0 是 schedule(self)（引擎调用 schedule()），main-verified 提交 1f486d96 增加了 throttle_prefills: bool = False（引擎调用 schedule(self._should_throttle_prefills())）。而 vllm-ascend CI 同时运行两者，因此重写不能硬编码任一签名——schedule(self) 在 1f486d96 上，当引擎传入参数时会引发 TypeError。正确的修复：重写签名采用**两个版本的并集** def schedule(self, throttle_prefills: bool = False)（一个带默认值的超集，两个引擎都能调用）；禁用路径委托给 super() 并使用**签名自省**（在导入时计算一次：_SUPER_SCHEDULE_HAS_THROTTLE = "throttle_prefills" in inspect.signature(Scheduler.schedule).parameters）来决定是否传递 throttle_prefills，而不是使用 vllm_version_is("0.23.0") 版本字符串——后者在开发检出分支（非 PEP 440 ValueError）上会引发 __version__，并且从根本上将事实性问题“super() 接受这个参数吗？”重构为脆弱的“版本字符串匹配吗？”。结论：当 CI 同时运行多个上游版本时，重写签名采用并集，版本相关的分支由签名自省驱动，而非版本字符串。 在单元测试层面，“签名等于已安装 vLLM 的签名”在双版本下是**错误**的（两条路径的已安装签名定义上就不同，因此相等性断言永远只能通过一条路径）；它被替换为“重写可被两个引擎的调用形态绑定，并且是已安装签名的超集”。

设计¶

原则¶

移除副本，而非功能。将收集逻辑拉入调度器本身（这样两个 EngineCore 副本就可以删除），并在未来通过一个最小的上游接缝注入门控（这样 schedule() 副本就可以删除）。balance_flag 语义严格保持不变。

实现状态与修正（本轮落地了阶段 1 + 2A + 3）¶

在实现过程中，我们验证了上游的实际结构，发现两个草案假设不成立；两者均已在此修正：

上游 Scheduler 没有 new_step_starts() 生命周期钩子（这是 kv_cache_manager 的方法），也没有可覆盖的“每步调度开始”接缝。因此 schedule() 不能放在 _has_global_unfinished_reqs 内部（参见步骤 1 — 将 gather 挂接到并删除 EngineCore 副本_has_global_unfinished_reqs中的死锁教训）；而是放在引擎核心的 dp_group 上。
调度器无法延迟获取 DP 组： _init_data_parallel 在 BalanceDPEngineCoreProc 中生成（早于调度器创建），且没有全局注册表。因此 _has_global_unfinished_reqs 不会被删除，而是精简为一个覆盖 dp_group 的子类（注入 balance_gather 并调用一次 run_engine_core）；DPEngineCoreProc 的副本被替换为修补模块级 run_engine_core 名称（上游的 VLLM_ASCEND_BALANCE_SCHEDULING 在调用时通过模块全局名称解析此类），且此替换**仅在启用 balance 时发生**（条件激活）。

因此，下面的阶段 1 描述和“重构后的文件结构”都已重写以匹配实际实现。阶段 3 将配置探测简化为两个回退（AscendConfig → additional_config），并**移除了直接的环境变量读取**——AscendConfig 仍由 additional_config 集中解析（作为 _balance_scheduling_enabled 的已弃用回退），但 AscendConfig 不再自行读取它，从而绕过了 ](#phased-rollout)。详情请参见[分阶段推出_has_global_unfinished_reqs。

步骤 1 — 将 gather 挂接到 `BalanceDPEngineCoreProc.run_busy_loop()` 并删除 EngineCore 副本¶

旧的 run_engine_core() 和 run_busy_loop 是上游方法的逐字副本，并且**已经过时偏离**：上游的 while self._handle_shutdown() 切换为 eep_scaling_state，添加了 is_sleeping / raise SystemExit 保护以及尾随的 while True，而补丁仍然使用 run_engine_core 和手写的信号处理器；SignalCallback 同样增加了 dp_group、numa、tracer 逻辑。目标是删除这两个副本。

实现过程中出现了三个约束条件：

调度器无法自行获取 DP 组。 parallel_config.stateless_init_dp_group() 由 DPEngineCoreProc._init_data_parallel 内部的 EngineCoreProc.__init__ 创建，并存储在引擎核心上；该方法在 super().__init__() 中运行，早于 dp_group（后者创建调度器），并且没有用于延迟查找的注册表。因此，引擎核心必须将 balance_gather 传递给调度器。
schedule() 必须在每个活跃（非空闲）wave 的每次迭代中的每个 rank 上运行（参见下面的死锁教训）。schedule() 不满足此要求——一个本地已排空的 rank 会获取一个虚拟批次，并且永远不会进入 AsyncRecomputeScheduler。
Balance 不得侵入未启用它的配置（例如，PD 分离的重计算 / BalanceDPEngineCoreProc），因此 balance_gather 的替换必须是有条件的。

因此，最终的落地方案是：

BalanceScheduler 被拉入 self.dp_group（无参数签名，使用 schedule()），但**不是从 BalanceDPEngineCoreProc 中调用**——引擎核心每步触发它（见下文）。
_has_global_unfinished_reqs 不会被删除，而是精简为一个覆盖：它挂接 super()._has_global_unfinished_reqs()，首先调用 dp_group（每 32 步运行一次跨 rank 的 all-reduce），然后在同一调用内部注入 balance_gather() 并调用一次 run_busy_loop。上游的 _has_global_unfinished_reqs 在每次迭代的每个非空闲路径上恰好调用一次 schedule()（包括已排空 rank 永远不会进入 run_busy_loop 的虚拟批次迭代），因此每个 rank 在每个活跃步骤都参与 gather。run_engine_core 的主体不再被复制。
run_engine_core 的副本被完全删除。 上游的 DPEngineCoreProc（一个静态方法）在其主体内部通过模块全局名称解析 engine_core = DPEngineCoreProc(*args, **kwargs)（](https://github.com/vllm-project/vllm)，参见 [vllm/v1/engine/core.pyrun_engine_core）。因此，一个薄包装器包装了 vllm_config：在其入口处（_balance_scheduling_enabled 可用时），它通过 DPEngineCoreProc 决定是将模块级的 BalanceDPEngineCoreProc 替换为 run_engine_core 还是恢复上游原始版本，然后调用原始的 SignalCallback。这是**条件激活**——关闭 balance 时，逐字使用上游的实现；信号处理、schedule()、numa 和 tracer 都保持上游正确。

教训A — 死锁（gather不能存在于balance_gather中）。 早期版本将BalanceScheduler.schedule()放在self.running的顶部，并论证"在两个步骤之间，schedule()只在update_from_output()/all_gather内部变化，因此快照是等价的，每个步骤仍然只有一个all_gather——安全"。这个论证**只对了一半**：gate看到的值确实等价，但它**忽略了has_requests()是一个集合操作，每个rank必须同步参与**。在DP MoE下，一个已经耗尽本地请求的rank（execute_dummy_batch()为False）会执行schedule()并且**永远不会进入all_gather，因此它会跳过那个_has_global_unfinished_reqs，而一个仍然繁忙的rank会调用它并永远等待——**集合不匹配，死锁。engines_running实际上每32步才进行一次真正的all-reduce，并且_process_engine_step在此期间是粘性的，这一事实扩大了此窗口。

教训B — 死锁（gather也不能存在于_has_global_unfinished_reqs中；它必须紧跟在_process_engine_step all-reduce之后）。 后续的一次迭代通过将gather移动到continue（每次迭代都调用，在同步之前和空闲_has_global_unfinished_reqs gate之前）"修复"了教训A。这**重新引入了一个不同的死锁**：all_gather是繁忙循环中**唯一**重新同步各rank关于波次/空闲状态的同步点。将每个步骤的continue放置在该同步**之前**（以及在空闲_process_input_queue之前）将gather与波次协调解耦。在波次边界——请求在不同时间在不同rank上完成，engines_running阻塞等待下一个波次/新请求，_process_input_queue粘性长达32步——一个rank可能到达gather，而另一个rank仍然阻塞在future.result()或all_gather中。然后sample_tokens死锁；卡住的EngineCore无法再排空其worker共享内存广播通道，worker的RPC call to sample_tokens timed out响应无处可去，60秒后引擎因balance_gather而死亡（间歇性观察到——"5次GPQA运行正常，第6次挂起"——因为触发条件取决于每次运行的完成时间）。这与专家并行是否跨越DP rank无关：故障机制是EngineCore↔worker shm耗尽，而不是worker前向传播。结论：super()._has_global_unfinished_reqs()必须紧跟在engines_running之后，这是唯一的每次迭代跨rank同步点，这样各rank在进入all-gather时刚刚就_has_global_unfinished_reqs达成一致。因为continue只在**没有**走空闲run_busy_loop的迭代中被调用，所以当所有rank都空闲时，gather会被所有rank一致地跳过（没有rank执行额外的gather）——与重构前复制的_has_global_unfinished_reqs相同，其中gather紧跟在all-reduce之后。此教训由单元测试中的_has_global_unfinished_reqs接缝保护锁定。

时序——与重构前逐位一致。 Gather现在运行在schedule()的尾部，即在update_from_output() + 执行 + schedule()之后，并紧跟在跨rank all-reduce之后——正是重构前由engine-core驱动的gather所在的位置。Gate在下一步的schedule()中消费该值；行为不变。

步骤2 — 用最小接缝替换`balance_flag`副本¶

schedule() gate被上游内联在schedule()中间；目前没有可覆盖的接缝。这分两个阶段解决。

阶段2A — 过渡期（无上游依赖）：

保留def schedule(self, throttle_prefills: bool = False)覆盖，但：

重写签名采用双版本联合： schedule()。vllm-ascend CI 同时运行 v0.23.0（引擎调用 schedule(throttle_prefills)）和 1f486d96（引擎调用 throttle_prefills），因此重写携带带有默认值的 super()，两个引擎均可调用。禁用路径通过签名内省委托给 _SUPER_SCHEDULE_HAS_THROTTLE — 在导入时计算一次 super().schedule(throttle_prefills)；如果为真则调用 super().schedule()，否则调用 vllm_version_is("0.23.0")。这取代了旧的 ValueError 版本字符串分支（该分支在开发检出上会引发 super()，并将“父类是否接受该参数？”重新表述为“版本是否匹配？”）。
将平衡变更压缩为3个带有清晰注释的差异点：(1) 禁用路径的早期返回委托给 balance_flag；(2) WAITING 循环内的 if request_queue is None: break 门控；(3) assert（上游使用 schedule()）。由于上游没有更细粒度的钩子，函数体仍需复制。
逐字比较现在可重现： .github/vllm-release-tag.commit 的副本与发布标签对齐（仅3个平衡差异点不同），因此固定的标签使得每次 CI 运行时“与上游的逐字比较”都能产生相同的基线。“意图锁定”测试（签名可被两个引擎调用，存在3个差异行，上游接缝仍然存在）继续作为 CPU 可及的护栏，并新增了一个“与发布标签的逐字比较（仅允许3个差异点）”的漂移测试（参见测试计划）。漂移测试**在运行时从 schedule() 读取标签**（与 CI 同源）——它不硬编码版本也不读取设计文档；当固定点前进时，测试自动与新标签比较并变红以指示“副本需要重新同步”。

“在固定点前进时重新对齐副本”现在是常规维护：每次发布标签前进时，将3个差异点重新应用到新标签的

# upstream vllm/v1/core/sched/scheduler.py
def _should_stop_admitting_waiting(self) -> bool:
    return len(self.running) >= self.max_num_running_reqs

一旦上游暴露了该接缝，Ascend 补丁将简化为：

class BalanceScheduler(Scheduler):
    def _should_stop_admitting_waiting(self) -> bool:
        if super()._should_stop_admitting_waiting():
            return True
        return self._balance_enabled and (
            max(t.item() for t in self.balance_queue) >= self.max_num_running_reqs
        )

（此处 >= 和 == 等价，因为任何 rank 的 len(running) 都不会超过 max_num_running_reqs；下面的契约将语义固定为 ==，以明确“领导者达到上限 ⇒ 冻结”，并拒绝“追赶领导者”的重新解释。）

结果： 约520行的 schedule() 副本被永久删除；该文件不再因上游对 schedule() 的编辑而产生漂移。这是 AGENTS.md 所要求的“向上游贡献的长期计划”。

关于 >= 与 ==，以及被拒绝的“临时降低上限”想法。 一个早期的想法是通过临时设置 self.max_num_running_reqs = min(cap, max(balance_queue)) 来重用上游现有的零拷贝中断条件。这将产生一个**不同**的语义（“让落后的 rank 追赶领导者”），并且被**明确拒绝**——请参见下面的契约。

步骤 3 — 规范化配置探测¶

_balance_scheduling_enabled() 简化为**两个回退选项（AscendConfig → additional_config）**。删除 run_engine_core 副本后，唯一的调用者是 BalanceScheduler.__init__，但此时 AscendConfig 是否已初始化仍无法保证（这是旧文件顶部 TODO 的起源），因此保留 additional_config 作为启动窗口期的回退，否则函数返回 False。本轮相对于旧实现收紧了一件事：

移除了直接的环境变量读取。 旧实现回退到裸的 os.getenv("VLLM_ASCEND_BALANCE_SCHEDULING")，违反了 AGENTS.md 的“无散落 os.getenv”原则。此函数不再自行读取环境变量——VLLM_ASCEND_BALANCE_SCHEDULING 由 AscendConfig 集中解析（作为 additional_config 的已弃用回退），并通过主路径 get_ascend_config().enable_balance_scheduling 生效，避免了多个入口点。
文件顶部的 TODO 更新为“一旦 AscendConfig 初始化被提前，这可以简化为单个 get_ascend_config().enable_balance_scheduling 读取”。

稍后（一旦 AscendConfig 的时机问题解决）：将两个回退合并为单个读取。

行为保持契约¶

此重构**必须**严格保持以下不变性；任何偏离都是错误。

领导者达到容量上限 ⇒ 全局冻结。 balance_flag 为 max(balance_queue) == max_num_running_reqs，根据上一步收集到的各 rank 的 len(running) 计算。当条件为真时，所有 rank 都不再接受新的 WAITING 请求。比较操作是 == 与配置的 max_num_running_reqs 进行比较——不是 >=，也不是“追赶领导者”。
相同输入 ⇒ 相同输出。 给定相同的 self.running、self.waiting、self.skipped_waiting、balance_queue 和 token 预算，重构后的 schedule() 产生的 SchedulerOutput 与当前实现完全相同（相同的已调度/已抢占/已恢复集合，相同的 num_scheduled_tokens，相同的连接器元数据）。
收集节奏不变。 每个活跃引擎步骤恰好执行一次 all_gather，在相同的 DP 组上，负载仍然是 len(self.running)，当所有 rank 都空闲时，所有 rank 一致地跳过。仅调用位置发生了变化。
禁用路径不变。 当 enable_balance_scheduling 为 false 时，_balance_run_engine_core 将模块级别的 DPEngineCoreProc 恢复为上游原始实现，并且引擎核心逐字运行上游的实现；BalanceScheduler 在 _balance_enabled=False 时将 schedule() 委托给 super().schedule()，不分配 balance_queue，并且不执行任何集合通信。即，当平衡功能关闭时，它不会触及任何配置（包括 PD 分离重计算 / AsyncRecomputeScheduler，后者已通过 platform.py 与平衡功能互斥；这是第二层防御）。
现有约束仍然适用。 profiling_chunk_config 互斥锁（参见 vllm_ascend/ascend_config.py）和 PD 混合模式限制（参见 vllm_ascend/platform.py）仍在原处强制执行。

重构后的文件结构¶

经过本轮（阶段 1 + 2A + 3）后，关键结构如下。schedule() 主体是发布标签 v0.23.0 的逐字副本（在阶段 2B 之前不能删除），包含三个已记录的差异（禁用路径的提前返回 + WAITING 循环中的 balance_flag 门控 +

# vllm_ascend/patch/platform/patch_balance_schedule.py
import inspect
import torch
import torch.distributed as dist
import vllm.v1.core.sched.scheduler as _sched_mod
import vllm.v1.engine.core as _engine_core_mod
from vllm.v1.core.sched.scheduler import Scheduler
from vllm.v1.engine.core import DPEngineCoreProc, EngineCoreProc
# ... other vllm imports ...


def _balance_scheduling_enabled(vllm_config) -> bool:
    try:
        from vllm_ascend.ascend_config import get_ascend_config
        return bool(get_ascend_config().enable_balance_scheduling)
    except Exception:
        pass
    additional_config = getattr(vllm_config, "additional_config", None) or {}
    if "enable_balance_scheduling" in additional_config:
        return bool(additional_config["enable_balance_scheduling"])
    return False  # no longer reads the env var itself; VLLM_ASCEND_BALANCE_SCHEDULING is parsed by AscendConfig


class BalanceScheduler(Scheduler):
    def __init__(self, ...):
        super().__init__(...)
        self._balance_enabled = _balance_scheduling_enabled(vllm_config)
        self.dp_group = None  # injected by BalanceDPEngineCoreProc before the first gather
        if self._balance_enabled:
            self.balance_queue = [torch.tensor([0], ...) for _ in range(dp_size)]

    def balance_gather(self):  # uses self.dp_group; no-op when disabled / not injected
        if not self._balance_enabled or self.dp_group is None:
            return
        running_tensor = torch.tensor([len(self.running)], dtype=torch.int, device="cpu")
        dist.all_gather(self.balance_queue, running_tensor, group=self.dp_group)

    def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:  # two-version union: both v0.23.0's schedule() and 1f486d96's schedule(throttle_prefills) bind
        if not self._balance_enabled:  # delta 1: disabled-path early return
            # Whether to forward throttle_prefills is decided by signature introspection
            # (_SUPER_SCHEDULE_HAS_THROTTLE), not a version string -- correct on both CI lanes
            # and unaffected by a dev checkout's bad __version__.
            if _SUPER_SCHEDULE_HAS_THROTTLE:
                return super().schedule(throttle_prefills)
            return super().schedule()
        # NOTE: balance_gather is NOT called here -- see BalanceDPEngineCoreProc.
        # ... upstream schedule() body (verbatim-aligned to the v0.23.0 tag) ...
        #   # inside the WAITING loop (deltas 2, 3):
        #   if max(t.item() for t in self.balance_queue) == self.max_num_running_reqs:  # delta 2: leader-at-cap => global freeze
        #       break
        #   request_queue = self._select_waiting_queue_for_scheduling()
        #   if request_queue is None:  # delta 3: keep if-break (upstream has assert)
        #       break
        # ...


class BalanceDPEngineCoreProc(DPEngineCoreProc):
    """Hook _has_global_unfinished_reqs: inject dp_group + one balance_gather per
    active step. Gather MUST sit immediately after super()._has_global_unfinished_reqs()
    (the only per-iteration cross-rank sync) -- NOT inside _process_engine_step
    (which runs before that sync and before the idle continue gate, and would
    deadlock at wave boundaries -> sample_tokens timeout), and NOT inside
    schedule() (drained ranks skip schedule() and would miss the all_gather)."""

    def _has_global_unfinished_reqs(self, local_unfinished: bool) -> bool:
        result = super()._has_global_unfinished_reqs(local_unfinished)
        self.scheduler.dp_group = self.dp_group
        self.scheduler.balance_gather()
        return result


_OriginalDPEngineCoreProc = _engine_core_mod.DPEngineCoreProc
_OriginalRunEngineCore = EngineCoreProc.run_engine_core


def _balance_run_engine_core(*args, dp_rank=0, local_dp_rank=0, **kwargs):
    # Conditional activation: swap the module-level DPEngineCoreProc only when balance is on.
    if _balance_scheduling_enabled(kwargs.get("vllm_config")):
        _engine_core_mod.DPEngineCoreProc = BalanceDPEngineCoreProc
    else:
        _engine_core_mod.DPEngineCoreProc = _OriginalDPEngineCoreProc
    return _OriginalRunEngineCore(*args, dp_rank=dp_rank, local_dp_rank=local_dp_rank, **kwargs)


# Scheduler is constructed by module-global name when scheduler_cls is unset
# (the PD-mixed balance path); recompute / dynamic-batch / profiling schedulers
# set scheduler_cls and bypass this name, which is correct.
_sched_mod.Scheduler = BalanceScheduler
EngineCoreProc.run_engine_core = staticmethod(_balance_run_engine_core)

本轮删除了约 95 行的 run_engine_core + run_busy_loop 副本及其无效导入，并添加了模块文档字符串、注释和 _balance_run_engine_core 条件激活包装器；净行数几乎没有减少，但**它消除的是陈旧漂移风险**（旧副本已落后于上游的 _handle_shutdown / eep_scaling_state / SignalCallback 演进），并修复了启用平衡时的死锁（先在 schedule() 内部收集，然后在 _process_engine_step 内部收集；现在在 _has_global_unfinished_reqs 之后）。

阶段 2B（上游提供 _should_stop_admitting_waiting）将删除约 520 行的 schedule() 副本，并将文件缩减至大约 60–80 行，且不包含任何逐字上游副本。

测试计划¶

签名 + 意图锁 + 逐字漂移测试（阶段2A）。 断言：（a） BalanceScheduler.schedule 的签名**可被两种引擎调用形状绑定**（schedule() 和 schedule(throttle_prefills=...)）并且是已安装 Scheduler.schedule 参数集的超集——注意，这不是“签名行等于已安装”，因为在双版本CI下，两条通道的已安装签名不同，相等性断言只能在一个通道上通过（这锁定了“双版本 ⇒ 取并集”的经验教训）；（b）3个平衡增量行必须存在于主体中（禁用路径 super().schedule() 委托、WAITING循环中的 balance_flag 门控、if request_queue is None: break）；（c）_balance_run_engine_core 包装器已安装，且 DPEngineCoreProc 未**在导入时被交换（延迟到包装器，在调用时条件性交换）；（d）上游 DPEngineCoreProc._has_global_unfinished_reqs 仍然存在（收集注入点——它必须在每个非空闲迭代中被调用，否则 all_gather 会死锁）；（e）上游 Scheduler 接缝方法（包括 _build_kv_connector_meta、_inflight_prefill_reserved_blocks）仍然存在；（f）**逐字漂移检测——首先从 .github/vllm-release-tag.commit 读取发布标签（与CI相同的来源，非硬编码，非从设计文档读取），然后 git show <tag>:vllm/v1/core/sched/scheduler.py 获取该标签的 schedule()，剥离相同的3个增量，并通过AST逐字比较 BalanceScheduler.schedule 的源代码；两者必须相同。读取锁定文件意味着锁定推进**自动**翻转测试以比较新标签并变红，表示“副本需要重新同步”——这正是我们想要的维护信号；如果锁定文件或标签不可达（例如，vLLM不是源代码检出，在 vllm-ascend 树外运行），则跳过测试，而非失败。
行为等价测试。 使用一个假的DP组和手动设置的 BalanceScheduler 构建 balance_queue，在几种代表性状态下驱动 schedule()（领导者达到容量冻结、落后排名未满、禁用、空等待），并断言 SchedulerOutput 相同（合同项2）。重用 tests/ut/test_platform.py 中的平衡测试框架。
收集节奏测试。 模拟 torch.distributed.all_gather（注意：在 balance_gather 内部，dist = torch.distributed，因此模拟目标必须是 torch.distributed.all_gather，而非 vllm.distributed.all_gather）；断言每个 balance_gather() 恰好执行一次 all_gather，负载为 len(self.running) 和注入的 dp_group（合同项3）。
禁用路径测试。 在标志关闭的情况下，断言 balance_queue 未被分配，all_gather 未被调用，且 schedule() 委托给 super().schedule()（合同项4）。
NPU性能检查。 根据 AGENTS.md 的 NPU 指南，max(t.item() for t in self.balance_queue) 每次步骤触发一次主机同步（不可避免，因为此值驱动主机端控制流）。通过性能分析确认重构**未**引入超出当前同步的额外同步。

分阶段发布¶

阶段	范围	风险	依赖项	状态
1	Hook gather onto `_has_global_unfinished_reqs` (after the cross-rank all-reduce — avoids both the schedule()-skip deadlock and the _process_engine_step wave-boundary deadlock); slim `BalanceDPEngineCoreProc` to that hook; delete the `run_engine_core`/`run_busy_loop` copies; `run_engine_core` wrapper conditionally activates `DPEngineCoreProc`; module-level `Scheduler` swap	Low	none	✅ Done
2A	覆盖签名采用双版本并集（`schedule(self, throttle_prefills=False)`，CI同时运行 v0.23.0 + 1f486d96）；主体与发布标签逐字对齐（仅3个平衡增量）；禁用路径通过签名内省（`super()`）委托给 `_SUPER_SCHEDULE_HAS_THROTTLE`；签名可调用性 + 意图锁 + 发布标签逐字漂移测试	低	无	✅ 已完成
3	Collapse config probing to two fallbacks (AscendConfig → additional_config); remove the direct env-var read (still parsed centrally by AscendConfig)	Low	Phase 1	✅ Done
2B	上游 `_should_stop_admitting_waiting` PR；删除 `schedule()` 副本	中	上游审查	⏳ 待办
测试	漂移回归 / 行为等价 / 收集节奏 / 禁用路径 / NPU性能检查	低	阶段1 + 2A	⏳ 待办（需要NPU）

每个阶段可以独立发布和回滚。阶段1、2A和3可以在同一个发布中落地；阶段2B在上游PR合并时落地。