minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Mick
2026-07-25 19:38:11 +08:00
committed by GitHub
parent d021990bf5
commit 1054060ef1
5 changed files with 175 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/gemm/marlin_moe/kernel.h
+3 -1
View File
@@ -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
+48 -31
View File
@@ -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<scalar_t*>(&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<scalar_t*>(&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<int4*>(&frag_bias)[0] = sh_bias[bias_sh_rd];
reinterpret_cast<int4*>(&frag_bias)[1] = sh_bias[bias_sh_rd + 4];
__syncthreads();
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]);
python/sglang/kernels/jit/csrc/gemm/marlin_moe/moe_wna16_marlin.cuh
+42 -25
View File
@@ -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;
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 <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
+5 -3
View File
@@ -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,
test/registered/kernels/ops/moe/test_moe_wna16_marlin.py
+77
View File
@@ -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",