# PR #30490 完整报告

- 仓库：`sgl-project/sglang`
- 标题：[refactor] Add the per-forward flags tier: ctx.forward
- 合并时间：2026-07-09 17:09
- 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/30490

---

# 执行摘要

- 一句话：统一 per-forward flags 层，替代分散的模块级状态
- 推荐动作：本 PR 是运行时上下文模块的重要重构，值得架构师和核心开发者精读。重点关注 ForwardFlags 的双备份设计、scoped 的事务语义、以及如何平衡 torch.compile 兼容性与 contextvar 的隔离性。review 中的讨论涵盖了 PEP8 权衡、API 设计哲学和类型注解策略，也值得参考。

# 功能与动机

Per-forward runtime state 分散在不同的模块 holder 中，存在不一致的并发选择（thread-local、裸 contextvar、混合 init-resolved 状态、sticky 类属性），缺乏统一的生命周期管理和线程隔离，且异常路径存在 flag 泄漏风险。

# 实现拆解

1. **在 runtime_context.py 中新增 ForwardFlags 类**：定义 `_DEFAULTS` 字典声明所有 per-forward 标志，通过 `_GRAPH_VISIBLE` 集合区分需要在 torch.compile 图中读写的标志，使用 contextvar（供 eager 读写）和 plain dict（供编译图读写）双备份存储。提供 `set()` 和 `scoped(**kw)` 两个写入路径，其中 `scoped()` 是事务性写入（保存旧值 → 写入新值 → 退出 / 异常时恢复）。
2. **将多流开关迁移到 forward 层**：修改 `python/sglang/srt/utils/multi_stream_utils.py`，删除 `do_multi_stream_local` 线程局部类和 `_local` 实例，将 `set_do_multi_stream`、`do_multi_stream`、`with_multi_stream` 全部重写为对 `get_forward().set/get/scoped` 的委托。
3. **将 MoE 输出缓冲上下文迁移到 forward 层**：修改 `python/sglang/srt/layers/moe/moe_runner/base.py`，删除模块级 contextvar `_moe_output_buf`，将 `moe_output_buffer_ctx` 重写为直接返回 `get_forward().scoped(moe_output_buffer=buf)`。
4. **将 AttnTpContext 的 per-forward 状态迁移到 forward 层**：修改 `python/sglang/srt/layers/communicator.py`，将 `input_scattered_` 和 `attn_inputs_` 实例属性替换为对 `forward.flags` 的读写，`maybe_input_scattered` 改为使用 `scoped()` 包裹，修复了异常时 flag 泄漏（旧代码无 try/finally）。
5. **将 DP 缓冲区的 `is_extend_in_batch` 迁移到 forward 层**：修改 `python/sglang/srt/layers/dp_attention.py`，删除 `_DpGatheredBufferWrapper._is_extend_in_batch` 类属性和对应的 classmethod，将 `set_is_extend_in_batch` / `get_is_extend_in_batch` 改为通过 `get_forward().set/is_extend_in_batch` 操作。
6. **补充完整单元测试**：在 `test/registered/unit/test_runtime_context.py` 中新增 `TestForwardFlags` 类，包含 scope 嵌套恢复异常验证、线程隔离、torch.compile 兼容性跨线程可见性等 45 个测试。

