平衡调度重构¶

摘要 现有 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.24.0，仅有 3 处 balance 差异），文件约 830 行。对齐到稳定发布标签（而不是移动的 main-verified commit hash）使“与上游逐字比较”成为可复现的漂移检查——固定标签在每次 CI 运行中指向相同源码。本轮实际移除的是 run_busy_loop / run_engine_core 副本的过时漂移风险，并修复一个启用 balance 后的死锁（最初是“在 schedule() 内 gather”，后来某次迭代中是“在 _process_engine_step 内 gather”）；现在 gather 位于 _has_global_unfinished_reqs 跨 rank all-reduce 之后立即执行。

背景¶

平衡调度的作用¶

当 data-parallel-size 较大且并发 ≈ DP × max-num-seqs 时，请求往往会堆积在部分 DP rank 上：饱和的 rank 同时承担 prefill 和 decode 而变慢，而其他 rank 继续接纳新请求，差距进一步扩大。平衡调度不会主动重新平衡每个 rank 的运行计数。相反，它提供一个全局准入门控：只要任意一个 rank 的运行计数达到上限，所有 rank 都停止从 WAITING 队列接纳新请求，从而给饱和 rank 一个排空其进行中请求的机会，使差距停止扩大。它不是“让落后 rank 追上 leader”（该语义在此被明确拒绝——见 行为保持契约，第 1 项）。

该功能通过 additional_config.scheduler_config.enable_balance_scheduling = true 启用。旧环境变量不再支持。它仅支持 PD-mixed 模式；校验位于 vllm_ascend/platform.py 和 vllm_ascend/ascend_config.py。

真正的逻辑（只有两处）¶

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

1. 运行计数的跨 rank 同步 — 每个引擎步骤一次 all_gather，收集每个 rank 的 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)

1. WAITING 调度循环内的准入门控 — 由于每个 rank 持有相同的 gather 向量，该检查在每个 rank 上产生相同结果。只要任意 rank 的运行计数在上一步结束时达到上限，所有 rank 在本步都停止接纳新的 WAITING 请求：

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

语义（必须逐位保持）： “leader 达到上限 ⇒ 全局冻结准入”。它不是“让落后 rank 追上 leader”。见 行为保持契约。

问题陈述¶

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

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

1. schedule() 副本现在已逐字对齐到发布标签（生产固定版本，当前为 v0.24.0）。 发布标签的唯一事实来源是 .github/vllm-release-tag.commit（CI 通过 tr -d '[:space:]' < .github/vllm-release-tag.commit 读取同一文件），当前为 v0.24.0；dev/CI 实际安装的是 .github/vllm-main-verified.commit 指向的 main-verified commit（它包含后续调度器演进）。旧补丁复制了来自早于 v0.24.0 的 vLLM 的 schedule()，因此整体已经过时。本轮将 schedule() 副本逐字对齐到发布标签的 Scheduler.schedule()，仅保留 3 处 balance 差异（禁用路径提前返回、balance_flag 门控、if request_queue is None: break）；run_busy_loop() / run_engine_core() 副本已在 Phase 1 删除。注意：本文档中任何具体的 v0.24.0 都只是 pin 文件当前值的快照——随着 pin 前进它会过时，并且不得用作版本权威；任何需要该标签的代码/测试都必须在运行时读取该文件。

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

成本与边界： 副本（v0.24.0 逻辑）与 main-verified 运行时在行为上略有差异，但 balance 的真实调度路径只会在 NPU + DP + MoE 组合下被触达，绝不会被 CPU UT（单元测试）触达（见测试计划）；这些差异不会影响 gate 自身的语义。两个受支持的修订版本——发布标签 v0.24.0 和 main-verified 提交 e5588e49——都暴露 schedule(self, throttle_prefills=False)，因此该重写实现匹配这一共享签名，并且禁用路径将 throttle_prefills 直接转发给 super()。即：方法体与发布标签对齐；签名与两个受支持的修订版本匹配；禁用路径直接委托。

1. 它违反了 AGENTS.md 补丁策略。 该策略要求补丁做到“最小且聚焦”，并有“长期向上游贡献的计划”。一个 500 行的逐字副本无法审查（不与上游做 diff 就看不到真正的改动），而且每次 vLLM 升级都必须手工重新同步。

