perf: speed up marlin moe with occupancy-aware launch specialization (#31552)

python/sglang/kernels/jit/csrc/gemm/marlin_moe/kernel.h

@@ -30,7 +30,9 @@ template <
                                          // fetch pipeline
     const int group_blocks,              // number of consecutive 16x16 blocks
                                          // with a separate quantization scale
-    const bool is_zp_float               // is zero point of float16 type?
+    const bool is_zp_float,              // is zero point of float16 type?
+    const bool kIsEP,                    // expert parallelism
+    const bool kHasBias                  // has per-expert bias
     >
 __global__ void Marlin(MARLIN_KERNEL_PARAMS);
 

python/sglang/kernels/jit/csrc/gemm/marlin_moe/marlin_template.h

@@ -51,7 +51,9 @@ template <
                                          // fetch pipeline
     const int group_blocks,              // number of consecutive 16x16 blocks
                                          // with a separate quantization scale
-    const bool is_zp_float               // is zero point of float16 type?
+    const bool is_zp_float,              // is zero point of float16 type?
+    const bool kIsEP,                    // expert parallelism
+    const bool kHasBias                  // has per-expert bias
     >
 __global__ void Marlin(
     const int4* __restrict__ A,                              // fp16 input matrix of shape mxk
@@ -292,7 +294,9 @@ template <
                                          // fetch pipeline
     const int group_blocks,              // number of consecutive 16x16 blocks
                                          // with a separate quantization scale
-    const bool is_zp_float               // is zero point of float16 type?
+    const bool is_zp_float,              // is zero point of float16 type?
+    const bool kIsEP,                    // expert parallelism
+    const bool kHasBias                  // has per-expert bias
     >
 __global__ void Marlin(
     const int4* __restrict__ A,  // fp16 input matrix of shape mxk
@@ -378,9 +382,11 @@ __global__ void Marlin(
   int num_tokens_past_padded = num_tokens_past_padded_ptr[0];
   int parallel = num_tokens_past_padded / moe_block_size;
   int num_valid_blocks = parallel;
+  if constexpr (kIsEP) {
     for (int i = 0; i < parallel; i++) {
       if (expert_ids_ptr[i] == -1) num_valid_blocks--;
     }
+  }
   int num_invalid_blocks = parallel - num_valid_blocks;
   parallel = num_valid_blocks;
 
@@ -510,6 +516,7 @@ __global__ void Marlin(
     if (par_id >= parallel) return;
 
     old_expert_id = expert_id;
+    if constexpr (kIsEP) {
       if (num_invalid_blocks > 0) {
         int skip_count = block_id == -1 ? par_id : 0;
         block_id++;
@@ -527,6 +534,10 @@ __global__ void Marlin(
         block_id = par_id;
         expert_id = expert_ids_ptr[block_id];
       }
+    } else {
+      block_id = par_id;
+      expert_id = expert_ids_ptr[block_id];
+    }
 
     if constexpr (w_type == host::kFE2M1f && s_type == host::kFE4M3fn) {
       uint16_t val = scale2_ptr[expert_id];
@@ -541,7 +552,7 @@ __global__ void Marlin(
     if constexpr (has_act_order) {
       g_idx += (expert_id - old_expert_id) * prob_k;
     }
-    if (has_bias) {
+    if constexpr (kHasBias) {
       b_bias_ptr += (expert_id - old_expert_id) * b_bias_expert_stride;
     }
 
@@ -1536,13 +1547,15 @@ __global__ void Marlin(
           res = __hmul2(res, global_scale);
         }
       }
-      if (has_bias && last) {
+      if constexpr (kHasBias) {
+        if (last) {
           scalar_t2 tmp_bias = b_bias[0];
           if constexpr (m_block_size_8) {
             tmp_bias = Dtype::num2num2(reinterpret_cast<scalar_t*>(&b_bias[0])[(threadIdx.x % 8) / 4]);
           }
           res = __hadd2(res, tmp_bias);
         }
+      }
 
       if constexpr (m_block_size_8) {
         ((scalar_t*)sh_red)[idx] = res.x;
@@ -1754,11 +1767,13 @@ __global__ void Marlin(
 
       thread_block_reduce();
 
-      if (has_bias && last) {
+      if constexpr (kHasBias) {
+        if (last) {
           __syncthreads();
           cp_async4_pred(&sh_bias[bias_sh_wr], &b_bias_ptr[bias_gl_rd], threadIdx.x < 16 * thread_n_blocks / 8);
           cp_async_fence();
         }
+      }
 
       if constexpr (!has_act_order && group_blocks == -1 && (has_zp && dequant_skip_flop || !has_zp)) {
         if (w_type.size_bits() == 8 || (last || use_atomic_add)) {
@@ -1813,13 +1828,15 @@ __global__ void Marlin(
         barrier_release(&locks[locks_off], last);
       }
 
-      if (has_bias && last) {
+      if constexpr (kHasBias) {
+        if (last) {
           cp_async_wait<0>();
           __syncthreads();
           reinterpret_cast<int4*>(&frag_bias)[0] = sh_bias[bias_sh_rd];
           reinterpret_cast<int4*>(&frag_bias)[1] = sh_bias[bias_sh_rd + 4];
           __syncthreads();
         }
+      }
 
       if (use_atomic_add && slice_count > 1 && slice_idx != 0) wait_negative_and_add(&locks[locks_off]);
       if (last || use_atomic_add)

python/sglang/kernels/jit/csrc/gemm/marlin_moe/moe_wna16_marlin.cuh

@@ -154,6 +154,9 @@ typedef struct {
   thread_config_t tb_cfg;
 } exec_config_t;
 
+constexpr int kSharedMemoryValidityMargin = 512;
+constexpr int kSharedMemoryLaunchReserve = 1024;
+
 int get_scales_cache_size(
     thread_config_t const& th_config,
     int prob_m,
@@ -285,7 +288,7 @@ bool is_valid_config(
       is_k_full,
       has_zp,
       is_zp_float);
-  return cache_size + 512 <= max_shared_mem;
+  return cache_size + kSharedMemoryValidityMargin <= max_shared_mem;
 }
 
 #define _GET_IF(                                                                                                       \
@@ -308,7 +311,9 @@ bool is_valid_config(
         M_BLOCK_SIZE_8,                                                                                                \
         pipe_stages,                                                                                                   \
         GROUP_BLOCKS,                                                                                                  \
-        IS_ZP_FLOAT>;                                                                                                  \
+        IS_ZP_FLOAT,                                                                                                   \
+        kIsEP,                                                                                                         \
+        kHasBias>;                                                                                                     \
   }
 
 // COMMON: cases for (group_blocks in [-1, 2, 4, 8] and is_zp_float == false)
@@ -432,7 +437,7 @@ bool is_valid_config(
   ACT_GET_IF_M234(W_TYPE, 16, 4, 256) \
   ACT_GET_IF_M234(W_TYPE, 8, 4, 128)
 
-template <typename scalar_t>
+template <typename scalar_t, bool kIsEP, bool kHasBias>
 MarlinFuncPtr get_marlin_kernel(
     const host::ScalarType q_type,
     int thread_m_blocks,
@@ -468,12 +473,13 @@ MarlinFuncPtr get_marlin_kernel(
   return kernel;
 }
 
-template <typename scalar_t>
+template <typename scalar_t, bool kIsEP, bool kHasBias>
 exec_config_t determine_exec_config(
     const host::ScalarType& q_type,
     int prob_m,
     int prob_n,
     int prob_k,
+    int top_k,
     int thread_m_blocks,
     bool m_block_size_8,
     int num_bits,
@@ -482,7 +488,8 @@ exec_config_t determine_exec_config(
     bool is_k_full,
     bool has_zp,
     bool is_zp_float,
-    int max_shared_mem) {
+    int max_shared_mem,
+    int sms) {
   exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}};
   thread_config_t* thread_configs = thread_m_blocks > 1 ? large_batch_thread_configs : small_batch_thread_configs;
   int thread_configs_size = thread_m_blocks > 1 ? sizeof(large_batch_thread_configs) / sizeof(thread_config_t)
@@ -529,7 +536,7 @@ exec_config_t determine_exec_config(
       group_blocks = group_size == -1 ? -1 : (group_size / 16);
     }
 
-    auto kernel = get_marlin_kernel<scalar_t>(
+    auto kernel = get_marlin_kernel<scalar_t, kIsEP, kHasBias>(
         q_type,
         thread_m_blocks,
         th_config.thread_n / 16,
@@ -543,26 +550,31 @@ exec_config_t determine_exec_config(
 
     if (kernel == MarlinDefault) continue;
 
-    if (thread_m_blocks > 1) {
-      exec_cfg = {1, th_config};
-      break;
-    } else {
     cudaFuncAttributes attr;
     cudaFuncGetAttributes(&attr, kernel);
     int reg_size = max(attr.numRegs, 1) * th_config.num_threads * 4;
-      int allow_count = min(device_max_reg_size / reg_size, max_shared_mem / (cache_size + 1024));
+    int allow_count =
-      allow_count = max(min(allow_count, 4), 1);
+        min(device_max_reg_size / reg_size,
+            max_shared_mem / (cache_size + kSharedMemoryValidityMargin + kSharedMemoryLaunchReserve));
+    allow_count = max(min(allow_count, thread_m_blocks == 1 ? 4 : 2), 1);
+
+    if (thread_m_blocks > 1) {
+      int problem_blocks = prob_n / th_config.thread_n * prob_m * top_k * 4;
+      if (problem_blocks < sms * allow_count) {
+        allow_count = max(problem_blocks / sms, 1);
+      }
+    }
+
     if (allow_count > count) {
       count = allow_count;
       exec_cfg = {count, th_config};
-      };
     }
   }
 
   return exec_cfg;
 }
 
-template <typename scalar_t>
+template <typename scalar_t, bool kIsEP, bool kHasBias>
 void marlin_mm(
     const void* A,
     const void* B,
@@ -702,11 +714,12 @@ void marlin_mm(
     host::RuntimeCheck(prob_k % thread_k == 0, "prob_k = ", prob_k, " is not divisible by thread_k = ", thread_k);
   } else {
     // Auto config
-    exec_cfg = determine_exec_config<scalar_t>(
+    exec_cfg = determine_exec_config<scalar_t, kIsEP, kHasBias>(
         q_type,
         prob_m,
         prob_n,
         prob_k,
+        top_k,
         thread_m_blocks,
         m_block_size_8,
         num_bits,
@@ -715,7 +728,8 @@ void marlin_mm(
         is_k_full,
         has_zp,
         is_zp_float,
-        max_shared_mem);
+        max_shared_mem,
+        sms);
     thread_tfg = exec_cfg.tb_cfg;
   }
 
@@ -723,7 +737,7 @@ void marlin_mm(
   thread_k = thread_tfg.thread_k;
   thread_n = thread_tfg.thread_n;
   int blocks = sms * exec_cfg.blocks_per_sm;
-  if (exec_cfg.blocks_per_sm > 1) max_shared_mem = max_shared_mem / exec_cfg.blocks_per_sm - 1024;
+  if (exec_cfg.blocks_per_sm > 1) max_shared_mem = max_shared_mem / exec_cfg.blocks_per_sm - kSharedMemoryLaunchReserve;
 
   int thread_k_blocks = thread_k / 16;
   int thread_n_blocks = thread_n / 16;
@@ -772,7 +786,7 @@ void marlin_mm(
       ", max_shared_mem = ",
       max_shared_mem);
 
-  auto kernel = get_marlin_kernel<scalar_t>(
+  auto kernel = get_marlin_kernel<scalar_t, kIsEP, kHasBias>(
       q_type,
       thread_m_blocks,
       thread_n_blocks,
@@ -823,7 +837,7 @@ void marlin_mm(
 
 }  // namespace device::marlin_moe
 
-template <typename scalar_t>
+template <typename scalar_t, bool kIsEP, bool kHasBias>
 void moe_wna16_marlin_gemm(
     tvm::ffi::TensorView a,
     tvm::ffi::TensorView c,
@@ -860,6 +874,9 @@ void moe_wna16_marlin_gemm(
     bool is_zp_float) {
   using namespace host;
 
+  RuntimeCheck(is_ep == kIsEP, "is_ep does not match the compiled Marlin MoE specialization");
+  RuntimeCheck(has_bias == kHasBias, "has_bias does not match the compiled Marlin MoE specialization");
+
   ScalarType const b_q_type = ScalarType::from_id(b_q_type_id);
   int pack_factor = 32 / b_q_type.size_bits();
 
@@ -1057,7 +1074,7 @@ void moe_wna16_marlin_gemm(
   // Early return for zero-size M (moved after all validation)
   if (size_m == 0) return;
 
-  device::marlin_moe::marlin_mm<scalar_t>(
+  device::marlin_moe::marlin_mm<scalar_t, kIsEP, kHasBias>(
       a.data_ptr(),
       b_q_weight.data_ptr(),
       c.data_ptr(),

python/sglang/kernels/ops/moe/moe_wna16_marlin.py

@@ -16,8 +16,10 @@ _MAX_THREAD_N = 256
 
 
 @cache_once
-def _jit_moe_wna16_marlin_module(dtype: torch.dtype) -> Module:
+def _jit_moe_wna16_marlin_module(
-    args = make_cpp_args(dtype)
+    dtype: torch.dtype, is_ep: bool, has_bias: bool
+) -> Module:
+    args = make_cpp_args(dtype, is_ep, has_bias)
     return load_jit(
         "moe_wna16_marlin",
         *args,
@@ -134,7 +136,7 @@ def moe_wna16_marlin_gemm(
     b_bias_t = _or_empty(b_bias_or_none, device, a.dtype)
     global_scale_t = _or_empty(global_scale_or_none, device, a.dtype)
 
-    module = _jit_moe_wna16_marlin_module(a.dtype)
+    module = _jit_moe_wna16_marlin_module(a.dtype, is_ep, has_bias)
     module.moe_wna16_marlin_gemm(
         a,
         c,

test/registered/kernels/ops/moe/test_moe_wna16_marlin.py

@@ -411,6 +411,83 @@ def test_fused_marlin_moe_non_gated_relu2():
     torch.testing.assert_close(output, output_ref, rtol=0.04, atol=0.04)
 
 
+@pytest.mark.parametrize("m", [123, 2304])
+@pytest.mark.parametrize("has_bias", [False, True])
+def test_fused_marlin_moe_large_non_ep_schedule(m, has_bias):
+    torch.manual_seed(0)
+
+    n = 1024
+    k = 512
+    e = 8
+    topk = 2
+    dtype = torch.bfloat16
+    group_size = 128
+    quant_type = scalar_types.uint4b8
+
+    hidden_states = torch.randn((m, k), device="cuda", dtype=dtype) / 10
+    w_ref1, qweight1, scales1, zeros1, g_idx1, sort_indices1 = _setup_moe_weights(
+        e, n, k, quant_type, group_size, False, dtype
+    )
+    w_ref2, qweight2, scales2, zeros2, g_idx2, sort_indices2 = _setup_moe_weights(
+        e, k, n, quant_type, group_size, False, dtype
+    )
+    w1_bias = (
+        torch.randn((e, n), device="cuda", dtype=dtype) / 100 if has_bias else None
+    )
+    w2_bias = (
+        torch.randn((e, k), device="cuda", dtype=dtype) / 100 if has_bias else None
+    )
+
+    router_logits = torch.randn((m, e), device="cuda", dtype=dtype)
+    score_softmax = torch.softmax(router_logits, dim=-1, dtype=torch.float32)
+    topk_weights, topk_ids = torch.topk(score_softmax, topk)
+
+    output = fused_marlin_moe(
+        hidden_states=hidden_states,
+        w1=qweight1,
+        w2=qweight2,
+        w1_scale=scales1,
+        w2_scale=scales2,
+        gating_output=router_logits,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        g_idx1=g_idx1,
+        g_idx2=g_idx2,
+        sort_indices1=sort_indices1,
+        sort_indices2=sort_indices2,
+        w1_zeros=zeros1,
+        w2_zeros=zeros2,
+        w1_bias=w1_bias,
+        w2_bias=w2_bias,
+        num_bits=4,
+        is_k_full=True,
+        routed_scaling_factor=1.0,
+        activation="relu2",
+        is_gated=False,
+    )
+
+    output_ref = torch.zeros_like(hidden_states, dtype=torch.float32)
+    for expert_id in range(e):
+        token_indices, route_indices = torch.where(topk_ids == expert_id)
+        intermediate = hidden_states[token_indices] @ w_ref1[expert_id].T
+        if w1_bias is not None:
+            intermediate += w1_bias[expert_id]
+        intermediate = torch.square(torch.relu(intermediate))
+        routed = intermediate @ w_ref2[expert_id].T
+        if w2_bias is not None:
+            routed += w2_bias[expert_id]
+        output_ref.index_add_(
+            0,
+            token_indices,
+            routed.float() * topk_weights[token_indices, route_indices, None],
+        )
+
+    torch.cuda.synchronize()
+    # The existing BF16/4-bit bias path has a few large-batch outliers just
+    # above 0.04 even without the compile-time specialization.
+    torch.testing.assert_close(output, output_ref.to(dtype), rtol=0.04, atol=0.06)
+
+
 @pytest.mark.skipif(
     not (is_sm80_supported() or is_sm90_supported()),
     reason="NVFP4 Marlin MoE padding test requires CUDA SM8X/SM9X",