关键文件：
- `python/sglang/srt/runtime_context.py`（模块 运行时上下文；类别 source；类型 core-logic；符号 ForwardFlags, __init__, __getattr__, __setattr__）: 核心变更文件，新增 ForwardFlags 类与 get_forward 访问器，实现 per-forward flags 双备份设计
- `test/registered/unit/test_runtime_context.py`（模块 测试；类别 test；类型 test-coverage；符号 TestForwardFlags, test_scoped_set_restore_and_nesting, test_scoped_restores_on_exception_and_validates_keys, test_threads_see_defaults）: 新增 TestForwardFlags 测试类，覆盖 scope 语义、线程隔离、torch.compile 兼容性等 45 个测试
- `python/sglang/srt/utils/multi_stream_utils.py`（模块 工具层；类别 source；类型 core-logic；符号 set_do_multi_stream, do_multi_stream, with_multi_stream）: 多流开关迁移到 forward 层，删除线程局部类，代之以对 get_forward().set/scoped 的委托
- `python/sglang/srt/layers/dp_attention.py`（模块 DP 注意力；类别 source；类型 core-logic；符号 set_is_extend_in_batch, get_is_extend_in_batch）: DP 缓冲区的 is_extend_in_batch 从类属性迁移到 forward 标志
- `python/sglang/srt/layers/moe/moe_runner/base.py`（模块 MoE 运行器；类别 source；类型 core-logic；符号 moe_output_buffer_ctx）: MoE 输出缓冲上下文从模块级 contextvar 迁移到 forward 标志
- `python/sglang/srt/layers/communicator.py`（模块 通信器；类别 source；类型 dependency-wiring；符号 AttnTpContext.input_scattered, AttnTpContext.set_attn_inputs, AttnTpContext.maybe_input_scattered）: AttnTpContext 的 per-forward 状态（input_scattered, attn_inputs）迁移到 forward 标志，maybe_input_scattered 使用 scoped 修复异常泄漏
- `python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py`（模块 MoE 运行器；类别 source；类型 dependency-wiring）: 适配 moe_output_buffer_ctx 的新签名（返回 scoped contextmanager 而非 yield），依赖更新

关键符号：ForwardFlags.__init__, ForwardFlags.__getattr__, ForwardFlags.__setattr__, ForwardFlags.set, ForwardFlags.scoped, get_forward, set_do_multi_stream, do_multi_stream, with_multi_stream, moe_output_buffer_ctx, AttnTpContext.input_scattered, AttnTpContext.set_attn_inputs, AttnTpContext.maybe_input_scattered, set_is_extend_in_batch, get_is_extend_in_batch

## 关键源码片段

### `test/registered/unit/test_runtime_context.py`

新增 TestForwardFlags 测试类，覆盖 scope 语义、线程隔离、torch.compile 兼容性等 45 个测试

```python
# test/registered/unit/test_runtime_context.py

class TestForwardFlags(_IsolatedServerArgs):
    """ctx.forward: contextvar-backed per-forward flags; scoped() restores,
    threads see defaults."""

    def test_scoped_set_restore_and_nesting(self):
        from sglang.srt.runtime_context import get_forward
        reset_context()
        fwd = get_forward()
        self.assertFalse(fwd.multi_stream)
        with fwd.scoped(multi_stream=True):
            self.assertTrue(fwd.multi_stream)
            with fwd.scoped(multi_stream=False):
                self.assertFalse(fwd.multi_stream)
            self.assertTrue(fwd.multi_stream)
        self.assertFalse(fwd.multi_stream)

    def test_graph_visible_flags_trace_under_torch_compile(self):
        # Regression test: 确保 _GRAPH_VISIBLE 标志在 torch.compile 下可用
        import torch
        from sglang.srt.runtime_context import get_forward
        reset_context()

        @torch.compile(fullgraph=True, backend="eager", dynamic=False)
        def probe(x):
            fwd = get_forward()
            if fwd.attn_input_scattered:
                x = x + 1
            if fwd.is_extend_in_batch:
                x = x + 2
            return x

        self.assertEqual(probe(torch.zeros(())).item(), 0)
        with get_forward().scoped(attn_input_scattered=True):
            self.assertEqual(probe(torch.zeros(())).item(), 1)
        get_forward().set("is_extend_in_batch", True)
        self.assertEqual(probe(torch.zeros(())).item(), 2)
        get_forward().set("is_extend_in_batch", False)

```

# 评论区精华