2. 静默且未记录在案的偏差会不断累积。 例如，该重写实现曾静默地把上游的 assert request_queue is not None 改成 if request_queue is None: break。这类偏差会让未来的 diff 变得不可信。

本轮的经验教训：在 v0.23 支持被移除后，两个受支持的修订版本都暴露相同的 Scheduler.schedule(self, throttle_prefills=False) 契约。因此，兼容性自省和版本字符串分支没有任何价值，应当移除。单元测试现在要求该重写实现的签名在每个 CI 通道中与已安装的上游签名一致。

设计¶

原则¶

删除副本，而不是删除功能。将 gather 拉入调度器自身（这样两个 EngineCore 副本就可以删除），未来再通过一个最小的上游接缝（seam）注入 gate（这样 schedule() 副本就可以删除）。balance_flag 的语义保持严格不变。

实现状态与修正（本轮完成阶段 1 + 2A + 3）¶

在实现过程中，我们核对了上游的真实结构，发现两个起草时的假设并不成立；这里对二者都做了修正：

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

相应地，下面的阶段 1 描述和“重构后的文件形态”都已重写，以匹配实际实现。阶段 3 将配置探测收敛为两个回退（AscendConfig → additional_config），并移除了对环境变量的直接读取。balance 调度现在只通过 additional_config 配置；_balance_scheduling_enabled 不再绕过 AscendConfig。详见分阶段推出。

步骤 1 — 将 gather 挂接到 _has_global_unfinished_reqs 并删除 EngineCore 副本¶

旧的 BalanceDPEngineCoreProc.run_busy_loop() 和 run_engine_core() 是上游方法的逐字副本，并且已经漂移过时：上游的 run_busy_loop 已切换到 while self._handle_shutdown()，加入了 eep_scaling_state / is_sleeping 守卫以及结尾的 raise SystemExit，而补丁里仍是 while True + 手写的信号处理器；run_engine_core 也同样新增了 SignalCallback、numa、tracer 逻辑。目标是删除这两个副本。

实现过程中浮现出三个约束：

1. 调度器自身无法获取 DP 组。 dp_group 由 parallel_config.stateless_init_dp_group() 在 DPEngineCoreProc._init_data_parallel 内部创建，并存储在引擎核心上；该方法在 EngineCoreProc.__init__ 中于 super().__init__()（它会创建调度器）之前运行，而且没有用于惰性查找的注册表。因此引擎核心必须把 dp_group 交给调度器。
2. balance_gather 必须在每个 rank 上、在每次迭代中都运行——只要该迭代属于活跃（非空闲）wave 的一部分（参见下面的死锁经验教训）。schedule() 无法满足这一点——本地已排空的 rank 会取走一个 dummy batch（空批次），并且永远不会进入 schedule()。
3. balance 不得侵入未启用它的配置（例如 PD 分离式重计算 / AsyncRecomputeScheduler），因此 BalanceDPEngineCoreProc 的替换必须是条件性的。

因此最终的落地方案是：

• balance_gather 被拉入 BalanceScheduler（无参签名，使用 self.dp_group），但不会从 schedule() 调用——引擎核心按步触发它（见下文）。
• BalanceDPEngineCoreProc 不会被删除，而是被瘦身为只做一处重写：它钩住 _has_global_unfinished_reqs，先调用 super()._has_global_unfinished_reqs()（每 32 步执行一次跨 rank 的 all-reduce），然后在同一次调用内部注入 dp_group 并调用一次 balance_gather()。上游的 run_busy_loop 在每条非空闲路径上每次迭代恰好调用一次 _has_global_unfinished_reqs（包括本地已排空的 rank 永远不会进入 schedule() 的 dummy-batch 迭代），因此每个活跃步骤中每个 rank 都会参与 gather。run_busy_loop 的方法体不再被复制。
• run_engine_core 副本被完全删除。 上游的 run_engine_core（一个 staticmethod）在方法体内按模块全局名解析 DPEngineCoreProc（engine_core = DPEngineCoreProc(*args, **kwargs)，参见 vllm/v1/engine/core.py）。因此用一个薄包装器包裹 run_engine_core：在其入口处（此时 vllm_config 可用）通过 _balance_scheduling_enabled 决定是将模块级 DPEngineCoreProc 替换为 BalanceDPEngineCoreProc，还是恢复为上游原始实现，然后调用原始的 run_engine_core。这就是条件激活——关闭 balance 时，上游实现被逐字使用；信号处理、SignalCallback、numa 和 tracer 全部保持与上游一致。

