FusedDeepMoE 通信模式 & SHMEM 编译机制深度剖析

代码深读报告 | 聚焦 FusedDeepMoE 的 SHMEM vs 非 SHMEM 通信路径、raw_distributed 通信原语、三条执行路径差异、Workspace 内存布局，以及编译系统的 SHMEM 注入机制。

1. 总体架构概览

FusedDeepMoE 是 CAM 中最复杂的融合算子，将 Dispatch → GMM1 → SwiGLU → GMM2 → Combine 五个阶段融合进单一 kernel。其核心挑战在于：

  • 通信与计算的融合：dispatch/combine 需要跨 EP rank 搬运 token，同时 GMM 需要密集 Cube 计算
  • 三条执行路径：非深度融合 / 深度融合 / 共享专家，各自通信模式不同
  • SHMEM vs HCCL Window：当前默认使用 HCCL Window 通信，SHMEM 路径为未来预留
┌─────────────────────────────────────────────────────────────────────┐
│                     FusedDeepMoE 融合流程                           │
│                                                                     │
│  x ──→ Dispatch ──→ GMM1 ──→ SwiGLU ──→ GMM2 ──→ Combine ──→ out  │
│       (HCCL Window)  (Cube)   (Vector)   (Cube)  (HCCL Window)     │
│                                                                     │
│  shared_x ──→ ShareGMM1 ──→ SwiGLU ──→ ShareGMM2 ──→ share_out    │
│               (Cube)        (Vector)   (Cube)                       │
└─────────────────────────────────────────────────────────────────────┘

2. FusedDeepMoE 通信模式实现

2.1 核心入口：fused_deep_moe.h

文件：op_kernel/fused_deep_moe.h

通信完全通过 raw_distributed 子模块实现，主 kernel 类 FusedDeepMoe 本身不包含任何通信逻辑：

// fused_deep_moe.h 第30行
#include "fused_deep_moe/raw_distributed/cam_moe_distribute_dispatch.h"
// 注意：combine 的 include 在 GMM2 kernel 内部（通过 combiner 指针）

Process() 函数揭示的三阶段执行模型（第354-448行）：

template <TemplateMC2TypeClass>
__aicore__ inline void FusedDeepMoe<TemplateMC2TypeFunc>::Process()
{
    // 阶段1: Dispatch（仅非深度融合路径）
    if constexpr ((EXEC_FLAG & EXEC_FLAG_DEEP_FUSE) == 0) {
        MoeDistributeDispatchImpl::CamMoeDistributeDispatch<...> dispatcher;
        dispatcher.Init(...);
        dispatcher.Process();  // ← HCCL Window AlltoAll
    }

    // 阶段2: GMM1 + SwiGLU（Cube 计算，无通信）
    GmmDeqSwigluQuant<...>(...);

    // 阶段3: GMM2 + Combine（融合通信计算）
    MoeDistributeCombineImpl::CamMoeDistributeCombine<...> combiner;
    GmmDeq<...>(..., &combiner);  // ← combiner 在 GMM2 的 AIV 侧执行
}

2.2 通信上下文获取：HCCL Window 模式

关键发现：当前代码中没有任何 SHMEM 分支。 所有通信通过 HcclOpResParam 结构体（HCCL Window 描述符）实现：

// cam_moe_distribute_dispatch.h 第236-237行
winContext_[COMM_EP_IDX] = (__gm__ HcclOpResParam *)AscendC::GetHcclContext<HCCL_GROUP_ID_0>();
winContext_[COMM_TP_IDX] = (__gm__ HcclOpResParam *)AscendC::GetHcclContext<1>();

// cam_moe_distribute_combine.h 第245-246行
auto contextGM0 = AscendC::GetHcclContext<HCCL_GROUP_ID_0>();
epWinContext_ = (__gm__ HcclOpResParam *)contextGM0;

HCCL Window 地址解析（dispatch 第83-91行）：

__aicore__ inline GM_ADDR GetWindAddrByRankId(uint8_t ctxIdx, const int32_t rankId)
{
    uint32_t curRankId = ctxIdx == COMM_EP_IDX ? epRankId_ : tpRankId_;
    if (curRankId == rankId) {
        // 本地 Window：直接访问 localWindowsIn
        return (GM_ADDR)(winContext_[ctxIdx]->localWindowsIn) + winDataSizeOffset_ + rankId * OPT_RANK_OFFSET;
    }
    // 远端 Window：通过 remoteRes 指针访问对端 windowsIn
    return (GM_ADDR)(((HcclRankRelationResV2 *)(winContext_[ctxIdx]->remoteRes[rankId].nextDevicePtr))->windowsIn)
           + winDataSizeOffset_ + rankId * OPT_RANK_OFFSET;
}

