[AMD] Enable Mori-EP on kimi-k3 (#35630)

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,11 +1719,55 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             )[0]
             return StandardCombineInput(hidden_states=trtllm_gen_output)
         if _use_aiter:
+            return self._apply_aiter(layer, dispatch_output)
+
+        backend = self.runner.runner_backend
+        if backend.is_triton_kernels():
+            from sglang.srt.layers.moe.moe_runner.triton_kernels import (
+                TritonKernelsQuantInfo,
+            )
+
+            assert (
+                layer.moe_ep_size == 1
+            ), "Expert parallel is not supported when using triton kernels"
+            quant_info = TritonKernelsQuantInfo(
+                w13_weight=(
+                    self.w13_weight_triton_tensor
+                    if self.w13_weight_triton_tensor is not None
+                    else layer.w13_weight
+                ),
+                w2_weight=(
+                    self.w2_weight_triton_tensor
+                    if self.w2_weight_triton_tensor is not None
+                    else layer.w2_weight
+                ),
+                w13_bias=getattr(layer, "w13_weight_bias", None),
+                w2_bias=getattr(layer, "w2_weight_bias", None),
+                w13_precision_config=getattr(self, "w13_precision_config", None),
+                w2_precision_config=getattr(self, "w2_precision_config", None),
+            )
+        else:
+            quant_info = TritonMoeQuantInfo(
+                w13_weight=layer.w13_weight,
+                w2_weight=layer.w2_weight,
+                b13=getattr(layer, "w13_weight_bias", None),
+                b2=getattr(layer, "w2_weight_bias", None),
+            )
+        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)
@@ -1768,40 +1821,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             dispatch_output._replace(hidden_states=x_padded), quant_info
         )
 
-        backend = self.runner.runner_backend
-        if backend.is_triton_kernels():
-            from sglang.srt.layers.moe.moe_runner.triton_kernels import (
-                TritonKernelsQuantInfo,
-            )
-
-            assert (
-                layer.moe_ep_size == 1
-            ), "Expert parallel is not supported when using triton kernels"
-            quant_info = TritonKernelsQuantInfo(
-                w13_weight=(
-                    self.w13_weight_triton_tensor
-                    if self.w13_weight_triton_tensor is not None
-                    else layer.w13_weight
-                ),
-                w2_weight=(
-                    self.w2_weight_triton_tensor
-                    if self.w2_weight_triton_tensor is not None
-                    else layer.w2_weight
-                ),
-                w13_bias=getattr(layer, "w13_weight_bias", None),
-                w2_bias=getattr(layer, "w2_weight_bias", None),
-                w13_precision_config=getattr(self, "w13_precision_config", None),
-                w2_precision_config=getattr(self, "w2_precision_config", None),
-            )
-        else:
-            quant_info = TritonMoeQuantInfo(
-                w13_weight=layer.w13_weight,
-                w2_weight=layer.w2_weight,
-                b13=getattr(layer, "w13_weight_bias", None),
-                b2=getattr(layer, "w2_weight_bias", None),
-            )
-        return self.runner.run(dispatch_output, quant_info)
-
 
 class Mxfp4DynamicQuantMoEMethod(FusedMoEMethodBase):
     def create_weights(

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
+        # EP a2a backends (megamoe / DeepEP / MoRI) move each row to its
-        # directly, so the MoE region can consume whatever rows this rank
+        # experts directly, so the MoE region can consume whatever rows this
-        # holds — an SP-MoE token shard (attn_tp > 1) or the DP-local batch
+        # rank holds — an SP-MoE token shard (attn_tp > 1) or the DP-local
-        # (DP attention) — with every global token dispatched exactly once.
+        # batch (DP attention) — with every global token dispatched exactly
-        # No DP gather and no TP reduce is needed anywhere in the region.
+        # 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