Prhub

#28740 refactor(runner): reuse a prepared static buffer for every dummy run

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

执行摘要

复用静态缓冲区替代逐次分配,优化 dummy run 内存管理

原始 flashinfer-autotune dummy forward 会分配一个与 req_to_token_pool.size 等大的临时 DecodeInputBuffers 集合。该 PR 改为复用 runner 已经分配的静态缓冲区,避免不必要的分配,同时修复旧版重填充可能超出复用缓冲区边界的溢出风险。

该 PR 值得精读,尤其是 _dummy_run 的缓冲区接管模式和 _autotune_buffers 的适配器设计。review 中提出的问题值得在后续 PR 中跟进修复。建议关注其与 #28387 等 runner 重构 PR 的关联,形成对 SRT 运行器架构演进的全局理解。

讨论亮点

review 中识别出三个关键问题:

  • gemini-code-assist 指出 buffers.mrope_positions 在非多模态模型下可能为 None,直接切片会引发 TypeError,建议添加 None 检查。
  • chatgpt-codex-connector 指出 eager registry 创建时 seq_len_fill_value=0,导致 autotune dummy decode 的 seq_lens 为零,可能影响 autotune 正确性(P1)。
  • chatgpt-codex-connector 指出 eager autotune 的 batch 未经 MLP sync 对齐,可能导致集合通信异常(P2)。
    以上评论均未在 PR 中收到作者回复,但 PR 已合并,可能已通过其他方式解决或确认不影响。

实现拆解

  1. 变更入口:BaseRunner.warmup 与 _flashinfer_autotune - 修改 warmup(),在触发 autotune 前通过 _autotune_buffers() 获取缓冲区集并断言非空;修改 _flashinfer_autotune 接受 buffersbatch_size 参数,直接传递给 _dummy_run,不再分配缓冲区。
  2. 子类适配器:_autotune_buffers 实现 - DecodeCudaGraphRunner 直接返回自身 buffersmax_bsEagerRunner_eager_registry 构建适配器,填补遗漏字段(logits 缓冲区、PP 代理张量、custom_mask 等),并保存 _eager_max_bs_eager_num_tokens_per_bs 供复用。
  3. DeepGEMM PP 并行预热优化 - pp_parallel_deep_gemm_warmup 通过 runner._alloc_dummy_decode_buffers 预分配一个最大缓冲区集,在整个 batch size 扫描中复用;新增内联 _align() 函数,使用 lcm(cp, attn_tp_size) 对齐,避免后续 MLP sync 对齐破坏 CP 倍数。
  4. 共享缓冲区分配函数 - 新增 BaseRunner._alloc_dummy_decode_buffers 作为统一入口,供非 autotune 预热场景(如 DeepGEMM 预热)使用。
  5. 配置与依赖调整 - 从 base_runner.py 移除 ceil_alignrequire_mlp_sync 的直接导入,改为局部导入;其他文件相应调整依赖。
文件 模块 状态 重要度
python/sglang/srt/model_executor/runner/base_runner.py 运行器 modified 8.48
python/sglang/srt/model_executor/runner/eager_runner.py 运行器 modified 7.75
python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py 编译工具 modified 6.65
python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py 运行器 modified 5.89

关键符号

_flashinfer_autotune _alloc_dummy_decode_buffers _autotune_buffers _slot _align

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

评论区精华

mrope_positions 切片缺少 None 检查 正确性

gemini-code-assist 指出如果 buffers.mrope_positions 为 None(非多模态模型),直接切片会引发 TypeError,建议添加 None 检查类似于 encoder_lens 的处理。

结论:未收到明确回复,但 PR 已合并,可能已修复或认为不影响。 · pending

eager registry 使用 seq_len_fill_value=0 导致 autotune dummy decode 的 seq_lens 为零 正确性

chatgpt-codex-connector 指出 eager registry 创建时 seq_len_fill_value=0,导致 ForwardBatch 的 seq_lens 为零,而旧版 allocate 使用非零值,可能影响 autotune 正确性。

结论:未回复,待作者确认。 · pending

eager autotune batch 未对齐 MLP sync 正确性

chatgpt-codex-connector 指出 max_bs 未经 MLP sync 对齐,可能导致集合通信异常。

结论:未回复,待作者确认。 · pending

风险与影响

  1. 类型错误风险:base_runner.py 中 buffers.mrope_positions[:, :num_tokens] 在非多模态模型下可能为 None,直接切片会引发 TypeError,需添加 None 检查。
  2. eager autotune 正确性风险:eager registry 的 seq_len_fill_value=0 导致 seq_lens 为零,与旧版行为不同,可能影响 autotune 的有效性。
  3. MLP sync 对齐风险:eager autotune 的 batch 未对齐 MLP sync,可能导致集合通信异常,需要对齐到 attn_tp_size。
  4. 回归风险:改动涉及多个 runner 基类和子类,若其他未覆盖的 forward mode(如预填充)也调用 _dummy_run,可能因缓冲区不匹配而断言失败。

影响范围:中等。直接影响所有使用 flashinfer autotune 和 DeepGEMM 预热的模型启动过程,包括 eager 和 cuda graph 模式。影响程度:正向性能改善(减少 warmup 分配),并修复潜在溢出 bug。对于非受影响模型(如 DLLM),逻辑不变。

缓冲区对齐风险 类型错误风险 eager autotune 未对齐 MLP sync seq_lens 为零风险

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论