执行摘要

• 一句话：Kimi-K3 支持 MoRI EP 后端，AITER 路径提前返回
• 推荐动作：值得精读。这个 PR 是理解 SGLang MoE 后端扩展模式的很好样例：kimi_k3.py 的 _ep_a2a / _sp_moe 谓词展示了 backend 能力抽象如何决定模型的执行布局，mxfp4.py 的 _apply_aiter() 提取则体现了“按 dispatch 格式契约在统一入口分流”的设计。建议阅读时重点看 DispatchOutputChecker.format_is_deepep() 与 use_deep_gemm 分支的对称关系，并留意后续是否有针对 MoRI 的专项测试补充。

功能与动机

PR body 的动机是 Enable Mori EP for Kimi-K3。提交信息进一步说明根因：models/kimi_k3.py 的 _ep_a2a 和 _sp_moe 只识别 megamoe / DeepEP，导致 --moe-a2a-backend mori 回退到普通 TP 布局（shared experts 保持 TP-sharded、MoE 区域保留 DP gather），而不是跑在 SP-MoE token shard 上；mxfp4.py 的 apply() 会解包 .topk_output，而 DeepEP 家族（MoRI 在内）的 dispatch 输出没有该字段，需要在解包前交给 AITER runner。

实现拆解

本 PR 的变更入口是 Kimi-K3 的 MoE 后端选择逻辑与 MXFP4 量化入口，共 2 个文件：

1. 模型层：扩展 EP-a2a 后端识别（python/sglang/srt/models/kimi_k3.py）
   在模型初始化中，_ep_a2a 与 _sp_moe 两个谓词的判断条件里追加 _a2a_backend.is_mori()，并同步更新相关注释。这样 --moe-a2a-backend mori 不再回退到普通 TP 布局，而是与 megamoe / DeepEP 一样进入 SP-MoE 布局：shared experts 按 tp1 复制、o_proj 延迟注意力 TP reduce 等下游分支都会被激活。_ep_a2a 决定 _shared_experts_tp1、_shared_experts_attn_tp_comm 与 MoE 区域的数据流，_sp_moe 决定 o_proj 的 reduce-scatter 时机，因此该谓词变更会连带改变整条 MoE 前向执行布局。

2. 量化层：AITER 路径提取与提前分流（python/sglang/srt/layers/quantization/mxfp4.py）
   在 apply() 中新增 DispatchOutputChecker.format_is_deepep() 检查，当启用 AITER 且 dispatch 输出属于 DeepEP 系列格式时，在解包 .topk_output 之前直接返回 _apply_aiter() 的结果；把原有 if _use_aiter: 分支的函数体原样提取为新的 _apply_aiter()，原分支改为委托调用。原因是 MoRI 属于 DeepEP 家族，dispatch 输出直接携带 topk_ids / topk_weights、没有 .topk_output 可解包，与 use_deep_gemm 分支面临同样的契约约束，所以必须在标准解包之前路由走。提取后 AITER 路径可由两个调用点（标准路径末尾与新 DeepEP 提前返回）共同复用，避免重复代码；同时保证 deepep_normal / deepep_ll / MoRI 格式都走 AITER runner。

3. 死代码清理
   依据 review 意见，移除 kimi_k3.py 中已废弃的 self._moe_front_needs_contiguouts 字段（main 分支已改用 self._moe_front_needs_dense_bf16）。

4. 测试与验证配套
   本次未新增单元测试；仅在 PR 描述中给出 gsm8k 准确度 0.951 与完整的启动命令（SGLANG_USE_AITER=1、AITER_SITUV2_A8W4=1、SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK=2048 等）。CI 中 AMD ROCm 7.2 与 Extra 任务标注为失败，需关注。

关键文件：

• python/sglang/srt/layers/quantization/mxfp4.py（模块 量化层；类别 source；类型 core-logic；符号 apply, _apply_aiter）: 核心量化入口：新增 DispatchOutputChecker.format_is_deepep() 提前分流，并提取 _apply_aiter()，是本次改动的主体。
• python/sglang/srt/models/kimi_k3.py（模块 K3模型；类别 source；类型 data-contract；符号 _ep_a2a, _sp_moe）: 模型层 backend 谓词扩展：_ep_a2a / _sp_moe 纳入 is_mori()，使 MoRI 走 SP-MoE 布局；同时清理死代码。

