From ce112c07fe117ae8bd06c1903035a72b82cb28b6 Mon Sep 17 00:00:00 2001
From: Qi Yuhang <45795032+HydraQYH@users.noreply.github.com>
Date: Mon, 27 Oct 2025 16:20:01 +0800
Subject: [PATCH] [sgl-kernel][4/N]Support Expert Specialization Grouped GEMM
 (#12080)

---
 .../csrc/expert_specialization/es_fp8_blockwise_functor.cuh | 6 +++---
 .../expert_specialization/es_fp8_blockwise_launcher.cuh     | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_functor.cuh b/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_functor.cuh
index c4af33447..db7f430f2 100644
--- a/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_functor.cuh
+++ b/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_functor.cuh
@@ -126,7 +126,7 @@ 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 <= 48) {
+    if (m < 64) {
       // Swap A/B
       problem_sizes[expert_id * 3 + 0] = n;
       problem_sizes[expert_id * 3 + 1] = m;
@@ -168,7 +168,7 @@ 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 > 48 && m <= 96) {
+    if (m >= 64 && m < 128) {
       problem_sizes[expert_id * 3 + 0] = m;
       problem_sizes[expert_id * 3 + 1] = n;
       problem_sizes[expert_id * 3 + 2] = k;
@@ -208,7 +208,7 @@ 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 > 96) {
+    if (m >= 128) {
       problem_sizes[expert_id * 3 + 0] = m;
       problem_sizes[expert_id * 3 + 1] = n;
       problem_sizes[expert_id * 3 + 2] = k;
diff --git a/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_launcher.cuh b/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_launcher.cuh
index f6ab33eb0..3ed98821d 100644
--- a/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_launcher.cuh
+++ b/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_launcher.cuh
@@ -232,7 +232,7 @@ void es_sm90_fp8_blockwise_scaled_group_mm_distpatch_out_dtype(
         workspace,
         stream);
   } else {
-    launch_sm90_fp8_blockwise_scaled_group_mm<HighMGemmHx00Traits>(
+    launch_sm90_fp8_blockwise_scaled_group_mm<MiddleMGemmH20Traits>(
         out_ptrs,
         a_ptrs,
         b_ptrs,