chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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

133
sgl-kernel/benchmark/bench_es_fp8_blockwise_grouped_gemm.py
View File

@@ -43,31 +43,23 @@ def per_block_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
)
def create_unbalanced_expert_token_distribution(max_num_experts):
ratios = [random.random() for _ in range(max_num_experts)]
def convert_to_tokens(ratio: float):
if ratio <= 0.7:
return random.randint(1, 32)
elif ratio > 0.7 and ratio <= 0.85:
return random.randint(32, 64)
elif ratio > 0.85 and ratio <= 0.95:
return random.randint(64, 128)
elif ratio > 0.95:
return random.randint(128, 1024)
else:
return 128
group_ms = [convert_to_tokens(ratio) for ratio in ratios]
def create_unbalanced_expert_token_distribution(
batch_size: int, topk: int, num_experts: int
):
expert_ids = np.random.randint(0, num_experts, size=(batch_size * topk,)).tolist()
expert_to_count = dict()
for expert_id in range(num_experts):
expert_to_count[expert_id] = 0
for expert_id in expert_ids:
expert_to_count[expert_id] += 1
group_ms = []
for expert_id in range(num_experts):
group_ms.append(expert_to_count[expert_id])
return group_ms
group_ms = create_unbalanced_expert_token_distribution(8192)
# group_ms = [128 for _ in range(8192)]
# group_ms = [128 if i % 2 == 0 else 64 for i in range(8192)]
def bench_es(
group_ms: List[int],
n: int,
k: int,
num_groups: int,
@@ -94,12 +86,13 @@ def bench_es(
m_g = group_ms[g]
expert_offsets[g + 1] = expert_offsets[g] + m_g
problem_sizes[g][:] = torch.tensor([m_g, n_g, k_g], device=device)
if m_g != 0:
a_g, a_scale = per_token_cast_to_fp8(torch.randn((m_g, k_g), device=device))
a_tensors.append(a_g)
a_scales_tensors.append(a_scale)
a_g, a_scale = per_token_cast_to_fp8(torch.randn((m_g, k_g), device=device))
b_g, b_scale = per_block_cast_to_fp8(torch.randn((n_g, k_g), device=device).t())
a_tensors.append(a_g)
b_tensors.append(b_g)
a_scales_tensors.append(a_scale)
b_scales_tensors.append(b_scale)
a_stack = torch.empty(
@@ -109,8 +102,11 @@ def bench_es(
(num_groups, n_g, k_g), device=device, dtype=torch.float8_e4m3fn
)
_aux_idx = 0
for g in range(num_groups):
a_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_tensors[g]
if group_ms[g] != 0:
a_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_tensors[_aux_idx]
_aux_idx += 1
b_stack[g] = b_tensors[g].t()
b_stack = b_stack.transpose(1, 2)
@@ -121,11 +117,17 @@ def bench_es(
(num_groups, n_g // 128, k_g // 128), device=device, dtype=torch.float32
)
_aux_idx = 0
for g in range(num_groups):
a_scale_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_scales_tensors[g]
if group_ms[g] != 0:
a_scale_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_scales_tensors[
_aux_idx
]
_aux_idx += 1
b_scale_stack[g] = b_scales_tensors[g].t()
b_scale_stack = b_scale_stack.transpose(1, 2)
workspace = torch.empty((1024 * 1024 * 1024), device=device, dtype=torch.uint8)
c_out = torch.empty((expert_offsets[-1], n_g), device=device, dtype=out_dtype)
a_strides = torch.full(
(num_groups,), a_stack.stride(0), device=device, dtype=torch.int64
@@ -133,7 +135,6 @@ def bench_es(
d_strides = torch.full(
(num_groups,), c_out.stride(0), device=device, dtype=torch.int64
)
workspace = torch.empty((1024 * 1024 * 1024), device=device, dtype=torch.uint8)
def run_cutlass():
es_fp8_blockwise_scaled_grouped_mm(
@@ -171,6 +172,7 @@ def bench_es(
def bench_sgl(
group_ms: List[int],
n: int,
k: int,
num_groups: int,
@@ -197,12 +199,13 @@ def bench_sgl(
m_g = group_ms[g]
expert_offsets[g + 1] = expert_offsets[g] + m_g
problem_sizes[g][:] = torch.tensor([m_g, n_g, k_g], device=device)
if m_g != 0:
a_g, a_scale = per_token_cast_to_fp8(torch.randn((m_g, k_g), device=device))
a_tensors.append(a_g)
a_scales_tensors.append(a_scale)
a_g, a_scale = per_token_cast_to_fp8(torch.randn((m_g, k_g), device=device))
b_g, b_scale = per_block_cast_to_fp8(torch.randn((n_g, k_g), device=device).t())
a_tensors.append(a_g)
b_tensors.append(b_g)
a_scales_tensors.append(a_scale)
b_scales_tensors.append(b_scale)
a_stack = torch.empty(
@@ -212,8 +215,11 @@ def bench_sgl(
(num_groups, n_g, k_g), device=device, dtype=torch.float8_e4m3fn
)
_aux_idx = 0
for g in range(num_groups):
a_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_tensors[g]
if group_ms[g] != 0:
a_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_tensors[_aux_idx]
_aux_idx += 1
b_stack[g] = b_tensors[g].t()
b_stack = b_stack.transpose(1, 2)
@@ -224,8 +230,13 @@ def bench_sgl(
(num_groups, n_g // 128, k_g // 128), device=device, dtype=torch.float32
)
_aux_idx = 0
for g in range(num_groups):
a_scale_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_scales_tensors[g]
if group_ms[g] != 0:
a_scale_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_scales_tensors[
_aux_idx
]
_aux_idx += 1
b_scale_stack[g] = b_scales_tensors[g].t()
b_scale_stack = b_scale_stack.transpose(1, 2)
@@ -300,16 +311,36 @@ def benchmark_one_shape(
num_run: int,
):
for shape in shape_args:
print(f"\nBenchmark: n={shape.n}, k={shape.k}, num_groups={shape.num_groups}")
for kernel_name, kernel_func in benchmark_kernels.items():
average_time, m = kernel_func(
shape.n,
shape.k,
shape.num_groups,
num_warmup,
num_run,
for batch_size in [
128,
256,
384,
512,
640,
768,
896,
1024,
1280,
1536,
2048,
3072,
]:
group_ms = create_unbalanced_expert_token_distribution(
batch_size, 8, shape.num_groups
)
print(f"{kernel_name}: {average_time} us")
print(
f"\nBenchmark: batch_size={batch_size}, n={shape.n}, k={shape.k}, num_groups={shape.num_groups}"
)
for kernel_name, kernel_func in benchmark_kernels.items():
average_time, m = kernel_func(
group_ms,
shape.n,
shape.k,
shape.num_groups,
num_warmup,
num_run,
)
print(f"{kernel_name}: {average_time} us")
def main():
@@ -317,18 +348,22 @@ def main():
parser.add_argument("--num-warmup", type=int, default=3)
parser.add_argument("--num-run", type=int, default=20)
shape_args = [
# Prefill, DeepSeek-R1, gateup, chunk_size = 4096, TP = 8
# DeepSeek-R1, gateup, TP = 8
ShapeArg(n=512, k=7168, num_groups=256),
# Prefill, DeepSeek-R1, down, chunk_size = 4096, TP = 8
# DeepSeek-R1, down, TP = 8
ShapeArg(n=7168, k=256, num_groups=256),
# Prefill, Qwen3-235B-A22B-FP8, gateup, TP = 4
# DeepSeek-R1, gateup, TP = 4
ShapeArg(n=1024, k=7168, num_groups=256),
# DeepSeek-R1, down, TP = 4
ShapeArg(n=7168, k=512, num_groups=256),
# Qwen3-235B-A22B-FP8, gateup, TP = 4
ShapeArg(n=768, k=4096, num_groups=128),
# Prefill, Qwen3-235B-A22B-FP8, down, TP = 4
# Qwen3-235B-A22B-FP8, down, TP = 4
ShapeArg(n=4096, k=384, num_groups=128),
# Decode, DeepSeek-R1, gateup, bs = 128, EP = 8
ShapeArg(n=4096, k=7168, num_groups=32),
# Decode, DeepSeek-R1, gateup, bs = 256, EP = 16
ShapeArg(n=4096, k=7168, num_groups=16),
# Qwen3-235B-A22B-FP8, gateup, TP = 2
ShapeArg(n=1536, k=4096, num_groups=128),
# Qwen3-235B-A22B-FP8, down, TP = 2
ShapeArg(n=4096, k=768, num_groups=128),
]
args = parser.parse_args()
benchmark_one_shape(shape_args, args.num_warmup, args.num_run)

38
sgl-kernel/csrc/expert_specialization/es_fp8_blockwise.cu
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);

21
sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_functor.cuh
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;

99
sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_launcher.cuh
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);
}
}

6
sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_traits.cuh
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());