diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index 953ba7b03..ce3f5d335 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -1412,7 +1412,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): dispatch_output: StandardDispatchOutput, ) -> CombineInput: - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + from sglang.srt.layers.moe.token_dispatcher import ( + DispatchOutputChecker, + StandardCombineInput, + ) from sglang.srt.layers.moe.topk import TopKOutputChecker if self.use_deep_gemm: @@ -1433,6 +1436,12 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): ) return self.runner.run(dispatch_output, quant_info) + # Same constraint as the deep_gemm branch above: the AITER runner also + # serves the DeepEP formats (deepep_normal / deepep_ll), which carry + # topk_ids/topk_weights directly and have no `.topk_output` to unpack. + if _use_aiter and DispatchOutputChecker.format_is_deepep(dispatch_output): + return self._apply_aiter(layer, dispatch_output) + x = dispatch_output.hidden_states topk_output = dispatch_output.topk_output if _is_cpu: @@ -1710,63 +1719,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): )[0] return StandardCombineInput(hidden_states=trtllm_gen_output) if _use_aiter: - from sglang.srt.layers.moe.moe_runner.aiter import ( - AiterMoeQuantInfo, - AiterQuantType, - ) - - if hasattr(torch, "float4_e2m1fn_x2"): - w13_weight = layer.w13_weight.view(torch.float4_e2m1fn_x2) - w2_weight = layer.w2_weight.view(torch.float4_e2m1fn_x2) - else: - w13_weight = layer.w13_weight - w2_weight = layer.w2_weight - - # `.view()` creates a fresh tensor that drops the `is_shuffled` - # marker we set in process_weights_after_loading. Re-tag it so the - # downstream aiter.fused_moe selects preshuffle_on kernels. - if getattr(layer.w13_weight, "is_shuffled", False): - w13_weight.is_shuffled = True - w2_weight.is_shuffled = True - - # Skip the explicit pad if x already arrives at the padded - # hidden_size (the upstream RMSNorm fused the pad into its - # output — see RMSNorm.x_pad_to_multiple). Saves a separate - # zero-pad kernel launch per layer. - if x.shape[-1] == self.hidden_size: - x_padded = x - else: - x_padded = torch.nn.functional.pad( - x, (0, self.hidden_pad), mode="constant", value=0.0 - ) - quant_info = AiterMoeQuantInfo( - w13_weight=w13_weight, - w2_weight=w2_weight, - quant_type=AiterQuantType.PER_1X32, - w13_scale=layer.w13_weight_scale, - w2_scale=layer.w2_weight_scale, - b13=layer.w13_weight_bias if self.with_bias else None, - b2=layer.w2_weight_bias if self.with_bias else None, - expert_mask=layer.dispatcher.expert_mask_gpu, - doweight_stage1=self.moe_runner_config.apply_router_weight_on_input, - hidden_pad=self.hidden_pad, - intermediate_pad=self.intermediate_pad, - # Applies swiglu clamp for GPT-OSS-style activations. K3 SiTU - # uses gemm1_clamp_limit as linear_beta, which is forwarded by - # the AITER runner and must not be treated as swiglu_limit. - swiglu_limit=( - 0.0 - if self.moe_runner_config.activation == "situ" - else ( - self.moe_runner_config.gemm1_clamp_limit - or self.moe_runner_config.swiglu_limit - or 0.0 - ) - ), - ) - return self.runner.run( - dispatch_output._replace(hidden_states=x_padded), quant_info - ) + return self._apply_aiter(layer, dispatch_output) backend = self.runner.runner_backend if backend.is_triton_kernels(): @@ -1802,6 +1755,72 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): ) return self.runner.run(dispatch_output, quant_info) + def _apply_aiter(self, layer, dispatch_output) -> CombineInput: + """MXFP4 MoE via the AITER runner. + + Reads only ``hidden_states`` off the dispatch output, so it serves the + standard and the DeepEP formats alike; the routing tensors are resolved + by the runner's registered permute hooks. + """ + from sglang.srt.layers.moe.moe_runner.aiter import ( + AiterMoeQuantInfo, + AiterQuantType, + ) + + x = dispatch_output.hidden_states + if hasattr(torch, "float4_e2m1fn_x2"): + w13_weight = layer.w13_weight.view(torch.float4_e2m1fn_x2) + w2_weight = layer.w2_weight.view(torch.float4_e2m1fn_x2) + else: + w13_weight = layer.w13_weight + w2_weight = layer.w2_weight + + # `.view()` creates a fresh tensor that drops the `is_shuffled` + # marker we set in process_weights_after_loading. Re-tag it so the + # downstream aiter.fused_moe selects preshuffle_on kernels. + if getattr(layer.w13_weight, "is_shuffled", False): + w13_weight.is_shuffled = True + w2_weight.is_shuffled = True + + # Skip the explicit pad if x already arrives at the padded + # hidden_size (the upstream RMSNorm fused the pad into its + # output — see RMSNorm.x_pad_to_multiple). Saves a separate + # zero-pad kernel launch per layer. + if x.shape[-1] == self.hidden_size: + x_padded = x + else: + x_padded = torch.nn.functional.pad( + x, (0, self.hidden_pad), mode="constant", value=0.0 + ) + quant_info = AiterMoeQuantInfo( + w13_weight=w13_weight, + w2_weight=w2_weight, + quant_type=AiterQuantType.PER_1X32, + w13_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + b13=layer.w13_weight_bias if self.with_bias else None, + b2=layer.w2_weight_bias if self.with_bias else None, + expert_mask=layer.dispatcher.expert_mask_gpu, + doweight_stage1=self.moe_runner_config.apply_router_weight_on_input, + hidden_pad=self.hidden_pad, + intermediate_pad=self.intermediate_pad, + # Applies swiglu clamp for GPT-OSS-style activations. K3 SiTU + # uses gemm1_clamp_limit as linear_beta, which is forwarded by + # the AITER runner and must not be treated as swiglu_limit. + swiglu_limit=( + 0.0 + if self.moe_runner_config.activation == "situ" + else ( + self.moe_runner_config.gemm1_clamp_limit + or self.moe_runner_config.swiglu_limit + or 0.0 + ) + ), + ) + return self.runner.run( + dispatch_output._replace(hidden_states=x_padded), quant_info + ) + class Mxfp4DynamicQuantMoEMethod(FusedMoEMethodBase): def create_weights( diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index b1d460c2c..45e83672d 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/python/sglang/srt/models/kimi_k3.py @@ -523,16 +523,17 @@ class KimiK3MoE(nn.Module): "got a checkpoint with different constants" ) - # EP a2a backends (megamoe / DeepEP) move each row to its experts - # directly, so the MoE region can consume whatever rows this rank - # holds — an SP-MoE token shard (attn_tp > 1) or the DP-local batch - # (DP attention) — with every global token dispatched exactly once. - # No DP gather and no TP reduce is needed anywhere in the region. + # EP a2a backends (megamoe / DeepEP / MoRI) move each row to its + # experts directly, so the MoE region can consume whatever rows this + # rank holds — an SP-MoE token shard (attn_tp > 1) or the DP-local + # batch (DP attention) — with every global token dispatched exactly + # once. No DP gather and no TP reduce is needed anywhere in the region. _a2a_backend = get_moe_a2a_backend() self._ep_a2a = ( _a2a_backend.is_megamoe() or _a2a_backend.is_deepep() or _a2a_backend.is_ascend_fuseep() + or _a2a_backend.is_mori() ) # Defer the trtllm-gen finalize (top-k weighted unpermute) out of the @@ -2081,7 +2082,7 @@ class KimiK3DecoderLayer(nn.Module): and layer_idx >= config.first_k_dense_replace and layer_idx % config.moe_layer_freq == 0 ) - # SP-MoE (EP a2a backend — megamoe or DeepEP): o_proj defers its + # SP-MoE (EP a2a backend — megamoe, DeepEP or MoRI): o_proj defers its # attention-TP reduction; this layer completes it as a reduce-scatter # so the whole MoE region (agg2, norms, gate, latent projs, tp1 # shared experts, EP a2a dispatch) runs on 1/attn_tp of the rows, @@ -2104,6 +2105,7 @@ class KimiK3DecoderLayer(nn.Module): _a2a_backend.is_megamoe() or _a2a_backend.is_deepep() or _a2a_backend.is_ascend_fuseep() + or _a2a_backend.is_mori() ) and self._is_moe_layer and get_parallel().attn_tp_group.world_size > 1