Delete sgl-kernel AOT bmm_fp8, use flashinfer.bmm_fp8 (#31202)

Co-authored-by: root <root@sgl-b300-inference.datacrunch.io>

python/sglang/kernels/ops/gemm/__init__.py

@@ -30,6 +30,19 @@ register_kernel(
         description="FP8 scaled matmul (sgl_kernel wheel).",
     )
 )
+register_kernel(
+    KernelSpec(
+        op="gemm.bmm_fp8",
+        backend=KernelBackend.FLASHINFER,
+        target="sglang.srt.layers.quantization.fp8_utils:bmm_fp8",
+        capabilities=_CUDA,
+        format_signature=FormatSignature(
+            supported_dtypes=("float8_e4m3fn", "float8_e5m2"),
+            description="batched (3D) per-tensor-scale FP8 matmul: D = A_fp8 @ B_fp8 * A_scale * B_scale",
+        ),
+        description="Batched FP8 matmul (flashinfer cuBLAS backend, torch.compile-safe wrapper).",
+    )
+)
 register_kernel(
     KernelSpec(
         op="gemm.dsv3_fused_a_gemm",
@@ -84,6 +97,20 @@ def fp8_scaled_mm(
     )
 
 
+def bmm_fp8(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    A_scale: torch.Tensor,
+    B_scale: torch.Tensor,
+    dtype: torch.dtype,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Batched (3D) per-tensor-scale FP8 matmul, via flashinfer's cuBLAS backend."""
+    return get_kernel("gemm.bmm_fp8", KernelBackend.FLASHINFER)(
+        A, B, A_scale, B_scale, dtype, out
+    )
+
+
 def dsv3_fused_a_gemm(
     mat_a: torch.Tensor,
     mat_b: torch.Tensor,
@@ -106,7 +133,7 @@ def dsv3_router_gemm(
     return impl(hidden_states, router_weights, out_dtype, output)
 
 
-__all__ = ["fp8_scaled_mm", "dsv3_fused_a_gemm", "dsv3_router_gemm"]
+__all__ = ["fp8_scaled_mm", "bmm_fp8", "dsv3_fused_a_gemm", "dsv3_router_gemm"]
 
 
 # LoRA SGMV Triton kernels migrated into this group (from lora/triton_ops);

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

@@ -173,6 +173,36 @@ if _is_cuda:
         N = mat_b.shape[-1]
         return mat_a.new_empty((M, N), dtype=out_dtype)
 
+    from flashinfer import bmm_fp8 as _raw_bmm_fp8_batched
+
+    @register_custom_op(op_name="flashinfer_bmm_fp8_batched", mutates_args=["out"])
+    def _bmm_fp8_batched_op(
+        A: torch.Tensor,
+        B: torch.Tensor,
+        out: torch.Tensor,
+        A_scale: torch.Tensor,
+        B_scale: torch.Tensor,
+    ) -> None:
+        _raw_bmm_fp8_batched(A, B, A_scale, B_scale, out.dtype, out)
+
+    def bmm_fp8(
+        A: torch.Tensor,
+        B: torch.Tensor,
+        A_scale: torch.Tensor,
+        B_scale: torch.Tensor,
+        dtype: torch.dtype,
+        out: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Batched (3D) per-tensor-scale FP8 matmul, via flashinfer's cuBLAS backend."""
+        if out is None:
+            out = torch.empty(
+                (A.shape[0], A.shape[1], B.shape[2]),
+                device=A.device,
+                dtype=dtype,
+            )
+        _bmm_fp8_batched_op(A, B, out, A_scale, B_scale)
+        return out
+
 
 use_triton_w8a8_fp8_kernel = get_bool_env_var("USE_TRITON_W8A8_FP8_KERNEL")
 

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py

@@ -88,30 +88,7 @@ class MlaBmmFusionPlan:
 
 
 if _is_cuda:
-    from sgl_kernel import bmm_fp8 as _raw_bmm_fp8
+    from sglang.kernels.ops.gemm import bmm_fp8
-
-    # TODO(yuwei): remove this wrapper after sgl-kernel registers its own fake/meta impl
-    # Wrap bmm_fp8 as a custom op so torch.compile does not trace into
-    # torch.cuda.current_blas_handle() (which returns a non-Tensor).
-    @register_custom_op(mutates_args=["out"])
-    def _bmm_fp8_op(
-        A: torch.Tensor,
-        B: torch.Tensor,
-        out: torch.Tensor,
-        A_scale: torch.Tensor,
-        B_scale: torch.Tensor,
-    ) -> None:
-        _raw_bmm_fp8(A, B, A_scale, B_scale, out.dtype, out)
-
-    def bmm_fp8(A, B, A_scale, B_scale, dtype, out=None):
-        if out is None:
-            out = torch.empty(
-                (A.shape[0], A.shape[1], B.shape[2]),
-                device=A.device,
-                dtype=dtype,
-            )
-        _bmm_fp8_op(A, B, out, A_scale, B_scale)
-        return out
 
 
 if _use_aiter:

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py

@@ -21,7 +21,7 @@ if TYPE_CHECKING:
     from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
 
 if _is_cuda:
-    from sgl_kernel import bmm_fp8
+    from sglang.kernels.ops.gemm import bmm_fp8
 
 if _is_hip:
     from sglang.kernels.ops.attention.rocm_mla_decode_rope import (

python/sglang/srt/models/minicpm3.py

@@ -43,32 +43,7 @@ from sglang.srt.utils import add_prefix, is_cuda
 from sglang.srt.utils.hf_transformers_utils import get_rope_config
 
 if is_cuda():
-    from sgl_kernel import bmm_fp8 as _raw_bmm_fp8
+    from sglang.kernels.ops.gemm import bmm_fp8
-
-    from sglang.srt.utils.custom_op import register_custom_op
-
-    # TODO(yuwei): remove this wrapper after sgl-kernel registers its own fake/meta impl
-    # Wrap bmm_fp8 as a custom op so torch.compile does not trace into
-    # torch.cuda.current_blas_handle() (which returns a non-Tensor).
-    @register_custom_op(mutates_args=["out"])
-    def _bmm_fp8_op(
-        A: torch.Tensor,
-        B: torch.Tensor,
-        out: torch.Tensor,
-        A_scale: torch.Tensor,
-        B_scale: torch.Tensor,
-    ) -> None:
-        _raw_bmm_fp8(A, B, A_scale, B_scale, out.dtype, out)
-
-    def bmm_fp8(A, B, A_scale, B_scale, dtype, out=None):
-        if out is None:
-            out = torch.empty(
-                (A.shape[0], A.shape[1], B.shape[2]),
-                device=A.device,
-                dtype=dtype,
-            )
-        _bmm_fp8_op(A, B, out, A_scale, B_scale)
-        return out
 
 
 class MiniCPM3MLP(nn.Module):

python/sglang/srt/models/sarvam_moe.py

@@ -81,9 +81,10 @@ _is_cublas_ge_129 = is_nvidia_cublas_version_ge_12_9()
 
 if _is_cuda:
     try:
-        from sgl_kernel import bmm_fp8, merge_state_v2
+        from sgl_kernel import merge_state_v2
 
         from sglang.jit_kernel.concat_mla import concat_mla_k
+        from sglang.kernels.ops.gemm import bmm_fp8
         from sglang.kernels.ops.quantization.fp8_kernel import per_tensor_quant_mla_fp8
 
         _has_fp8_support = True

sgl-kernel/CMakeLists.txt

@@ -263,7 +263,6 @@ set(SOURCES
     "csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cu"
 
     "csrc/gemm/awq_kernel.cu"
-    "csrc/gemm/bmm_fp8.cu"
     "csrc/gemm/dsv3_fused_a_gemm.cu"
     "csrc/gemm/fp8_gemm_kernel.cu"
     "csrc/gemm/int8_gemm_kernel.cu"

sgl-kernel/csrc/common_extension.cc

@@ -344,12 +344,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
   /*
    * From FlashInfer
    */
-  m.def(
-      "bmm_fp8(Tensor A, Tensor B, Tensor! D, Tensor A_scale, Tensor B_scale, Tensor workspace_buffer, "
-      "int cublas_handle) -> ()",
-      {at::Tag::needs_fixed_stride_order});
-  m.impl("bmm_fp8", torch::kCUDA, &bmm_fp8);
-
   m.def("top_k_renorm_probs(Tensor probs, Tensor! renorm_probs, Tensor? maybe_top_k_arr, int top_k_val) -> ()");
   m.impl("top_k_renorm_probs", torch::kCUDA, &top_k_renorm_probs);
 

sgl-kernel/csrc/common_extension_musa.cc

@@ -256,12 +256,6 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
   /*
    * From FlashInfer
    */
-  m.def(
-      "bmm_fp8(Tensor A, Tensor B, Tensor! D, Tensor A_scale, Tensor B_scale, Tensor workspace_buffer, "
-      "int cublas_handle) -> ()",
-      {at::Tag::needs_fixed_stride_order});
-  m.impl("bmm_fp8", torch::kMUSA, &bmm_fp8);
-
   m.def("top_k_renorm_probs(Tensor probs, Tensor! renorm_probs, Tensor? maybe_top_k_arr, int top_k_val) -> ()");
   m.impl("top_k_renorm_probs", torch::kMUSA, &top_k_renorm_probs);
 

sgl-kernel/csrc/gemm/bmm_fp8.cu

@@ -1,75 +0,0 @@
-/*
- * Copyright (c) 2024 by FlashInfer team.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *   http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include <driver_types.h>
-
-#include <flashinfer/gemm/bmm_fp8.cuh>
-
-#include "pytorch_extension_utils.h"
-
-void bmm_fp8(
-    at::Tensor A,
-    at::Tensor B,
-    at::Tensor D,
-    at::Tensor A_scale,
-    at::Tensor B_scale,
-    at::Tensor workspace_buffer,
-    int64_t cublas_handle) {
-  TORCH_CHECK(A.is_cuda(), "A must be a CUDA tensor");
-  TORCH_CHECK(B.is_cuda(), "B must be a CUDA tensor");
-  TORCH_CHECK(D.is_cuda(), "D must be a CUDA tensor");
-  TORCH_CHECK(A.dim() == 3, "Expected 3D tensor for A");
-  TORCH_CHECK(B.dim() == 3, "Expected 3D tensor for B");
-  TORCH_CHECK(D.dim() == 3, "Expected 3D tensor for D");
-  TORCH_CHECK(A.size(0) == B.size(0) && A.size(0) == D.size(0), "Batch sizes must match");
-  TORCH_CHECK(A.size(2) == B.size(1), "Incompatible matrix sizes");
-  TORCH_CHECK(A.size(1) == D.size(1) && B.size(2) == D.size(2), "Result tensor has incorrect shape");
-
-  // PyTorch is row major by default. cuBLASLt is column major by default.
-  // We need row major D as expected.
-  // A ^ T * B = D, so D ^ T = B ^ T * A
-  DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP8(B.scalar_type(), b_type, [&] {
-    return DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP8(A.scalar_type(), a_type, [&] {
-      return DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(D.scalar_type(), d_type, [&] {
-        auto batch_size = A.size(0);
-        auto m = A.size(1);
-        auto k = A.size(2);
-        auto n = B.size(2);
-
-        auto lt_handle = reinterpret_cast<cublasLtHandle_t>(cublas_handle);
-        auto stream = at::cuda::getCurrentCUDAStream();
-
-        auto status = flashinfer::bmm_fp8::bmm_fp8_internal_cublaslt(
-            workspace_buffer.data_ptr(),
-            workspace_buffer.numel(),
-            static_cast<b_type*>(B.data_ptr()),
-            static_cast<a_type*>(A.data_ptr()),
-            static_cast<d_type*>(D.data_ptr()),
-            batch_size,
-            n,
-            m,
-            k,
-            static_cast<float*>(B_scale.data_ptr()),
-            static_cast<float*>(A_scale.data_ptr()),
-            lt_handle,
-            stream);
-        TORCH_CHECK(
-            status == CUBLAS_STATUS_SUCCESS, "bmm_fp8_internal_cublaslt failed: ", cublasGetStatusString(status));
-        return true;
-      });
-    });
-  });
-}

sgl-kernel/include/sgl_kernel_ops.h

@@ -256,14 +256,6 @@ void sgl_per_token_group_quant_8bit_v2(
     bool fuse_silu_and_mul,
     const std::optional<torch::Tensor>& masked_m);
 void sgl_per_token_quant_fp8(at::Tensor input, at::Tensor output_q, at::Tensor output_s);
-void bmm_fp8(
-    at::Tensor A,
-    at::Tensor B,
-    at::Tensor D,
-    at::Tensor A_scale,
-    at::Tensor B_scale,
-    at::Tensor workspace_buffer,
-    int64_t cublas_handle);
 void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a, torch::Tensor const& mat_b);
 
 torch::Tensor gptq_gemm(

sgl-kernel/python/sgl_kernel/__init__.py

@@ -55,7 +55,6 @@ else:
     )
     from sgl_kernel.gemm import (
         awq_dequantize,
-        bmm_fp8,
         dsv3_fused_a_gemm,
         fp8_scaled_mm,
         gptq_gemm,
@@ -153,7 +152,6 @@ else:
         "apply_shuffle_mul_sum",
         "apply_token_bitmask_inplace_cuda",
         "awq_dequantize",
-        "bmm_fp8",
         "build_tree_kernel_efficient",
         "causal_conv1d_fwd",
         "causal_conv1d_update",

sgl-kernel/python/sgl_kernel/gemm.py

@@ -1,7 +1,6 @@
 from typing import Optional
 
 import torch
-from sgl_kernel.utils import _get_cache_buf
 
 
 def awq_dequantize(
@@ -32,45 +31,6 @@ def fp8_scaled_mm(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None):
     )
 
 
-def _bmm_fp8_internal(
-    workspace_buffer: torch.Tensor,
-    A: torch.Tensor,
-    B: torch.Tensor,
-    D: torch.Tensor,
-    A_scale: torch.Tensor,
-    B_scale: torch.Tensor,
-) -> None:
-    cublas_handle = torch.cuda.current_blas_handle()
-    torch.ops.sgl_kernel.bmm_fp8.default(
-        A,
-        B,
-        D,
-        A_scale,
-        B_scale,
-        workspace_buffer,
-        cublas_handle,
-    )
-
-
-def bmm_fp8(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    A_scale: torch.Tensor,
-    B_scale: torch.Tensor,
-    dtype: torch.dtype,
-    out: Optional[torch.Tensor] = None,
-) -> torch.Tensor:
-    if out is None:
-        out = torch.empty(
-            (A.shape[0], A.shape[1], B.shape[2]),
-            device=A.device,
-            dtype=dtype,
-        )
-    workspace_buffer = _get_cache_buf("bmm_fp8_workspace", 32 * 1024 * 1024, A.device)
-    _bmm_fp8_internal(workspace_buffer, A, B, out, A_scale, B_scale)
-    return out
-
-
 def dsv3_fused_a_gemm(
     mat_a: torch.Tensor,
     mat_b: torch.Tensor,

sgl-kernel/setup_musa.py

@@ -95,7 +95,6 @@ sources = [
     "csrc/speculative/speculative_sampling.cu",
     "csrc/kvcacheio/transfer.cu",
     "csrc/gemm/awq_kernel.cu",
-    "csrc/gemm/bmm_fp8.cu",
     "csrc/gemm/dsv3_fused_a_gemm.cu",
     "csrc/gemm/dsv3_router_gemm_bf16_out.cu",
     "csrc/gemm/dsv3_router_gemm_entry.cu",

sgl-kernel/tests/test_bmm_fp8.py

@@ -1,45 +0,0 @@
-# Adapted from https://github.com/flashinfer-ai/flashinfer/blob/4e8eb1879f9c3ba6d75511e5893183bf8f289a62/tests/test_bmm_fp8.py
-
-import sys
-
-import pytest
-import torch
-import torch.nn.functional as F
-from sgl_kernel import bmm_fp8
-
-
-def to_float8(x, dtype=torch.float8_e4m3fn):
-    finfo = torch.finfo(dtype)
-    min_val, max_val = x.aminmax()
-    amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12)
-    scale = finfo.max / amax
-    x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max)
-    return x_scl_sat.to(dtype), scale.float().reciprocal()
-
-
-@pytest.mark.parametrize("input_dtype", [torch.float8_e4m3fn, torch.float8_e5m2])
-@pytest.mark.parametrize("mat2_dtype", [torch.float8_e4m3fn, torch.float8_e5m2])
-@pytest.mark.parametrize("res_dtype", [torch.bfloat16, torch.float16])
-def test_bmm_fp8(input_dtype, mat2_dtype, res_dtype):
-    if input_dtype == torch.float8_e5m2 and mat2_dtype == torch.float8_e5m2:
-        pytest.skip("Invalid combination: both input and mat2 are e5m2")
-
-    input = torch.randn([16, 48, 64], device="cuda", dtype=torch.bfloat16)
-    input_fp8, input_inv_s = to_float8(input, dtype=input_dtype)
-
-    # mat2 row  major -> column major
-    mat2 = torch.randn([16, 80, 64], device="cuda", dtype=torch.bfloat16).transpose(
-        -2, -1
-    )
-    mat2_fp8, mat2_inv_s = to_float8(mat2, dtype=mat2_dtype)
-
-    res = torch.empty([16, 48, 80], device="cuda", dtype=res_dtype)
-    bmm_fp8(input_fp8, mat2_fp8, input_inv_s, mat2_inv_s, res_dtype, res)
-
-    reference = torch.bmm(input, mat2)
-    cos_sim = F.cosine_similarity(reference.reshape(-1), res.reshape(-1), dim=0)
-    assert cos_sim > 0.99
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__]))