# PR #30347 完整报告

- 仓库：`sgl-project/sglang`
- 标题：[refactor] Collect MoE and DP-attention runtime state into typed flag groups
- 合并时间：2026-07-08 12:29
- 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/30347

---

# 执行摘要

- 一句话：迁移 MoE/DP 运行时状态到类型化标志组
- 推荐动作：值得仔细阅读 runtime_context.py 和 test_module_state_ratchet.py，了解如何通过类型化标志组和 AST 回归测试管理运行时状态。这是清理遗留技术债的典范，类似模块（如 speculative decoding 中的状态）可借鉴此模式。

# 功能与动机

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

# 实现拆解

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`（模块 运行时上下文；类别 source；类型 core-logic；符号 MoeFlags, DpFlags）: 定义 MoeFlags 和 DpFlags 数据类，扩展 Flags 根类，是重构的核心
- `python/sglang/srt/layers/moe/utils.py`（模块 MoE 配置；类别 source；类型 dependency-wiring；符号 initialize_moe_config, get_moe_a2a_backend, get_moe_runner_backend, is_tbo_enabled）: 迁移 12 个模块级全局变量到 flags.moe，调整所有访问器
- `python/sglang/srt/layers/dp_attention.py`（模块 DP 注意力；类别 source；类型 dependency-wiring；符号 initialize_dp_attention, is_dp_attention_enabled, get_dp_padding_mode）: 迁移 DP 注意力运行时布尔值到 flags.dp，更新所有引用
- `test/registered/unit/test_runtime_context.py`（模块 运行时测试；类别 test；类型 test-coverage；符号 TestMoeFlagsGroup, TestDpFlagsGroup）: 新增针对 MoeFlags 和 DpFlags 的单元测试，覆盖懒加载、materialize、speculative 上下文交换和异常恢复
- `test/registered/unit/test_module_state_ratchet.py`（模块 回归测试；类别 test；类型 test-coverage；符号 TestModuleStateRatchet）: 新增基于 AST 的模块状态回归测试，锁定允许的 global 语句，防止重新引入模块级运行时状态
- `test/registered/ops/test_aiter_greedy_sample_amd.py`（模块 AMD 采样测试；类别 test；类型 test-coverage）: 修改为通过 get_flags().dp 直接设置启用标志，替代模块绑定注入

关键符号：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`

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

```python
@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`

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

```python
_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.",
            )

```

# 评论区精华

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

- 测试中应使用 override 上下文管理器代替直接赋值 (testing): 最终提交未采用该建议，仍使用直接赋值。可能作者认为当前测试框架已通过 reset_context 或其他机制保证隔离，但 review 指出的问题仍然存在。

# 风险与影响

- 风险：主要风险：若外部代码直接引用被移除的模块级全局变量（如 layers.moe.utils.MOE_A2A_BACKEND），将导致 AttributeError。但该仓库内部已通过访问器封装，外部调用者很少直接使用。测试覆盖了主要路径。新引入的 AST 回归测试可能过于严格，阻碍合法的新模块级全局变量（如引入新的资源），但已设计了 pin 更新机制。性能影响可忽略（增加一次 get_flags() 调用）。
- 影响：对用户无行为变化。对系统：测试更健壮，状态管理更清晰。对团队：降低了新引入状态泄漏的风险，提升了代码可维护性。
- 风险标记：模块级状态迁移 , 测试隔离 , AST 回归测试 , 调用点兼容性

# 关联脉络

- PR #30348 [refactor] ctx.resources: named slots, stream leases, and workspace buffer leases: 同属 runtime context 重构系列，本 PR 构建在 #30348 之上（PR body 提及 'Stacked on the read-convergence PR.'，即 #30348）。