执行摘要

此 PR 在 sgl-kernel 模块中暴露 get_scheduler_metadata torch op，预计算 Flash Attention v3 的调度元数据，以优化解码性能。通过避免每层重复内核调用，提升效率，变更向后兼容，是两部分优化的基础工作。

功能与动机

核心动机是优化 FA3 解码阶段的性能。PR body 明确指出："Precomputes FA3 tile scheduling metadata so that the prepare_varlen_num_blocks kernel does not need to run per-layer during decode." 这意味着通过一次预计算减少冗余操作，提升整体推理速度。

实现拆解

变更涉及三个文件，按层次拆解：

• C++ 头文件 (sgl_flash_kernel_ops.h)：声明 mha_fwd_get_scheduler_metadata 函数，提供参数接口。
• C++ 源文件 (flash_extension.cc)：注册 torch op sgl_kernel.get_scheduler_metadata，集成到 PyTorch 生态。代码片段展示了参数列表：
  

  m.def(
    "get_scheduler_metadata(" +
    " int batch_size," +
    " ..." +
    ") -> Tensor");
  

• Python 包装器 (flash_attn.py)：添加 get_scheduler_metadata() 函数，提供用户友好接口，但初始版本缺失参数，后经补全。

评论区精华

review 中主要讨论围绕 Python 包装器的完整性展开。gemini-code-assist[bot] 指出：

"The Python wrapper for get_scheduler_metadata is missing several parameters... could cause incorrect behavior when used with features like left padding."
这引发了正确性担忧，但通过后续提交（refactor: complete Python wrapper with all C++ op parameters）解决了问题，体现了代码审查中的质量把关。

风险与影响

• 风险：Python 包装器初始参数不完整可能导致使用错误，但已修复；新函数依赖 flash_ops.so 中的 C++ 实现，需确保参数传递正确。
• 影响：解码性能提升，减少内核调用开销；为团队后续集成（PR #21104）提供基础，推动整体系统优化。

关联脉络

此 PR 是两阶段优化的第1部分，直接关联 PR #21104（第2部分），后者将集成此 op 到 sglang Python 层。从近期历史 PR 看，仓库注重性能优化（如 PR #21318、#21253），本 PR 延续了这一趋势，聚焦内核层调度优化。