通信模式总结： - 数据通道：windowsIn — 存放实际 token 数据（int8/half/bf16） - 状态通道：windowsExp — 存放同步状态标志（flag/sumTarget） - 同步机制：通过 DataCopy 在 Window 上写入 flag，轮询对端 flag 实现握手

3. SHMEM vs 非 SHMEM 代码路径

3.1 当前状态：无 SHMEM 分支

对整个 fused_deep_moe/ 目录进行全量 grep：

grep -rn "aclshmem\|SHMEM\|shmem_" src/cam/comm_operator/ascend_kernels/fused_deep_moe/
# 结果：0 命中

结论：FusedDeepMoE kernel 代码中不存在任何 SHMEM 相关代码路径。 所有通信均通过 HCCL Window（HcclOpResParam）实现。

3.2 编译系统的 SHMEM 预留

虽然 kernel 代码没有 SHMEM 分支，但编译系统已完整预留 SHMEM 注入点：

层级 文件 SHMEM 机制
算子选择 select_ops.py requires_shmem() 过滤
编译脚本 compile_ascend_proj.sh SHMEM_HOME_PATH 检测 + include 注入
CMake 工具 ascendc_impl_build.py SHMEM include 路径硬编码注入

3.3 抽象层设计

当前通信原语通过 HcclOpResParam 抽象，没有引入统一的 CommunicationInterface。dispatch/combine 直接依赖 HCCL 的 Window 结构体：

// 通信上下文是 HCCL 强绑定的
__gm__ HcclOpResParam *winContext_[COMM_NUM];  // dispatch
__gm__ HcclOpResParam *epWinContext_;           // combine
__gm__ HcclOpResParam *tpWinContext_;           // combine (TP domain)

如果要支持 SHMEM，需要： 1. 引入 ICommContext 抽象基类 2. 派生 HcclWindowContext 和 ShmemContext 3. dispatch/combine 通过基类指针访问通信原语

4. raw_distributed 子模块：通信原语详解

4.1 cam_moe_distribute_dispatch.h

文件：op_kernel/fused_deep_moe/raw_distributed/cam_moe_distribute_dispatch.h

类模板参数（第19-24行 fused_deep_moe_base.h）：

#define TemplateDispatchTypeClass
    typename XType, typename ExpandXOutType, bool StaticQuant,
    bool DynamicQuant, bool IsSmoothScaleExist, bool IsNeedAllgater, uint32_t EXEC_FLAG

核心通信流程（第1023-1039行）：

void Process()
{
    if ASCEND_IS_AIV {
        if constexpr (EXEC_FLAG & EXEC_FLAG_SHARED_EXPERT) {
            QuantForShareExpert();     // 共享专家量化（本地操作）
        }
        AlltoAllDispatch();            // ← 核心：AlltoAll 数据搬运
        SetStatus();                   // 写入本 rank 状态到各远端 Window
        WaitDispatch();                // 轮询等待所有远端状态就绪
        LocalWindowCopy();             // 从本地 Window 读取数据到 workspace
        if constexpr (IsNeedAllgater) {
            AllGatherSetStatusAndWait();
            AllgatherProcessOut();
        }
        UpdataTokenNumsOut();          // 更新 expertTokenNums 计数
    }
}

AlltoAllDispatch 关键路径（第575-602行）：

void AlltoAllDispatch()
{
    // 1. 加载 expertIds（决定每个 token 去哪个 expert/rank）
    DataCopyPad(expertIdsTensor_, expertIdsGMTensor_, expertIdsCntParams, copyPadParams);

    // 2. 计算每个 expert 的 token 计数 + 状态更新
    if (enableAivOpt_) {
        CountAndPrepareStatusWithAivOpt(expertIdsCnt);  // 优化路径
    } else {
        for (uint32_t tokenIndex = 0; tokenIndex < expertIdsCnt; ++tokenIndex) {
            expertCountTensor_(tokenIndex) = statusTensor_(expertId * INT32_NUM_PER_BLOCK + 1);
            statusTensor_(expertId * INT32_NUM_PER_BLOCK + 1)++;  // 原子计数
        }
    }

    // 3. 将 token 数据写入目标 rank 的 Window
    SendToMoeExpert();
}

SendToMoeExpert 数据搬运（第433-487行）：

