执行摘要
此 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 生态。代码片段展示了参数列表：
cpp
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 延续了这一趋势，聚焦内核层调度优化。
参与讨论