Prhub

#34052 [Scheduler] Unify WAR read-done gating behind shared-read boundary declarations

原始 PR 作者 hnyls2002 合并时间 2026-08-08 18:26 文件变更 11 提交数 7 评论 0 代码增减 +105 / -70

执行摘要

统一 WAR 栅栏为共享读边界声明,删除策略枚举

PR body 明确指出:overlap scheduler 的 WAR barrier 每个 forward 只回答一个问题——该 forward 最后一次读取 scheduler 共享缓冲的位置在哪,以便在那里记录 read-done 事件。此前这个问题由四个互不关联的片段回答(prefill_shared_reads_end_at_metadata_init 类布尔、supports_target_verify_war_read_done 布尔方法、use_captured_forward_metadata_for_breakable_cuda_graph 借用标志、WarReadDonePolicy + if 链),"Each consumer held its own fragment, so supporting a new (backend, mode) combination meant adding another flag and threading it through the chain — and asymmetries crept in silently (e.g. the BCG check existed on the verify branch but not decode)"。

值得精读,尤其是对计划接入新 attention backend 或需要理解 sglang WAR overlap 机制的工程师。重点学习三点设计:SharedReadBoundary 用枚举表达相对 replay 的边界位置、UNKNOWN 作为安全底值的降级策略、以及默认值从 out-graph/in-graph 契约推导而不是从布尔标志拼接。配合两个 CPU 单测的平移方式可快速掌握整个决策矩阵。

讨论亮点

本 PR 无 review 评论与 issue 讨论。从 7 个 commit 的演进序列可还原设计迭代:先完成初步统一(a9bb5072),随即放弃 WarReadDonePolicy 枚举、改用单一 SharedReadBoundary 类型(a4e882eb),再调整默认边界推导来源、由 base 类持有契约默认值、decode runner 保留直接 case 表并删除 resolve helper(d27c5753、2bc85c65、ede7e369)。最终形态的核心收敛决策是:UNKNOWN 作为安全底值、IN_REPLAY 作为契约上界默认、后端只对审计过的偏差做 override。

实现拆解

实现按 5 步拆解:

  1. 声明层(base_attn_backend.py):新增 SharedReadBoundary 四值枚举(PRE_REPLAY / IN_REPLAY / POST_REPLAY / UNKNOWN),并在 AttentionBackend 上新增 shared_read_boundary(forward_mode) 方法。默认实现依据 out-graph/in-graph 元数据初始化契约返回 decode/verify → IN_REPLAY、其余 → UNKNOWN;同时删除旧的 prefill_shared_reads_end_at_metadata_init 类布尔,UNKNOWN 作为安全底值保证未审计后端退回粗粒度整轮围栏。

  2. 算法层(spec_info.py / spec_registry.py):SpeculativeAlgorithm.is_war_publish_phase(forward_mode) 取代 supports_target_verify_war_read_done,把"哪个阶段拥有发布权"从后端布尔提升为算法语义:dflash/dspark 家族在 target verify 阶段发布,EAGLE 等其余算法在 draft extend v2 阶段发布。

  3. 消费层(decode_cuda_graph_runner.py):_war_read_done_policy 重写为 _war_read_done_record,变成"算法门控 → 后端声明 → 图内锚点降级"的机械查表;execute 中按 PRE_REPLAY / POST_REPLAY / IN_REPLAY 三档落点调用 _publish_war_read_done。eagle_draft_extend_cuda_graph_runner.py 补注释说明 draft extend 是 EAGLE 家族的发布阶段。

  4. 后端迁移(trtllm_mha_backend.py / deepseek_v4_backend.py):trtllm_mha 用 EXTEND → PRE_REPLAY 声明取代 prefill 布尔;deepseek_v4 用 TARGET_VERIFY → POST_REPLAY 声明取代对 use_captured_forward_metadata_for_breakable_cuda_graph 的借用,该标志回归纯 BCG 职责。

  5. 清理与测试(war_event.py / runner_utils/init.py / 两个单测):war_event.py 删除 WarReadDonePolicy,maybe_publish_prefill_war_read_done 改用 boundary is PRE_REPLAY 判断;runner_utils/init.py 移除对应导出;test_decode_cuda_graph_war_fence.py 与 test_prefill_war_read_done.py 按新 API 平移同一决策矩阵,断言一一对应。

文件 模块 状态 重要度
python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py 解码图执行 modified 7.79
python/sglang/srt/layers/attention/base_attn_backend.py 注意力后端 modified 7.28
test/registered/unit/model_executor/runner/test_decode_cuda_graph_war_fence.py 单元测试 modified 7.13
python/sglang/srt/model_executor/runner_utils/war_event.py 事件工具 modified 7.03
python/sglang/srt/layers/attention/deepseek_v4_backend.py 注意力后端 modified 6.48
python/sglang/srt/speculative/spec_info.py 投机解码 modified 6.4
python/sglang/srt/layers/attention/trtllm_mha_backend.py 注意力后端 modified 6.38
python/sglang/srt/speculative/spec_registry.py 投机解码 modified 5.99
test/registered/unit/model_executor/runner/test_prefill_war_read_done.py 单元测试 modified 5.05
python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py 投机解码 modified 4.62
python/sglang/srt/model_executor/runner_utils/__init__.py 事件工具 modified 4.59

关键符号

shared_read_boundary _war_read_done_record is_war_publish_phase _publish_war_read_done maybe_publish_prefill_war_read_done _war_read_done_policy

关键源码片段

python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py data-contract

WAR read-done 事件的最终消费与发布点,_war_read_done_policy 重写为 _war_read_done_record,execute 按 SharedReadBoundary 三档落点发布,是本次重构的机械消费核心。

# python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
# Runner 变为 " 机械消费者 ":先查算法发布阶段,再查后端边界声明,
# 最后按边界位置决定 read-done 事件落点。UNKNOWN 表示本次 replay
# 不参与 WAR 快速路径发布。def _war_read_done_record(self, attn_backend, forward_mode) -> SharedReadBoundary:
    """决定本次 replay 在何处记录 WAR read-done 事件;UNKNOWN 表示不记录。"""
    if forward_mode.is_target_verify():
        # 算法门控优先:非发布阶段(如 EAGLE 的 verify)不在此 runner 发布
        if not self.model_runner.spec_algorithm.is_war_publish_phase(forward_mode):
            return SharedReadBoundary.UNKNOWN
    elif not forward_mode.is_decode():
        # decode 之外的模式由其他 runner(如 EAGLE draft extend)负责发布
        return SharedReadBoundary.UNKNOWN
    boundary = attn_backend.shared_read_boundary(forward_mode)
    if (
        boundary is SharedReadBoundary.IN_REPLAY
        and not self._war_read_done_node_planted
    ):
        # 非捕获运行 / 不支持外部事件时图内锚点不存在,回退到 replay 前记录
        return SharedReadBoundary.PRE_REPLAY
    return boundary
​
​
def execute(self, forward_batch, pp_proxy_tensors=None):
    # ... 前置逻辑省略 ...
    war_record = self._war_read_done_record(
        self.attn_backend, forward_batch.forward_mode
    )
    with timer_ctx, self.backend.replay_session():
        self.load_batch(forward_batch, pp_proxy_tensors)
        if war_record is SharedReadBoundary.PRE_REPLAY:
            self._publish_war_read_done(in_graph=False)
        output = self.backend.replay(self._replay_graph_key, forward_batch)
        if war_record is SharedReadBoundary.POST_REPLAY:
            self._publish_war_read_done(in_graph=False)
        elif war_record is SharedReadBoundary.IN_REPLAY:
            self._publish_war_read_done(in_graph=True)
python/sglang/srt/layers/attention/base_attn_backend.py core-logic

新增 SharedReadBoundary 枚举与 shared_read_boundary 声明方法,取代旧 bool 标志,是整个统一方案的契约定义处。