教训 A —— 死锁（gather 不能放在 schedule() 中）。 早期版本把 balance_gather 放在 BalanceScheduler.schedule() 的开头，并声称“两步之间 self.running 只会在 schedule() / update_from_output() 内部变化，所以快照是等价的，并且每步仍然只有一个 all_gather —— 安全”。这个论证只说对了一半：门控看到的值确实等价，但它忽略了 all_gather 是集合操作，每个 rank 都必须同步参与。在 DP MoE 下，某个 rank 如果已经耗尽本地请求（has_requests() 为 False），就会执行 execute_dummy_batch() 并永远不会进入 schedule()，因此它会跳过那个 all_gather；与此同时仍在忙碌的 rank 调用它并永远等待下去 —— 集合通信不匹配，死锁。_has_global_unfinished_reqs 实际上每 32 步才真正做一次 all-reduce，并且 engines_running 在这期间保持粘滞，这进一步扩大了这个时间窗口。

教训 B —— 死锁（gather 也不能放在 _process_engine_step 中；它必须紧跟在 _has_global_unfinished_reqs 的 all-reduce 之后）。 后续迭代通过把 gather 移入 _process_engine_step（每次迭代都会调用，位于同步之前、空闲 continue 门控之前）“修复”了教训 A。这重新引入了一个不同的死锁：_has_global_unfinished_reqs 是繁忙循环中唯一一个在 wave/空闲状态上重新同步各 rank 的点。把每步的 all_gather 放在该同步之前（以及空闲 continue 之前），会使 gather 与 wave 协调脱钩。在 wave 边界处——各 rank 上的请求在不同时间完成、_process_input_queue 在下一个 wave / 新请求上阻塞、engines_running 最多 32 步保持粘滞——一个 rank 可能已经到达 gather，而另一个 rank 仍阻塞在 _process_input_queue 或 future.result() 中。此时 all_gather 就会死锁；卡住的 EngineCore 无法再排空其 worker 的共享内存广播通道，worker 的 sample_tokens 响应无处落地，60 秒后引擎以 RPC call to sample_tokens timed out 退出（已间歇性观察到——“5 次 GPQA 运行正常，第 6 次挂起”——因为触发条件取决于每次运行的完成时机）。这与专家并行是否跨越 DP rank 无关：故障机制是 EngineCore↔worker 的共享内存耗尽，而不是 worker 前向传播。结论：balance_gather 必须紧跟在 super()._has_global_unfinished_reqs() 之后——这是每次迭代唯一的跨 rank 同步点——这样各 rank 在进入 all-gather 时刚刚就 engines_running 达成一致。 由于 _has_global_unfinished_reqs 只在未走空闲 continue 分支的迭代上被调用，所以当所有 rank 都空闲时，每个 rank 都会一致地跳过 gather（没有任何 rank 会多做一次 gather）——这与重构前复制的 run_busy_loop 相同，当时 gather 就紧跟在 all-reduce 之后。这一教训已由单元测试中的 _has_global_unfinished_reqs seam 守卫锁定。

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

步骤 2 —— 用最小 seam 替换 schedule() 副本¶

上游在 schedule() 中间内联了 balance_flag 门控；目前没有可覆盖的 seam。这个问题分两个阶段解决。

阶段 2A —— 过渡方案（不依赖上游）：

保留 schedule() 覆盖，但：

