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

sgl-kernel/benchmark/bench_es_fp8_blockwise_grouped_gemm.py

@@ -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):
+def create_unbalanced_expert_token_distribution(
-    ratios = [random.random() for _ in range(max_num_experts)]
+    batch_size: int, topk: int, num_experts: int
-
+):
-    def convert_to_tokens(ratio: float):
+    expert_ids = np.random.randint(0, num_experts, size=(batch_size * topk,)).tolist()
-        if ratio <= 0.7:
+    expert_to_count = dict()
-            return random.randint(1, 32)
+    for expert_id in range(num_experts):
-        elif ratio > 0.7 and ratio <= 0.85:
+        expert_to_count[expert_id] = 0
-            return random.randint(32, 64)
+    for expert_id in expert_ids:
-        elif ratio > 0.85 and ratio <= 0.95:
+        expert_to_count[expert_id] += 1
-            return random.randint(64, 128)
+    group_ms = []
-        elif ratio > 0.95:
+    for expert_id in range(num_experts):
-            return random.randint(128, 1024)
+        group_ms.append(expert_to_count[expert_id])
-        else:
-            return 128
-
-    group_ms = [convert_to_tokens(ratio) for ratio in ratios]
     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))
-        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_g, b_scale = per_block_cast_to_fp8(torch.randn((n_g, k_g), device=device).t())
+        b_tensors.append(b_g)
         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))
-        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_g, b_scale = per_block_cast_to_fp8(torch.randn((n_g, k_g), device=device).t())
+        b_tensors.append(b_g)
         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,9 +311,29 @@ 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 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"\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,
@@ -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
+        # Qwen3-235B-A22B-FP8, gateup, TP = 2
-        ShapeArg(n=4096, k=7168, num_groups=32),
+        ShapeArg(n=1536, k=4096, num_groups=128),
-        # Decode, DeepSeek-R1, gateup, bs = 256, EP = 16
+        # Qwen3-235B-A22B-FP8, down, TP = 2
-        ShapeArg(n=4096, k=7168, num_groups=16),
+        ShapeArg(n=4096, k=768, num_groups=128),
     ]
     args = parser.parse_args()
     benchmark_one_shape(shape_args, args.num_warmup, args.num_run)

sgl-kernel/csrc/expert_specialization/es_fp8_blockwise.cu

@@ -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

@@ -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

@@ -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,
+        backup_workspace_1,
-        stream);
+        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,
+        backup_workspace_1,
-        stream);
+        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,
+        backup_workspace_0,
-        stream);
+        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,
+        backup_workspace_0,
-        stream);
+        backup_stream_0,
-  }
+        78);
-
-  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);
   }
 }
 

sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_traits.cuh

@@ -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 KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperativeFP8Blockwise;
-  using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecializedPingpong;
+  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());