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

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

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,