From 1054060ef1192bceab0e844b8a51fd083cab1b39 Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 25 Jul 2026 19:38:11 +0800 Subject: [PATCH] perf: speed up marlin moe with occupancy-aware launch specialization (#31552) --- .../kernels/jit/csrc/gemm/marlin_moe/kernel.h | 4 +- .../csrc/gemm/marlin_moe/marlin_template.h | 79 +++++++++++-------- .../csrc/gemm/marlin_moe/moe_wna16_marlin.cuh | 67 ++++++++++------ .../kernels/ops/moe/moe_wna16_marlin.py | 8 +- .../kernels/ops/moe/test_moe_wna16_marlin.py | 77 ++++++++++++++++++ 5 files changed, 175 insertions(+), 60 deletions(-) diff --git a/python/sglang/kernels/jit/csrc/gemm/marlin_moe/kernel.h b/python/sglang/kernels/jit/csrc/gemm/marlin_moe/kernel.h index 522a77d40..caed4de3b 100644 --- a/python/sglang/kernels/jit/csrc/gemm/marlin_moe/kernel.h +++ b/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); diff --git a/python/sglang/kernels/jit/csrc/gemm/marlin_moe/marlin_template.h b/python/sglang/kernels/jit/csrc/gemm/marlin_moe/marlin_template.h index cc2d15986..f5bda01dd 100644 --- a/python/sglang/kernels/jit/csrc/gemm/marlin_moe/marlin_template.h +++ b/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,8 +382,10 @@ __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; - for (int i = 0; i < parallel; i++) { - if (expert_ids_ptr[i] == -1) num_valid_blocks--; + 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,18 +516,23 @@ __global__ void Marlin( if (par_id >= parallel) return; old_expert_id = expert_id; - if (num_invalid_blocks > 0) { - int skip_count = block_id == -1 ? par_id : 0; - block_id++; - for (int i = block_id; i < num_tokens_past_padded / moe_block_size; i++) { - expert_id = expert_ids_ptr[i]; - if (expert_id != -1) { - if (skip_count == 0) { - block_id = i; - break; + if constexpr (kIsEP) { + if (num_invalid_blocks > 0) { + int skip_count = block_id == -1 ? par_id : 0; + block_id++; + for (int i = block_id; i < num_tokens_past_padded / moe_block_size; i++) { + expert_id = expert_ids_ptr[i]; + if (expert_id != -1) { + if (skip_count == 0) { + block_id = i; + break; + }; + skip_count--; }; - skip_count--; - }; + } + } else { + block_id = par_id; + expert_id = expert_ids_ptr[block_id]; } } else { block_id = par_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,12 +1547,14 @@ __global__ void Marlin( res = __hmul2(res, global_scale); } } - if (has_bias && last) { - scalar_t2 tmp_bias = b_bias[0]; - if constexpr (m_block_size_8) { - tmp_bias = Dtype::num2num2(reinterpret_cast(&b_bias[0])[(threadIdx.x % 8) / 4]); + if constexpr (kHasBias) { + if (last) { + scalar_t2 tmp_bias = b_bias[0]; + if constexpr (m_block_size_8) { + tmp_bias = Dtype::num2num2(reinterpret_cast(&b_bias[0])[(threadIdx.x % 8) / 4]); + } + res = __hadd2(res, tmp_bias); } - res = __hadd2(res, tmp_bias); } if constexpr (m_block_size_8) { @@ -1754,10 +1767,12 @@ __global__ void Marlin( thread_block_reduce(); - if (has_bias && 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 (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)) { @@ -1813,12 +1828,14 @@ __global__ void Marlin( barrier_release(&locks[locks_off], last); } - if (has_bias && last) { - cp_async_wait<0>(); - __syncthreads(); - reinterpret_cast(&frag_bias)[0] = sh_bias[bias_sh_rd]; - reinterpret_cast(&frag_bias)[1] = sh_bias[bias_sh_rd + 4]; - __syncthreads(); + if constexpr (kHasBias) { + if (last) { + cp_async_wait<0>(); + __syncthreads(); + reinterpret_cast(&frag_bias)[0] = sh_bias[bias_sh_rd]; + reinterpret_cast(&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]); diff --git a/python/sglang/kernels/jit/csrc/gemm/marlin_moe/moe_wna16_marlin.cuh b/python/sglang/kernels/jit/csrc/gemm/marlin_moe/moe_wna16_marlin.cuh index a3b563151..b824c9858 100644 --- a/python/sglang/kernels/jit/csrc/gemm/marlin_moe/moe_wna16_marlin.cuh +++ b/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 +template MarlinFuncPtr get_marlin_kernel( const host::ScalarType q_type, int thread_m_blocks, @@ -468,12 +473,13 @@ MarlinFuncPtr get_marlin_kernel( return kernel; } -template +template 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( + auto kernel = get_marlin_kernel( q_type, thread_m_blocks, th_config.thread_n / 16, @@ -543,26 +550,31 @@ exec_config_t determine_exec_config( if (kernel == MarlinDefault) continue; + 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 + kSharedMemoryValidityMargin + kSharedMemoryLaunchReserve)); + allow_count = max(min(allow_count, thread_m_blocks == 1 ? 4 : 2), 1); + 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)); - allow_count = max(min(allow_count, 4), 1); - if (allow_count > count) { - count = allow_count; - exec_cfg = {count, th_config}; - }; + 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 +template 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( + exec_cfg = determine_exec_config( 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( + auto kernel = get_marlin_kernel( q_type, thread_m_blocks, thread_n_blocks, @@ -823,7 +837,7 @@ void marlin_mm( } // namespace device::marlin_moe -template +template 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( + device::marlin_moe::marlin_mm( a.data_ptr(), b_q_weight.data_ptr(), c.data_ptr(), diff --git a/python/sglang/kernels/ops/moe/moe_wna16_marlin.py b/python/sglang/kernels/ops/moe/moe_wna16_marlin.py index bbac46d10..3e3e25d0d 100644 --- a/python/sglang/kernels/ops/moe/moe_wna16_marlin.py +++ b/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: - args = make_cpp_args(dtype) +def _jit_moe_wna16_marlin_module( + 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, diff --git a/test/registered/kernels/ops/moe/test_moe_wna16_marlin.py b/test/registered/kernels/ops/moe/test_moe_wna16_marlin.py index 4667a1cc8..005e9bd37 100644 --- a/test/registered/kernels/ops/moe/test_moe_wna16_marlin.py +++ b/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",