关键符号：apply, _apply_aiter

关键源码片段

python/sglang/srt/models/kimi_k3.py

模型层 backend 谓词扩展：_ep_a2a / _sp_moe 纳入 is_mori()，使 MoRI 走 SP-MoE 布局；同时清理死代码。

# python/sglang/srt/models/kimi_k3.py
# EP a2a 后端（Megamoe / DeepEP / MoRI）会把每行 token 直接送到其
# experts，因此 MoE 区域可以消费本 rank 持有的任意行 —— SP-MoE token
# shard（attn_tp > 1）或 DP-local batch（DP attention）—— 每个全局
# token 恰好被 dispatch 一次，区域内部无需任何 DP gather 或 TP reduce。
_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()
)

评论区精华

核心 review 交锋集中在两点：

• Duyi-Wang 指出 kimi_k3.py 中 self._moe_front_needs_contiguouts 是死代码，与本次 PR 无关，main 分支已经改用 self._moe_front_needs_dense_bf16。作者 RolaoDenthu 回复致谢并确认已移除，该线程已解决。

• 设计层面：mxfp4.py 中新增的 DeepEP 提前返回分支与既有 use_deep_gemm 分支在注释中形成对称说明，核心权衡是“dispatch 格式契约不同，必须在解包 .topk_output 之前按格式分流”；AITER 与 DeepGEMM runner 各自通过注册的 pre/post-permute 或专用 quant_info 处理 raw topk 格式，避免在统一入口里做格式特判扩散。

• kimi_k3.py 中遗留死代码字段 (style): 作者 RolaoDenthu 回复感谢并已移除该字段。

风险与影响

• 风险：
  • 核心路径重构回归：mxfp4.py 的 apply() 是所有 MXFP4 MoE 的公共入口，AITER 分支从末尾提前到开头，任何 format_is_deepep() 判定边界（如新格式未登记）都可能导致行为偏差；_apply_aiter() 提取后需确保所有调用点语义一致。
  • 后端布局切换：kimi_k3.py 将 MoRI 纳入 _ep_a2a / _sp_moe 后，shared experts 复制、o_proj reduce-scatter 等路径从 TP 布局切换为 SP-MoE 布局，若 MoRI 的实际 dispatch 契约与 DeepEP 存在细微差异，可能引入精度或性能回归。
  • 测试覆盖缺失：PR 没有配套单元测试，仅依赖 gsm8k 0.951 与手工启动验证；AMD ROCm 7.2 与 CI Extra 任务标红，需要人工确认是否为环境问题。
  • 配置面广：启用 MoRI 依赖 SGLANG_USE_AITER、AITER_FLYDSL_FORCE、AITER_SITUV2_A8W4、SGLANG_MLA_DECODE_TUNE、MORI_SHMEM_MODE 等多个环境变量组合，排障面较大。
• 影响：
  • 用户：AMD 平台上运行 Kimi-K3 的用户新增 --moe-a2a-backend mori 可选路径，MoE 区域切换到 SP-MoE 布局并走 AITER runner，预期可获得更优的 EP 通信与计算重叠。
  • 系统：默认路径（非 MoRI、非 AITER）完全不受影响；只有显式启用 MoRI 或 DeepEP 系列格式并开启 AITER 时才进入新分支。
  • 团队：改动集中在两个文件，规模小，但 mxfp4.py 的 apply 入口是量化 MoE 的公共路径，后续新增 backend 时需继续遵守“格式契约先行拆分”的模式；缺少测试意味着回归风险靠人工兜底。
  • 风险标记：核心路径重构, 缺少测试覆盖, AMD 专用路径, 后端布局切换

关联脉络

• PR #35672 [AMD] Enable draft_extend CUDA graph for HIP DSA backend: 同为 AMD 后端能力使能类改动，且都围绕 K3/DeepSeek 家族的 HIP 路径。
• PR #33021 [AMD] Drop redundant FP8 bpreshuffle scale transpose via fused AR kernel: AMD + 量化层（layernorm/通信）的相邻优化，与本次 MXFP4 量化路径调整同属 AMD 量化执行链路。