diff --git a/python/sglang/kernels/ops/gemm/__init__.py b/python/sglang/kernels/ops/gemm/__init__.py index 600203826..1851a9831 100644 --- a/python/sglang/kernels/ops/gemm/__init__.py +++ b/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); diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 5093f8602..de940c480 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/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") diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index c399cf39f..bc26c7522 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/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 - - # 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 + from sglang.kernels.ops.gemm import bmm_fp8 if _use_aiter: diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py index e3ee27a48..6112b565b 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py +++ b/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 ( diff --git a/python/sglang/srt/models/minicpm3.py b/python/sglang/srt/models/minicpm3.py index fbcfab5fd..b9374dfe2 100644 --- a/python/sglang/srt/models/minicpm3.py +++ b/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.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 + from sglang.kernels.ops.gemm import bmm_fp8 class MiniCPM3MLP(nn.Module): diff --git a/python/sglang/srt/models/sarvam_moe.py b/python/sglang/srt/models/sarvam_moe.py index 4d184e11b..8add4bfe2 100644 --- a/python/sglang/srt/models/sarvam_moe.py +++ b/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 diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt index 85916ca7f..d4382e073 100644 --- a/sgl-kernel/CMakeLists.txt +++ b/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" diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc index 3998c9d67..bd8420659 100644 --- a/sgl-kernel/csrc/common_extension.cc +++ b/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); diff --git a/sgl-kernel/csrc/common_extension_musa.cc b/sgl-kernel/csrc/common_extension_musa.cc index 52b139def..bf83cba3c 100644 --- a/sgl-kernel/csrc/common_extension_musa.cc +++ b/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); diff --git a/sgl-kernel/csrc/gemm/bmm_fp8.cu b/sgl-kernel/csrc/gemm/bmm_fp8.cu deleted file mode 100644 index cef85a7de..000000000 --- a/sgl-kernel/csrc/gemm/bmm_fp8.cu +++ /dev/null @@ -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 - -#include - -#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(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.data_ptr()), - static_cast(A.data_ptr()), - static_cast(D.data_ptr()), - batch_size, - n, - m, - k, - static_cast(B_scale.data_ptr()), - static_cast(A_scale.data_ptr()), - lt_handle, - stream); - TORCH_CHECK( - status == CUBLAS_STATUS_SUCCESS, "bmm_fp8_internal_cublaslt failed: ", cublasGetStatusString(status)); - return true; - }); - }); - }); -} diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h index bd86c0c90..2e63f5fec 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/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& 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( diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py index 1632d64e3..f4874cd73 100644 --- a/sgl-kernel/python/sgl_kernel/__init__.py +++ b/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", diff --git a/sgl-kernel/python/sgl_kernel/gemm.py b/sgl-kernel/python/sgl_kernel/gemm.py index 97ff82273..e6a9b50d1 100644 --- a/sgl-kernel/python/sgl_kernel/gemm.py +++ b/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, diff --git a/sgl-kernel/setup_musa.py b/sgl-kernel/setup_musa.py index 92bb6962e..db8cb3fd7 100644 --- a/sgl-kernel/setup_musa.py +++ b/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", diff --git a/sgl-kernel/tests/test_bmm_fp8.py b/sgl-kernel/tests/test_bmm_fp8.py deleted file mode 100644 index c6c463d9b..000000000 --- a/sgl-kernel/tests/test_bmm_fp8.py +++ /dev/null @@ -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__]))