• 覆盖方法匹配共享的受支持签名： def schedule(self, throttle_prefills: bool = False)。v0.24.0 和 e5588e49 都暴露了这个签名，禁用路径直接委托给 super().schedule(throttle_prefills)。不再需要旧的 v0.23 兼容分支和签名内省。
• 将 balance 改动收敛为 3 个注释清晰的 delta：（1）禁用路径中委托给 super() 的提前返回；（2）WAITING 循环内的 balance_flag 门控；（3）if request_queue is None: break（上游是 assert）。因为上游没有更细粒度的钩子，函数体仍然必须复制。
• 逐字对比现在可复现： schedule() 副本对齐到发布 tag（只有 3 处 balance delta 不同），因此固定的 tag 使得“与上游逐字对比”在每次 CI 运行中都产生相同的基线。“意图锁定”测试（签名相等、3 处 delta 行存在、上游 seam 仍然存在）仍然作为可在 CPU 上运行的护栏，同时新增了一个“与发布 tag 逐字对比（仅允许这 3 处 delta）”的漂移测试（见测试计划）。该漂移测试在运行时从 .github/vllm-release-tag.commit 读取 tag（与 CI 同源）——它不硬编码版本，也不读取设计文档；当 pin 更新时，测试会自动对比新 tag 并变红，以提示“副本需要重新同步”。
• “在 pin 更新时重新对齐副本”现在属于常规维护：每次发布 tag 前进时，将 3 处 delta 重新应用到新 tag 的 schedule() 上（持续到阶段 2B 删除该副本为止）。

阶段 2B —— 目标（随上游贡献一起落地）：

向上游贡献一个最小重构，将 WAITING 循环的停止条件提取为可覆盖的方法：

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

一旦上游暴露了该 seam，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；下面的契约将语义钉死为 ==，以明确 "leader-at-cap ⇒ 全局冻结"，并拒绝"追赶 leader"的重解释。）

结果: 约 520 行的 schedule() 副本被彻底删除；该文件不再随上游对 schedule() 的修改而漂移。这就是 AGENTS.md 所要求的"长期上游贡献计划"。

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

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

_balance_scheduling_enabled() 收敛为两个回退（AscendConfig → additional_config）。删除 run_engine_core 副本后，唯一调用方是 BalanceScheduler.__init__，但此刻 AscendConfig 是否已初始化仍无法保证（旧文件顶部 TODO 的由来），因此 additional_config 保留为启动窗口期的回退，否则函数返回 False。与旧实现相比，本轮收紧了一件事：

• 环境变量路径被移除。 旧实现从环境变量读取该设置。平衡调度现在仅通过主路径 get_ascend_config().enable_balance_scheduling 生效，避免多个入口点。
• 文件顶部 TODO 更新为"一旦 AscendConfig 初始化被提前，这里可以收敛为单次 get_ascend_config().enable_balance_scheduling 读取"。

之后（一旦 AscendConfig 时序确定）：将两个回退收敛为单次读取。

行为保持契约¶

本次重构必须严格保持以下不变量；任何偏离都是 bug。

1. Leader 达到 cap ⇒ 全局冻结。 balance_flag 为 max(balance_queue) == max_num_running_reqs，在每个 rank 上基于上一步 gather 的 len(running) 计算。当为真时，任何 rank 都不接纳新的 WAITING 请求。比较使用 == 与配置的 max_num_running_reqs 比较——不是 >=，不是"追赶 leader"。
2. 相同输入 ⇒ 相同输出。 给定相同的 self.running、self.waiting、 self.skipped_waiting、balance_queue 和 token 预算，重构后的 schedule() 产生的 SchedulerOutput 与当前实现完全一致（相同的 scheduled / preempted / resumed 集合，相同的 num_scheduled_tokens，相同的 connector 元数据）。
3. Gather 节奏不变。 每个活跃 engine step 恰好一次 all_gather，在相同的 DP group 上，payload 仍为 len(self.running)，所有 rank 全部空闲时一致地跳过。只是调用位置移动了。
4. 禁用路径不变。 当 enable_balance_scheduling 为 false 时， _balance_run_engine_core 将模块级的 DPEngineCoreProc 恢复为上游原始版本，engine core 逐字运行上游实现；_balance_enabled=False 的 BalanceScheduler 将 schedule(throttle_prefills) 委托给 super().schedule(throttle_prefills)，不分配 balance_queue，并且 不执行任何集合通信。即关闭时 balance 不触碰任何配置（包括 PD 分离式 recompute / AsyncRecomputeScheduler， 后者已通过 platform.py 与 balance 互斥；这是第二层防御）。
5. 已有约束仍然适用。 profiling_chunk_config 互斥（参见 vllm_ascend/ascend_config.py）和 PD 混合模式限制（参见 vllm_ascend/platform.py）仍在原有位置强制生效。

