执行摘要

• 一句话：移除 warmup 阶段 inference_mode，修复懒分配 buffer 更新崩溃
• 推荐动作：值得快速精读：PR 本体仅 2 行变更，但关联 Issue #33470 对 PyTorch grad mode 边界的分析质量很高，揭示了 inference_mode 与 no_grad() 的语义差异以及 tensor 生命周期跨模式边界的陷阱。对模型后端与 runner 相关开发者，应关注'懒分配 buffer 必须跨 grad 模式安全'的约定，以及本 PR 采用的'不主动覆盖调用方 grad mode'的设计原则。后续可跟踪 Issue 提出的 workspace 显式初始化方案是否落地。

功能与动机

Issue #33470 指出：SGLang 的 FlashInfer autotune/warmup forward 运行在 torch.inference_mode() 下，而真实 eager serving、CUDA graph capture 与 replay 均不在 inference_mode 下，模式不一致会让懒分配的可复用 tensors 成为 inference tensors，后续 Python 侧 in-place 更新失败，报错 Inplace update to inference tensor outside InferenceMode is not allowed。实测案例包括 FlashInfer MegaMOE 的 topk_idx_out[num_tokens:prev_n].fill_(-1) 与 final_hidden_states += shared_output。Issue 被标记为 high priority。

实现拆解

1. 定位模式边界：依据 Issue #33470 的对比表，确认 FlashInfer autotune/warmup（inference_mode 开启）与真实 eager forward、CUDA graph capture/replay（inference_mode 关闭）之间的 grad 模式不一致，是懒分配 inference tensor 的根源。
2. 修改 flashinfer_autotune.py：在 flashinfer_autotune_context() 中移除 torch.inference_mode() 包裹，仅保留 autotune()、maybe_skip_logits 与 forward_stream 上下文，使 autotune 阶段的 tensor 分配不再带 inference 语义。
3. 修改 base_runner.py：shared dummy warmup 的 run_once() 调用处不再叠加 torch.inference_mode()，仅使用调用方传入的 run_ctx（可能为 torch.no_grad() 等）或空上下文，避免覆盖调用方的 grad mode 意图——这是第 3 个提交 fix: avoid overriding warmup grad mode 的核心考量。
4. 演进与收敛：4 个提交依次为初始修复、触发 CI、修正 grad mode 覆盖问题、按 review 反馈删除新增注释与 98 行单元测试，最终收敛为 2 文件各 1 行净变更。
5. 测试与配套：最终未保留自动化测试（新增的 test/registered/unit/model_executor/test_flashinfer_autotune.py 在 review 中被作者自审删除），回归保障依赖 run-ci 标签触发的集成测试与实际模型运行。

关键文件：

• python/sglang/srt/model_executor/runner/flashinfer_autotune.py（模块 执行器；类别 source；类型 core-logic；符号 flashinfer_autotune_context）: 核心修复点之一：flashinfer_autotune_context() 移除 torch.inference_mode() 包裹，使 autotune 阶段懒分配的 workspace 不再成为 inference tensor，这是 Issue #33470 指出的首要问题来源。
• python/sglang/srt/model_executor/runner/base_runner.py（模块 执行器；类别 source；类型 core-logic；符号 run_once）: 共享 dummy warmup 路径同样移除 torch.inference_mode()，与 FlashInfer autotune 修复配套，覆盖所有复用该 warmup 后端的懒分配场景。

关键符号：flashinfer_autotune_context, run_once

关键源码片段

python/sglang/srt/model_executor/runner/flashinfer_autotune.py

核心修复点之一：flashinfer_autotune_context() 移除 torch.inference_mode() 包裹，使 autotune 阶段懒分配的 workspace 不再成为 inference tensor，这是 Issue #33470 指出的首要问题来源。

@contextlib.contextmanager
def flashinfer_autotune_context(model_runner: ModelRunner, *, skip_logits: bool):
    from flashinfer.autotuner import autotune
​
    mr = model_runner
    cache_path = flashinfer_autotune_cache_path(mr)
    if envs.SGLANG_FLASHINFER_AUTOTUNE_CACHE.get():
        autotune_cache = cache_path
        logger.info("Running FlashInfer autotune with cache: %s", autotune_cache)
    else:
        # 不启用缓存时，每次写入带时间戳的独立结果文件，便于回溯
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        runs_dir = cache_path.parent / "runs"
        runs_dir.mkdir(parents=True, exist_ok=True)
        autotune_cache = runs_dir / f"{cache_path.stem}.{timestamp}{cache_path.suffix}"
        logger.info(
            "Running FlashInfer autotune (cache reuse DISABLED via "
            "SGLANG_FLASHINFER_AUTOTUNE_CACHE=0); writing fresh result to: %s",
            autotune_cache,
        )