# python/sglang/srt/layers/attention/base_attn_backend.py
# WAR read-done 边界的统一声明点:SharedReadBoundary 描述 " 后端对
# scheduler 共享缓冲的最后一次读 " 相对 CUDA graph replay 的位置。
# UNKNOWN 是安全底值——不发布事件,scheduler 退回整轮 forward
# 的粗粒度围栏(coarse whole-forward fence)。class SharedReadBoundary(Enum):
    PRE_REPLAY = auto() # 共享读在 replay 启动前已结束(快照式后端)
    IN_REPLAY = auto() # 共享读在图内元数据初始化处结束(图内锚点)
    POST_REPLAY = auto() # 共享读贯穿整个 graph replay(如 DSV4 BCG verify)
    UNKNOWN = auto() # 未审计 -> 不发布,保留整轮粗粒度围栏
​
​
class AttentionBackend(ABC):
​
    def shared_read_boundary(self, forward_mode: ForwardMode) -> SharedReadBoundary:
        """按 forward 模式声明共享读结束位置。        decode/verify 默认 IN_REPLAY:out-graph/in-graph 元数据初始化契约
        保证任何遵守该契约的后端,其共享读不会越过图内元数据锚点,
        因此 IN_REPLAY 是安全上界。有审计结论的后端才需要 override。
        """
        if forward_mode.is_decode() or forward_mode.is_target_verify():
            return SharedReadBoundary.IN_REPLAY
        return SharedReadBoundary.UNKNOWN
python/sglang/srt/speculative/spec_info.py core-logic

is_war_publish_phase 取代 supports_target_verify_war_read_done,把发布权从后端布尔提升为算法语义。

# python/sglang/srt/speculative/spec_info.py
# 每个 spec 步骤的 " 最后读共享缓冲的阶段 " 负责发布 read-done 事件。
# dflash/dspark 系列在 target verify 阶段发布;EAGLE 等其余算法在
# draft extend 阶段发布(由 EAGLE draft extend runner 执行记录)。def is_war_publish_phase(self, forward_mode) -> bool:
    if self.is_dflash_family():
        return forward_mode.is_target_verify()
    return forward_mode.is_draft_extend_v2()

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

  1. 行为等价性依赖审计:PR 声称逐分支严格等价,但 verify 分支存在一个隐性变化点——原来通过 use_captured_forward_metadata_for_breakable_cuda_graph 判断 POST_REPLAY,现在只有显式声明 POST_REPLAY 的 deepseek_v4_backend 保留该行为;若未来出现第三个同时设置 BCG 标志并支持 dflash verify 的后端而未声明边界,将从 POST_REPLAY 静默降级为 IN_REPLAY,可能引入 WAR race。base_attn_backend.py 的默认值是建立在"后端遵守 out-graph/in-graph 元数据初始化契约"这一前提上的安全上界,契约一旦被违反(如在 replay 前向中读 scheduler 共享张量),IN_REPLAY 默认即失效。
  2. 测试覆盖:两个 CPU 单测覆盖了完整决策矩阵,但没有 GPU 端集成测试(如 BCG + dflash verify + deepseek_v4 的真实 replay)验证事件落点与实际 barrier 行为。
  3. 无 review 复核:PR 由作者自行 merge,讨论环节缺失,契约迁移的审计结论缺少第二人确认。

对最终用户透明,无 API 或行为变化。对系统:decode/verify 每次 replay 的 read-done 事件位置改由声明驱动,新增 (backend, mode) 组合时只需实现一个 shared_read_boundary override,接入成本显著下降;BCG 检查此前只存在于 verify 分支的不对称也被修复。对团队:删除 WarReadDonePolicy 与两个 bool 方法后,WAR 机制只剩"声明 → 记录 → 消费"三层,认知负担大幅降低,后续 deepseek_v4、trtllm_mha 等后端的边界语义集中在各自文件中。

核心路径变更 行为等价依赖审计 默认上界依赖契约 缺少 GPU 集成测试 无 review 复核

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论