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

[Perf][Moe]improve cutlass_moe_fp4 performance by using apply_router_weight_on_i… (#19493)

Browse Source
This commit is contained in:
chengchao23
2026-05-26 00:07:49 -07:00
committed by GitHub
parent 3f5e2c7688
commit 137168539a
1 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/cutlass_moe.py
+6 -4
View File
@@ -491,10 +491,12 @@ def cutlass_moe_fp4(
params.to_gemm2_args(), params.to_gemm2_args(),
) )
del int_fp4, int_blockscale del int_fp4, int_blockscale
if no_combine:
c2 = shuffle_rows(c2, c_map, (m_a * num_topk, params.hidden_size)) c2 = shuffle_rows(c2, c_map, (m_a * num_topk, params.hidden_size))
c2 = c2.view(m_a, num_topk, params.hidden_size) c2 = c2.view(m_a, num_topk, params.hidden_size)
if no_combine:
return c2.to(out_dtype) return c2.to(out_dtype)
if not apply_router_weight_on_input: output = torch.empty((m_a, k_a), device=device, dtype=out_dtype)
c2 = c2 * topk_weights.view(m_a, num_topk, 1).to(out_dtype) weights = topk_weights.to(out_dtype) if not apply_router_weight_on_input else None
return c2.sum(dim=1).to(out_dtype) apply_shuffle_mul_sum(c2, output, c_map, weights)
return output