void SendToMoeExpert()
{
    for (uint32_t tokenIndex = startTokenId; tokenIndex < endTokenId; ++tokenIndex) {
        int32_t dstExpertId = expertIdsTensor_(tokenIndex);
        uint32_t tempRankId = dstExpertId / moeExpertNumPerRank_;

        // 计算目标 Window 地址
        GM_ADDR rankGM = GetWindAddrByRankId(COMM_EP_IDX, tempRankId)
                         + expertPerSizeOnWin_ * (epRankId_ * moeExpertNumPerRank_ + dstExpertId % moeExpertNumPerRank_)
                         + hCommuSize_ * expertCountTensor_(tokenIndex);

        // 量化 + 写入目标 Window
        if constexpr (DynamicQuant || StaticQuant) {
            QuantProcess(expertIndex, scalesGMTensor_);
            DataCopy(dstWinGMTensor, xOutTensor_, axisHCommu_);
        } else {
            DataCopy(dstWinGMTensor, xTmpTensor_, axisHCommu_);
        }
    }
}

WaitDispatch 同步机制（第684-726行）：

void WaitDispatch()
{
    // 轮询 windowInstatusFp32Tensor_ 中的 flag
    while ((sumOfFlag < minTarget) || (sumOfFlag > maxTarget)) {
        DataCopy(statusFp32Tensor_, windowInstatusFp32Tensor_[startStatusIndex * stateOffset_ / sizeof(float)], intriParams);
        GatherMask(gatherMaskOutTensor, statusFp32Tensor_, gatherTmpTensor, true, mask, ...);
        Sum(statusSumOutTensor, gatherMaskOutTensor, sumParams);
        sumOfFlag = statusSumOutTensor.GetValue(0);
    }
    SyncAll<true>;  // 全核同步
}

4.2 cam_moe_distribute_combine.h

文件：op_kernel/fused_deep_moe/raw_distributed/cam_moe_distribute_combine.h

类模板参数（与主 kernel 共享 TemplateMC2TypeClass）：

#define TemplateMC2TypeClass
    typename ExpandXType, typename W1ScaleType, typename W2ScaleType,
    typename ExpandIdxType, bool IsNeedReduceScatter, uint32_t EXEC_FLAG

CombineCalcInfo 结构体（第46-56行）— 传递给 GMM2 AIV 侧的关键信息：

struct CombineCalcInfo {
    uint64_t expertPerSizeOnWin_;   // 每个 expert 在 Window 上占用的字节数
    uint32_t epRankId_;
    uint32_t epWorldSize_;
    uint32_t moeExpertPerRankNum_;
    uint32_t axisH_;
    uint32_t moeSendNum_;
    GM_ADDR epSendCount_;
    __gm__ HcclOpResParam *epWinContext_;  // ← HCCL 强绑定
    uint64_t winDataSizeOffset_;
};

Process 流程（第719-734行）：

void Process()
{
    SyncAll<true>;
    if constexpr (IsNeedReduceScatter) {
        ReduceScatterTrans();          // TP 域 ReduceScatter
    }
    if constexpr ((EXEC_FLAG & EXEC_FLAG_DEEP_FUSE) == 0) {
        BuffInit();
        SetWaitTpStatusAndDisPatch();  // 非融合路径的 AlltoAll
    }
    AlltoAllBuffInit();
    SetStatus();                       // 写入 EP 状态
    WaitDispatch();                    // 等待 EP 域数据就绪
    LocalWindowCopy();                 // 从 Window 读取 + 加权求和
}

LocalWindowCopy 核心逻辑（第627-716行）：

