| [Feature] Support fused moe
Co-authored-by: betta18<jiangmengyu1@huawei.com>
# message auto-generated for no-merge-commit merge:
!310 merge support_fused_moe_new into dev
[Feature] Support fused moe
Created-by: betta18
Commit-by: betta18
Merged-by: ascend-robot
Description: <!--
PR描述模板更新日期：20251224
-->
# Which issue(s) this PR fixes or accomplishes
Fix part of #ISSUE ID
# Purpose
本 PR 新增 MindIE-SD 的 fused_moe 融合算子接口，用于支持开源框架在 NPU 上执行 MoE 前向计算。当前 PR 接入的是 fused_moe 的 fallback 路径，已基于 vllm-omni 中 HunyuanImage-3.0 模型的 MindIE-SD MoE 组网路径完成适配验证。
当前 fused_moe 对外提供统一入口，并区分两条执行路径：
- **融合算子路径**：预留融合 MoE 算子入口，当前版本暂不支持。
- **fallback 路径**：本 PR 接入的执行路径，用于完成 MoE 路由、专家计算和结果合并流程。
fallback 路径将 MoE 前向过程拆分为多个阶段：
- prepare：整理输入激活和 router logits，并根据输入 layout 准备后续计算所需的数据。
- select_experts：根据 router 输出，为每个 token 选择对应的 top-k experts，并生成 routing weights。
- dispatch：根据专家选择结果，将 token 按 expert 路由并重排，生成专家计算所需的输入。
- mlp：执行专家侧 grouped MLP 计算，完成 routed experts 的前馈计算。
- combine：将专家输出按原 token 顺序合并，恢复 routed MoE 的输出结果。
- finalize：完成输出恢复和通信后的收尾处理。
通过 fallback 路径的阶段化封装，外部调用方只需要使用 fused_moe 统一入口，即可完成当前版本支持的 MoE 前向计算流程。
主要修改包括：
- 新增并导出 fused_moe 融合算子接口。
- 预留融合 MoE 算子路径，并接入 fallback 路径作为当前执行实现。
- 新增 mindiesd.layers.moe 阶段化 MoE fallback 实现。
- 新增 runtime context 处理逻辑，用于集中校验外部入参、生成通信上下文并封装各阶段输入对象。
- 新增专家选择逻辑，支持 top-k routing、routing weight 归一化和自定义 routing function。
- 新增 static dispatcher，用于完成静态 token 分发、expert token 统计和结果恢复。
- 新增 dynamic dispatcher，用于完成动态 token 分发、all-to-all token 交换和结果恢复。
- 新增 grouped expert MLP 计算逻辑，基于 torch_npu.npu_grouped_matmul 完成专家前馈计算。
- 新增 MoE 通信上下文和基础通信算子封装。
- 扩展 NPU 平台识别，新增 A3 平台，并使相关 MoE/attention 路径能够识别 A3。
- 新增 fused_moe 中文特性文档，补充接口说明、参数约束、融合算子路径、fallback 路径、通信配置和使用示例。
后续计划：
- 性能优化。
- 支持量化。
# Test Plan
测试重点包括：
- fused_moe 入口：
- 融合算子路径当前版本暂不支持时回退到 fallback 路径。
- fallback 路径调用阶段化 MoE 实现。
- fallback moe 主流程：
- static dispatcher 路径。
- dynamic dispatcher 路径。
- 无通信组、TP 通信组、EP 通信组场景。
- runtime context：
- 外部输入参数校验。
- prepare、routing、dispatch、MLP compute 阶段输入对象封装。
- MoE 通信上下文选择。
- 专家选择：
- top-k expert 选择。
- top-k weight 归一化。
- 自定义 routing function。
- token dispatch/combine：
- static dispatch 的 token 排序、expert token 统计和结果恢复。
- dynamic dispatch 的 all-to-all token 交换和顺序恢复。
- grouped expert MLP：
- w13 / w2 两阶段 grouped matmul。
- bias 路径。
- swiglu activation。
- 通信算子：
- all-gather。
- reduce-scatter。
- all-reduce。
- all-to-all single。
- 精度测试：
- 对比 static MoE 输出与参考实现结果的一致性。
- 对比 grouped expert MLP 输出与参考实现结果的一致性。
- 验证 HunyuanImage-3.0 MindIE-SD MoE 组网路径下的推理结果正确性。
# Test Report
在 MindIE-SD 中新增 MoE 相关单元测试，主要覆盖 fallback 路径中的通信算子、专家选择、runtime context、主流程、专家 MLP 计算和 token dispatch/combine 等核心逻辑：
- tests/layers/moe/test_comm_ops.py
- tests/layers/moe/test_experts_selector.py
- tests/layers/moe/test_runtime_context.py
- tests/layers/moe/test_moe.py
- tests/layers/moe/test_moe_mlp.py
- tests/layers/moe/test_token_dispatcher.py
在 vllm-omni 框架中，基于 HunyuanImage-3.0 的 MindIE-SD MoE 组网路径完成端到端推理验证，结果正确。验证覆盖：
- EP / TP 通信场景。
- static / dynamic dispatcher 路径。
- tokens_full 不同输入 layout。
- reduce_results 不同结果规约配置。
See merge request: Ascend/MindIE-SD!310 | 4 天前 |