​
    # 在非默认流上跑 warmup，规避 NCCL 2.29+ 在默认流上调用
    # cudaMemcpyBatchAsync 的问题（--enable-symm-mem 场景）
    mr.forward_stream.wait_stream(torch.cuda.current_stream())
    with torch.get_device_module(mr.device).stream(mr.forward_stream):
        maybe_skip_logits = contextlib.nullcontext()
        if skip_logits:
            from sglang.srt.layers.logits_processor import autotune_dummy_run_mode
            maybe_skip_logits = autotune_dummy_run_mode()
        skip_ops = get_flashinfer_autotune_skip_ops(mr)
        # 关键修复：不再包裹 torch.inference_mode()。
        # 若此处创建 inference tensor，真实 eager serving / CUDA graph capture
        # 后续对懒分配 workspace 做 in-place 更新时会抛
        # "Inplace update to inference tensor outside InferenceMode is not allowed"。
        # grad 模式改由调用方决定，避免覆盖 warmup 的 mode 意图。
        with autotune(
            True,
            cache=str(autotune_cache),
            skip_ops=skip_ops,
        ), maybe_skip_logits:
            yield
    torch.cuda.current_stream().wait_stream(mr.forward_stream)
    logger.info("FlashInfer autotune completed.")

python/sglang/srt/model_executor/runner/base_runner.py

共享 dummy warmup 路径同样移除 torch.inference_mode()，与 FlashInfer autotune 修复配套，覆盖所有复用该 warmup 后端的懒分配场景。

# 共享 dummy warmup：与真实 eager forward、CUDA graph capture 的
# grad 模式边界对齐，避免懒分配 workspace 变成 inference tensor。
# run_ctx 由调用方（如 torch.no_grad()）显式控制，本层不再叠加
# torch.inference_mode()，防止覆盖调用方的 grad mode 意图。
torch.get_device_module(mr.device).synchronize()
mr.tp_group.barrier()
with forward_context(ForwardContext(attn_backend=mr.attn_backend)):
    with run_ctx or empty_context():
        run_once()

评论区精华

作者 nvpohanh 对自己的 PR 进行了 3 轮 COMMENTED 自审，均已在最终提交中解决：

• 要求在 base_runner.py 中删除新增的说明注释（代码自解释，无需注释）。
• 要求在 flashinfer_autotune.py 中删除同类注释。

• 要求删除新增的 98 行单元测试 test/registered/unit/model_executor/test_flashinfer_autotune.py，说明该 inference_mode 场景难以在单元测试中稳定复现，回归保障依赖真实集成路径。
  值得注意的设计权衡：移除 inference_mode() 后不显式补 torch.no_grad()，是因为 warmup 路径可能自带 grad 上下文（run_ctx），主动叠加会改变调用方语义；且 no_grad() 创建的普通 tensor 本身可安全跨模式 in-place 更新，真正需要规避的只有 inference tensor。

• 删除 base_runner.py 中新增的说明注释 (style): 已在后续提交中删除该注释。

• 删除 flashinfer_autotune.py 中新增的说明注释 (style): 已在后续提交中删除该注释。
• 删除新增的 98 行单元测试 (testing): 测试文件已被删除，回归保障依赖 run-ci 集成测试与真实模型路径。

风险与影响

• 风险：
  1. 性能/内存回归风险（低-中）：移除 torch.inference_mode() 后，若调用方未通过 run_ctx 传入 no_grad()，dummy warmup forward 理论上可能构建 autograd 图，增加 warmup 阶段显存与时间开销。需要确认所有调用点（如 EAGLE draft、speculative decode 的 warmup 路径）都正确传入 run_ctx。
  2. 回归检测缺口（中）：最终没有保留自动化测试，而该 bug 只在特定后端（如 FlashInfer MegaMOE workspace 懒分配）上触发，回归可能静默重现。
  3. 影响面（中）：变更作用于所有使用 FlashInfer autotune 的模型 warmup 路径，属于核心执行路径；但行为是'放宽'而非'收紧'，对已正常运行的模型无副作用。
  4. CI 状态：Extra CI 运行曾失败（Run #31099758373 标记 x），过程中触发过 /rerun-failed-ci，最终合并，说明集成验证通过但曾有不稳定迹象。
     - 影响：用户/服务影响：消除 FlashInfer MegaMOE 等后端在 warmup 后首个请求阶段因 inference tensor 跨模式更新而崩溃的问题；对正常模型无行为变化。系统影响：warmup/autotune 与真实 forward、CUDA graph capture 的 grad 模式边界统一为'由调用方决定'，后端开发者不再需要为懒分配 workspace 手动包 torch.inference_mode(False) 绕过。团队影响：确立了 runner 层不强制 grad mode 的约定，与 Issue #33470 中建议的'方案 2（避免在 warmup 下用 inference_mode）'一致；但 Issue 提出的'显式初始化阶段'或'workspace 分配 API'（方案 3/4）仍是更彻底的长期方向。
     - 风险标记：核心执行路径变更, 缺少测试覆盖, warmup 行为影响面广, 依赖调用方传 run_ctx 保性能

关联脉络

• PR #24370 Profiling Enhancements [1/3]: cuda graph profile traces: 同属 model_executor/runner 模块的 warmup/捕获阶段路径改动，涉及 CUDA graph capture 期间的执行上下文行为；Issue #33470 的论证明确将 CUDA graph capture/replay 路径列为 grad 模式边界的一部分。