Prhub

#30347 [refactor] Collect MoE and DP-attention runtime state into typed flag groups

原始 PR 作者 ch-wan 合并时间 2026-07-08 12:29 文件变更 7 提交数 4 评论 2 代码增减 +292 / -110

执行摘要

迁移 MoE/DP 运行时状态到类型化标志组

模块级全局变量生命周期管理混乱,跨测试 teardown 泄漏,测试中只能通过重新绑定 import 覆盖。通过集中到类型化标志组,自动获得生命周期重置、类型安全、事务性测试覆盖。

值得仔细阅读 runtime_context.py 和 test_module_state_ratchet.py,了解如何通过类型化标志组和 AST 回归测试管理运行时状态。这是清理遗留技术债的典范,类似模块(如 speculative decoding 中的状态)可借鉴此模式。

讨论亮点

review 指出直接设置 get_flags().dp.enabled = False 会永久影响进程内其他测试,建议使用 get_flags().dp.override(enabled=False) 上下文管理器以确保测试隔离。但最终提交未采用此建议,可能认为当前测试隔离已通过其他机制保证。后续可考虑改进。

实现拆解

  1. 定义标志组类型:在 runtime_context.py 中新增 MoeFlags 和 DpFlags 数据类,继承 _FlagGroupBase,声明与旧全局变量一一对应的字段。将 Flags 根类扩展为包含 moe 和 dp 子组。

  2. 迁移 MoE 配置:修改 layers/moe/utils.py 中的 initialize_moe_config,将写入 12 个全局变量的代码改为写入 get_flags().moe 对应的字段。调整所有访问器(如 get_moe_a2a_backend)从标志组读取,保持懒加载默认值。speculative 上下文仍通过 override 机制交换叶子值。

  3. 迁移 DP 注意力标志:修改 layers/dp_attention.py,将 _ENABLE_DP_ATTENTION_FLAG 和 _DP_MAX_LEN_WITH_IDLE 迁移到 get_flags().dp;is_dp_attention_enabled 变为薄封装;get_dp_padding_mode 直接读取 max_len_with_idle 标志。

  4. 引入模块状态回归测试:新增 test_module_state_ratchet.py,AST 解析 moe/utils.py 和 dp_attention.py 的 global 语句,与预设的 pin 集合对比。任何新增的模块级全局变量都会导致测试失败。

  5. 清理测试注入:将两个 AMD 测试文件(test_aiter_greedy_sample_amd.py、test_aiter_allreduce_fusion_amd.py)中通过模块 import 覆盖 is_dp_attention_enabled 的做法改为直接设置 get_flags().dp.enabled。

文件 模块 状态 重要度
python/sglang/srt/runtime_context.py 运行时上下文 modified 7.4
python/sglang/srt/layers/moe/utils.py MoE 配置 modified 7.31
python/sglang/srt/layers/dp_attention.py DP 注意力 modified 6.21
test/registered/unit/test_runtime_context.py 运行时测试 modified 6.93
test/registered/unit/test_module_state_ratchet.py 回归测试 added 6.59
test/registered/ops/test_aiter_greedy_sample_amd.py AMD 采样测试 modified 4.16

关键符号

MoeFlags DpFlags initialize_moe_config get_moe_a2a_backend is_tbo_enabled initialize_dp_attention is_dp_attention_enabled

关键源码片段

python/sglang/srt/runtime_context.py core-logic

定义 MoeFlags 和 DpFlags 数据类,扩展 Flags 根类,是重构的核心

@dataclasses.dataclass
class MoeFlags(_FlagGroupBase):
    """MoE 运行时标志,由 initialize_moe_config 在调度器初始化时填充。    a2a_backend / runner_backend / disable_fp4_allgather 是“活跃”值:
    speculative 上下文会在 draft-model 前向时交换它们。
    """
​
    a2a_backend: Any = None
    runner_backend: Any = None
    speculative_runner_backend: Any = None
    speculative_a2a_backend: Any = None
    deepep_mode: Any = None
    deepep_config: str | None = None
    tbo_enabled: bool | None = None
    sbo_enabled: bool | None = None
    tbo_token_distribution_threshold: float | None = None
    disable_fp4_allgather: bool | None = None
    quantization: str | None = None
​
​
@dataclasses.dataclass
class DpFlags(_FlagGroupBase):
    """DP 注意力运行时标志,由 initialize_dp_attention 在分布式设置后填充。
    Topology 值(size/rank)暂留在 layers.dp_attention 模块中。
    """
​
    enabled: bool = False
    max_len_with_idle: bool = False
​
​
@dataclasses.dataclass
class Flags(_FlagGroupBase):
    capture: CaptureFlags = dataclasses.field(default_factory=CaptureFlags)
    moe: MoeFlags = dataclasses.field(default_factory=MoeFlags) # 新增 MoE 子组
    dp: DpFlags = dataclasses.field(default_factory=DpFlags) # 新增 DP 子组
test/registered/unit/test_module_state_ratchet.py test-coverage

新增基于 AST 的模块状态回归测试,锁定允许的 global 语句,防止重新引入模块级运行时状态

_PINNED_GLOBALS = {
    "layers/moe/utils.py": frozenset(),
    "layers/dp_attention.py": frozenset(
        {
            # DP-attention topology (parallel vertical scope).
            "_ATTN_DP_RANK",
            "_ATTN_DP_SIZE",
            "_LOCAL_ATTN_DP_SIZE",
            "_LOCAL_ATTN_DP_RANK",
            # Comm stream resource (resources vertical scope).
            "_DP_TBO_COMM_STREAM",
        }
    ),
}
​
​
class TestModuleStateRatchet(CustomTestCase):
    def test_global_statements_match_the_pins(self):
        for rel, pinned in _PINNED_GLOBALS.items():
            tree = ast.parse((_SRT_ROOT / rel).read_text())
            declared = {
                name
                for node in ast.walk(tree)
                if isinstance(node, ast.Global)
                for name in node.names
            }
            grown = declared - pinned
            self.assertFalse(
                grown,
                f"{rel} declares new module-level runtime state {sorted(grown)}; "
                "put runtime flags on a get_flags() group instead "
                "(see runtime_context.MoeFlags / DpFlags).",
            )
            shrunk = pinned - declared
            self.assertFalse(
                shrunk,
                f"{rel} no longer declares {sorted(shrunk)}; "
                "shrink the pin in this file to lock in the progress.",
            )

评论区精华

测试中应使用 override 上下文管理器代替直接赋值 测试

reviewer(chatgpt-codex-connector)建议在 test_aiter_greedy_sample_amd.py 中使用 get_flags().dp.override(enabled=False) 而不是直接设置 get_flags().dp.enabled = False,以确保测试隔离,避免影响其他测试。

结论:最终提交未采用该建议,仍使用直接赋值。可能作者认为当前测试框架已通过 reset_context 或其他机制保证隔离,但 review 指出的问题仍然存在。 · unresolved

风险与影响

主要风险:若外部代码直接引用被移除的模块级全局变量(如 layers.moe.utils.MOE_A2A_BACKEND),将导致 AttributeError。但该仓库内部已通过访问器封装,外部调用者很少直接使用。测试覆盖了主要路径。新引入的 AST 回归测试可能过于严格,阻碍合法的新模块级全局变量(如引入新的资源),但已设计了 pin 更新机制。性能影响可忽略(增加一次 get_flags() 调用)。

对用户无行为变化。对系统:测试更健壮,状态管理更清晰。对团队:降低了新引入状态泄漏的风险,提升了代码可维护性。

模块级状态迁移 测试隔离 AST 回归测试 调用点兼容性

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论