[Fix] Route the Mooncake MoE A2A backend through Kimi K3's EP-A2A / SP-MoE fast path (#36862)

python/sglang/srt/models/kimi_k3.py

@@ -531,15 +531,17 @@ class KimiK3MoE(nn.Module):
                 "got a checkpoint with different constants"
             )
 
-        # EP a2a backends (megamoe / DeepEP / MoRI) move each row to its
+        # EP a2a backends (megamoe / DeepEP / Mooncake / Ascend-FuseEP / MoRI)
-        # experts directly, so the MoE region can consume whatever rows this
+        # move each row to its experts directly, so the MoE region can consume
-        # rank holds — an SP-MoE token shard (attn_tp > 1) or the DP-local
+        # whatever rows this rank holds — an SP-MoE token shard (attn_tp > 1) or
-        # batch (DP attention) — with every global token dispatched exactly
+        # the DP-local batch (DP attention) — with every global token dispatched
-        # once. No DP gather and no TP reduce is needed anywhere in the region.
+        # 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_mooncake()
             or _a2a_backend.is_ascend_fuseep()
             or _a2a_backend.is_mori()
         )
@@ -2162,8 +2164,9 @@ 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, DeepEP or MoRI): o_proj defers its
+        # SP-MoE (EP a2a backend — megamoe, DeepEP, Mooncake, Ascend-FuseEP or
-        # attention-TP reduction; this layer completes it as a reduce-scatter
+        # 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,
         # then all-gathers rows back after the MoE tail add. RS+AG moves the
@@ -2184,6 +2187,7 @@ class KimiK3DecoderLayer(nn.Module):
             (
                 _a2a_backend.is_megamoe()
                 or _a2a_backend.is_deepep()
+                or _a2a_backend.is_mooncake()
                 or _a2a_backend.is_ascend_fuseep()
                 or _a2a_backend.is_mori()
             )