### 讨论线程
- **torch.compile 兼容性修复**：CI（stack-review 载体 #30349）发现 `ContextVar.get` 在 Dynamo 中不可追踪，导致 `vocab_parallel_embedding.forward` 编译失败。作者新增 `_GRAPH_VISIBLE` 集合，将需要在编译图中读写的标志放在 plain dict 备份中，解决了问题。
- **导入位置与模块约定**：bot 建议将 `import contextvars` 移到文件顶部；作者回应函数级导入遵循该模块的零导入依赖约定，已在模块 docstring 说明。
- **set 方法欠缺验证**：bot 建议对无效键抛出友好 `AttributeError`；作者回复 `set()` 是 legacy shim，未知键已由 `KeyError` 触发且报错信息友好。
- **返回类型注释缺失**：bot 建议对返回 contextmanager 的函数添加 `ContextManager[Any]` 注解；作者认为 shim 返回类型保持最小化有意。

- PEP8 imports: import contextvars 应移至文件顶部 (style): 作者坚持模块约定，使用函数局部导入，拒绝修改。
- ForwardFlags.set 方法欠缺 key 验证 (correctness): 作者保留当前方式，认为 KeyError 已足够友好，拒绝添加验证。
- 缺少显式返回类型注解 (style): 作者保留最小化类型注解，拒绝添加。
- torch.compile 兼容性：ContextVar.get 不可追踪 (correctness): 作者通过双备份设计解决了兼容问题，并在测试中添加了 fullgraph=True 的回归测试。

# 风险与影响

- 风险：
 1. **torch.compile 回归风险**：`ForwardFlags` 对 `_GRAPH_VISIBLE` 集合外的标志使用 contextvar，若将来有标志被移到编译图中但未加入该集合，会引发 Dynamo 错误。需在 `_GRAPH_VISIBLE` 集合处设置强制检查机制。
 2. **迁移完整性风险**：可能存在其他模块残留的 per-forward 状态（如 `FlashInferTrtllmRunner` 中直接读写 `moe_output_buffer_ctx`）未迁移，导致不一致。已迁移的 shim 函数（如 `set_do_multi_stream`）必须确保所有调用点都已切换。
 3. **双备份复杂性**：两个备份存储的读写逻辑通过 `__getattr__` / `set` 统一，但 `scoped` 内部通过 `getattr(self, key)` 获取旧值，可能因预期存储不同而产生问题（例如 `attn_inputs` 是 plain 存储，`scoped` 的保存 / 恢复逻辑一致工作）。
 4. **线程隔离语义分化**：`_GRAPH_VISIBLE` 标志是跨线程可见（process-global），而 contextvar 标志是线程隔离。开发者若不理解此差异，可能错误假设语义一致。测试用例已涵盖此差异。
 - 影响：对使用者（模型开发者）无直接影响，因为 API 兼容（`set_do_multi_stream` 等 shim 保持签名不变）。对系统内部影响大：统一了 per-forward 状态管理，消除了异常泄漏，为未来新增 per-forward 状态提供标准位置。对团队影响：后续所有 per-forward 标志都应通过 `ctx.forward` 添加，废弃在模块级定义 thread-local / contextvar 的做法。测试覆盖率达到 45 个单元测试，降低重构风险。
 - 风险标记：torch.compile 兼容需特别注意 _GRAPH_VISIBLE 集合 , 双备份设计中 plain 与 contextvar 语义差异可能引起误解 , 需确保所有 per-forward 状态调用点已迁移到 forward 层 , scoped 内 getattr(self, key) 的正确性依赖 __getattr__统一接口

# 关联脉络

- PR #30489 [refactor] Add the per-forward flags tier: ctx.forward (stack base): 本 PR 声称 stacked on #30489，是该 PR 的直接基础依赖
- PR #30493 [refactor] Retire the legacy config accessor and the remaining process singletons: 同一作者（ch-wan）的运行时上下文重构系列，涉及统一的配置访问和进程单例迁移，共享 runtime_context.py 和测试文件
- PR #30492 [refactor] Adopt get_parallel() everywhere and close out the parallel wrapper surface: 同一重构系列，全局统一并行拓扑访问模式，与本 PR 共同推动 RuntimeContext 成为中心化状态容器