Prhub

#28739 refactor(runner): move kernel warmup into the shared runner lifecycle (warmup())

原始 PR 作者 ch-wan 合并时间 2026-06-20 04:13 文件变更 16 提交数 1 评论 4 代码增减 +678 / -606

执行摘要

统一 kernel warmup 至 BaseRunner 生命周期

Kernel warmup / flashinfer autotune 属于 runner 生命周期,应在 init 时执行,而不是作为 ModelRunner 的独立步骤或在 forward 路径中。随着 EagerRunner 成为 BaseRunner 的子类,需要共享 warmup 方法,因此将其上提至 BaseRunner。

精读并理解 warmup 统一的设计;部署前优先修复 pp_proxy_topk_size 缺失问题。

讨论亮点

Review 中 gemini-code-assist[bot] 和 chatgpt-codex-connector[bot] 共同指出 base_runner.py 中 _allocate_decode_buffers 函数缺少 pp_proxy_topk_size 参数和对应的 topk_indices 分配,会导致 pipeline parallel MoE(DSA)模型在 dummy run 或 flashinfer autotune 时崩溃。该问题标记为高优先级,但 PR 合并时未修复。

实现拆解

  1. 在 base_runner.py 中定义 warmup() 方法,包含 run_once 守卫、设备上下文、flashinfer allreduce 预初始化、autotune、dummy run 和 PP-DeepGEMM warmup。
  2. 将 _allocate_decode_buffers 从 decode_cuda_graph_runner.py 移至 base_runner.py 模块级,同时迁移其他 autotune 相关函数。
  3. 调整 ModelRunner.init_backends 顺序:先初始化 attention backends,然后构建 EagerRunner(其 init 调用 warmup()),再构建 prefill 和 decode CUDA graph runners(capture() 中调用 warmup())。
  4. 修改 pp_parallel_deep_gemm_warmup 签名,接收 runner 而非 model_runner。
  5. 在 attention 测试 mock runner 中设置 _kernel_warmed_up = True,避免重复 warmup。
文件 模块 状态 重要度
python/sglang/srt/model_executor/runner/base_runner.py 执行器 modified 8.86
python/sglang/srt/model_executor/model_runner.py 模型运行器 modified 8.86
python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py 解码图运行器 modified 7.7
python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py 编译工具 modified 5.77
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py 预填图运行器 modified 4.63
python/sglang/srt/model_executor/runner/eager_runner.py 即时运行器 modified 4.53
python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py 注意力测试 modified 3.66

关键符号

BaseRunner.warmup BaseRunner._dummy_run BaseRunner._flashinfer_autotune BaseRunner._pre_initialize_flashinfer_allreduce_workspace BaseRunner._should_run_flashinfer_autotune BaseRunner._allocate_decode_buffers ModelRunner.init_backends PrefillCudaGraphRunner.capture pp_parallel_deep_gemm_warmup

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

评论区精华

Missing pp_proxy_topk_size in _allocate_decode_buffers 正确性

gemini-code-assist[bot] 指出现有 _allocate_decode_buffers 签名缺少 pp_proxy_topk_size 参数,导致 topk_indices 未分配,pipeline parallel MoE(DSA)会在 dummy run 时崩溃。chatgpt-codex-connector[bot] 也提出同样问题。

结论:PR 已合并但该问题未修复。 · unresolved

风险与影响

  1. 深度流水线 MoE 兼容性风险:_allocate_decode_buffers 缺失 pp_proxy_topk_size,DSA 模型 warmup 可能崩溃。
  2. 初始化顺序依赖:warmup 后移,若外部代码依赖原时间点可能冲突。
  3. 测试覆盖不足:warmup 行为未有单元测试覆盖。

用户影响:非 MoE 模型无影响;使用 DeepSeek V2/V3 DSA 等 pipeline parallel MoE 需注意启动时崩溃。系统影响:forward 路径不再包含 autotune,消除运行时抖动。团队影响:代码结构更清晰。

缺失 DSA topk 代理缓冲区 pipeline parallel MoE 兼容风险 warmup 测试覆盖不足 初始化顺序潜在冲突

关联 Issue

未识别关联 Issue

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

参与讨论