[feat] A5 deep fused moe supported by 666syh · Pull Request #661 · sgl-project/sgl-kernel-npu

666syh · 2026-08-06T08:30:50Z

Motivation

Deep Fused MoE is now supported on Ascend 950. This operator fuses the Dispatch, expert computation (GMM), and Combine stages into a single large operator so that communication and computation can be overlapped, improving end-to-end MoE execution efficiency.
This PR introduces the corresponding implementation updates in csrc/ and synchronizes the related examples.

Modifications

Host
- Updated the host-side operator registration and launch flow for Deep Fused MoE.
- Added the Ascend 950 fused execution path for Dispatch + GMM + Combine.
- Aligned parameter preparation and invocation flow with the new fused operator behavior.
Kernel
- Updated the Ascend 950 Deep Fused MoE kernel implementation.
- Integrated Dispatch, GMM, and Combine into a unified execution flow.
- Improved communication/computation overlap to reduce stage-by-stage synchronization overhead.
Tiling
- Updated tiling calculation and configuration logic for the fused path.
- Adjusted tile partitioning, buffer layout, and execution granularity for Ascend 950.
- Ensured consistency between tiling outputs and kernel-side execution requirements.

Testing

Original Fused MoE operator based on Ascend 910C:

[Rank 0] Difference between base and fused recv_count -> max: 0, mean: 0
[Rank 1] Difference between base and fused recv_count -> max: 0, mean: 0
[Rank 5] Difference between base and fused recv_count -> max: 0, mean: 0
[Rank 2] Difference between base and fused recv_count -> max: 0, mean: 0
[Rank 3] Difference between base and fused recv_count -> max: 0, mean: 0
[Rank 7] Difference between base and fused recv_count -> max: 0, mean: 0
[Rank 4] Difference between base and fused recv_count -> max: 0, mean: 0
[Rank 6] Difference between base and fused recv_count -> max: 0, mean: 0
[Rank 4] baseline_time= 267.56 us
[Rank 4] fused_moe_time= 218.31 us
[Rank 5] baseline_time= 268.56 us
[Rank 5] fused_moe_time= 218.16 us
[Rank 1] baseline_time= 268.03 us
[Rank 1] fused_moe_time= 218.09 us
[Rank 3] baseline_time= 268.43 us
[Rank 3] fused_moe_time= 217.98 us
[Rank 0] baseline_time= 268.18 us
[Rank 0] fused_moe_time= 218.01 us
[Rank 6] baseline_time= 268.16 us
[Rank 6] fused_moe_time= 218.42 us
[Rank 2] baseline_time= 267.80 us
[Rank 2] fused_moe_time= 218.01 us
[Rank 7] baseline_time= 268.22 us
[Rank 7] fused_moe_time= 218.20 us

Fused MoE operator based on Ascend 950:

| Rank | Local Experts | Fused Counts         | Sum |
|-----:|--------------:|----------------------|----:|
|    0 |             8 | [24, 24, 24, 24, 24, 24, 24, 24] | 192 |
|    1 |             8 | [24, 24, 24, 24, 24, 24, 24, 24] | 192 |
|    2 |             8 | [24, 24, 24, 24, 24, 24, 24, 24] | 192 |
|    3 |             8 | [24, 24, 24, 24, 24, 24, 24, 24] | 192 |
|    4 |             8 | [24, 24, 24, 24, 24, 24, 24, 24] | 192 |
|    5 |             8 | [24, 24, 24, 24, 24, 24, 24, 24] | 192 |
|    6 |             8 | [24, 24, 24, 24, 24, 24, 24, 24] | 192 |
|    7 |             8 | [24, 24, 24, 24, 24, 24, 24, 24] | 192 |
Accuracy check passed. avg_diff=0.000000, max_diff=0.000000, calc_diff=0.000000
rank skew summary:
| Rank | Dispatch (us) | GMM1 (us) | SwiGLU (us) | Requant (us) | GMM2 (us) | Combine (us) | Small Total (us) | Fused (us) |
|-----:|--------------:|----------:|------------:|-------------:|----------:|-------------:|-----------------:|-----------:|
|    0 |         18.68 |    235.81 |       10.92 |         6.29 |    121.91 |        23.61 |           417.22 |     401.65 |
|    1 |         18.20 |    236.05 |       10.92 |         6.35 |    121.89 |        23.88 |           417.28 |     401.68 |
|    2 |         18.52 |    236.92 |       10.83 |         6.18 |    122.14 |        22.61 |           417.20 |     401.69 |
|    3 |         18.50 |    236.09 |       10.99 |         6.34 |    121.81 |        23.64 |           417.37 |     401.70 |
|    4 |         18.23 |    236.32 |       10.77 |         6.17 |    121.68 |        24.01 |           417.19 |     401.70 |
|    5 |         18.69 |    236.44 |       11.01 |         6.37 |    122.12 |        22.51 |           417.14 |     401.56 |
|    6 |         18.03 |    236.45 |       11.20 |         6.16 |    121.85 |        23.51 |           417.20 |     401.68 |
|    7 |         18.97 |    236.25 |       11.04 |         6.21 |    121.83 |        23.02 |           417.32 |     401.68 |
| mean |         18.48 |    236.29 |       10.96 |         6.26 |    121.90 |        23.35 |           417.24 |     401.67 |

Benchmarking and Profiling

Use single operators to simulate framework-side model invocation.

premise:
- FP8 quant
- no share experts
- Topk distribution is uniform.
case:

Dimension value

EP 8

hidden 7168

moe-intermediate-size 3072

topk 6

performance:

num-tokens-per-device	num-experts-per-device	Dispatch + GMM + Combine (us)	Fused Moe (us)	improvement
32	6	323.59	302.68	6.4% ↑
64	6	366.38	320.99	12.4% ↑
96	6	431.81	341.94	20.8% ↑
128	6	493.37	360.92	26.8% ↑

Checklist

Format your code.
Add unit tests. (N/A for docs)
Update documentation.
Provide accuracy and speed benchmark results. (N/A for docs)

kaniel-outis · 2026-08-07T09:26:58Z

@@ -0,0 +1,53 @@
+## DeepEP-DeepFusedMoE


这块融合算子是有相关的文档，a5只需要在a3的基础修改即可，无需新启文件。

kaniel-outis · 2026-08-13T09:20:10Z

+    uint64_t shareX1MxScaleNum = x1MxScaleNum;
+    uint64_t x2MxScaleNum = CeilUp(Ceil(gmm2HLen, 32), 2);
+    uint64_t shareX2MxScaleNum = CeilUp(Ceil(shareGmm2HLen, 32), 2);
+    ;


有孤立分号；同一函数还保留了注释掉的 workspace 计算。

kaniel-outis · 2026-08-13T09:20:40Z

+    size_t x1TokenSize = MxActStorageBytes(shareExpertTokenNum * h + maxTokenNum * h, isMxFp4);
+    size_t x2TokenSize = MxActStorageBytes(shareExpertTokenNum * shareGmm2HLen + maxTokenNum * gmm2HLen, isMxFp4);
+    size_t maxTokenSize = CeilUp(x1TokenSize < x2TokenSize ? x2TokenSize : x1TokenSize, GM_ALIGN_SIZE);
+    // size_t tokenScaleSize = CeilUp((shareExpertTokenNum + maxTokenNum) * sizeof(float), GM_ALIGN_SIZE);


删除死代码、孤立语句。

kaniel-outis · 2026-08-13T09:24:49Z

@@ -0,0 +1,295 @@
+#ifndef DEEPEP_OP_KERNEL_PROFILING_COMMON_PROFILE_PROTOCOL_COMMON_H


profiling 协议被复制两份，已出现漂移
Host 和 kernel 分别维护近乎相同的协议头：
csrc/deepep/profiling/common/profile_protocol_common.h
csrc/deepep/ops/op_kernel/profiling/common/profile_protocol_common.h
当前两者已经存在 API 差异，后续字段布局或常量变更很容易导致 host/kernel 静默不兼容。
建议：抽取唯一共享协议头；若受编译环境限制，至少加入构建期一致性校验或由同一模板生成。

kaniel-outis · 2026-08-13T09:30:10Z

+                                                                static_cast<bool>(EXEC_FLAG & EXEC_FLAG_SMOOTH_QUANT),
+                                                                false, EXEC_FLAG>
+                dispatcher;
+            dispatcher.Init(gmX_, gmexpertIds_, gmSmoothScales_, gmShareSmoothScales_, xActiveMask_, gmShareX1, gmX1,


Init 签名（cam_moe_distribute_dispatch.h）第 5/6 参数为 shareX1Token / xActiveMask，但调用处写反：
dispatcher.Init(gmX_, gmexpertIds_, gmSmoothScales_, gmShareSmoothScales_, xActiveMask_, // ← 实参5 传入 shareX1Token 位 gmShareX1, // ← 实参6 传入 xActiveMask 位 gmX1, ...);
QuantForShareExpert() 把量化后的 shared-expert token 写进 xActiveMask 缓冲（破坏掩码）；TokenActiveMaskCal() 从 gmShareX1 读掩码得到垃圾activeMaskBsCnt_
触发条件: moeExpertNumPerRank == 1（tiling 仅在 != 1 时置 EXEC_FLAG_DEEP_FUSE）。PR 测试用 64 experts / EP8 = 8 experts/rank，未覆盖此路径。
修复: 交换两个实参为 gmShareX1, xActiveMask_。

kaniel-outis · 2026-08-13T11:29:27Z

+    if constexpr ((EXEC_FLAG & EXEC_FLAG_DEEP_FUSE) == 0) {
+        if constexpr (g_coreType == AscendC::AIV) {
+            AscendC::TPipe tpipe;
+            MoeDistributeDispatchImpl::CamMoeDistributeDispatch<ExpandXType, int8_t, false, true,


非 deep-fuse 分支实例化 CamMoeDistributeDispatch<ExpandXType, int8_t, false, true, ...>，其 QuantProcess() 产出普通 int8 对称量化 + float per-token scale。但下游 GMM1 kernel（dispatch_mx_gmm1_swiglu.h:70-76）的 static_assert 强制 ElementA 为 float8_e4m3_t/ float8_e5m2_t，scale 为 fp8_e8m0_t（per-32-group MX 格式）。int8_t 被重解释为 float8_e4m3_t 数值含义完全不同。
后果: expert/rank 配置下 GMM1 计算结果错误。
修复: 要么让独立 dispatcher 也走 MX 量化（QuantDynamicMx），要么在 tiling 阶段对 moeExpertNumPerRank == 1 报错拒绝。

kaniel-outis · 2026-08-13T11:33:07Z

+    return ge::GRAPH_SUCCESS;
+}
+
+static ge::graphStatus GetAttrAndSetTilingData(const gert::TilingContext &context, const char *nodeName,


fused_deep_moe_tiling.cpp 的GetAttrAndSetTilingData 从未给这两个字段赋值，而 kernel 侧 cam_moe_distribute_combine.h:251/259 与 cam_moe_distribute_dispatch.h:235 会读取。当前零初始化且暂未用于计算/边界检查，属潜在 0 值隐患。建议填充或删除字段。

kaniel-outis · 2026-08-13T11:33:59Z

+        }
+
+        AscendC::PipeBarrier<PIPE_ALL>();
+        Arch::CrossCoreFlag gmm1AivFinished{0};


fused_deep_moe_a5.h:425 和 :445 均用 flag ID {0}。非 deep-fuse 路径下两次 SetFlag/WaitFlag 中间虽有 PipeBarrier<PIPE_ALL>()，但若 flag 需显式 clear 则第二次等待可能受残留状态影响。建议用不同 ID 或确认自动清除语义。

kaniel-outis · 2026-08-13T11:35:37Z

+        return;
+    }
+    const auto *calibration = GetTimeCalibration(rank);
+    exporter::ExportAggregatedTrace(sources, rank, manager.profileTraceDir, calibration, manager.numRanks);


profile_session.cpp：ExportAggregatedTrace() 返回值未检查，随后 manager.Reset() 清空所有 buffer。磁盘满/权限错→整轮 profiling 静默丢失。建议检查返回值，失败时保留 session 可重试。

profling风险不大，无需dfx逻辑

kaniel-outis · 2026-08-13T11:41:24Z

+    }
+
+    CATLASS_DEVICE
+    uint32_t SendToMoeExprt(GM_ADDR gmX, GM_ADDR gmExpandIdx, GM_ADDR gmMoeSmoothScales)


拼写 SendToMoeExprt → Expert

kaniel-outis · 2026-08-13T11:42:05Z

+    // MM1/GMM1-Swiglu output
+    int64_t shareSwigluOffset;
+    int64_t swigluOffset;
+    // MM2/GMM2 input


建议统一 UTF-8

kaniel-outis · 2026-08-13T11:43:11Z

+
+        uint64_t profWeightSumStart = 0;
+
+        do {


do{...}while(false) 无 break，等价普通 {}，建议简化

kaniel-outis · 2026-08-13T11:43:43Z

+    int64_t launchCountCaptured{0};
+};
+
+void ExportBufferToTrace(const at::Tensor &profileBuffer, int64_t rank, const std::string &profileTraceDir,


ExportBufferToTrace 完整实现但无调用方，dead code

kaniel-outis · 2026-08-13T11:46:43Z

    num_experts: int,
    quant_mode: int = 1,
    fuse_mode: FuseMode = FuseMode.FUSED_DEEP_MOE,
 ) -> Tuple[torch.Tensor, torch.Tensor]


Python API 签名漏掉 profile_enable

kaniel-outis · 2026-08-13T11:56:00Z

+    auto profileBufferBytesPtr = attrs->GetAttrPointer<int64_t>(ATTR_PROFILE_BUFFER_BYTES_INDEX);
+    auto profileLaunchIdPtr = attrs->GetAttrPointer<int64_t>(ATTR_PROFILE_LAUNCH_ID_INDEX);
+
+    uint32_t epRankSize = static_cast<uint32_t>(*epRankSizePtr);


Tiling 在校验前解引用属性并做除法，非法输入会崩溃而不是报错在检查 epRankSizePtr、epRankIdPtr、moeExpertNumPtr 是否为空、epRankSize 是否大于零之前，已经解引用并执行 moeExpertNum / epRankSize。
建议：先完整检查全部指针和 epRankSize > 0，再转换类型与相除。

* origin/main: release A5 image (sgl-project#683) add minimax m3 attention ops on A3 (sgl-project#668) [Attentions][CI][Fix] Add attentions to build-and-release CI and build module bugfix (sgl-project#571)

Y0Hang added 8 commits August 6, 2026 11:16

fused_moe init

45037c6

remove relwithdebug

990a279

remove ia_a5_build

d821187

check profiling symbol

7c29e73

merge main

7508ee0

rm lisence head

c47c5f9

rm unused code

cf88312

fix compile cann 9.0

ee33ec3

666syh force-pushed the fused_moe_a5_rebuild branch from 3972f03 to ee33ec3 Compare August 7, 2026 01:19

Y0Hang added 2 commits August 7, 2026 14:16

cleancode

499f6e4

merge main

e984969

kaniel-outis reviewed Aug 7, 2026

View reviewed changes

Y0Hang added 3 commits August 12, 2026 14:41

profling ts align && fix fp8 nz

16ea4c9

Merge remote-tracking branch 'origin' into fused_moe_a5_rebuild

14b7e6f

modify docs

1110cc0

kaniel-outis reviewed Aug 13, 2026

View reviewed changes

cleancode

39d3933

Y0Hang added 2 commits August 14, 2026 11:35

unity profiling common

30a2cc5

Merge remote-tracking branch 'origin/main' into fused_moe_a5_rebuild

bbb1170

* origin/main: release A5 image (sgl-project#683) add minimax m3 attention ops on A3 (sgl-project#668) [Attentions][CI][Fix] Add attentions to build-and-release CI and build module bugfix (sgl-project#571)

666syh force-pushed the fused_moe_a5_rebuild branch from 2b38637 to bbb1170 Compare August 14, 2026 08:31

		@@ -0,0 +1,295 @@
		#ifndef DEEPEP_OP_KERNEL_PROFILING_COMMON_PROFILE_PROTOCOL_COMMON_H

Dimension	value
EP	8
hidden	7168
moe-intermediate-size	3072
topk	6

Uh oh!

Conversation

666syh commented Aug 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Testing

Benchmarking and Profiling

Checklist

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kaniel-outis Aug 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kaniel-outis Aug 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kaniel-outis Aug 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

666syh commented Aug 6, 2026 •

edited

Loading

kaniel-outis Aug 13, 2026 •

edited

Loading

kaniel-outis Aug 13, 2026 •

edited

Loading

kaniel-outis Aug 13, 2026 •

edited

Loading