void LocalWindowCopy()
{
    for (uint32_t tokenIndex = beginIndex; tokenIndex < endIndex; tokenIndex++) {
        Duplicate(sumFloatBufLocal, (float)0, axisH_);  // 清零累加器
        for (uint32_t i = 0; i < axisK_; i++) {
            int32_t moeExpert = expertIdsLocal.GetValue(index);
            float scaleVal = expandScalesLocal.GetValue(index);

            // 从 Window 读取该 expert 的 token 数据
            GM_ADDR wAddr = epWindowGM_ + expertPerSizeOnWin_ * moeExpert
                            + indexCountsLocal.GetValue(index) * axisHExpandXTypeSize_ + tokenOffset;
            rowTmpGlobal_.SetGlobalBuffer((__gm__ ExpandXType *)wAddr);

            // 量化反量化 + 缩放 + 累加
            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. 三条执行路径的通信差异

5.1 路径判定机制

EXEC_FLAG 位图定义（fused_deep_moe_tiling.h 第104-108行）：

constexpr uint32_t EXEC_FLAG_DEEP_FUSE     = (1U << 0);  // 深度融合
constexpr uint32_t EXEC_FLAG_TENSOR_LIST   = (1U << 1);  // 权重 tensor list
constexpr uint32_t EXEC_FLAG_X_ACTIVE_MASK = (1U << 2);  // 激活 mask
constexpr uint32_t EXEC_FLAG_SHARED_EXPERT = (1U << 3);  // 共享专家
constexpr uint32_t EXEC_FLAG_SMOOTH_QUANT  = (1U << 4);  // 平滑量化

Tiling 阶段计算 EXEC_FLAG（fused_deep_moe_tiling.cpp 第35-40行）：

constexpr uint32_t EXEC_FLAG =
    (TPL_IS_DEEP_FUSE     ? EXEC_FLAG_DEEP_FUSE     : 0) |
    (TPL_IS_TENSOR_LIST   ? EXEC_FLAG_TENSOR_LIST   : 0) |
    (TPL_IS_X_ACTIVE_MASK ? EXEC_FLAG_X_ACTIVE_MASK : 0) |
    (TPL_IS_SHARED_EXPERT ? EXEC_FLAG_SHARED_EXPERT : 0) |
    (TPL_IS_SMOOTH_QUANT  ? EXEC_FLAG_SMOOTH_QUANT  : 0);

5.2 路径对比

特性 非深度融合 深度融合 (DEEP_FUSE) 共享专家 (SHARED_EXPERT)
触发条件 moeExpertNumPerRank == 1 moeExpertNumPerRank > 1 share_gmm1_weight != nullptr
Dispatch 独立执行（AIV 上） 跳过（由 GMM2 AIV 侧融合） 独立执行
Combine 独立执行（AIV 上） 融合进 GMM2 AIV 侧 融合进 GMM2 AIV 侧
GMM1 后同步 PipeBarrier + CrossCoreFlag 同左 同左
GMM2 AIV 角色 仅做 epilogue 执行 AllToAllSend + ReducePermute 执行 ProcessCombine
共享专家 GMM 无 无 有（额外 GMM1+GMM2）
核分配 AIC + AIV 1:1 AIC + AIV 1:2 (1C2V) AIC + AIV 1:2

5.3 深度融合路径的通信融合

GMM2 AIV 侧执行流程（grouped_matmul_slice_m_per_token_dequant_multistage_workspace.h 第291-458行）：

template <>
CATLASS_DEVICE void operator()<AscendC::AIV>(Params const &params)
{
    auto *combiner = (CamMoeDistributeCombine *)params.combiner;

    // 1. 先执行 GMM2 的 epilogue（反量化 + SwiGLU 后处理）
    for (uint32_t groupIdx = 0; groupIdx < params.problemCount; ++groupIdx) {
        for (uint32_t loopIdx = startLoopIdx; loopIdx < coreLoops; loopIdx += coreNum) {
            CheckSyncFlag(syncGmAddr + SOFT_SYNC_OFFSET, coreIdx, target);  // 等待 AIC 完成
            blockEpilogue(gmGroupOffsetD, groupIdx, ...);  // epilogue 计算
        }
    }

    // 2. 共享专家 GMM2（如有）
    if constexpr (EXEC_FLAG & EXEC_FLAG_SHARED_EXPERT) {
        // ... 执行共享专家 GMM2 epilogue
    }

    // 3. 融合 Combine 通信
    if constexpr (EXEC_FLAG & EXEC_FLAG_SHARED_EXPERT) {
        if (GetSubBlockIdx() == 1) {
            combiner->ProcessCombine();  // ← 在 GMM2 AIV 上执行 Combine
        }
    } else if constexpr (EXEC_FLAG & EXEC_FLAG_DEEP_FUSE) {
        if (GetSubBlockIdx() == 0) {
            combiner->AllToAllSend();    // ← 发送阶段
        } else {
            combiner->ReducePermute();   // ← 接收+累加阶段
        }
    } else {
        combiner->Process();             // ← 非融合：完整 Combine
    }
}

关键同步机制： - SEND_SYNC_EVENT_ID = 6 — GMM2 epilogue 完成 → Combine 开始 - RECV_SYNC_EVENT_ID = 7 — Combine 接收完成 → 主 kernel 继续 - CheckSyncFlag — AIV 等待 AIC 的 Cube 计算完成（soft sync）

6. Workspace 内存布局

6.1 WorkSpaceOffset 结构体

文件：fused_deep_moe_tiling.h 第19-40行

struct WorkSpaceOffset {
    // MM1/GMM1-Swiglu input
    int64_t shareX1TokenOffset;     // 共享专家 GMM1 输入
    int64_t x1TokenOffset;          // MoE GMM1 输入（dispatch 后）
    int64_t shareX1ScaleOffset;     // 共享专家量化 scale
    int64_t x1ScaleOffset;          // MoE 量化 scale

    // MM1/GMM1-Swiglu output
    int64_t shareSwigluOffset;      // 共享专家 SwiGLU 输出
    int64_t swigluOffset;           // MoE SwiGLU 输出

    // MM2/GMM2 input
    int64_t shareX2TokenOffset;     // 共享专家 GMM2 输入
    int64_t x2TokenOffset;          // MoE GMM2 输入
    int64_t shareX2ScaleOffset;     // 共享专家 GMM2 scale
    int64_t x2ScaleOffset;          // MoE GMM2 scale

    int64_t swapSpaceOffset;        // C↔V 数据交换区（GMM1 输出/GMM2 输入）
    int64_t y2TokenOffset;          // 浅融合路径 GMM2 输出
    int64_t groupListOffset;        // expert token 计数前缀和
    int64_t expandIdxOffset;        // dispatch 远端 token 索引
    int64_t epSendCountOffset;      // 每个 expert 从每个 rank 接收的 token 数
    int64_t reservedOffset;         // 保留区
};

6.2 内存复用策略（ENABLE_REUSE_MEMORY）

文件：fused_deep_moe_tiling.cpp 第671-711行

#ifdef ENABLE_REUSE_MEMORY
    // shareX1 和 shareX2 复用同一块内存（因为 SwiGLU 后不再需要 X1）
    tilingData.workSpaceOffset.shareX1TokenOffset = offset;
    tilingData.workSpaceOffset.shareX2TokenOffset = offset;
    tilingData.workSpaceOffset.x1TokenOffset = offset + shareExpertTokenNum * h * sizeof(int8_t);
    tilingData.workSpaceOffset.x2TokenOffset = offset + shareExpertTokenNum * shareGmm2HLen * sizeof(int8_t);
    offset += maxTokenSize;  // 取 x1/x2 中较大者

    // scale 同样复用
    tilingData.workSpaceOffset.shareX1ScaleOffset = offset;
    tilingData.workSpaceOffset.shareX2ScaleOffset = offset;
    tilingData.workSpaceOffset.x1ScaleOffset = offset + shareExpertTokenNum * sizeof(float);
    tilingData.workSpaceOffset.x2ScaleOffset = offset + shareExpertTokenNum * sizeof(float);
    offset += tokenScaleSize;

    // swapSpace 和 shareSwiglu 复用
    tilingData.workSpaceOffset.swapSpaceOffset = offset;
    tilingData.workSpaceOffset.shareSwigluOffset = offset;
    tilingData.workSpaceOffset.swigluOffset = offset + shareExpertTokenNum * shareGmm1HLen * sizeof(float);
    offset += maxSwapSwigluSize;
    tilingData.workSpaceOffset.y2TokenOffset = offset;
#else
    // 非复用模式：各区域独立分配
    // ... 每个区域单独 CeilUp 对齐
#endif

6.3 Workspace 总大小计算

// 第716-717行
size_t usrSize = tilingData.workSpaceOffset.reservedOffset + reservedSize;
workSpaces[0] = SYSTEM_NEED_WORKSPACE + usrSize;  // 16MB 系统预留 + 用户区

各区域大小公式：

区域 大小公式 说明
x1/x2 token max(shareBs*h, maxTokenNum*h) int8 存储
scale (shareBs + maxTokenNum) * sizeof(float) per-token scale
swap/swiglu max(maxTokenNum*gmm1HLen, maxTokenNum*h) * sizeof(int32/float) GMM1 输出
y2 maxTokenNum * h * 2 (if moeExpertNumPerRank==1) GMM2 输出
groupList moeExpertNumPerRank * sizeof(int64_t) 前缀和
expandIdx bs * topK * sizeof(int32_t) 远端索引
epSendCount epRankSize * moeExpertNumPerRank * sizeof(int32_t) 接收计数
reserved 256KB 保留

6.4 数据流转图

HBM (GM)
  │
  ├── x [bs, h] ─────────────────────────────────────────────────────────────┐
  ├── expert_ids [bs, k] ─────────────────────────────────────────────────────┤
  ├── gmm1_weight [moeExpertNumPerRank, h, gmm1HLen] ─────────────────────────┤
  ├── gmm2_weight [moeExpertNumPerRank, gmm1HLen/2, h] ───────────────────────┤
  └── workspace ──────────────────────────────────────────────────────────────┘
       │
       ├── shareX1Token (Window Copy ← Dispatch)
       │     ↓ GMM1 (Cube) + SwiGLU (Vector)
       ├── shareSwigluOut
       │     ↓
       ├── shareX2Token ──────────────────────────→ ShareGMM2 → share_output
       │
       ├── x1Token (Window Copy ← Dispatch)
       │     ↓ GMM1 (Cube) + SwiGLU (Vector)
       ├── swigluOut
       │     ↓
       ├── x2Token
       │     ↓ GMM2 (Cube, 输出到 swapSpace)
       │     ↓ epilogue (Vector, 反量化 + 写回 Window)
       │     ↓ Combine (AllToAllSend → Window → ReducePermute → output)
       │
       ├── swapSpace (GMM2 AIC 输出 → AIV epilogue 输入)
       ├── y2Token (浅融合路径 GMM2 最终输出)
       ├── groupList (expert token 前缀和)
       ├── expandIdx (远端 token 索引)
       ├── epSendCount (各 rank 接收计数)
       └── reserved

7. 编译系统 SHMEM 机制

7.1 select_ops.py — SHMEM 算子过滤

文件：build/cam/comm_operator/select_ops.py

requires_shmem() 函数（第55-56行）：

def requires_shmem(registry, op):
    return bool(get_meta(registry, op).get("requires_shmem", False))

SHMEM 过滤逻辑（第116-121行）：

# 3. SHMEM filter: drop requires_shmem operators when SHMEM is not installed.
if not shmem:
    dropped = [o for o in candidates if requires_shmem(registry, o)]
    for o in dropped:
        sys.stderr.write(f"NOTE: dropping {o} (requires SHMEM, which is not installed)\n")
    candidates = [o for o in candidates if not requires_shmem(registry, o)]

完整选择管线：

1. 确定候选集（全 SOC 列表 或 用户 -a 列表）
        ↓
2. fused_deep_moe base/w4a8 互斥规则
   -q → 选 w4a8，剔除 base
   非 -q → 选 base，剔除 w4a8
   fwk 不受影响
        ↓
3. SHMEM 过滤
   shmem=0 → 剔除 requires_shmem=true 的算子
        ↓
4. 按 registry 顺序排列输出

7.2 compile_ascend_proj.sh — SHMEM 路径注入

文件：build/cam/comm_operator/compile_ascend_proj.sh

SHMEM_HOME_PATH 检测（第19-22行）：

# Whether SHMEM is installed (determined by the SHMEM_HOME_PATH env variable)
shmem_installed=1
if [ -z "${SHMEM_HOME_PATH}" ]; then
    shmem_installed=0
fi

SHMEM include 注入（第179-192行）：

KERNEL_COMPILE_OPTS=""
if [ -n "${SHMEM_HOME_PATH}" ]; then
    KERNEL_COMPILE_OPTS="-I${SHMEM_HOME_PATH}/shmem/include -I${SHMEM_HOME_PATH}/shmem/src/device"
fi
if [ -n "${CATLASS_HOME_PATH}" ]; then
    KERNEL_COMPILE_OPTS="${KERNEL_COMPILE_OPTS} -I${CATLASS_HOME_PATH}/include"
fi
if [ -n "${KERNEL_COMPILE_OPTS}" ]; then
    echo "Patching kernel compile options: ${KERNEL_COMPILE_OPTS}"
    find build_out/op_kernel/CMakeFiles -name "build.make" -exec sed -i \
        "s|--compile-options=\"\"|--compile-options=\\\"${KERNEL_COMPILE_OPTS}\\\"|g" {} +
fi

关键机制：通过 sed 替换 .make 文件中的 --compile-options="" 为实际的 include 路径。这是因为 npu_op_* 构建系统的 simple_kernel_compile 不读取 custom_compile_options.ini。

7.3 ascendc_impl_build.py — SHMEM 编译选项生成

文件：src/cam/comm_operator/ascend_kernels/cmake_files/cmake/util/ascendc_impl_build.py

SHMEM include 路径硬编码（第167-179行）：

SHMEM_HOME_PATH = os.environ.get('SHMEM_HOME_PATH')
if SHMEM_HOME_PATH is None:
    SHMEM_HOME_PATH = os.path.realpath("/usr/local/Ascend/shmem/latest")

options.append("-I" + SHMEM_HOME_PATH + "/shmem/include")
options.append("-I" + SHMEM_HOME_PATH + "shmem/include/device")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/include/host_device")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/include/internal")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/src")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/src/device")
options.append("-I" + SHMEM_HOME_PATH + "/memfabric_hybrid/include/smem/device")

注意：这里有一个 bug — 第172行缺少 / 分隔符：

options.append("-I" + SHMEM_HOME_PATH + "shmem/include/device")  # 缺少 /
# 应该是：
options.append("-I" + SHMEM_HOME_PATH + "/shmem/include/device")

7.4 operator_registry.json — SHMEM 标记

文件：src/cam/comm_operator/ascend_kernels/operator_registry.json

{
  "operator_meta": {
    "dispatch_layout": { "requires_shmem": false },
    "fused_deep_moe": {
      "requires_shmem": false,
      "family": "fused_deep_moe",
      "variant": "base"
    },
    "moe_combine_normal": { "requires_shmem": false },
    "moe_dispatch_normal": { "requires_shmem": false },
    "notify_dispatch": { "requires_shmem": false }
  }
}

当前所有算子的 requires_shmem 均为 false，即 SHMEM 过滤当前不影响任何算子。

8. SHMEM_HOME_PATH 环境变量

8.1 作用域

阶段 使用方式 默认值
编译时 compile_ascend_proj.sh 检测是否设置 未设置 → shmem_installed=0
编译时 ascendc_impl_build.py 读取路径 /usr/local/Ascend/shmem/latest
运行时 kernel 代码中 未使用 N/A

8.2 编译时行为

SHMEM_HOME_PATH 未设置:
  ├── select_ops.py: shmem=0 → 剔除 requires_shmem=true 的算子
  ├── compile_ascend_proj.sh: 不注入 SHMEM include
  └── ascendc_impl_build.py: 使用默认路径 /usr/local/Ascend/shmem/latest

SHMEM_HOME_PATH 已设置:
  ├── select_ops.py: shmem=1 → 保留所有算子
  ├── compile_ascend_proj.sh: 注入 -I${SHMEM_HOME_PATH}/shmem/include 等
  └── ascendc_impl_build.py: 使用 SHMEM_HOME_PATH 的值

8.3 运行时行为

当前 kernel 代码不读取 SHMEM_HOME_PATH。SHMEM 的运行时行为（如 aclshmem_* API 的链接）由 CANN 的 SHMEM 库在算子加载时动态解析。

9. 未来 SHMEM 集成路径

9.1 需要修改的代码

要让 FusedDeepMoE 支持 SHMEM，需要以下改动：

9.1.1 通信抽象层（新增）

新增文件：
  op_kernel/fused_deep_moe/comm/icomm_context.h        # 抽象基类
  op_kernel/fused_deep_moe/comm/hccl_window_context.h   # HCCL 实现
  op_kernel/fused_deep_moe/comm/shmem_context.h         # SHMEM 实现

// icomm_context.h (概念设计)
class ICommContext {
public:
    virtual GM_ADDR GetDataAddr(int rankId) = 0;
    virtual GM_ADDR GetStateAddr(int rankId) = 0;
    virtual void SetFlag(int rankId, int value) = 0;
    virtual int WaitFlag(int rankId, int target) = 0;
    virtual void DataCopy(void* dst, void* src, size_t len) = 0;
};

9.1.2 dispatch/combine 改造

// cam_moe_distribute_dispatch.h 改造方向
template <typename CommContext>
class CamMoeDistributeDispatch {
    CommContext* commCtx_;  // 替代 winContext_[COMM_NUM]
    // ...
};

9.1.3 编译系统改动

  1. operator_registry.json — 新增 SHMEM 变体：
{
  "fused_deep_moe_shmem": {
    "requires_shmem": true,
    "family": "fused_deep_moe",
    "variant": "shmem"
  }
}
  1. select_ops.py — 新增 SHMEM 变体互斥规则：
SHMEM = "fused_deep_moe_shmem"
# base/w4a8/shmem 三者互斥
  1. compile_ascend_proj.sh — 已支持，无需改动（通过 SHMEM_HOME_PATH 自动注入）

9.1.4 运行时 SHMEM API 调用

SHMEM 路径的核心变化：

// HCCL Window 模式（当前）
GM_ADDR dstAddr = GetWindAddrByRankId(COMM_EP_IDX, rankId);
DataCopy(dstWinGMTensor, xOutTensor_, axisHCommu_);

// SHMEM 模式（未来）
// 使用 aclshmem_* API 直接 RDMA
aclshmem_put(dstAddr, srcAddr, len, dstRankId);
aclshmem_fence();
aclshmem_wait_flag(flagAddr, target);

9.2 SHMEM 集成挑战

挑战 说明
同步语义差异 HCCL Window 是双向对称的（localWindowsIn + remoteRes），SHMEM 是单边 RDMA
状态通道设计 当前 windowsExp 用于 flag 同步，SHMEM 需要额外设计 fence/notify 机制
多域通信 dispatch 使用 COMM_EP_IDX + COMM_TP_IDX 两个域，SHMEM 需要对应两个通信组
地址计算 expertPerSizeOnWin_、winDataSizeOffset_ 等偏移量需要重新映射到 SHMEM 内存池
量化融合 dispatch 中量化+Window 拷贝是融合的，SHMEM 路径需要保持同样的融合度

9.3 推荐集成策略

  1. Phase 1：引入 ICommContext 抽象，HCCL 路径保持不变
  2. Phase 2：实现 ShmemContext，在 notify_dispatch 等简单算子上验证
  3. Phase 3：将 SHMEM 路径扩展到 FusedDeepMoE 的 dispatch/combine
  4. Phase 4：性能调优，利用 SHMEM 的单边 RDMA 降低同步开销

10. 关键代码索引

文件 关键函数/结构 行号
fused_deep_moe.h FusedDeepMoe::Process() 354-448
fused_deep_moe.h GmmDeqSwigluQuant() 56-152
fused_deep_moe.h GmmDeq() 154-238
fused_deep_moe_tiling.h WorkSpaceOffset 19-40
fused_deep_moe_tiling.h EXEC_FLAG_* 定义 104-108
fused_deep_moe_tiling.cpp SetWorkSpace() 636-719
fused_deep_moe_tiling.cpp SetHcommCfg() 605-616
fused_deep_moe_tiling_key.h ASCENDC_TPL_SEL 34-55
cam_moe_distribute_dispatch.h GetWindAddrByRankId() 83-91
cam_moe_distribute_dispatch.h AlltoAllDispatch() 575-602
cam_moe_distribute_dispatch.h SendToMoeExpert() 433-487
cam_moe_distribute_dispatch.h WaitDispatch() 684-726
cam_moe_distribute_combine.h CombineCalcInfo 46-56
cam_moe_distribute_combine.h GetWinAddrByRankId() 96-111
cam_moe_distribute_combine.h LocalWindowCopy() 627-716
cam_moe_distribute_combine.h Process() 719-734
grouped_matmul_slice_m_per_token_dequant_multistage_workspace.h operator()<AIV>() 291-458
select_ops.py requires_shmem() 55-56
select_ops.py SHMEM 过滤逻辑 116-121
compile_ascend_proj.sh SHMEM_HOME_PATH 检测 19-22
compile_ascend_proj.sh SHMEM include 注入 179-192
ascendc_impl_build.py SHMEM include 路径 167-179
operator_registry.json operator_meta 12-30

11. 结论

  1. FusedDeepMoE 当前完全基于 HCCL Window 通信，kernel 代码中不存在任何 SHMEM 分支。

  2. 通信原语通过 HcclOpResParam 强绑定，dispatch/combine 直接操作 windowsIn（数据）和 windowsExp（状态）。

  3. 三条执行路径的通信差异： - 非深度融合：dispatch 和 combine 独立执行 - 深度融合：combine 融合进 GMM2 AIV 侧，通过 AllToAllSend + ReducePermute 实现 - 共享专家：额外执行共享专家 GMM1+GMM2，combine 通过 ProcessCombine 执行

  4. Workspace 内存布局采用复用策略（ENABLE_REUSE_MEMORY），shareX1/X2 共用、shareX1Scale/X2Scale 共用、swapSpace/shareSwiglu 共用。

  5. 编译系统已完整预留 SHMEM 注入点： - select_ops.py 通过 requires_shmem 过滤 - compile_ascend_proj.sh 通过 SHMEM_HOME_PATH 注入 include - ascendc_impl_build.py 硬编码 SHMEM include 路径（含一个路径 bug）

  6. SHMEM 集成需要引入通信抽象层，将 HcclOpResParam 替换为 ICommContext 接口，并在 dispatch/combine 中通过模板参数化通信后端。

  7. ascendc_impl_build.py 第172行存在路径 bug：SHMEM_HOME_PATH + "shmem/include/device" 缺少 / 分隔符。

← notify_dispatch + dispatch_layout 📑 索引 ZB Dispatch 算子族 →