执行摘要

本PR优化了SM120 GPU上的FP8块状GEMM内核，通过引入pingpong schedule并基于M大小动态选择，显著提升了小M场景下的性能，推理速度提升约53%，同时保持了准确性。

功能与动机

原SM120 fp8 blockwise GEMM kernel使用KernelScheduleAuto，在SM120上仅选择cooperative schedule，导致小M时性能不足。pingpong schedule相比cooperative对于小M快约2倍，因此本PR旨在利用这一性能机会，提升整体效率。

实现拆解

关键改动集中在sgl-kernel/csrc/gemm/fp8_blockwise_gemm_kernel.cu文件的launch_sm120_fp8_blockwise_scaled_mm函数：

• 添加kCanUsePingpong常量检查。
• 根据M大小（M≤64）选择pingpong路径使用KernelTmaWarpSpecializedBlockwisePingpongSm120，否则使用cooperative路径。
• 重构kernel setup代码，提高可读性和维护性。

代码示例：

constexpr bool kCanUsePingpong = (64 % ScaleGranularityM == 0);
int m = a.size(0);
// ... 基于 m 选择 schedule

评论区精华

review过程中无实质性技术讨论，仅由BBuf批准。作者在issue评论中提供了NCU性能报告，进一步验证了变更的有效性，例如pingpong schedule相比cooperative schedule性能提升约一倍。

风险与影响

风险：

• 对于M>64使用cooperative路径以避免CUTLASS库问题，但可能存在未解决的准确性隐患。
• 新添加的pingpong路径可能引入回归或兼容性问题，尤其是在边缘case或不同硬件配置下。
• 代码变更集中在核心路径，测试覆盖可能不足。

影响：

• 用户受益：在SM120 GPU上运行小M形状的FP8 GEMM时，推理速度显著提升，基准测试显示速度从34.14 token/s提升至52.11 token/s。
• 系统层面：优化了内核调度策略，可能减少延迟并提高资源利用率。

关联脉络

本PR性能测试中与FlashInfer比较，关联到PR #20214（添加FlashInfer集成），表明团队在持续优化性能并集成第三方库。此外，无直接关联Issue，但体现了对硬件特定调优的关注。