重构后的文件形态¶

本轮（阶段 1 + 2A + 3）之后，关键结构如下。schedule() 函数体是发布标签 v0.24.0 的逐字副本（在阶段 2B 之前不能删除），带有三处记录的差异（禁用路径提前返回 + WAITING 循环中的 balance_flag 门控 + if request_queue is None: break）：

# vllm_ascend/patch/platform/patch_balance_schedule.py
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 environment-variable fallback


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:  # shared by v0.24.0 and e5588e49
        if not self._balance_enabled:  # delta 1: disabled-path early return
            return super().schedule(throttle_prefills)
        # NOTE: balance_gather is NOT called here -- see BalanceDPEngineCoreProc.
        # ... upstream schedule() body (verbatim-aligned to the v0.24.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 副本及其无用导入，并添加了模块 docstring、注释和 _balance_run_engine_core 条件激活包装器；净行数几乎没有下降，但它消除的是陈旧漂移风险（旧副本已落后于上游的 _handle_shutdown / eep_scaling_state / SignalCallback 演进），同时修复了开启 balance 时的死锁（先前分别在 schedule() 和 _process_engine_step 内 gather；现在在 _has_global_unfinished_reqs 之后）。

第 2B 阶段（上游提供 _should_stop_admitting_waiting）将删除约 520 行的 schedule() 副本，并把文件缩减到约 60–80 行，且完全没有逐字复制的上游代码。

测试计划¶

1. 签名 + 意图锁 + 逐字漂移测试（阶段 2A）。 断言：(a) BalanceScheduler.schedule 的签名与已安装的 Scheduler.schedule 签名相同；两个受支持的修订版本都遵循这一契约，因此每条 CI 通道都校验同一个不变量；(b) 方法体中必须存在这 3 处 balance 增量（禁用路径的 super().schedule(throttle_prefills) 委托、WAITING 循环中的 balance_flag 闸门、if request_queue is None: break）；(c) _balance_run_engine_core 包装器已安装，且 DPEngineCoreProc 在导入时不会被替换（延迟到包装器中，在调用时按条件替换）；(d) 上游的 DPEngineCoreProc._has_global_unfinished_reqs 仍然存在（gather 注入点——它必须在每次非空闲迭代中被调用，否则 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 的源码逐字比较；两者必须完全相同。读取 pin 文件意味着 pin 前进时测试会自动切换为与新标签比较并变红，发出“副本需要重新同步”的信号——这正是我们想要的维护信号；如果 pin 文件或标签不可达（例如 vLLM 不是源码检出，或在 vllm-ascend 树之外运行），则跳过测试而不是失败。
2. 行为等价测试。 使用一个模拟的 DP 组和手工设置的 balance_queue 构建 BalanceScheduler，在若干代表性状态下驱动 schedule()（leader 达上限冻结、落后的 rank 未满、禁用、空等待队列），并断言 SchedulerOutput 相同（契约项 2）。复用 tests/ut/test_platform.py 中的 balance 测试脚手架。
3. Gather 节奏测试。 模拟 torch.distributed.all_gather（注意：在 balance_gather 内部，dist = torch.distributed，因此 mock 目标必须是 torch.distributed.all_gather，而不是 vllm.distributed.all_gather）；断言每次 balance_gather() 恰好执行一次 all_gather，payload 为 len(self.running) 且使用注入的 dp_group（契约项 3）。
4. 禁用路径测试。 关闭该标志后，断言 balance_queue 未被分配、all_gather 未被调用，且 schedule(throttle_prefills) 委托给 super().schedule(throttle_prefills)（契约项 4）。
5. NPU 性能检查。 按照 AGENTS.md 的 NPU 指引，max(t.item() for t in self.balance_queue) 每步会触发一次 host 同步（这是不可避免的，因为该值驱动宿主机侧的控制流）。进行 profile 以确认本次重构不会引入超出当前同步的额外同步。

分阶段推出¶

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