Skip to content

feat(npu): support CANN 9.1 MegaMoe for Qwen3.5#1989

Draft
pjgao wants to merge 6 commits into
xLLM-AI:mainfrom
pjgao:feat/cann91-megamoe-qwen35
Draft

feat(npu): support CANN 9.1 MegaMoe for Qwen3.5#1989
pjgao wants to merge 6 commits into
xLLM-AI:mainfrom
pjgao:feat/cann91-megamoe-qwen35

Conversation

@pjgao

@pjgao pjgao commented Jul 21, 2026

Copy link
Copy Markdown
Collaborator

背景与目标

Qwen3.5-35B-A3B 是 MoE 模型。原有 NPU FusedMoE 路径将专家路由、跨 rank token dispatch、专家 FFN、激活和 combine 拆成多个阶段,存在额外的 kernel launch、中间张量读写以及通信与计算衔接开销。

CANN 9.1 ops-transformer 提供的 MegaMoe 将 MC2 通信与专家 FFN 组合到统一算子路径中,可以减少阶段切换和中间数据搬运,并为后续在 A3 上优化 Qwen3.5 的 decode 吞吐与时延提供基础。本 PR 的目标是把该算子安全接入 xLLM,同时不影响不满足条件的模型和原有 FusedMoE fallback。

接入方案

1. 配置与能力判断

  • 增加 MegaMoe 配置开关和 kernel 配置解析。
  • mega_moe_policy 中集中检查平台、模型、dtype、shape、EP 拓扑和算子可用性。
  • EP size 不再硬编码为 16;只要 EP 为正数、能够整除专家数且满足算子约束,即可进入 MegaMoe 路径,因此 EP8/EP16 均可验证。
  • 未通过能力检查时继续使用原有 FusedMoE 实现。

2. ProcessGroup 级通信资源

  • 新增 MegaMoeCommResource,由 NPU ProcessGroup 持有并管理 KFC/HCCL context、远端 window 和拓扑信息。
  • 资源按通信域复用,避免每层重复初始化 communicator。
  • 初始化时校验 rank、server 拓扑、buffer span、对齐和边界;销毁前完成未结束任务的 drain。

3. ACL 算子适配

  • 增加 MegaMoe ACL contract,并调用 aclnnMegaMoeGetWorkspaceSizeaclnnMegaMoe
  • 明确可选输入、输出顺序、top-k dtype、权重 shape 和 topology 参数,避免 xLLM wrapper 与外置算子 ABI 错位。
  • 对 MegaMoe vendor 做符号来源校验:workspace 和 execute 两个符号必须来自同一份已验证的 libcust_opapi;标准 OPP 可作为其他算子的 companion vendor,但不能覆盖 MegaMoe ABI。

4. FusedMoE 运行时接入

  • fused_moe 根据 policy 选择 MegaMoe 或原有 fallback。
  • mega_moe_runtime 负责整理本地专家权重、输入、top-k 结果、通信 context 和输出张量,并调用 NPU kernel adapter。
  • 支持 attention data parallelism 与 expert parallelism 组合;空 DP/EP shard 使用一致的 synthetic-token metadata 参与 collective,避免各 rank tensor length 不一致。

5. 异步 launch 生命周期

  • aclnnMegaMoe 在设备任务入队后即可返回,而同一 rank 的多个 FusedMoE 层会复用 ProcessGroup 级 KFC context 和 HCCL window。
  • 本 PR 增加资源级 launch fence:在复用同一通信资源前等待上一次 launch 完成,在实际 launch stream 上记录完成事件,并覆盖异常退出和资源销毁路径。
  • fence 只约束同一 MegaMoe 通信资源,不对无关算子或不同资源做全局同步。

定位过程中遇到的问题与解决方法

问题 根因 解决方法
Qwen3.5 shape 被 tiling 拒绝 Qwen3.5 使用对齐后的 intermediate_hidden=512,A2/A3 共用检查采用了 A2 的 1024 下界 ops-transformer !8966 按 SoC 选择下界:A3 支持 512,A2 仍保持 1024
xLLM wrapper 与 CANN 9.1 MC2 接口不一致 xllm-ops 固定 header 仍描述旧版 HCCL/MC2 context 布局 xllm-ops #21 同步 CANN 9.1 接口,同时保留打包所需 include fallback
NPU registry 启动存在非确定性 多 worker 线程并发读写进程级 KernelRegistry torch-npu-ops #2 对读取、注册和清理统一串行化
只能使用 EP16 policy 把一次验证配置错误固化成能力限制 改为按专家数、EP size 和算子约束做通用能力判断,支持 EP8/EP16 等合法配置
attention DP + EP 下 collective shape 不一致 空 shard 仍需参与 collective,但 synthetic token 和线性注意力 metadata 不一致 统一空 shard token、序列长度、state id 和全局 token count
远端通信 buffer 越界风险 window span 的计算和校验没有完整覆盖所有 rank 与对齐条件 按实际布局计算 span,并在资源初始化阶段校验容量、偏移和边界
EP8 输出随机且跨轮次不一致 standalone 算子数值正确,但异步 launch 后共享 KFC/HCCL context 被下一层提前复用 对同一通信资源增加 stream-aware launch fence,并在 communicator teardown 前 drain
标准算子与外置 MegaMoe ABI 混用 动态符号可能从不同 vendor 解析,workspace/execute ABI 不成对 校验两个 ACL 符号的实际来源,并要求 MegaMoe vendor 位于 vendor 链首位

