执行摘要

• 一句话：修复 FlashInfer bmm_fp8 后端选择错误
• 推荐动作：建议合入，这是一次经过调研的低风险性能修复，改动极小，收益明确。无需深入精读，但值得关注后续是否有更多 backend 硬编码调整。

功能与动机

PR body 指出 FlashInfer 的 autotune 在 auto 后端下会错误选择 CUTLASS（实际 cuBLAS 更优），导致 Nemotron 模型 GEMM 性能损失，M <= 256 时可提升 10%。替代方案 #29071 不如本 PR 直接有效。

实现拆解

1. 修改文件 python/sglang/srt/layers/quantization/fp8_utils.py 中 flashinfer_bmm_fp8 函数调用 _raw_flashinfer_bmm_fp8 时的 backend 参数，从 "auto" 改为 "cublas"。
2. 该修改仅影响 Blackwell 架构且 FlashInfer 可用时的路径，不影响其他后端（如 groupwise 的 CUTLASS/TRTLLM 选择逻辑）。

关键文件：

• python/sglang/srt/layers/quantization/fp8_utils.py（模块 量化；类别 source；类型 core-logic；符号 flashinfer_bmm_fp8）: 唯一修改的文件，将 flashinfer_bmm_fp8 的 backend 参数从 auto 改为 cublas，修复性能退化。

关键符号：flashinfer_bmm_fp8

关键源码片段

python/sglang/srt/layers/quantization/fp8_utils.py

唯一修改的文件，将 flashinfer_bmm_fp8 的 backend 参数从 auto 改为 cublas，修复性能退化。

# python/sglang/srt/layers/quantization/fp8_utils.py
​
@register_custom_op(
    op_name="flashinfer_bmm_fp8",
    mutates_args=[],
    fake_impl=lambda q_input, weight, x_scale, weight_scale, out_dtype: (
        q_input.new_empty((q_input.shape[0], weight.shape[1]), dtype=out_dtype)
    ),
)
def flashinfer_bmm_fp8(
    q_input: torch.Tensor, # [M, K] fp8 e4m3
    weight: torch.Tensor, # [K, N] fp8 e4m3, column-major
    x_scale: torch.Tensor, # per-tensor scalar
    weight_scale: torch.Tensor, # per-tensor scalar
    out_dtype: torch.dtype,
) -> torch.Tensor:
    m, n = q_input.shape[0], weight.shape[1]
    return _raw_flashinfer_bmm_fp8(
        q_input.unsqueeze(0),
        weight.unsqueeze(0),
        x_scale.reshape(1),
        weight_scale.reshape(1),
        out_dtype,
        # 修复：强制使用 cuBLAS 后端替代 auto，避免 autotune 错误选择 CUTLASS
        backend="cublas",
    ).view(m, n)

评论区精华

PR 无 review 评论，但 body 中与 #29071 关联，表明这是经过调研后的更优方案。

• 暂无高价值评论线程

风险与影响

• 风险：风险较低：仅修改一行 backend 参数，从 auto 改为 cuBLAS，cuBLAS 是 NVIDIA 官方库，稳定可靠。但需确认 cuBLAS 在所有支持的 Blackwell 硬件上均可用且表现最优，否则可能在其他场景退化。
• 影响：影响范围小：只影响 Blackwell + FlashInfer 可用时的 bmm_fp8 调用路径。受益场景：Nemotron 等模型 M <= 256 的 GEMM，预期性能提升 10%。不会影响非 Blackwell 架构或未使用 FlashInfer 的场景。
• 风险标记：hardware-specific, 缺少测试覆盖

关联脉络

• PR #29071 Alternative fix for CuDNN failure on bmm_fp8 when two libcudart.so exists: PR body 指出本 PR 是该替代方案的更优选择，关联同一问题的修复尝试。