From 137168539a732ea5688d09998eb0c49ef5a3fad7 Mon Sep 17 00:00:00 2001 From: chengchao23 <55905375+chengchao23@users.noreply.github.com> Date: Tue, 26 May 2026 15:07:49 +0800 Subject: [PATCH] =?UTF-8?q?[Perf][Moe]improve=20cutlass=5Fmoe=5Ffp4=20perf?= =?UTF-8?q?ormance=20by=20using=20apply=5Frouter=5Fweight=5Fon=5Fi?= =?UTF-8?q?=E2=80=A6=20(#19493)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- python/sglang/srt/layers/moe/cutlass_moe.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/layers/moe/cutlass_moe.py b/python/sglang/srt/layers/moe/cutlass_moe.py index fd02d6718..05cfe00fc 100755 --- a/python/sglang/srt/layers/moe/cutlass_moe.py +++ b/python/sglang/srt/layers/moe/cutlass_moe.py @@ -491,10 +491,12 @@ def cutlass_moe_fp4( params.to_gemm2_args(), ) del int_fp4, int_blockscale - c2 = shuffle_rows(c2, c_map, (m_a * num_topk, params.hidden_size)) - c2 = c2.view(m_a, num_topk, params.hidden_size) + if no_combine: + c2 = shuffle_rows(c2, c_map, (m_a * num_topk, params.hidden_size)) + c2 = c2.view(m_a, num_topk, params.hidden_size) return c2.to(out_dtype) - if not apply_router_weight_on_input: - c2 = c2 * topk_weights.view(m_a, num_topk, 1).to(out_dtype) - return c2.sum(dim=1).to(out_dtype) + output = torch.empty((m_a, k_a), device=device, dtype=out_dtype) + weights = topk_weights.to(out_dtype) if not apply_router_weight_on_input else None + apply_shuffle_mul_sum(c2, output, c_map, weights) + return output