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

[sgl-kernel][6/7]Support Expert Specialization Grouped GEMM (#15471)

Browse Source
This commit is contained in:
Qi Yuhang
2026-03-19 15:39:52 +08:00
committed by GitHub
parent 574572b21b
commit cb8105fe28
5 changed files with 192 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/expert_specialization/es_fp8_blockwise.cu
+32 -6
View File
@@ -1,3 +1,4 @@
#include <ATen/cuda/CUDAEvent.h>
#include <torch/all.h>
#include <tuple>
@@ -70,10 +71,18 @@ void es_fp8_blockwise_scaled_grouped_mm(
torch::Tensor mm_problem_sizes = torch::empty({num_experts, 3}, options_int32);
torch::Tensor hm_problem_sizes = torch::empty({num_experts, 3}, options_int32);
torch::Tensor backup_workspace_0 = torch::empty_like(workspace);
torch::Tensor backup_workspace_1 = torch::empty_like(workspace);
const std::string H20_device_type_str("NVIDIA H20");
bool is_h20_device = std::string(at::cuda::getCurrentDeviceProperties()->name) == H20_device_type_str;
at::cuda::CUDAGuard device_guard{(char)a.get_device()};
cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.get_device());
auto stream = at::cuda::getCurrentCUDAStream();
static auto backup_stream_0 = at::cuda::getStreamFromPool();
static auto backup_stream_1 = at::cuda::getStreamFromPool();
at::cuda::CUDAEvent start_event;
at::cuda::CUDAEvent end_event_0;
at::cuda::CUDAEvent end_event_1;
if (output.dtype() == torch::kBFloat16) {
expert_specialization::es_sm90_fp8_blockwise_scaled_group_mm_pre_compute<cutlass::bfloat16_t>(
@@ -95,7 +104,7 @@ void es_fp8_blockwise_scaled_grouped_mm(
problem_sizes,
expert_offsets,
is_h20_device,
stream);
stream.stream());
} else if (output.dtype() == torch::kFloat16) {
expert_specialization::es_sm90_fp8_blockwise_scaled_group_mm_pre_compute<cutlass::half_t>(
out_ptrs,
@@ -116,11 +125,15 @@ void es_fp8_blockwise_scaled_grouped_mm(
problem_sizes,
expert_offsets,
is_h20_device,
stream);
stream.stream());
} else {
TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
}
start_event.recordOnce(stream);
start_event.block(backup_stream_0);
start_event.block(backup_stream_1);
if (output.dtype() == torch::kBFloat16) {
expert_specialization::es_sm90_fp8_blockwise_scaled_group_mm_distpatch_out_dtype<cutlass::bfloat16_t>(
out_ptrs,
@@ -137,8 +150,12 @@ void es_fp8_blockwise_scaled_grouped_mm(
mm_problem_sizes,
hm_problem_sizes,
workspace,
backup_workspace_0,
backup_workspace_1,
is_h20_device,
stream);
stream.stream(),
backup_stream_0.stream(),
backup_stream_1.stream());
} else if (output.dtype() == torch::kFloat16) {
expert_specialization::es_sm90_fp8_blockwise_scaled_group_mm_distpatch_out_dtype<cutlass::half_t>(
out_ptrs,
@@ -155,11 +172,20 @@ void es_fp8_blockwise_scaled_grouped_mm(
mm_problem_sizes,
hm_problem_sizes,
workspace,
backup_workspace_0,
backup_workspace_1,
is_h20_device,
stream);
stream.stream(),
backup_stream_0.stream(),
backup_stream_1.stream());
} else {
TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
}
end_event_0.recordOnce(backup_stream_0);
end_event_1.recordOnce(backup_stream_1);
end_event_0.block(stream);
end_event_1.block(stream);
#else
TORCH_CHECK_NOT_IMPLEMENTED(
can_implement, "No implemented fp8_blockwise_scaled_grouped_mm for current compute capability: ", sm_version);
sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_functor.cuh
+18 -3
View File
@@ -126,7 +126,12 @@ struct Fp8BlockwiseGroupedGemmProblemSizeFilterFunctor<PerfConfigLowMH20> {
Fp8BlockwiseGroupedGemmProblemSizeFilterFunctor(int* _problem_sizes) : problem_sizes(_problem_sizes) {}
void CUTE_DEVICE operator()(int64_t expert_id, int m, int n, int k) {
if (m < 64) {
float m_f = __int2float_rn(m);
float n_f = __int2float_rn(n);
float k_f = __int2float_rn(k);
float arithmetic_intensity = 2.0f * m_f * n_f * k_f / (m_f * k_f + k_f * n_f + 2.0f * m_f * n_f);
if (m <= 32 || arithmetic_intensity < 70.0f) {
// Swap A/B
problem_sizes[expert_id * 3 + 0] = n;
problem_sizes[expert_id * 3 + 1] = m;
@@ -168,7 +173,12 @@ struct Fp8BlockwiseGroupedGemmProblemSizeFilterFunctor<PerfConfigMiddleMH20> {
Fp8BlockwiseGroupedGemmProblemSizeFilterFunctor(int* _problem_sizes) : problem_sizes(_problem_sizes) {}
void CUTE_DEVICE operator()(int64_t expert_id, int m, int n, int k) {
if (m >= 64 && m < 128) {
float m_f = __int2float_rn(m);
float n_f = __int2float_rn(n);
float k_f = __int2float_rn(k);
float arithmetic_intensity = 2.0f * m_f * n_f * k_f / (m_f * k_f + k_f * n_f + 2.0f * m_f * n_f);
if ((!(m <= 32 || arithmetic_intensity < 70.0f)) && m <= 64) {
problem_sizes[expert_id * 3 + 0] = m;
problem_sizes[expert_id * 3 + 1] = n;
problem_sizes[expert_id * 3 + 2] = k;
@@ -208,7 +218,12 @@ struct Fp8BlockwiseGroupedGemmProblemSizeFilterFunctor<PerfConfigHighMH20> {
Fp8BlockwiseGroupedGemmProblemSizeFilterFunctor(int* _problem_sizes) : problem_sizes(_problem_sizes) {}
void CUTE_DEVICE operator()(int64_t expert_id, int m, int n, int k) {
if (m >= 128) {
float m_f = __int2float_rn(m);
float n_f = __int2float_rn(n);
float k_f = __int2float_rn(k);
float arithmetic_intensity = 2.0f * m_f * n_f * k_f / (m_f * k_f + k_f * n_f + 2.0f * m_f * n_f);
if ((!(m <= 32 || arithmetic_intensity < 70.0f)) && m > 64) {
problem_sizes[expert_id * 3 + 0] = m;
problem_sizes[expert_id * 3 + 1] = n;
problem_sizes[expert_id * 3 + 2] = k;
sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_launcher.cuh
+55 -44
View File
@@ -99,7 +99,8 @@ void launch_sm90_fp8_blockwise_scaled_group_mm(
const torch::Tensor& layout_sfb,
const torch::Tensor& problem_sizes,
const torch::Tensor& workspace,
cudaStream_t stream) {
cudaStream_t stream,
int sm_count) {
using ElementA = typename GemmTraits::ElementA;
using StrideA = typename GemmTraits::StrideA;
using ElementB = typename GemmTraits::ElementB;
@@ -128,7 +129,7 @@ void launch_sm90_fp8_blockwise_scaled_group_mm(
cutlass::KernelHardwareInfo hw_info;
hw_info.device_id = c10::cuda::current_device();
hw_info.sm_count = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
hw_info.sm_count = sm_count;
typename GemmKernel::EpilogueArguments epilogue_args{
{}, nullptr, nullptr, static_cast<ElementD**>(out_ptrs.data_ptr()), static_cast<StrideD*>(stride_d.data_ptr())};
@@ -147,7 +148,7 @@ void launch_sm90_fp8_blockwise_scaled_group_mm(
auto status = gemm_op.initialize(args, workspace.data_ptr(), stream);
TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to initialize GEMM");
status = gemm_op.run(stream, nullptr, true); // Enable PDL
status = gemm_op.run(stream, nullptr);
TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
}
@@ -167,8 +168,12 @@ void es_sm90_fp8_blockwise_scaled_group_mm_distpatch_out_dtype(
const torch::Tensor& mm_problem_sizes,
const torch::Tensor& hm_problem_sizes,
const torch::Tensor& workspace,
const torch::Tensor& backup_workspace_0,
const torch::Tensor& backup_workspace_1,
bool is_h20_device,
cudaStream_t stream) {
cudaStream_t stream,
cudaStream_t backup_stream_0,
cudaStream_t backup_stream_1) {
using LowMGemmH20Traits =
ExpertSpecializationSm90FP8BlockwiseGroupedGemmTraits<OutType, cutlass::layout::ColumnMajor, PerfConfigLowMH20>;
using LowMGemmHx00Traits =
@@ -184,6 +189,40 @@ void es_sm90_fp8_blockwise_scaled_group_mm_distpatch_out_dtype(
using HighMGemmHx00Traits =
ExpertSpecializationSm90FP8BlockwiseGroupedGemmTraits<OutType, cutlass::layout::RowMajor, PerfConfigHighMHx00>;
if (!is_h20_device) {
launch_sm90_fp8_blockwise_scaled_group_mm<HighMGemmHx00Traits>(
out_ptrs,
a_ptrs,
b_ptrs,
a_scales_ptrs,
b_scales_ptrs,
stride_a,
stride_b,
stride_d,
layout_sfa,
layout_sfb,
hm_problem_sizes,
workspace,
stream,
132);
} else {
launch_sm90_fp8_blockwise_scaled_group_mm<HighMGemmH20Traits>(
out_ptrs,
a_ptrs,
b_ptrs,
a_scales_ptrs,
b_scales_ptrs,
stride_a,
stride_b,
stride_d,
layout_sfa,
layout_sfb,
hm_problem_sizes,
workspace,
stream,
78);
}
if (!is_h20_device) {
launch_sm90_fp8_blockwise_scaled_group_mm<LowMGemmHx00Traits>(
out_ptrs,
@@ -197,8 +236,9 @@ void es_sm90_fp8_blockwise_scaled_group_mm_distpatch_out_dtype(
layout_sfb,
layout_sfa,
lm_problem_sizes,
workspace,
stream);
backup_workspace_1,
backup_stream_1,
132);
} else {
launch_sm90_fp8_blockwise_scaled_group_mm<LowMGemmH20Traits>(
out_ptrs,
@@ -212,8 +252,9 @@ void es_sm90_fp8_blockwise_scaled_group_mm_distpatch_out_dtype(
layout_sfb,
layout_sfa,
lm_problem_sizes,
workspace,
stream);
backup_workspace_1,
backup_stream_1,
78);
}
if (!is_h20_device) {
@@ -229,8 +270,9 @@ void es_sm90_fp8_blockwise_scaled_group_mm_distpatch_out_dtype(
layout_sfb,
layout_sfa,
mm_problem_sizes,
workspace,
stream);
backup_workspace_0,
backup_stream_0,
132);
} else {
launch_sm90_fp8_blockwise_scaled_group_mm<MiddleMGemmH20Traits>(
out_ptrs,
@@ -244,40 +286,9 @@ void es_sm90_fp8_blockwise_scaled_group_mm_distpatch_out_dtype(
layout_sfa,
layout_sfb,
mm_problem_sizes,
workspace,
stream);
}
if (!is_h20_device) {
launch_sm90_fp8_blockwise_scaled_group_mm<HighMGemmHx00Traits>(
out_ptrs,
a_ptrs,
b_ptrs,
a_scales_ptrs,
b_scales_ptrs,
stride_a,
stride_b,
stride_d,
layout_sfa,
layout_sfb,
hm_problem_sizes,
workspace,
stream);
} else {
launch_sm90_fp8_blockwise_scaled_group_mm<HighMGemmH20Traits>(
out_ptrs,
a_ptrs,
b_ptrs,
a_scales_ptrs,
b_scales_ptrs,
stride_a,
stride_b,
stride_d,
layout_sfa,
layout_sfb,
hm_problem_sizes,
workspace,
stream);
backup_workspace_0,
backup_stream_0,
78);
}
}
sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_traits.cuh
+3 -3
View File
@@ -30,10 +30,10 @@ using namespace cute;
struct PerfConfigLowMH20 {
// Swap A/B
using ElementA = cutlass::float_e4m3_t;
using MmaTileShape = Shape<_128, _32, _128>;
using MmaTileShape = Shape<_256, _32, _128>;
using ClusterShape = Shape<_2, _1, _1>;
using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpongFP8Blockwise;
using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecializedPingpong;
using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperativeFP8Blockwise;
using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecializedCooperative;
using ScaleConfig =
cutlass::detail::Sm90BlockwiseScaleConfig<128, 1, 128, cute::GMMA::Major::K, cute::GMMA::Major::K>;
using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA());