执行摘要

• 一句话：统一 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 个测试

# 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 成为中心化状态容器