代码改动规模与分布

当前 PR 相对目标分支共修改 38 个文件,新增 3589 行、删除 16 行,净增 3573 行

  • 生产代码与构建描述:27 个文件,+2432/-12
  • 聚焦测试:8 个文件,+1153/-0
  • submodule 地址与指针:3 个文件,+4/-4
模块 文件数 新增 删除 主要内容
测试 8 1153 0 配置解析、通信资源、ACL contract、policy 和 runtime 聚焦测试
通信资源 7 927 3 MegaMoeCommResource、ProcessGroup 生命周期、buffer span 与 launch fence
FusedMoE 运行时 6 685 8 MegaMoe runtime、FusedMoE 分流、空 shard metadata 与 worker 接入
ACL 与 kernel 适配 8 519 0 ACL contract、动态符号来源校验、workspace/execute 调用与参数定义
配置与策略 5 299 0 配置开关、能力判断、EP/DP 约束与 fallback policy
模型兼容 1 2 1 Qwen3.5 hybrid attention metadata 的 device 传递
submodule 地址与指针 3 4 4 切换到 GitHub 官方仓库并更新 xllm-ops 与 torch-npu-ops 依赖提交
合计 38 3589 16 净增 3573 行

新增代码主要集中在通信资源、运行时和测试三个部分,合计 2765 行,占全部新增行约 77%;模型层只做最小兼容修改,核心接入逻辑保持在 NPU 通信、kernel adapter 和 FusedMoE runtime 内。

依赖关系

建议先合入三个算子、wrapper 依赖 PR,再刷新 xLLM submodule gitlink,最后合入本 PR。

验证环境

项目 版本或配置
硬件 Atlas A3,Ascend 910,SoC target ascend910_93,EP8 使用 8 张 NPU,每张 64 GiB HBM
NPU 驱动工具 npu-smi 26.0.rc1
操作系统 openEuler 24.03 LTS-SP2,aarch64
CANN 9.1.0-beta.3,OPP 构建时间戳 20260616_185108448
ATB 9.0.0,提交 7a8ae13c16daf1a2179d7520e71e6b21445f3b2f
Python 3.11.15
PyTorch 2.9.0+cpu
torch_npu 2.9.0.post2
transformers 5.10.1
tokenizers 0.22.2
safetensors 0.7.0
NumPy 1.26.4
GCC 12.3.1
CMake 3.27.9
模型 Qwen3.5-35B-A3B
并行配置 DP8 / EP8
xLLM PR HEAD 0a838b8c2c55d932fbdcdbba3c6cda1f88be8164
验证 binary SHA256 73a676ea9d95051e0be6481aecad9231b921e66dfa1b0f34c72915f308f0ca18
MegaMoe OPP 标准 CANN OPP + 独立外置 MegaMoe vendor;MegaMoe vendor 位于 vendor 链首位

验证结果

聚焦测试与构建

  • MegaMoe policy、runtime、ACL contract、通信资源和空 shard metadata 聚焦测试通过。
  • ops-transformer MegaMoe op-host RED→GREEN:修复前 A3 512 shape 用例失败;修复后 12/12 通过,且 A2 仍拒绝 512。
  • 受影响 NPU translation unit 使用 -Werror 编译通过。
  • 正式 xLLM binary 增量链接通过;未重复执行无关的大范围产品构建。

EP8 服务级验证

immutable attempt:fix-launch-fence-refactor-ep8-20260721T1215CST

检查项 结果
服务初始化 通过:8 个 worker 均完成初始化,API ready
真实生成 通过:提示“只回答最终数字:123加456等于多少?”,temperature=0,最终输出 579
并发确定性 通过:连续 3 轮 × 每轮 8 请求,共 24/24 输出非空且字节一致
运行时错误扫描 通过:未发现 HCCL、ACL、AICore、非法内存访问或进程异常
进程状态 通过:未出现 mmap/GUP D-state
清理 通过:停止后 API 不再接受连接,使用的 NPU 均已释放

该结果属于针对本次 MegaMoe 跨 rank 随机 token 问题的 L2 集成回归,不等同于完整 CEval/GSM8K 分数评测。PR 历史清理仅移除了不随代码发布的机器适配,MegaMoe 接入和精度修复 diff 与上述验证路径一致。

PR 范围

本 PR 只提交 MegaMoe 接入、EP/DP 适配、通信资源管理、ABI/shape 校验和异步生命周期修复;机器相关适配不进入上游代码。

georgezhanglei85-eng pushed a commit to georgezhanglei85-eng/xllm that referenced this pull request Jul 21, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant