FusedDeepMoE Kernel 深度分析报告

CAM (Communication Acceleration for Matrix) 皇冠上的明珠 — 将 Dispatch + GMM1 + Swiglu + GMM2 + Combine 融合为单一算子的终极优化

1. 整体架构与融合策略

1.1 算子定义

FusedDeepMoE 是 CAM 中最复杂的融合算子，将 MoE Decode 阶段的完整计算流水线融合进单一 kernel：

标准小算子组合 (5次kernel launch + 多次GM读写):
  Dispatch → GMM1 → Swiglu → GMM2 → Combine

FusedDeepMoE (1次kernel launch, 全在片上完成):
  Dispatch → GMM1 → Swiglu → GMM2 → Combine (全部融合)

1.2 输入/输出签名

输入 (14个): | 索引 | 名称 | 类型 | 说明 | |------|------|------|------| | 0 | x | bf16/fp16 | 输入 token [bs, h] | | 1 | expert_ids | int32 | 专家索引 [bs, topk] | | 2 | gmm1_weight | int8 (TensorList) | 路由专家 GMM1 权重 | | 3 | gmm1_weight_scale | fp32/bf16/fp16 | GMM1 权重 scale | | 4 | gmm2_weight | int8 (TensorList) | 路由专家 GMM2 权重 | | 5 | gmm2_weight_scale | fp32/bf16/fp16 | GMM2 权重 scale | | 6 | expert_scales | fp32 | token级专家权重 [bs, topk] | | 7-8 | share_gmm1_weight/scale | int8/fp32 | 共享专家 GMM1 权重 (可选) | | 9-10 | share_gmm2_weight/scale | int8/fp32 | 共享专家 GMM2 权重 (可选) | | 11 | expert_smooth_scales | fp32 | 平滑 scale (可选) | | 12 | share_smooth_scales | fp32 | 共享专家平滑 scale (可选) | | 13 | x_active_mask | bool | token 激活掩码 (可选) |

输出 (3个): - output: [bs, h] — MoE 路由输出 - share_output: [bs, h] — 共享专家输出 - expert_token_nums: [moeExpertNumPerRank] — 每个专家接收的 token 数

属性 (6个): - group_ep, ep_rank_size, ep_rank_id, moe_expert_num, quant_mode, global_bs

1.3 为什么要融合？

核心动机：消除中间数据的 Global Memory 往返

在标准小算子组合中： 1. Dispatch 写出 expand_x, dynamic_scales 到 GM → GMM1 读回 (读写 2次) 2. GMM1 写出 int32 累加结果到 GM → Swiglu 读回 (读写 2次) 3. Swiglu 写出 int8 + scale 到 GM → GMM2 读回 (读写 2次) 4. GMM2 写出 结果到 GM → Combine 读回 (读写 2次)

总计至少 8 次 GM 访问，对于 Decode 阶段（batch小、延迟敏感），GM 带宽是瓶颈。

FusedDeepMoE 的融合策略： - X1 (GMM1输入)：AIV 在 UB 中完成 Dispatch 量化后，直接通过 workspace 偏移量传递给 AIC 做 GMM1 - Swiglu 输出：GMM1 的 int32 结果在 L0C/UB 中直接做 Swiglu + 动态量化，写出 float 到 workspace - X2 (GMM2输入)：Swiglu 后立即做动态量化为 int8，通过 workspace 传递给 GMM2 - GMM2 输出：直接在 UB 中完成 Combine 的 AlltoAll 发送

关键收益：中间数据 X1、SwigluOut、X2 全部在片上 workspace 中流转，不访问 GM。

2. op_host Tiling 策略

2.1 核心 Tiling 数据结构

// fused_deep_moe_tiling.h
struct FusedDeepMoeInfo {
    uint32_t epRankSize;           // EP 并行度
    uint32_t epRankId;             // 当前 rank
    uint32_t moeExpertNum;         // 总专家数
    uint32_t moeExpertNumPerRank;  // 每 rank 专家数
    uint32_t quantMode;            // 量化模式
    uint32_t globalBs;             // 全局 batch size
    uint32_t bs;                   // 本地 batch size
    uint32_t k;                    // topK
    uint32_t h;                    // hidden size
    uint32_t aicNum;               // AIC 核数
    uint32_t aivNum;               // AIV 核数
    uint64_t gmm1HLen;             // GMM1 隐藏维度 (intermediate*2)
    uint64_t shareGmm1HLen;        // 共享专家 GMM1 隐藏维度
    bool isTensorList;             // 权重是否为 TensorList
};

2.2 Workspace 布局计算

Workspace 是 FusedDeepMoE 的核心设计，所有中间数据在此共享：

// 各段偏移量计算 (启用 ENABLE_REUSE_MEMORY 时复用内存)
WorkSpaceOffset:
  ├── shareX1TokenOffset / x1TokenOffset      // GMM1 输入 (int8)
  ├── shareX1ScaleOffset / x1ScaleOffset      // GMM1 输入 per-token scale
  ├── swapSpaceOffset / shareSwigluOffset     // GMM1 输出 (int32) / Swiglu 输出 (fp32)
  ├── shareX2TokenOffset / x2TokenOffset      // GMM2 输入 (int8)
  ├── shareX2ScaleOffset / x2ScaleOffset      // GMM2 输入 per-token scale
  ├── y2TokenOffset                           // GMM2 输出 (浅融合使用)
  ├── groupListOffset                         // 专家 token 数前缀和
  ├── expandIdxOffset                         // Dispatch 索引
  ├── epSendCountOffset                       // AlltoAll 发送计数
  └── reservedOffset                          // 保留空间 (256KB)

内存复用策略 (ENABLE_REUSE_MEMORY): - X1 和 X2 共享同一块 workspace 空间（因为 X1 和 X2 生命周期不重叠） - Scale 空间同理复用 - SwapSpace 和 Swiglu 输出复用

2.3 Tile Shape 常量

// GMM1 Tile Shape
GMM1_L1M = 256, GMM1_L1N = 128, GMM1_L1K = 512
GMM1_L0K = 128, GMM1_EPIM = 64
GMM1_SWIZZLE_OFFSET = 3, GMM1_SWIZZLE_DIRECTION = 0

// GMM2 Tile Shape
GMM2_L1M = 128, GMM2_L1N = 256, GMM2_L1K = 512
GMM2_L0K = 128, GMM2_EPIM = 32
GMM2_L1A_STAGES = 4, GMM2_L1B_STAGES = 2  // GMM2 使用 ResidentA 优化

// 通用
WORKSPACE_STAGES = 4
CUSTOM_PRELOAD_STAGES = 1
CUSTOM_L1_STAGES = 2, CUSTOM_L0A_STAGES = 2, CUSTOM_L0B_STAGES = 2

2.4 Tiling Key 路由机制

使用华为 ASCENDC_TPL 机制，通过 5 个 bool 标志 + 1 个 ARCH 参数组合出多种执行路径：

// 关键执行标志
EXEC_FLAG_DEEP_FUSE      = 1 << 0   // 是否深度融合 (moeExpertNumPerRank != 1)
EXEC_FLAG_TENSOR_LIST    = 1 << 1   // 权重是否为 TensorList
EXEC_FLAG_X_ACTIVE_MASK  = 1 << 2   // 是否启用 x_active_mask
EXEC_FLAG_SHARED_EXPERT  = 1 << 3   // 是否有共享专家
EXEC_FLAG_SMOOTH_QUANT   = 1 << 4   // 是否启用 smooth quant

Tiling Key 决定： - 是否跳过独立 Dispatch（深度融合时 Dispatch 在 kernel 内完成） - 是否走共享专家路径 - 是否启用 smooth scale 处理

3. op_kernel 执行流程

3.1 Kernel 入口

// fused_deep_moe.cpp
template <bool TPL_IS_DEEP_FUSE, bool TPL_IS_TENSOR_LIST, ...>
__global__ __aicore__ void fused_deep_moe(...) {
    KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2);  // 1C2V 模式
    // 重组 EXEC_FLAG 位图
    constexpr uint32_t EXEC_FLAG = (TPL_IS_DEEP_FUSE ? EXEC_FLAG_DEEP_FUSE : 0) | ...;
    // 根据 TPL_IS_SHARED_EXPERT 选择 tiling struct
    // 实例化 FusedDeepMoe<...> op; op.Init(...); op.Process();
}

1C2V 模式：每个 AI Core 组包含 1 个 AIC (AI Core Compute) + 2 个 AIV (AI Core Vector)。AIC 负责 GMM 计算，AIV 负责通信调度。

3.2 FusedDeepMoe::Process() 主流程

void Process() {
    // 1. 非深度融合时，先执行独立 Dispatch
    if constexpr ((EXEC_FLAG & EXEC_FLAG_DEEP_FUSE) == 0) {
        if (g_coreType == AIV) {
            CamMoeDistributeDispatch dispatcher;
            dispatcher.Init(...);
            dispatcher.Process();  // AlltoAll Dispatch
        }
        PipeBarrier<PIPE_ALL>();
        CrossCoreBarrier/Wait();  // AIV→AIC 同步
    }

    // 2. GMM1 + Swiglu + 动态量化 (融合执行)
    GmmDeqSwigluQuant<Gmm1L1TileShape, ...>(...);
    PipeBarrier<PIPE_ALL>();
    CrossCoreBarrier/Wait();  // AIC→AIV 同步

    // 3. 初始化 Combine (AIV 端)
    CamMoeDistributeCombine combiner;
    combiner.Init(...);

    // 4. GMM2 + Combine (融合执行)
    GmmDeq<Gmm2L1TileShape, ...>(..., &combiner);
}

3.3 GmmDeqSwigluQuant 函数

这是 GMM1 的模板化入口，定义了完整的 GMM1 计算流水线：

template <..., class L1TileShape_, class L0TileShape_, class EpilogueTileShape_, class BlockScheduler_>
CATLASS_DEVICE void GmmDeqSwigluQuant(...) {
    // 数据类型定义
    using AType = Gemm::GemmType<int8_t, layout::RowMajor>;   // 输入 X1
    using BType = Gemm::GemmType<int8_t, layout::zN>;          // 权重 (NZ格式)
    using CType = Gemm::GemmType<int32_t, layout::RowMajor>;   // 累加器

    // BlockMmad: L1→L0 多级流水
    using BlockMmad = Gemm::Block::BlockMmad<DispatchPolicy, L1TileShape, L0TileShape, AType, BType, CType>;

    // Epilogue: PerTokenDequantSwiglu — 核心融合点!
    using EpilogueDispatchPolicy = Epilogue::EpilogueAtlasA2PerTokenDequantSwiglu<ubStages, 0>;
    using BlockEpilogue = Epilogue::Block::BlockEpilogue<...>;

    // Kernel: GroupedMatmulSliceMPerTokenDequantSwigluQuantMultiStageWorkspace
    using GemmKernel = Gemm::Kernel::GroupedMatmulSliceMPerTokenDequantSwigluQuantMultiStageWorkspace<...>;
    GemmKernel gemm;
    gemm(params);
}

3.4 GmmDeq 函数

GMM2 的模板化入口，结构类似但使用不同的 Epilogue：

// GMM2 使用 PerTokenDequantCombine epilogue
using EpilogueDispatchPolicy = Epilogue::EpilogueAtlasA2PerTokenDequantCombine<ubStages, EXEC_FLAG>;
// 注意 GMM2 使用 ResidentA Dispatch Policy (A矩阵常驻 L1)
using Gmm2DispatchPolicy = Gemm::MmadAtlasA2PreloadAsyncWithCallbackResidentA<...>;

4. raw_distributed 子模块

4.1 CamMoeDistributeDispatch

文件: raw_distributed/cam_moe_distribute_dispatch.h

负责将 token 按 expert_id 路由到对应 rank 的窗口内存中。

核心流程 (Process):

void Process() {
    if ASCEND_IS_AIV {
        if constexpr (EXEC_FLAG & EXEC_FLAG_SHARED_EXPERT)
            QuantForShareExpert();   // 共享专家量化
        AlltoAllDispatch();          // 路由专家 AlltoAll 分发
        SetStatus();                 // 通知对端数据就绪
        WaitDispatch();              // 等待所有 rank 数据到达
        LocalWindowCopy();           // 从窗口拷贝到本地 tensor
        if constexpr (IsNeedAllgater) {
            AllGatherSetStatusAndWait();
            AllgatherProcessOut();
        }
        UpdataTokenNumsOut();        // 更新 expert_token_nums
    }
}

关键优化 — AIV 循环优化 (enableAivOpt):

// 当 batch_size <= 64 且 expert_num 对齐后 <= 256 时启用
// 使用 table 法替代逐 token 计数，减少 UB 访问
if (axisBS_ <= LOOP_OPT_MAX_BS && moeExpertRankNumAligned_ <= LOOP_OPT_MAX_MOE_RANK) {
    enableAivOpt_ = true;
    // 构建 (bs+1) × expert_num 的 table，用前缀和计算 expand_idx
}

量化融合：Dispatch 阶段同时完成动态量化（bf16→int8），将 scale 附加在 token 后面传输：

// hCommuSize = hOutSize + scaleParamPad (128B for scale)
hCommuSize_ = axisH_ * sizeof(ExpandXOutType) + scaleParamPad_;

4.2 CamMoeDistributeCombine

文件: raw_distributed/cam_moe_distribute_combine.h

负责将 GMM2 计算结果路由回原始 rank，并做加权求和。

核心流程:

void Process() {
    SyncAll<true>();
    if constexpr (IsNeedReduceScatter) ReduceScatterTrans();
    if constexpr ((EXEC_FLAG & EXEC_FLAG_DEEP_FUSE) == 0) {
        BuffInit();
        SetWaitTpStatusAndDisPatch();  // 非融合模式
    }
    AlltoAllBuffInit();
    SetStatus();
    WaitDispatch();
    LocalWindowCopy();  // 从窗口读取 + 加权求和 + 写出
}

深度融合时的特殊路径:

// 融合模式下，Combine 被拆分为两部分：
void AllToAllSend() {
    // AIV subblock 0: AlltoAll 发送
    // AIV subblock 1: ReducePermute
    if constexpr (EXEC_FLAG & EXEC_FLAG_DEEP_FUSE) {
        CrossCoreSetFlag<0x0, PIPE_MTE3>(SEND_SYNC_EVENT_ID);
        CrossCoreWaitFlag(SEND_SYNC_EVENT_ID);
    }
}

void ReducePermute() {
    // 在 GMM2 的 AIV 端执行
    AlltoAllBuffInit();
    WaitDispatch();
    LocalWindowCopy();
}

加权求和 (LocalWindowCopy):

// 对每个 token，遍历 topK 个专家的结果，加权求和
for (uint32_t tokenIndex = beginIndex; tokenIndex < endIndex; ++tokenIndex) {
    Duplicate(sumFloatBufLocal, 0, axisH_);
    for (uint32_t i = 0; i < axisK_; i++) {
        float scaleVal = expandScalesLocal.GetValue(index);
        // 从对应专家的窗口位置读取
        rowTmpGlobal_.SetGlobalBuffer(wAddr + expertPerSizeOnWin_ * moeExpert + ...);
        Cast(rowTmpFloatLocal, tmpUb, CAST_NONE, processLen);
        Muls(mulBufLocal, rowTmpFloatLocal, scaleVal, processLen);
        Add(sumFloatBufLocal, sumFloatBufLocal, mulBufLocal, processLen);
    }
    Cast(sumBufLocal, sumFloatBufLocal, CAST_RINT, processLen);
    DataCopy(expandOutGlobal_[tokenIndex * axisH_ + tokenOffset], sumBufLocal, processLen);
}

5. gemm 子模块深度分析

5.1 Catlass 集成

FusedDeepMoE 使用 Catlass (CUDA Template Library analogue for Ascend) 作为 GMM 计算框架。Catlass 提供：

类型化的矩阵类型: GemmType<Element, Layout>
Block 级矩阵乘法: BlockMmad<DispatchPolicy, L1Tile, L0Tile, AType, BType, CType>
Epilogue 融合: BlockEpilogue<DispatchPolicy, CType, ScaleType, PerTokenScaleType, DType, ...>
Kernel 级调度: GroupedMatmul<...>

5.2 GroupedMatmul 变体

5.2.1 GroupedMatmulSliceMPerTokenDequantSwigluQuantMultiStageWorkspace

用途: GMM1 — 完成 int8 GMM + PerTokenDequant + Swiglu + 动态量化

文件: gemm/kernel/grouped_matmul_slice_m_per_token_dequant_swiglu_quant_multistage_workspace.h

核心特点: - 1C2V 架构: AIC 做 GMM 计算，AIV 做通信 + Swiglu + 量化 - 三角色分离: - isSendCore (aivIdx % 2 == 1): 负责将量化后的 token 发送给目标 rank - isRecvCore (aivIdx % 2 == 0): 负责接收其他 rank 的 token - isCompCore: 负责 Swiglu + 动态量化计算

AIC 端 (operator()):

// 1. 先处理共享专家 (如果启用)
if constexpr (EXEC_FLAG & EXEC_FLAG_SHARED_EXPERT) {
    // 等待 AIV 完成共享专家量化
    WaitGroupTokenNumReady(shareQuantTokenStateTensor, expected);
    // 执行共享专家 GMM1
    blockMmad(gmA[gmOffsetA], layoutA, gmB[gmOffsetB], layoutB, ...);
}

// 2. 处理路由专家 (按 group 循环)
for (uint32_t groupIdx = 0; groupIdx < localExpertNum; ++groupIdx) {
    // 等待 AIV 完成该专家的 token 接收
    WaitGroupTokenNumReady(groupTokenNumStateTensor, expected);
    currentM = groupTokenNumStateTensor.GetValue(GROUP_TOKEN_COUNT);
    // 执行该专家的 GMM1
    blockMmad(gmA[gmGroupOffsetA + gmOffsetA], ...);
}

AIV 端 (operator()):

// 1. 共享专家量化 (shareQuantCoreFunc)
if (isShareQuantCore) {
    // 对每个 token 做动态量化: bf16 → int8 + per-token scale
    shareQuantCoreFunc(gmX, gmShareSmoothScales, gmShareX1, gmShareX1Scale);
    // 通知 AIC 量化完成
}

// 2. 路由专家 Dispatch (SendCoreFunc)
if (isSendCore) {
    TokenActiveMaskCal();       // 计算有效 token 数
    CalAndSendTokenCount();     // 计算每个专家的 token 数并通知对端
    SendToMoeExprt();           // 量化 + 发送到目标 rank 窗口
}

// 3. 路由专家接收 (RecvCoreFunc)
if (isRecvCore) {
    RecvCount();                // 等待所有 rank 的 token 计数到达
    for (each expert) {
        GetCumSum();            // 计算前缀和索引
        RecvToken();            // 从各 rank 窗口接收 token
    }
    // 通知 AIC 接收完成 (通过 GM flag)
}

// 4. Swiglu + 动态量化 (CompCoreFunc)
if (isCompCore) {
    CompCoreFunc();  // 调用 BlockEpilogue 做 Swiglu
}

// 5. 后处理: Swiglu 输出的动态量化
// 共享专家 Swiglu 输出量化
BlockQuant<ArchTag> blockQuant(resource, quantParams);
// 路由专家 Swiglu 输出量化
BlockQuant<ArchTag> blockQuant(resource, quantParams);

5.2.2 GroupedMatmulSliceMPerTokenDequantMultiStageWorkspace

用途: GMM2 — 完成 int8 GMM + PerTokenDequant + Combine 发送

文件: gemm/kernel/grouped_matmul_slice_m_per_token_dequant_multistage_workspace.h

核心特点: - 使用 ResidentA 策略（GMM2 的 A 矩阵是 Swiglu 输出，可常驻 L1） - Epilogue 中融合 Combine 的 AlltoAll 发送逻辑 - AIV 端直接调用 combiner->ProcessCombine() 或 combiner->AllToAllSend()

AIV 端特殊处理:

// 共享专家路径
if constexpr (EXEC_FLAG & EXEC_FLAG_SHARED_EXPERT) {
    if (AscendC::GetSubBlockIdx() == 1) {
        combiner->TPipeSet(&resource.pipe);
        combiner->ProcessCombine();  // 执行 Combine
    }
}
// 深度融合路径
else if constexpr (EXEC_FLAG & EXEC_FLAG_DEEP_FUSE) {
    if (AscendC::GetSubBlockIdx() == 0) {
        combiner->AllToAllSend();   // subblock 0: 发送
    } else {
        combiner->ReducePermute();  // subblock 1: 接收+permute
    }
}

5.3 BlockMmad — L1/L0 多级流水

文件: gemm/block/block_mmad_preload_async_with_callback_resident_a.h

// 使用 ResidentA 策略的 BlockMmad
// 关键优化: 当 A 矩阵不变时，跳过 GM→L1 的加载
bool useResidentA = (kTileCount == L1A_STAGES) && (!isFirstLoad) 
                    && (gmBlockA.GetPhyAddr() == lastGmBlockA.GetPhyAddr());
if (!useResidentA) {
    copyGmToL1A(l1ATensorList[l1AListId], gmTileA, L1A_LAYOUT, layoutTileA);
}

多级流水阶段: 1. PRELOAD_STAGES: 预加载指令提前发射 2. L1A_STAGES / L1B_STAGES: L1 上 A/B 矩阵的 double buffer 3. L0A_STAGES / L0B_STAGES: L0 上 A/B 矩阵的多级流水 4. L0C_STAGES: L0C 累加器 5. ENABLE_SHUFFLE_K: K 维度循环打乱，避免 bank conflict

5.4 同步机制

AIC→AIV 同步 (Soft Sync):

// AIC 完成一个 L1 tile 的计算后，设置 flag
aicSetFunc = {syncGmAddr + SOFT_SYNC_OFFSET, static_cast<uint8_t>(AscendC::GetBlockIdx())};
// callbackAfterFixpipe = MakeCallback(&aicSetFunc);
// AIV 等待 flag 满足条件
CheckSyncFlag(syncGmAddr + SOFT_SYNC_OFFSET, static_cast<uint8_t>(coreIdx), target);

C/V 数据空间交替 (cvDataState):

// 使用 dataState 实现双缓冲，避免读写冲突
cvDataState = FlushAndSpinValue<int32_t>(selfDataStatusTensor, aivStateGlobalCoreIdx * UB_ALIGN);
vToCFlag = (cvDataState == 0) ? V_TO_C_FLAG_1 : V_TO_C_FLAG_2;

6. epilogue 子模块

6.1 EpilogueAtlasA2PerTokenDequantSwiglu

用途: GMM1 的 Epilogue — int32→float 反量化 + Swiglu 激活

文件: epilogue/block/block_epilogue_per_token_dequant_swiglu.h

计算流程:

// 1. 从 GM 加载 int32 累加结果到 UB
copyGmToUbC(ubC, gmTileC, layoutUbC, layoutGmTileC);

// 2. int32 → float 转换
Cast(ubTmpMxN, ubC, CAST_RINT, TileShape::COUNT);

// 3. 乘以 weight scale (行广播: 每列共享一个 scale)
tileRowBroadcastMul(ubTmpMxN, ubTmpMxN, ubFp32Scale);

// 4. 乘以 per-token scale (行广播: 每行共享一个 scale)
tileBroadcastOneBlk(ubTmpMx32B, ubPerTokenScale);
tileOneBlkColumnBroadcastMul(ubTmpMxN, ubTmpMxN, ubTmpMx32B);

// 5. Swiglu 激活 (左右两半)
// 左半: x / (1 + exp(-x))  — sigmoid 门控
// 右半: 保持不变
if (isLeft) {
    Muls(ubDenominatorMxN, ubTmpMxN, -1.0f, TileShape::COUNT);
    Exp(ubDenominatorMxN, ubDenominatorMxN, TileShape::COUNT);
    Adds(ubDenominatorMxN, ubDenominatorMxN, 1.0f, TileShape::COUNT);
    Div(ubD, ubTmpMxN, ubDenominatorMxN, TileShape::COUNT);
} else {
    tileOneBlkColumnBroadcastMul(ubD, ubTmpMxN, ubTmpMx32B);
}

// 6. 写出 float 结果到 workspace
copyUbToGmD(gmTileD, ubD, layoutGmTileD, layoutUbD);

Swiglu 数学公式:

Swiglu(x) = [x_left / (1 + exp(-x_left)), x_right]
          = [x_left * sigmoid(x_left), x_right]

6.2 EpilogueAtlasA2PerTokenDequantCombine

用途: GMM2 的 Epilogue — int32→bf16/fp16 反量化 + Combine 发送

文件: epilogue/block/block_epilogue_per_token_dequant.hpp

计算流程:

// 1-4: 与 Swiglu Epilogue 相同 (int32 → float → scale → per-token scale)
Cast(ubCFp32, ubC, CAST_RINT, TileShape::COUNT);
tileRowBroadcastMul(ubMul, ubCFp32, ubFp32Scale);
tileBroadcastOneBlk(ubPerTokenScaleBrcb, ubPerTokenScale);
tileOneBlkColumnBroadcastMul(ubPerTokenMul, ubMul, ubPerTokenScaleBrcb);

// 5. float → bf16/fp16 转换
Cast(ubD, ubPerTokenMul, CAST_RINT, TileShape::COUNT);

// 6. 写出或 Combine 发送
if constexpr (EXEC_FLAG & EXEC_FLAG_DEEP_FUSE) {
    if (expertIdx == UINT32_MAX) {
        copyUbToGmD(gmTileD, ubD, layoutGmTileD, layoutUbD);  // 共享专家
    } else {
        DoCombineSend(ubD, layoutGmTileD, layoutUbD, ...);    // 路由专家: 直接发送到对端窗口
    }
}

DoCombineSend — 融合 Combine 的关键:

// 根据 epSendCount 计算每个 rank 的 token 范围
// 直接将 UB 中的结果 DataCopyPad 到对端窗口
for (uint32_t epRank = epRankStart; epRank < calcInfo.epWorldSize_; ++epRank) {
    GM_ADDR rankGM = GetWinAddrByRankId(remoteEpRank, expertIdx) + ...;
    rankWindow.SetGlobalBuffer((__gm__ ElementD *)rankGM);
    DataCopyPad(rankWindow[(itToken - prevSendCount) * axisH_ + tokenOffset],
                ubD[(itToken - startToken) * layoutUbD.stride(0)], dataCopyParams);
}

6.3 Tile 级操作

TileStrideBinary (tile_stride_binary.h): - 支持 Mul/Div 操作，带自定义 stride - 使用 BinaryRepeatParams 实现跨行跨列的灵活访问

TileStrideMuls (tile_stride_muls.h): - 标量乘法，支持不同 DstTileShape 和 SrcTileShape - 用于 per-token scale 广播

7. Swiglu + Quantization 融合机制

7.1 完整数据流

GMM1 输入: int8 → 累加: int32
     ↓
PerTokenDequant: int32 × weight_scale × per_token_scale → float32
     ↓
Swiglu: [x_left * sigmoid(x_left), x_right] → float32
     ↓
动态量化: float32 → int8 + per_token_scale → workspace
     ↓
GMM2 输入: int8 → 累加: int32
     ↓
PerTokenDequant: int32 × weight_scale × per_token_scale → float32
     ↓
类型转换: float32 → bf16/fp16
     ↓
Combine 发送: 直接到对端窗口

7.2 BlockQuant — 动态量化实现

文件: gemm/kernel/post_swiglu_dynamic_quant.h

// BlockQuant 完成 Swiglu 输出的动态量化
// 输入: float32 Swiglu 输出 [token_count, n]
// 输出: int8 [token_count, n/2] + per_token_scale [token_count]

// 关键步骤:
// 1. 计算 Swiglu 右半部分与左半部分相乘 (完成 Swiglu)
Mul(ubInput, ubInput, ubInputRightHalf, tileCount);

// 2. 取绝对值
Abs(ubAbs, ubInput, tileCount);

// 3. 对每行取 max (分半处理)
for (uint32_t rowIdx = 0; rowIdx < tileRow; ++rowIdx) {
    Max(ubMax[rowIdx * halfTileColumn], ubAbs[rowIdx * tileColumn],
        ubAbs[rowIdx * tileColumn + halfTileColumn], halfTileColumn);
}

// 4. 计算量化 scale = max / 127
Muls(ubDequantScale, ubReduceMax, 1.0f / 127.0f, tileRow);

// 5. 量化: x / scale → int8
Muls(ubInputTmp, ubInput, 127.f, tileCount);
for (uint32_t rowIdx = 0; rowIdx < tileRow; ++rowIdx) {
    Muls(ubQuantF32[rowIdx * tileColumn], ubInputTmp[rowIdx * tileColumn],
          1.f / ubReduceMax.GetValue(rowIdx), tileColumn);
}
Cast(ubQuantS32, ubQuantF32, CAST_RINT, tileCount);
Cast(ubQuantF16, ubQuantS32, CAST_RINT, tileCount);
Cast(ubOutput, ubQuantF16, CAST_RINT, tileCount);

7.3 量化参数传递

Per-token scale 的存储:

// 每个 token 附带 512B 的额外空间存储 scale
constexpr uint32_t TOKEN_EXTRA_SPACE = 512;
scaleParamPad = TOKEN_EXTRA_SPACE;
hCommuSize = hOutSize + scaleParamPad;

通信时的 scale 传递:

// 发送端: 将 scale 附加在 token 后面
yFp32Tensor.SetValue(tokenLength / sizeof(float), float(1.0) / dynamicQuantScale);
DataCopy(dstWinGMTensor, yInt8Tensor[index], tokenLength);  // token
DataCopy(dstWinGMTensor[tokenLength], yInt8Tensor[index][tokenLength], scaleParamPad);  // scale

// 接收端: 从 token 后面读取 scale
DataCopyPad(dynamicScalesOutGMTensor_[beginIdx + j],
            xOutFp32Tensor_[tokenLength / sizeof(float)], dataCopyParamsFloat);

8. 性能优化关键点

8.1 内存带宽优化

优化点	收益	实现方式
中间数据片上消除	减少 8 次 GM 访问	Workspace 复用 + 流水线传递
X1/X2 Scale 压缩	减少 1/3 通信量	Per-token scale 替代 per-channel
权重 NZ 格式	提升 Cube 效率	layout::zN 布局
L1 常驻 A 矩阵	减少 GM→L1 加载	ResidentA Dispatch Policy

8.2 计算通信重叠

1C2V 架构的天然优势: - AIC 在做 GMM 计算时，AIV 可以并行做： - 下一组专家的 token 接收 - 当前组 Swiglu 输出的量化 - Combine 的 AlltoAll 发送

流水线示意:

时间 →
AIC: [GMM1_Expert0] [GMM1_Expert1] [GMM1_Expert2] [GMM2_Expert0] ...
AIV: [Dispatch_Send] [Recv_Expert0] [Swiglu_Quant_0] [Combine_Send_0] ...

8.3 同步优化

Soft Sync 替代 Hard Sync:

// 使用 GM 上的 flag 做轻量级同步
// 避免 PipeBarrier<PIPE_ALL>() 的全局阻塞
CheckSyncFlag(syncGmAddr + SOFT_SYNC_OFFSET, coreIdx, target);

CV 双缓冲:

// dataState 在 0/1 之间交替
// C 端写 state0 时 V 端读 state1，反之亦然
winDataSizeOffset = dataState_ * epWorldSize_ * expertPerSizeOnWin_ * moeExpertNumPerRank_;

8.4 循环优化

AIV Dispatch 优化 (enableAivOpt): - 小 batch (≤64) 时，用 table 法替代逐 token 计数 - 将 expert_id 映射到 (bs+1) × expert_num 的 table - 用前缀和指令快速计算 expand_idx

K 维度 Shuffle:

// 打乱 K 维度的循环顺序，避免 bank conflict
if constexpr (ENABLE_SHUFFLE_K) {
    startTileIdx = AscendC::GetBlockIdx() % kTileCount;
}

8.5 Workspace 内存复用

#ifdef ENABLE_REUSE_MEMORY
// X1 和 X2 共享空间 (生命周期不重叠)
shareX1TokenOffset == shareX2TokenOffset
// Scale 空间复用
shareX1ScaleOffset == shareX2ScaleOffset
// SwapSpace 和 Swiglu 输出复用
swapSpaceOffset == shareSwigluOffset
#endif

9. 与标准 MoE (DeepSeek-style) 的对应关系

9.1 数学对应

FusedDeepMoe 实现了 DeepSeek-V3 风格的 MoE Decode:

标准 MoE 计算:
  1. Routing: expert_ids = TopK(router(x), k)
  2. Dispatch: 将 token 发送到对应 expert 的 rank
  3. GMM1: y = x @ W1  (int8 量化)
  4. Swiglu: y = Swiglu(y)  (GLU 变体)
  5. GMM2: z = y @ W2  (int8 量化)
  6. Combine: output = Σ(expert_scale_i * z_i)

FusedDeepMoE 对应:
  1. expert_ids 由外部 router 计算后传入
  2. Dispatch → CamMoeDistributeDispatch (融合在 kernel 内)
  3. GMM1 → GroupedMatmulSliceMPerTokenDequantSwigluQuantMultiStageWorkspace
  4. Swiglu → BlockEpilogue 中的 Swiglu 激活
  5. GMM2 → GroupedMatmulSliceMPerTokenDequantMultiStageWorkspace
  6. Combine → CamMoeDistributeCombine (融合在 kernel 内)

9.2 量化策略对应

标准 MoE	FusedDeepMoE	说明
BF16 输入	Dynamic Quant → int8	Dispatch 阶段完成
W1 int8 per-channel	W1 int8 + per-channel scale	GMM1 权重
Swiglu 输出 BF16	Dynamic Quant → int8	Swiglu 后立即量化
W2 int8 per-channel	W2 int8 + per-channel scale	GMM2 权重
输出 BF16	PerTokenDequant → BF16	GMM2 Epilogue

9.3 与 torch_npu 标准算子对比

标准小算子组合 (SmallOps):

# 5 次 kernel launch
expand_x, scales, assist_info, token_nums, ep_send_counts, _, expand_scales = \
    torch_npu.npu_moe_distribute_dispatch_v2(...)
y1_int32 = torch_npu.npu_grouped_matmul(x=[expand_x], weight=[gmm1_weight], ...)
y1, y1_scale = torch_npu.npu_dequant_swiglu_quant(x=y1_int32, ...)
y2 = torch_npu.npu_grouped_matmul(x=[y1], weight=[gmm2_weight], ...)
output = torch_npu.npu_moe_distribute_combine_v2(expand_x=y2, ...)

FusedDeepMoE (FusionOp):

# 1 次 kernel launch
output, share_output, expert_token_nums = torch.ops.umdk_cam_op_lib.fused_deep_moe(
    x, expert_ids, gmm1_weight, gmm1_weight_scale, gmm2_weight, gmm2_weight_scale,
    expert_scales, share_gmm1_weight, ..., group_ep=..., ep_rank_size=..., ...)

9.4 支持的配置范围

参数	范围	说明
batchSize	0-256	受 AIV 循环优化限制
tokenLength (h)	512-7168	常见 7168 (DeepSeek-V3)
gmm1HLen	1048-6144, 对齐 1024	中间维度
moeExpertNum	1-512, 整除 epRankSize	总专家数
topK	1-12	路由 topK
epRankSize	任意	EP 并行度

9.5 执行路径总结

                        ┌─────────────────────────────────────┐
                        │         FusedDeepMoe Entry          │
                        └──────────────┬──────────────────────┘
                                       │
                    ┌──────────────────┼──────────────────┐
                    │                  │                  │
              moeExpertNumPerRank==1   │   moeExpertNumPerRank>1
              (非深度融合)              │   (深度融合)
                    │                  │                  │
                    ▼                  │                  ▼
          独立 Dispatch kernel         │         GMM1 内完成 Dispatch
                    │                  │                  │
                    └──────────────────┼──────────────────┘
                                       │
                                       ▼
                        ┌──────────────────────────────┐
                        │  GmmDeqSwigluQuant (GMM1)    │
                        │  AIC: int8 GMM → int32       │
                        │  AIV: Swiglu + Dynamic Quant │
                        │  → workspace (int8 + scale)  │
                        └──────────────┬───────────────┘
                                       │
                                       ▼
                        ┌──────────────────────────────┐
                        │  GmmDeq (GMM2 + Combine)     │
                        │  AIC: int8 GMM → int32       │
                        │  AIV: PerTokenDequant +      │
                        │       Combine AlltoAll Send  │
                        └──────────────────────────────┘

附录：关键文件索引

文件	功能
`op_kernel/fused_deep_moe.h`	主 kernel 类定义，Process() 入口
`op_kernel/fused_deep_moe.cpp`	kernel 入口，模板实例化
`op_kernel/fused_deep_moe_tiling.h`	Tiling 数据结构，Tile Shape 常量
`op_kernel/fused_deep_moe_tiling_key.h`	Tiling Key 路由机制
`op_host/fused_deep_moe_tiling.cpp`	Host 端 tiling 计算
`op_host/fused_deep_moe.cpp`	算子定义 (输入/输出/属性)
`op_host/fused_deep_moe_infer.cpp`	Shape/DataType 推断
`raw_distributed/cam_moe_distribute_dispatch.h`	Dispatch 实现
`raw_distributed/cam_moe_distribute_combine.h`	Combine 实现
`gemm/dispatch_policy.h`	GMM Dispatch Policy 定义
`gemm/kernel/grouped_matmul_slice_m_per_token_dequant_swiglu_quant_multistage_workspace.h`	GMM1 完整实现
`gemm/kernel/grouped_matmul_slice_m_per_token_dequant_multistage_workspace.h`	GMM2 完整实现
`gemm/kernel/post_swiglu_dynamic_quant.h`	Swiglu 后动态量化
`gemm/block/block_mmad.h`	BlockMmad 入口
`gemm/block/block_mmad_preload_async_with_callback_resident_a.h`	ResidentA BlockMmad
`epilogue/dispatch_policy.h`	Epilogue Dispatch Policy
`epilogue/block/block_epilogue.h`	BlockEpilogue 入口
`epilogue/block/block_epilogue_per_token_dequant_swiglu.h`	Swiglu Epilogue
`epilogue/block/block_epilogue_per_token_dequant.hpp`	Combine Epilogue
`epilogue/tile/tile_stride_binary.h`	Tile 级 Mul/Div
`epilogue/tile/tile_stride_muls.h`	Tile 级 Muls
`fused_deep_moe_utils.h`	同步工具函数
`op_api/aclnn_fused_deep_moe.cpp`	aclnn API 入口
`pybind/fused_deep_moe.cpp`	PyTorch 绑定
`examples/fused_deep_moe_sample.py`	使用示例

总结

FusedDeepMoE 是 CAM 库的皇冠级算子，其核心价值在于：

极致融合: 将 5 个独立算子 (Dispatch + GMM1 + Swiglu + GMM2 + Combine) 融合为 1 个 kernel
中间数据全片上化: X1、SwigluOut、X2 等中间数据在 workspace 中流转，零 GM 访问
1C2V 异构并行: AIC 做 GMM 计算，AIV 做通信调度，天然计算通信重叠
全链路量化: 从输入动态量化到 Swiglu 后动态量化，全程 int8 计算
精细同步: Soft Sync + CV 双缓冲 + 多级 Flag 同步，最小化同步开销

这是华为昇腾 Ascend 910B 平台上 MoE Decode 阶段的终极优化实现，代表了 NPU 上 MoE 算子融合的最高水平。

← EP 通信算子深度分析 📑 索引测试体系与回归测试 →

FusedDeepMoE 融合算子 — UMDK/CAM