diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index 6064795f3..c164b4c0a 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/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 - # 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 / Mooncake / Ascend-FuseEP / 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_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 - # attention-TP reduction; this layer completes it as a reduce-scatter + # SP-MoE (EP a2a backend — megamoe, DeepEP, Mooncake, Ascend-FuseEP 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, # 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() )