minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: root <root@sgl-b300-inference.datacrunch.io>
This commit is contained in:
Brayden Zhong
2026-07-22 07:44:47 +08:00
committed by GitHub
co-authored by root
parent 1b4cb6b8c1
commit 2f4f2362fb
15 changed files with 63 additions and 237 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/CMakeLists.txt
-1
View File
@@ -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
-6
View File
@@ -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
-6
View File
@@ -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
-75
View File
@@ -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
-8
View File
@@ -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
-2
View File
@@ -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
-40
View File
@@ -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
-1
View File
@@ -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
-45
View File
@@ -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__]))