From 745abd6cc00d51bf4536c30fa7dca80a48df0ee3 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Mon, 18 May 2026 15:59:55 -0700 Subject: [PATCH] Add no_combine support to cutlass_moe_fp4 (#25688) Co-authored-by: Hanming Lu <69857889+hanming-lu@users.noreply.github.com> --- python/sglang/srt/layers/moe/cutlass_moe.py | 3 +++ python/sglang/srt/layers/quantization/modelopt_quant.py | 1 + 2 files changed, 4 insertions(+) diff --git a/python/sglang/srt/layers/moe/cutlass_moe.py b/python/sglang/srt/layers/moe/cutlass_moe.py index 3716fc6a6..16dbc74c9 100755 --- a/python/sglang/srt/layers/moe/cutlass_moe.py +++ b/python/sglang/srt/layers/moe/cutlass_moe.py @@ -359,6 +359,7 @@ def cutlass_moe_fp4( topk_ids: torch.Tensor, params: CutlassMoEParams, apply_router_weight_on_input: bool = False, + no_combine: bool = False, ): """ MoE implementation for FP4 Inputs @@ -492,6 +493,8 @@ def cutlass_moe_fp4( 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: + 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) diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index bab2e6f59..c0d9b300f 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -2175,6 +2175,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): topk_ids=topk_ids, params=layer.cutlass_moe_params, apply_router_weight_on_input=moe_runner_config.apply_router_weight_on_input, + no_combine=moe_runner_config.no_combine, ).to(x.dtype) # Scale by routed_scaling_factor is fused into select_experts. return StandardCombineInput(hidden_states=output)