UMDK/CAM 通信加速库深度分析 v2

1. 架构总览

1.1 算子分类全景

1.2 核心设计哲学

1.3 芯片存储层次

2. EP Normal 通信算子

2.1 总体架构

EP Normal 通信算子构成 MoE 模型 Expert Parallelism 的完整 dispatch/combine 流水线，采用 zero-buffer 优化策略。

EP Normal Pipeline (4 Stages)

Stage 1              Stage 2              Stage 3              Stage 4
┌──────────┐         ┌──────────┐         ┌──────────┐         ┌──────────┐
│ Dispatch │         │  Notify  │         │ Dispatch │         │ Combine  │
│  Layout  │────────▶│ Dispatch │────────▶│  Normal  │────────▶│  Normal  │
│   (zb)   │         │   (zb)   │         │   (zb)   │         │   (zb)   │
└──────────┘         └──────────┘         └──────────┘         └──────────┘

本地计算             AllGather            shmem 写入           shmem 读取
token count         tokenPerExpert       token data           + weighted sum
+ routing info      + recvCount          to remote rank       + dequant

Sync:              Sync:                Sync:                Sync:
SyncAll             BarrierAll           BarrierAll           BarrierAll
(intra-rank)        (inter-rank)         (inter-rank)         (inter-rank)
                    SyncAll              SyncAll              SyncAll

2.2 共享基础设施

2.2.1 comm_args.h — 通信参数定义

namespace Moe {
constexpr int CAM_MAX_RANK_SIZE = 384;        // 最大支持 384 张 NPU 卡
constexpr int64_t IPC_BUFF_MAX_SIZE = 100 * 1024 * 1024;  // 100MB IPC buffer
constexpr int64_t IPC_DATA_OFFSET = 2 * 1024 * 1024;      // 前 2MB 为 flag 区
constexpr int64_t PING_PONG_SIZE = 2;
constexpr int64_t UB_SINGLE_DMA_SIZE_MAX = 190 * 1024;    // 单次 DMA 最大 190KB
constexpr int64_t UB_ALIGN_SIZE = 32;              // 32B 对齐
constexpr uint8_t COMM_NUM = 2;                // EP + TP 两个通信域

struct CommArgs {
    int rank;                                   // 全局 rank ID
    int rankSize;                               // 全局 rank 数
    GM_ADDR peerMems[CAM_MAX_RANK_SIZE];        // 初始化获取的对端 buffer 地址
    int64_t sendCountMatrix[CAM_MAX_RANK_SIZE * CAM_MAX_RANK_SIZE];  // All2AllV 发送计数矩阵
    int64_t batchSize, hiddenSize, topk;
    int64_t expertNumPerRank;                   // 每 rank 专家数
};
}

2.2.2 zero_buffer_sync_flag.h — 核心同步机制

class ZeroBufferSyncFlag {
    constexpr uint32_t FLAG_SLOT_SIZE = 32U;        // 每个 flag slot 32B
    constexpr uint32_t FLAG_ELEM_NUM = 4U;          // 4 × int64_t
    constexpr uint64_t FLAG_AREA_BASE = 100UL * 1024UL;  // flag 区起始: 100KB
    constexpr uint32_t MAGIC_SLOT_ALIGN = 512U;     // 每 core 512B 对齐
    constexpr uint32_t MAX_CORE_NUM_SYNC = 48U;     // 最大 48 核

    // 每个 kernel 入口调用一次，获取唯一 magic
    __aicore__ inline void IncrementMagic() {
        GM_ADDR magicAddr = flagBaseGM_ + coreIdx_ * MAGIC_SLOT_ALIGN;
        DataCacheCleanAndInvalid(magicGT);          // 读前 cache flush
        magic_ = static_cast(magicGT(0));
        magicGT(0) = static_cast(magic_) + 1ULL;  // 自增
        DataCacheCleanAndInvalid(magicGT);          // 写后 cache flush
        barrierSeq_ = 0;
    }

    // All-to-all barrier: SyncAll → Write flags → Wait flags → SyncAll
    __aicore__ inline void BarrierAll() {
        SyncAll<true>();  // Step 1: 卡内同步
        // Step 2a: 向所有对端写 barrier flag
        for (uint32_t r = startRank; r < endRank; r++) {
            DataCopy(remoteSlot, ubFlag, FLAG_ELEM_NUM);
        }
        // Step 2b: 等待所有对端的 barrier flag
        for (uint32_t r = startRank; r < endRank; r++) {
            do { DataCopy(ubFlag, localSlot, FLAG_ELEM_NUM); } 
            while (ubFlag.GetValue(0) != barrierFlag);
        }
        SyncAll<true>();  // Step 3: 卡内同步
    }
};

2.3 moe_dispatch_layout_zb — Dispatch 布局计算

根据 topkIdx[bs, k] 计算每个 rank/expert 需要接收的 token 数。

// Phase 1: 统计 token 分布 (需要跨核原子操作)
for (round = 0; round < maxNumRounds; ++round) {
    CountTokensInRound(roundStart, roundTokens, expertsPerRank);
    SyncAll<true>();  // 每轮结束同步，保证原子加可见
}

// Phase 2: 计算发送序号 (无跨核依赖)
for (round = 0; round < numRounds; ++round) {
    CalcSendTokenIdxInRound(roundStart, roundTokens);
}

2.4 moe_notify_dispatch_zb — Notify 数据交换

通过 AllGather 交换各 rank 的 tokenPerExpert 数据，计算 putOffset（前缀和）。

__aicore__ inline void Process() {
    syncFlag_.IncrementMagic();     // Step 1: 获取 magic M
    syncFlag_.BarrierAll();         // Step 2: 确保所有 rank 就绪
    AllGatherSendData();            // Step 3: shmem 读取所有 rank 的 tokenPerExpert
    SyncAll<true>();
    ReloadRecvData();
    BuildTotalRecvCount();          // 计算 putOffset (前缀和)
    BuildMaxBs();
    BuildTotalRecvTokens();
    BuildRecvTokenPerExp();
    SyncAll<true>();
}

2.5 moe_dispatch_normal_zb — Normal Dispatch 数据写入

根据 layout 结果将 token 数据通过 shmem 写入对端 expandXOut buffer。支持可选的动态量化 (BF16/FP16 → INT8)。

__aicore__ inline void QuantProcess() {
    // 1. BF16/FP16 → FP32
    Cast(floatLocalTemp, xInTensor, RoundMode::CAST_NONE, h);
    // 2. 计算动态 scale: scale = 127.0 / max(|x|)
    Abs(floatLocalAbsTemp, floatLocalTemp, h);
    ReduceMaxInplace(floatLocalAbsTemp, h);
    dynamicScale = float(127.0) / (floatLocalAbsTemp.GetValue(0) + 1e-12f);
    // 3. 量化: x * scale → int32 → int16 → int8
    Muls(floatLocalTemp, floatLocalTemp, dynamicScale, h);
    Cast(int32LocalTemp, floatLocalTemp, RoundMode::CAST_RINT, h);
    Cast(halfLocalTemp, int32LocalTemp, RoundMode::CAST_ROUND, h);
    Cast(xOutTensor, halfLocalTemp, RoundMode::CAST_TRUNC, h);
    // 4. 保存 scale 到 token 末尾
    floatLocalTemp.SetValue(hUBAlignSize / sizeof(float), float(1.0) / dynamicScale);
}

2.6 moe_combine_normal_zb — Normal Combine 加权归约

__aicore__ inline void ReadBufferAndWeightedSum(uint32_t tokenIndex, ...) {
    Duplicate(sumFloatBufLocal, static_cast<float>(0), axisH_);  // 初始化累加器
    for (uint32_t topkId = 0; topkId < axisK_; topkId++) {
        float scale = topkWeightsLocal.GetValue(topkId);
        int32_t expertId = topkIdxLocal.GetValue(topkId);
        int32_t remoteReadOffset = sendTokenIdxLocal(topkId);
        int32_t remoteReadBase = allRecvCountLocal(expertId * epWorldSize_ + epRankId_);
        uint64_t remoteReadAddr = static_cast<uint64_t>(remoteReadBase + remoteReadOffset) * hRecvXTypeLen_;
        int32_t dstRankId = expertId / moeExpertPerRankNum_;
        auto ptr = reinterpret_cast<__gm__ uint8_t *>(shmem_ptr(recvXGM_, dstRankId));
        dstGT.SetGlobalBuffer((__gm__ XType *)(ptr + remoteReadAddr));
        // 读取 → 类型转换 → 加权累加
        LocalTensor<XType> tmpToken = weightedSumQueue_.AllocTensor<XType>();
        DataCopyPad(tmpToken, dstGT, xOutCopyParams);
        Cast(tokenFloatLocal, tmpToken, AscendC::RoundMode::CAST_NONE, axisH_);
        Muls(weightedMulBufLocal, tokenFloatLocal, scale, axisH_);
        Add(sumFloatBufLocal, sumFloatBufLocal, weightedMulBufLocal, axisH_);
    }
    Cast(xOutLocal, sumFloatBufLocal, AscendC::RoundMode::CAST_RINT, axisH_);
    DataCopyPad(xOutGlobal_[tokenIndex * axisH_], xOutLocal, xOutCopyParams);
}

2.7 全 Pipeline 同步协议

Magic 递增序列:

Notify ──────▶ Dispatch ──────▶ Combine ──────▶ Notify ──────▶ ...
magic=M       magic=M+1       magic=M+2       magic=M+3

BarrierAll    BarrierAll      BarrierAll      BarrierAll
(确保所有      (确保 Notify    (确保 GMM       (确保所有
 rank 的       results 可见)  完成)           rank 的
 tokenPerExpert                                              tokenPerExpert
 就绪)                                                       就绪)

2.8 Zero-Buffer 内存布局

Zero-Buffer Memory Layout (Per Rank):

+0x00000 ┌──────────────────────────────────────────────────────────────┐
         │  User Data Region (100KB)                                    │
         │  - tokenPerExpertData, recvData, putOffset, maxBs, etc.     │
+0x19000 ├──────────────────────────────────────────────────────────────┤
         │  Flag Area Base (100KB offset from gva_gm)                   │
+0x19000 │  ┌────────────────────────────────────────────────────────┐  │
         │  │ Magic Area (24KB = 48 cores × 512B)                   │  │
         │  │   slot[coreIdx] at coreIdx × 512B                     │  │
+0x1F000 │  ├────────────────────────────────────────────────────────┤  │
         │  │ Barrier PingPong Buffer 0 (epWorldSize × 32B)          │  │
         │  ├────────────────────────────────────────────────────────┤  │
         │  │ Barrier PingPong Buffer 1 (same layout)                │  │
         │  ├────────────────────────────────────────────────────────┤  │
         │  │ Flag PingPong Buffer 0                                 │  │
         │  │   flag[srcRank][eventID] (32B each)                    │  │
         │  ├────────────────────────────────────────────────────────┤  │
         │  │ Flag PingPong Buffer 1                                 │  │
         │  └────────────────────────────────────────────────────────┘  │
         ├──────────────────────────────────────────────────────────────┤
         │  Token Data Region (Dispatch/Combine)                       │
         │  - expandXOut (Dispatch 写入, Combine 读取)                 │
         │  - recvX (GMM 输出, Combine 读取)                           │
         └──────────────────────────────────────────────────────────────┘

3. EP Low-Latency 算子

3.1 设计哲学对比

3.2 moe_dispatch_lowlatency_zb — 8 阶段流水线

template <TemplateMC2TypeClass>
__aicore__ inline void MoeDispatchLowlatencyZeroBuffer<...>::Process()
{
    if ASCEND_IS_AIV {
        LocalLayout();          // Phase 1: 本地 expert count 统计
        SyncAll<true>();        // 本地全核同步
        SetLayoutStatus();      // Phase 2: 通知各 rank "我已写完 layout"
        SendCountNotify();      // Phase 3: 等待并读取各 rank 的 layout
        SyncAll<true>();
        WaitNotify();           // Phase 4: 聚合全局 count 矩阵, 计算前缀和
        SyncAll<true>();
        InputToDstOutput();     // Phase 5: 执行 token → remote window 写入
        SetSyncFlag();          // Phase 6: 通知各 rank "我的数据已就绪"
        WaitSyncFlag();         // Phase 7: 等待所有 rank 数据就绪
        SyncAll<true>();
        CleanUp();              // Phase 8: 清理元数据
    }
}

3.2.1 LocalLayout — 本地 Expert Count 统计

__aicore__ inline void LocalLayout() {
    LocalTensor<int32_t> sendCountLt = sendCountLocalBuf_.Get<int32_t>();
    Duplicate<int32_t>(sendCountLt, 0, moeExpertNum_);
    DataCopyPad(expertIdsTensor_, expertIdsGMTensor_, expertIdsCntParams, ...);
    SyncFunc<AscendC::HardEvent::MTE2_S>();
    // 按 AIV 核分片统计
    SplitToCore(sendToMoeExpTokenCnt_, moeUsedAivNum_, startTokenId, endTokenId, sendTokenNum);
    for (int32_t i = startTokenId; i < endTokenId; ++i) {
        int32_t expertIdx = expertIdsTensor_(i);
        int32_t curCnt = sendCountLt.GetValue(expertIdx) + 1;
        sendCountLt.SetValue(expertIdx, curCnt);
    }
    // 原子写入共享内存
    AscendC::SetAtomicAdd<int32_t>();
    DataCopyPad(LocalNotifyDataTensor_, sendCountLt, sendCountDataCopyParams);
    AscendC::SetAtomicNone();
}

3.2.2 SetSyncFlag / WaitSyncFlag — 细粒度跨卡同步

__aicore__ inline void SetSyncFlag() {
    // 1. 准备 flag 值 (1.0f 的 IEEE754 表示 = 0x3F800000)
    LocalTensor<int32_t> statusTensor = syncStatusBuf_.Get<int32_t>();
    Duplicate<int32_t>(statusTensor, 0, rankNumPerBlock_ * 8);
    uint64_t mask[2] = {0x101010101010101, 0};
    Duplicate<int32_t>(statusTensor, 0x3F800000, mask, statusCntAlign / 8, 1, 8);
    // 2. 向每个目标 rank 的元数据区写入 flag
    for (uint32_t targetRankId = curBlockStartRankId_; targetRankId < curBlockEndRankId_; ++targetRankId) {
        auto ptr = GetSyncFlagAddrByRankId(targetRankId) + epRankId_ * DISPATCH_SYNC_STATE_OFFSET;
        gmRemoteStatusGT.SetGlobalBuffer((__gm__ int32_t *)(ptr));
        DataCopy<int32_t>(gmRemoteStatusGT, statusTensor[...], 8UL);
    }
    SyncFunc<AscendC::HardEvent::MTE3_S>();
}

__aicore__ inline void WaitSyncFlag() {
    float compareTarget = 1.0f * rankNumPerBlock_;
    while ((sumOfFlag < minTarget) || (sumOfFlag > maxTarget)) {
        DataCopy(statusFp32Tensor, statusFp32TensorGT[...], intriParams);
        SyncFunc<AscendC::HardEvent::MTE2_V>();
        ReduceSum(statusSumOutTensor, statusFp32Tensor, gatherMaskOutTensor, mask, rankNumPerBlock_, 1);
        sumOfFlag = statusSumOutTensor.GetValue(0);
    }
    // 清理 flag 为下一轮做准备
    Duplicate<int32_t>(cleanStateTensor, 0, duplicateMask, Ceil(rankNumPerBlock_, 8), 1, 8);
    DataCopy(statusFp32TensorGT[...], cleanStateTensor.ReinterpretCast<float>(), intriOutParams);
}

3.3 moe_combine_lowlatency_zb — 4 阶段流水线

__aicore__ inline void MoeCombineLowlatencyZeroBuffer<...>::Process()
{
    if ASCEND_IS_AIV {
        CopyValidExpandXToRemote();  // Phase 1: 本地 expert 数据拷贝到 window
        SyncAll<true>();
        tpipe_->Reset();             // 释放 UB 供后续使用
        SetSyncFlag();               // Phase 2: 通知各 rank "我的数据已就绪"
        WaitSyncFlag();              // Phase 3: 等待所有 rank 数据就绪
        SyncAll<true>();
        InputToDstOutput();          // Phase 4: 读取远程数据 + 加权求和
    }
}

3.3.1 CopyValidExpandXToRemote — 本地 Expert 数据拷贝

__aicore__ inline void CopyValidExpandXToRemote() {
    if (!needCopyExpandX_) return;
    SplitToCore(validTokenNum_, aivNum_, startTokenId, endTokenId, sendTokenNum);
    constexpr uint32_t COPY_BUFFER_NUM = 2U;
    tpipe_->InitBuffer(copyInBuf_, COPY_BUFFER_NUM * GM2GM_COPY_BLOCK_SIZE);
    for (uint64_t i = 0; dataSizeRemain > 0U; ++i) {
        uint32_t copySize = min(dataSizeRemain, GM2GM_COPY_BLOCK_SIZE);
        uint32_t bufferOffset = (i & 1U) * GM2GM_COPY_BLOCK_SIZE;
        event_t eventId = (i & 1U) ? EVENT_ID0 : EVENT_ID1;
        // 乒乓流水线: MTE2 读 → MTE3 写
        AscendC::WaitFlag<HardEvent::MTE3_MTE2>(eventId);
        DataCopyPad(copyLocal[bufferOffset], oriXByteGT[srcOffset], copyParams, copyPadParams);
        AscendC::SetFlag<HardEvent::MTE2_MTE3>(eventId);
        AscendC::WaitFlag<HardEvent::MTE2_MTE3>(eventId);
        DataCopyPad(expandXByteGT[dstOffset], copyLocal[bufferOffset], copyParams);
        AscendC::SetFlag<HardEvent::MTE3_MTE2>(eventId);
        srcOffset += copySize; dstOffset += copySize; dataSizeRemain -= copySize;
    }
}

3.4 同步标志内存布局

ZeroBuffer Metadata Layout (per rank):
┌──────────────────────────────────────────────────────────────┐
│ FLAG_AREA_BASE (100KB offset from zeroBufferPtr)             │
├──────────────────────────────────────────────────────────────┤
│ Magic Area (24KB)                                            │
├──────────────────────────────────────────────────────────────┤
│ Barrier PingPong Buffer 0 (epWorldSize × 32B)               │
├──────────────────────────────────────────────────────────────┤
│ Barrier PingPong Buffer 1 (epWorldSize × 32B)               │
├──────────────────────────────────────────────────────────────┤
│ Flag PingPong Buffer 0 (epWorldSize × slotsPerRank × 32B)   │
├──────────────────────────────────────────────────────────────┤
│ Flag PingPong Buffer 1 (epWorldSize × slotsPerRank × 32B)   │
└──────────────────────────────────────────────────────────────┘

Dispatch 使用 DISPATCH_SYNC_FLAG_BASE_OFFSET = FLAG_AREA_BASE + 64KB
Combine 使用 COMBINE_SYNC_FLAG_BASE_OFFSET = FLAG_AREA_BASE

3.5 适用场景决策

4. Fused Deep MoE 算子家族

4.1 核心动机

Fused Deep MoE 将传统上分散在多个算子中的操作融合为单一算子：

传统方式：Dispatch → GMM1 → SwiGLU → Quant → GMM2 → Combine
                    ↑________ 跨设备通信 ________↓

融合方式：[Dispatch + GMM1 + DeqSwigluQuant + GMM2 + DeqCombine] → Output
          ↑________ 全在 NPU 内部完成，无额外同步点 ________↓

4.2 三变体横向对比

4.3 三阶段流水线架构

┌─────────────────────────────────────────────────────────────────────────────┐
│                         Fused Deep MoE 三阶段流水线                          │
│                                                                             │
│  Stage 1: Dispatch (AIV)                                                    │
│  ┌─────────────────────────────────────────────────────────────────────┐   │
│  │  x [bs, h] ──→ Read expert_ids ──→ Quantize (BF16→INT8)           │   │
│  │                    │                    │                            │   │
│  │                    ▼                    ▼                            │   │
│  │              expertTokenNums    expand_x [M, h] (INT8)              │   │
│  │              epSendCount        dynamic_scale [M]                    │   │
│  └─────────────────────────────────────────────────────────────────────┘   │
│                              │                                              │
│                              ▼                                              │
│  Stage 2: GMM1 + DeqSwigluQuant (AIC + AIV 协作)                           │
│  ┌─────────────────────────────────────────────────────────────────────┐   │
│  │  AIC: GMM1 (INT8×INT8→INT32) → workspace                          │   │
│  │       ↓ callback (CrossCoreFlag)                                    │   │
│  │  AIV: Read INT32 → Cast FP32 → ×weight_scale → ×per_token_scale   │   │
│  │       → SwiGLU (left: x*sigmoid(x), right: passthrough)            │   │
│  │       → Quantize FP32→INT8 → gm_x2                                 │   │
│  └─────────────────────────────────────────────────────────────────────┘   │
│                              │                                              │
│                              ▼                                              │
│  Stage 3: GMM2 + DeqCombine (AIC + AIV 协作)                               │
│  ┌─────────────────────────────────────────────────────────────────────┐   │
│  │  AIC: GMM2 (INT8×INT8→INT32) → workspace                          │   │
│  │       ↓ callback                                                    │   │
│  │  AIV: Read INT32 → Cast FP32 → ×weight_scale → ×per_token_scale   │   │
│  │       → AlltoAll Send → ReducePermute → Scale×ExpertScale → Add   │   │
│  │       → output [bs, h]                                             │   │
│  └─────────────────────────────────────────────────────────────────────┘   │
└─────────────────────────────────────────────────────────────────────────────┘

4.4 AIC/AIV 异构执行模型

┌──────────────────────────────────────────────────────────────┐
│                     Single NPU Core (1C2V)                   │
│                                                              │
│  ┌──────────────────────┐    ┌──────────────────────────┐   │
│  │  AIC (Cube Engine)   │    │  AIV (Vector Engine)     │   │
│  │  - GMM 矩阵乘        │    │  - Element-wise 运算     │   │
│  │  - L0A/L0B/L0C      │    │  - SwiGLU/Quant/Dequant  │   │
│  │  - MTE2/MTE3 DMA    │    │  - Dispatch/Combine 逻辑 │   │
│  └──────────┬───────────┘    └──────────┬───────────────┘   │
│             │    CrossCoreFlag/Callback  │                   │
│             └────────────────────────────┘                   │
│                          │                                   │
│                     GM (Global Memory)                       │
└──────────────────────────────────────────────────────────────┘

4.5 Tiling Key 位掩码

constexpr uint32_t EXEC_FLAG_DEEP_FUSE      = (1U << 0);  // 多 expert 融合
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);  // shared expert
constexpr uint32_t EXEC_FLAG_SMOOTH_QUANT   = (1U << 4);  // smooth scale
constexpr uint32_t EXEC_FLAG_ZERO_BUFFER    = (1U << 5);  // SHMEM 零拷贝

4.6 Workspace 内存布局

Workspace 布局 (GM):
┌─────────────────────────────────────────────────────────────────┐
│ 0x00000: System Workspace (16MB)                               │
├─────────────────────────────────────────────────────────────────┤
│ gmShareX1 (shared expert tokens, INT8)                          │
│ gmX1 (MoE tokens, INT8)                                         │
├─────────────────────────────────────────────────────────────────┤
│ gmShareX1Scale (FP32 per-token scale)                           │
│ gmX1Scale (FP32 per-token scale)                                │
├─────────────────────────────────────────────────────────────────┤
│ gmTokenFlag (INT32, token 到达标志)                              │
├─────────────────────────────────────────────────────────────────┤
│ gmCVSwap (AIC/AIV 交换缓冲, blockDim * L1M * L1N * 4B * stages)│
├─────────────────────────────────────────────────────────────────┤
│ gmShareSwigluOut (FP32, shared expert SwiGLU 输出)              │
│ gmSwigluOut (FP32, MoE SwiGLU 输出)                             │
├─────────────────────────────────────────────────────────────────┤
│ gmGroupList (INT64, expert token 累积计数)                      │
│ gmExpandIdx (INT32, 扩展索引)                                    │
│ gmEpSendCount (INT32, EP 发送计数)                               │
├─────────────────────────────────────────────────────────────────┤
│ gmAllExpertTokenNums (INT64, 全局 expert token 计数)            │
│ gmCombineSend (FP32, combine 发送缓冲)                           │
│ gmAllEpRecvCount (INT32, 全局 EP 接收计数, ZeroBuffer 模式)     │
└─────────────────────────────────────────────────────────────────┘

4.7 GEMM 引擎：Catlass TPL 深度分析

4.7.1 Catlass 架构层次

Catlass (CUDA-like Template Library for Ascend)
├── Arch::AtlasA2          ← 架构标签（Atlas A2）
├── Gemm::Kernel           ← Kernel 级别（问题级调度）
├── Gemm::Block            ← Block 级别（单核内分块）
├── Gemm::Tile             ← Tile 级别（指令级）
└── Epilogue::Block        ← Epilogue 后处理

4.7.2 Dispatch Policy 层次

// 基础版 GMM1：MmadAtlasA2Custom
using MmadAtlasA2Custom =
    Gemm::MmadAtlasA2PreloadAsyncWithCallback<
        CUSTOM_PRELOAD_STAGES,   // 1
        CUSTOM_L1_STAGES,        // 2
        CUSTOM_L0A_STAGES,       // 2
        CUSTOM_L0B_STAGES,       // 2
        CUSTOM_L0C_STAGES,       // 1
        CUSTOM_ENABLE_UNIT_FLAG, // true
        CUSTOM_ENABLE_SHUFFLE_K  // true
    >;

// 基础版 GMM2：MmadAtlasA2PreloadAsyncWithCallbackResidentA
using Gmm2DispatchPolicy =
    Gemm::MmadAtlasA2PreloadAsyncWithCallbackResidentA<
        CUSTOM_PRELOAD_STAGES,   // 1
        GMM2_L1A_STAGES,         // 4
        GMM2_L1B_STAGES,         // 2
        GMM2_L0A_STAGES,         // 4
        GMM2_L0B_STAGES,         // 2
        CUSTOM_L0C_STAGES,       // 1
        CUSTOM_ENABLE_UNIT_FLAG, // true
        CUSTOM_ENABLE_SHUFFLE_K  // true
    >;

4.7.3 BlockMmad 核心执行流程

void L1TileMmad(L1TileMmadParams const &params) {
    uint32_t mPartLoop = CeilDiv<L0TileShape::M>(params.mRound);
    uint32_t nPartLoop = CeilDiv<L0TileShape::N>(params.nRound);
    uint32_t kPartLoop = CeilDiv<L0TileShape::K>(params.kActual);

    for (mPartIdx = 0; mPartIdx < mPartLoop; ++mPartIdx) {
        for (kPartIdx = 0; kPartIdx < kPartLoop; ++kPartIdx) {
            // 1. L1 → L0A (A 矩阵分块)
            copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, L1A_LAYOUT);
            for (nPartIdx = 0; nPartIdx < nPartLoop; ++nPartIdx) {
                // 2. L1 → L0B (B 矩阵分块)
                copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, L1B_LAYOUT);
                // 3. L0C = L0A × L0B (Cube 引擎矩阵乘)
                tileMmad(l0CTile, l0ATile, l0BTile, 
                         mPartActual, nPartActual, kPartActual, initC, unitFlag);
            }
        }
    }
    // 4. K 循环结束后：L0C → GM (写回)
    if (params.isKLoopLast) {
        params.callbackBeforeFixpipe();  // 通知 AIV 可以开始 epilogue
        copyL0CToGm(params.gmBlockC, l0CTensor, layoutCInGm, layoutCInL0, 0b11);
        params.callbackAfterFixpipe();   // 通知 AIV epilogue 完成
    }
}

4.7.4 Workspace 多阶段流水线

// 4-Stage Workspace 流水线（WORKSPACE_STAGES = 4）
// 时间 →
// AIC: [GMM1_blk0] [GMM1_blk1] [GMM1_blk2] [GMM1_blk3] [GMM1_blk4]
//       ↓callback   ↓callback   ↓callback   ↓callback   ↓callback
// AIV:              [EPI_blk0]  [EPI_blk1]  [EPI_blk2]  [EPI_blk3]
//
// 每个 stage 大小 = L1M × L1N × sizeof(int32) = 256 × 128 × 4 = 128KB

4.8 Epilogue：后处理流水线

4.8.1 GMM1 Epilogue（PerTokenDequantSwiGLU）

void operator()(...) {
    // 1. 加载 INT32 accumulator
    copyGmToUbC(ubC, gmTileC, layoutUbC, layoutGmTileC);
    // 2. 加载 weight_scale (per-column)
    copyGmToUbScale(ubFp32Scale, gmTileScale, ...);
    // 3. 加载 per_token_scale (per-token)
    copyGmToUbPerTokenScale(ubPerTokenScale, gmTilePerTokenScale, ...);
    // 4. INT32 → FP32
    Cast(ubTmpMxN, ubC, CAST_RINT, TileShape::COUNT);
    // 5. × weight_scale (row broadcast: [1,N] × [M,N])
    tileRowBroadcastMul(ubTmpMxN, ubTmpMxN, ubFp32Scale);
    // 6. × per_token_scale (broadcast: [M,1] × [M,N])
    tileBroadcastOneBlk(ubTmpMx32B, ubPerTokenScale);
    tileOneBlkColumnBroadcastMul(ubTmpMxN, ubTmpMxN, ubTmpMx32B);
    // 7. SwiGLU 激活（仅左半部分）
    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);
    }
    // 8. 写回 GM
    copyUbToGmD(gmTileD, ubD, layoutGmTileD, layoutUbD);
}

4.9 W4A8 量化方案

// W4A8 Epilogue 额外步骤：
// 1. 高低 4-bit 分离
high = (int8_weight >> 4) & 0xF;   // 高 4 位
low  = int8_weight & 0xF;          // 低 4 位
// 2. 有符号扩展
if (high >= 8) high -= 16;
if (low >= 8) low -= 16;
// 3. 反量化合并
AscendC::Muls(ubTmpMxN, ubTmpMxN, 16.0f, ...);  // high * 16
AscendC::Add(ubTmpMxN, ubTmpMxN, ubTmpMxN[TileShape::COLUMN/2], ...);  // + low
AscendC::Add(ubTmpMxN, ubTmpMxN, ubweighAux, ...);  // + assistMatrix

4.10 通信层：Raw Distributed Dispatch/Combine

// Dispatch 流程
void Process() {
    if ASCEND_IS_AIV {
        // 1. Shared Expert 量化（可选）
        if constexpr (EXEC_FLAG & EXEC_FLAG_SHARED_EXPERT) {
            QuantForShareExpert();
        }
        // 2. AlltoAll Dispatch
        AlltoAllDispatch();
        // 3. SetStatus：通知目标 rank 数据已到达
        SetStatus();
        // 4. WaitDispatch：等待所有 rank 数据到达
        WaitDispatch();
        // 5. LocalWindowCopy：从 window 读取数据到本地
        LocalWindowCopy();
        // 6. AllGather（可选，用于 TP）
        if constexpr (IsNeedAllgater) {
            AllGatherSetStatusAndWait();
            AllgatherProcessOut();
        }
        // 7. 更新 expertTokenNums
        UpdataTokenNumsOut();
    }
}

4.11 执行流程与跨核同步

时间 →
──────────────────────────────────────────────────────────────────────────────

AIV Core:
  ┌──────────────┐  ┌──────────────────────────────────────────────────────┐
  │   Dispatch   │  │              Epilogue (GMM1 + GMM2)                  │
  │ AlltoAll +   │  │  [EPI_blk0] [EPI_blk1] [EPI_blk2] ... [Combine]    │
  │ LocalCopy    │  │       ↑          ↑          ↑                        │
  └──────┬───────┘  └───────│──────────│──────────│────────────────────────┘
         │                  │          │          │
         │ CrossCoreFlag    │          │          │
         │ (gmm1AivFinished)│          │          │
         │                  │          │          │
AIC Core:│                  │          │          │
  ───────┼──────────────────┼──────────┼──────────┼─────────────────────────
         │    ┌─────────────▼──────────▼──────────▼─────────────────────┐
         │    │              GMM1 + GMM2 Compute                         │
         │    │  [GMM1_blk0] [GMM1_blk1] ... [GMM2_blk0] [GMM2_blk1]   │
         │    │       ↓callback      ↓callback      ↓callback           │
         │    └──────────────────────────────────────────────────────────┘
         │
HCCL:    │    ┌─────────────┐              ┌─────────────┐
         │    │ AlltoAll    │              │ AlltoAll    │
         │    │ (Dispatch)  │              │ (Combine)   │
         │    └─────────────┘              └─────────────┘

4.12 通信隐藏策略

传统方式: [====Dispatch====] [GMM1] [SwiGLU] [GMM2] [====Combine====]
          ↑ 通信暴露 ↑                              ↑ 通信暴露 ↑

融合方式: [Dispatch+GMM1] [SwiGLU+GMM2] [Combine]
          ↑ 通信隐藏 ↑    ↑ 计算重叠 ↑    ↑ 通信隐藏 ↑
          
          AlltoAll 发送在 GMM1 计算期间完成
          AlltoAll 接收在 GMM2 计算期间完成

4.13 演进方向

fused_deep_moe (2025-07)  →  fused_deep_moe_fwk (2026-01)  →  fused_deep_moe_w4a8 (2026-06)
       BF16 基础版               TPL 框架验证                    W4A8 量化推理
         ↓                          ↓                              ↓
    功能完整                    简洁接口                      内存节省 50%
    64 tiling key               4 tiling key                  支持 bias
    Zero Buffer                 无 Zero Buffer                 INT4 补偿矩阵

5. 辅助通信算子

5.1 执行摘要

CAM 辅助通信算子是一组绕过 HCCL 标准集合通信调度、直接在 AICore 上通过 IPC 共享内存窗口和 MTE 引擎完成数据搬运的底层通信原语。

5.2 all_to_all_detour — 全互打 Detour 算子

5.2.1 算法核心

// 共享内存偏移计算
int64_t GetShareMemOffset(int srcRank, int dstRank, int memRank) {
    return (rankSize * srcRank + dstRank) * memLen 
           + memRank * OPT_RANK_OFFSET  // 512B 对齐
           + (magic % PING_PONG_SIZE) * (memLen / 2); // ping-pong
}

// 获取第 order 个非通信 rank 作为 detour
int GetDstRankId(int order) {
    for (int i = 0; i < rankSize; ++i) {
        if (!IsCommRank(i)) {  // 跳过通信 rank
            if (count == order) return i;
            ++count;
        }
    }
}

5.2.2 数据流示意

4 rank，commRank=[0,1]，detourRank=[2,3]:

Rank 0 需要发送给 Rank 1 的数据：
  ┌─────────────────────────────────────────────────────────┐
  │ 分片 0 (直接)  │ 分片 1 (经 Rank 2) │ 分片 2 (经 Rank 3) │
  │ → shareAddrs[1] │ → shareAddrs[2]    │ → shareAddrs[3]    │
  └─────────────────────────────────────────────────────────┘

同步协议：
  发送方：CpGM2GM → SetFlag(S_MTE3) → CpUB2GM(flag)
  接收方：WaitSyncFlag(magic, 1, srcCommRank, detourRank) → CpGM2GM

5.3 reduce_scatter_detour — 归约散射 Detour 算子

// 核心参数计算
detourRankSize = rankSize - commRankSize;  // 中转 rank 数
in2IpcCoreNum = (detourRankSize + 1) * (commRankSize - 1) + 1;
reduceCoreNum = detourRankSize + 2;
perDataNumForDetourRank = (len + detourRankSize) / (detourRankSize + 1);

// 归约流水线
for (int i = 0; i < piplineNum; i++) {
    int srcCommRank = commRankIds[(commRankIdx + 1 + i) % commRankSize];
    int detourRank = offsetIdx == 0 ? rank : GetDetourRankId(offsetIdx - 1);
    // 等待数据就绪
    sync.WaitSyncFlag(magic, 1, rank, rank, 1, shareMemOffsetInRank);
    sync.WaitSyncFlag(magic, 1, srcCommRank, detourRank, 1, shareMemOffsetInDstRank);
    // 原子归约
    CpGM2GMPingPong(curReduceDataNum * sizeof(U), reduceSrcGt, reduceDstGt, atomOp);
    // 通知下一级
    sync.SetOuterFlagDetour(magic, i + 1, rank, rankSize + rank, shareMemOffsetInRank);
}

5.4 a2e / e2a — MoE 跨域通信

5.4.1 通信模式

Attention 域 (rank ≥ expertRankSize)          Expert 域 (rank < expertRankSize)
         │                                              │
         │  A2E: x, expert_ids, scales ──────────────▶ │
         │                                              │ Expert Compute
         │  E2A: expert output ◀─────────────────────── │
         │                                              │

5.4.2 A2E 共享内存布局

shareAddrs[sendRank] + IPC_DATA_OFFSET:
┌──────────────┬───────────────────┬─────────────────────┬─────────────────┐
│ batch_size   │ expert_ids        │ expert_scales       │ x (token data)  │
│ (4B)         │ (batch*topk*4B)   │ (batch*topk*4B)     │ (batch*hidden*2B)│
└──────────────┴───────────────────┴─────────────────────┴─────────────────┘

5.4.3 多 Flag 同步

// 使用不同偏移的 flag 区分数据类型
WAIT_FLAG_OFFSET_2 = 2;  // expert_scales
WAIT_FLAG_OFFSET_3 = 3;  // x (token)
WAIT_FLAG_OFFSET_4 = 4;  // batch_size

// 发送 expert_ids 后的 flag
shareFlagGt = shareAddrs[rank % expertRankSize] 
    + (1 * attnToMoeRatio + sendOffset) * flagUnitInt64Num;

// 发送 expert_scales 后的 flag
shareFlagGt = shareAddrs[rank % expertRankSize] 
    + (WAIT_FLAG_OFFSET_2 * attnToMoeRatio + sendOffset) * flagUnitInt64Num;

// 发送 x 后的 flag
shareFlagGt = shareAddrs[rank % expertRankSize] 
    + (WAIT_FLAG_OFFSET_3 * attnToMoeRatio + sendOffset) * flagUnitInt64Num;

5.5 四大算子对比

5.6 使用场景决策树

需要集合通信？
├─ 通信 rank 数 < 总 rank 数？
│  ├─ 是 → 数据量 < 1MB？
│  │  ├─ 是 → 使用 Detour 算子
│  │  │  ├─ All-to-All → all_to_all_detour
│  │  │  └─ Reduce-Scatter → reduce_scatter_detour
│  │  └─ 否 → 使用 HCCL 标准接口
│  └─ 否 → 使用 HCCL 标准接口
└─ MoE 跨域通信？
   ├─ Attention → Expert 分发 → a2e
   └─ Expert → Attention 回传 → e2a

6. 框架与基础设施

6.1 架构总览

Python (torch.ops.umdk_cam_op_lib.*)
    ↓ pybind11
ext_utils.cpp (g_comms[] 全局状态管理)
    ↓ C API
cam_api.cpp (CamCreateComm / CamGetCommArgs / CamDestroyComm)
    ↓
cam_comm.cpp (CamComm: IPC 内存管理 + 拓扑发现)
    ↓
cam_socket_exchange.cpp (Bootstrap AllGather)
    ↓
aclrtMalloc / rtIpc* / aclrtDeviceEnablePeerAccess (华为 ACL/Runtime API)

6.2 CamComm：通信器核心

6.2.1 初始化流程

CamComm::Init() 执行步骤：
1. 参数校验 (rank 范围检查)
2. 创建 CamSocketExchange 实例
3. GetMyExchangeIds() → 收集本 rank 的 devId/pid/sdid/IPC name
4. IdsExchangeAndProcess() → AllGather 交换所有 rank 元数据
5. InitCommon() → EnablePeerAccess() 启用 P2P
6. InitCommMem() → SetIpcPidSdid + GetName + OpenIpcMem
7. SyncCommArgs() → 将 CommArgs 拷贝到 Device HBM
8. 销毁 socketExchange（初始化完成后不再需要）

6.2.2 芯片拓扑与物理链路

enum class ChipName {
    CHIP_310P3 = 0, CHIP_910B1, CHIP_910B2, CHIP_910B3, CHIP_910B4,
    CHIP_910B41, CHIP_910B2C, CHIP_910_9391, CHIP_910_9381,
    CHIP_910_9392, CHIP_910_9382, CHIP_910_9372, CHIP_910_9361,
    CHIP_910A5, RESERVED
};

enum class PhysicalLink { HCCS = 0, PCIE = 1, RESERVED };

6.2.3 IPC 内存权限管理

// 910B 系列：仅设置 PID
if (physicalInfo_.chipName < ChipName::CHIP_910_9391) {
    rtSetIpcMemPid(name.c_str(), &pidInt32, HCCL_IPC_PID_ARRAY_SIZE);
} 
// 910A3 系列：设置 SuperPod PID + SDID
else {
    rtSetIpcMemorySuperPodPodPid(name.c_str(), exchangeIds[i].sdid, 
                                  &pidInt32, HCCL_IPC_PID_ARRAY_SIZE);
}

6.3 Socket Exchange：Bootstrap 握手协议

6.3.1 网络地址解析优先级

void CamSocketExchange::GetIpAndPort()
{
    // 1. 优先使用构造函数传入的 serverIpPort 参数
    if (!serverIpPort_.empty() && ParseIpAndPort(...) == CAM_SUCCESS) { ... }
    // 2. 回退到环境变量 CAM_COMM_ID
    else if (Cam::GetEnv("CAM_COMM_ID") != nullptr && ...) { ... }
    // 3. 默认值：127.0.0.1:10067
    if (!isParseOk) {
        ip_ = "127.0.0.1";   // 仅监听 localhost (安全设计)
        port_ = 10067;
    }
    port_ += serverRank;  // 多 communicator 端口偏移
}

6.3.2 节点发现

GetNodeNum() 通过 /proc/sys/kernel/random/boot_id 读取节点 UUID，实现自动节点计数，无需用户手动配置 nodeNum。

6.4 CommArgs 结构体：Host-Device 契约

constexpr int CAM_MAX_RANK_SIZE = 384;  // 最大支持 384 rank

struct CommArgs {
    int rank = 0;                    // 全局 rank ID
    int localRank = -1;              // 节点内 rank ID
    int rankSize = 0;                // 全局 rank 总数
    int localRankSize = -1;          // 节点内 rank 总数
    uint32_t extraFlag = 0;          // 功能标志位
    int testFlag = 0;                // 测试标志
    
    // 共享内存地址表：每个 rank 的 HBM buffer 地址
    GM_ADDR peerMems[CAM_MAX_RANK_SIZE] = {};
    
    // 全连接通信矩阵
    int64_t sendCountMatrix[CAM_MAX_RANK_SIZE * CAM_MAX_RANK_SIZE] = {};
    int64_t sendCounts[CAM_MAX_RANK_SIZE] = {};
    int64_t sdispls[CAM_MAX_RANK_SIZE] = {};
    int64_t recvCounts[CAM_MAX_RANK_SIZE] = {};
    int64_t rdispls[CAM_MAX_RANK_SIZE] = {};
    
    // MoE 模型参数
    int64_t batchSize, hiddenSize, topk;
    int64_t sharedExpertRankNum, expertNumPerRank;
    int64_t dfx[50] = {};            // DFX 调试字段
    int64_t memLen;                  // buffer 总长度
};

6.4.1 ExtraFlag 位域

struct ExtraFlag {
    static constexpr uint32_t RDMA = 1;              // 启用 RDMA
    static constexpr uint32_t TOPO_910B2C = 1 << 1;  // 910B2C 拓扑
    static constexpr uint32_t TOPO_910_93 = 1 << 2;  // 910A3 拓扑
    static constexpr uint32_t DETERMINISTIC = 1 << 3;// 确定性计算
    static constexpr uint32_t QUANT_FP16 = 1 << 4;   // FP16 量化
    static constexpr uint32_t QUANT_FP32 = 1 << 5;   // FP32 量化
    static constexpr uint32_t TOPO_910A5 = 1 << 6;   // 910A5 拓扑
    static constexpr uint32_t QUANT_DELAY = 1 << 7;  // 延迟量化
    static constexpr uint32_t QUANT_CURRENT = 1 << 8;// 当前量化
    static constexpr uint32_t TOPO_PCIE = 1 << 9;    // PCIE 拓扑
    static constexpr uint32_t ATOMIC_ENABLE = 1 << 15;// 原子操作使能
};

6.5 Buffer 管理

6.5.1 HBM 分配策略

int CamComm::InitMem() {
    void *peerMem = nullptr;
    aclError ret = aclrtMalloc(&peerMem, memLen_,
        (GetChipName() == ChipName::CHIP_310P3) ? 
            ACL_MEM_MALLOC_HUGE_FIRST_P2P :    // 310P 使用 P2P 大页
            ACL_MEM_MALLOC_HUGE_FIRST);        // 其他使用普通大页
    peerMem_[rank_] = static_cast<GM_ADDR>(peerMem);
    aclrtMemset(peerMem_[rank_], memLen_, 0, memLen_);
    return CAM_SUCCESS;
}

// Buffer 大小由环境变量 CAM_BUFFSIZE 控制
static uint64_t GetCamMaxWindowSize() {
    auto defaultWindowSize = DEFAULT_CAM_BUFF_SIZE;  // 204 MB
    if (getenv(CAM_BUFFSIZE) != nullptr) {
        defaultWindowSize = std::stoi(getenv(CAM_BUFFSIZE));
        if (defaultWindowSize > MAX_CAM_BUFF_SIZE) {  // 上限 2000 MB
            CAM_LOG(ERROR) << "CAM_BUFFSIZE larger than MAX";
        }
    }
    return static_cast<uint64_t>(defaultWindowSize) * 1024UL * 1024UL;
}

6.6 Pybind 层

6.6.1 双重注册机制

// 1. pybind11 直接绑定（Python 直接调用）
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("fused_deep_moe", &FusedDeepMoeImplAutograd, "fused_deep_moe");
    m.def("all2_all_detour", &all2_all_detour_impl_autograd, "all2_all_detour");
    m.def("reduce_scatter_detour", &reduce_scatter_detour_impl_autograd, "reduce_scatter_detour");
    m.def("cam_get_comm", &cam_get_comm, "cam_get_comm");
    m.def("a2e", &A2eImplAutograd, "a2e");
    m.def("e2a", &E2aImplAutograd, "e2a");
    pybind11::class_<fused_deep_moe::Buffer>(m, "FusedDeepMoeBuffer")
        .def(pybind11::init<>())
        .def("init", &fused_deep_moe::Buffer::init)
        .def("zb_fused_deep_moe", &fused_deep_moe::Buffer::zb_fused_deep_moe);
}

// 2. PyTorch Operator Schema 注册（支持 TorchScript/编译优化）
TORCH_LIBRARY(umdk_cam_op_lib, m) {
    m.def("fused_deep_moe(Tensor x, Tensor expert_ids, ...) -> Tensor[]");
    m.def("all2_all_detour(Tensor sendData, Tensor commRankIds, ...) -> Tensor");
    m.def("a2e(Tensor x, ...) -> Tensor[]");
    m.def("e2a(Tensor expand_x, ...) -> Tensor");
}

6.7 EXEC_NPU_CMD 宏

#define EXEC_NPU_CMD(aclnn_api, ...) do {
    // 1. 获取 OpApi 函数地址（libopapi.so / libcust_opapi.so）
    static const auto getWorkspaceSizeFuncAddr = GetOpApiFuncAddr(#aclnn_api "GetWorkspaceSize");
    static const auto opApiFuncAddr = GetOpApiFuncAddr(#aclnn_api);
    // 2. 初始化大内存线程局部存储
    InitHugeMemThreadLocal initMemFunc = ...;
    if (initMemFunc) { initMemFunc(nullptr, false); }
    // 3. 类型转换：ATen Tensor → aclTensor
    auto converted_params = ConvertTypes(__VA_ARGS__, workspace_size_addr, executor_addr);
    // 4. 获取 workspace 大小
    auto getWorkspaceSizeFunc = ConvertToOpApiFunc(converted_params, getWorkspaceSizeFuncAddr);
    call(getWorkspaceSizeFunc, converted_params);
    // 5. 分配 workspace
    auto workspace_tensor = at::empty({workspace_size}, options.dtype(c10::kByte));
    // 6. 提交到 NPU Stream
    auto acl_call = [...]() { opApiFunc(workspace_addr, workspace_size, executor, acl_stream); };
    at_npu::native::OpCommand cmd;
    cmd.Name(#aclnn_api);
    cmd.SetCustomHandler(acl_call);
    cmd.Run();
    // 7. 清理
    UnInitHugeMemThreadLocal unInitMemFunc = ...;
    if (unInitMemFunc) { unInitMemFunc(nullptr, false); }
} while (false)

6.8 Buffer 类（ZeroBubble 流水线支持）

void Buffer::init(int rank, int num_ranks, uint64_t memsize, const std::string &ip_port) {
    // 初始化 SHMEM
    aclshmemx_init_attr(ACLSHMEMX_INIT_WITH_DEFAULT, &attributes);
    // 分配 ext_info 区（2MB）
    ext_info_ptr_ = aclshmem_malloc(EXT_INFO_SIZE);  // 2MB
    // 分配 workspace 区（剩余空间）
    shmem_workspace_size_ = local_mem_size - EXT_INFO_SIZE;
    shmem_workspace_ptr_ = aclshmem_malloc(shmem_workspace_size_);
}

6.9 与 NCCL 的对比

7. 构建系统 / HFA / OPX

7.1 构建系统

7.1.1 顶层构建脚本

7.1.2 双层 CMake 架构

第一层：comm_operator/CMakeLists.txt — 通信库
  add_library(cam SHARED ${CAM_SOURCE_FILE})
  add_library(cam_static STATIC ${CAM_SOURCE_FILE})
  target_link_libraries(cam ascendcl runtime c_sec)

第二层：ascend_kernels/CMakeLists.txt — 算子包
  find_package(ASC REQUIRED)
  npu_op_package(${package_name} TYPE RUN CONFIG INSTALL_PATH ${CMAKE_BINARY_DIR}/)

7.1.3 cmake_files 构建基础设施

7.2 算子注册机制

{
  "soc_versions": {
    "ascend910_93": [
      "a2e", "all_to_all_detour", "e2a", "fused_deep_moe",
      "fused_deep_moe_fwk", "fused_deep_moe_w4a8",
      "moe_combine_lowlatency_zb", "moe_combine_normal_zb",
      "moe_dispatch_layout_zb", "moe_dispatch_lowlatency_zb",
      "moe_dispatch_normal_zb", "moe_notify_dispatch_zb",
      "reduce_scatter_detour"
    ]
  },
  "operator_meta": {
    "fused_deep_moe": { "requires_shmem": true, "family": "fused_deep_moe", "variant": "base" },
    "fused_deep_moe_w4a8": { "requires_shmem": false, "family": "fused_deep_moe", "variant": "w4a8" }
  }
}

7.3 HFA — HBM Fast Allocator

7.3.1 架构概览

HFA 是一个高性能 HBM 内存分配器，专为 Ascend NPU 设计：

• ~99.7% HBM 利用率（跨所有对象大小）
• ~0.5 µs 稳定 alloc/free 延迟
• 支持最大 128GB 虚拟内存、32GB 物理内存

7.3.2 三层架构

┌─────────────────────────────────────────┐
│  HbmFastAlloc (顶层 API, thread-safe)   │
│  - alloc/free/free_all                  │
│  - 互斥锁 m_alloc_map_lock              │
├─────────────────────────────────────────┤
│  FlexAlloc (逻辑层, 非线程安全)          │
│  - 双索引: StartMap + SizeMap           │
│  - Best-fit 分配 + 邻居合并             │
├─────────────────────────────────────────┤
│  AclVirtualMemory (物理层)              │
│  - 虚拟内存预留 + 物理页映射            │
│  - defrag 碎片整理                      │
└─────────────────────────────────────────┘

7.3.3 FlexAlloc — 灵活分配器

// 核心数据结构 — 双索引
StartMap m_free_ranges_by_offset;  // map<offset_t, Range> — 按起始地址排序
SizeMap  m_free_ranges_by_size;    // multimap<length_t, StartMap::iterator> — 按大小排序

// 分配算法 (Best-fit)
FlexRetCode FlexAlloc::alloc(length_t req_size, size_t &offset, size_t &alloc_size) {
    alloc_size = round_up_bits(req_size, OBJ_SIZE_BITS);  // 64 字节对齐
    auto sizeIt = m_free_ranges_by_size.lower_bound(alloc_size);  // 找第一个 >= 请求的
    if (sizeIt == end) {
        if (get_total_free_mem() >= alloc_size)
            return FLEX_RC_ERR_FRAGMENTATION;  // 总量够但碎片化了
        else
            return FLEX_RC_ERR_NO_MEM;         // 真的没内存了
    }
    // 分裂剩余空间
    if (blkSz > alloc_size) {
        m_free_ranges_by_offset.emplace(offset + alloc_size, Range{blkSz - alloc_size});
    }
}

// 释放算法 — 邻居合并
bool FlexAlloc::free(offset_t offset, length_t size) {
    // 1. 查找前驱，若相邻则合并
    if (prevEnd == offset) { newStart = prevStart; newSize += prevSize; }
    // 2. 查找后继，循环合并所有相邻范围
    while (newStart + newSize == rangeStart) { newSize += rangeSize; }
    // 3. 插入合并后的范围到双索引
}

7.3.4 AclVirtualMemory — 虚拟内存抽象

7.4 OPX — Operator Extension Library

7.4.1 架构概览

OPX 是一个设备端算子编程框架，提供完整的 AscendC 内核开发原语：

opx.h (统一入口)
├── type_traits.h      — 元编程基础设施（自实现 tuple/invoke/apply）
├── ub_queue.h         — UBQue 硬件队列抽象
├── comm_def.h         — 管线同步原语 + ProcessParam
├── swizzle.h          — 数据分块策略
├── data_looper.h      — 数据遍历主循环（CRTP + callback）
├── data_looper_ctx.h  — 遍历上下文（config + status）
├── matrix_data_ctx.h  — 矩阵运算上下文
├── vector_data_ctx.h  — 向量运算上下文
├── matrix_data_reader.h — GM→UB 数据读取（封装 Catlass CopyGm2Ub）
├── matrix_data_writer.h — UB→GM 数据写出（封装 Catlass CopyUb2Gm）
├── matrix_op_start.h  — 矩阵算子起点（DataLooper + Reader）
├── matrix_op_end.h    — 矩阵算子终点（Writer）
├── vector_op_start.h  — 向量算子起点
├── vector_coord.h     — 一维坐标
├── fused.h            — 算子融合（FusedVV/Parallel/Sequential）
└── dep_catlass.h      — Catlass 依赖适配

7.4.2 UBQue 硬件队列

template <auto src, auto dst, uint32_t depth_ = 1>
class UBQue {
    AscendC::TBufType* bufStart;
    AscendC::TBufHandle que[depth_];
    uint8_t bufNum, usedCount, head, tail, bufUsedCount, bufCursor;
};

// 核心管线同步
GetHardEventByPipe<src, dst>()  // 编译期计算硬件事件 ID
AllocTensor<T>()               // 分配 LocalTensor，等待生产者完成
EnQue()                        // 入队，设置生产者事件标志
DeQue<T>()                    // 出队，等待消费者事件完成
FreeTensor()                   // 释放 tensor，设置消费者完成标志

7.4.3 DataLooper — 数据遍历主循环

// 四级循环结构
for group_loop:       // group 列表（如 expert 列表）
  for block_loop:     // block 级分块（多核分配）
    for tile_loop:    // tile 级分块（sub-core 分配）
      callback(ctx)   // 用户回调

// CRTP 模式 + 8 个回调点
// PRE/POST × LOOP/GROUP/BLOCK/TILE

7.4.4 算子融合框架

7.5 测试覆盖分析

测试覆盖缺口

8. 关键设计洞察与跨章节主题

8.1 Flag-Based Sync vs. 传统 Barrier

8.2 Detour 模式的核心思想

传统方式（HCCL）：
  Rank 0 ──RDMA──▶ Rank 1（需全局 barrier，所有 rank 等待）

Detour 方式：
  Rank 0 ──GM write──▶ Rank 2（detour）──GM read──▶ Rank 1
  （仅需子集同步，Rank 3 不受影响）

8.3 MC2 子通道的利用

e2a/a2e 虽然绕过 HCCL 的标准集合通信 API，但仍然使用 MC2（Memory-Core Communication） 子通道的 tiling 配置：

AscendC::Mc2CcTilingConfig mc2CcTilingConfig(groupEp, opType1, algConfigAllToAllStr);
mc2CcTilingConfig.GetTiling(tiling->mc2InitTiling);

这表明 MC2 是 HCCL 内部的硬件通信调度层，提供 fullmesh/pairwise 等拓扑算法。CAM 算子直接配置 MC2 tiling，利用其底层通信能力，但自主管理数据搬运和同步。这是一种"半绕过"策略：复用硬件能力，但规避软件调度开销。

8.4 编译期多态

// 大量使用 C++ 模板元编程实现编译期优化

// 1. Tiling Key 编译期分支
if constexpr (EXEC_FLAG & EXEC_FLAG_DEEP_FUSE) {
    // Deep Fuse 路径
} else {
    // Shallow Dispatch 路径
}

// 2. 数据类型编译期选择
using ScaleType = Gemm::GemmType<W1ScaleType, layout::VectorLayout>;

// 3. 编译期常量 + static_assert
constexpr uint32_t L1A_TILE_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA);
static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES <= ArchTag::L1_SIZE);

8.5 通信隐藏策略


Fused Deep MoE 的核心创新：将 AlltoAll 通信与 GMM 计算重叠。

传统方式: [====Dispatch====] [GMM1] [SwiGLU] [GMM2] [====Combine====]
          ↑ 通信暴露 ↑                              ↑ 通信暴露 ↑

融合方式: [Dispatch+GMM1] [SwiGLU+GMM2] [Combine]
          ↑ 通信隐藏 ↑    ↑ 计算重叠 ↑    ↑ 通信隐藏 ↑
          
          AlltoAll 发送在 GMM1 计算期间完成
          AlltoAll 接收在 GMM2 计算期间完成


8.6 内存层次优化


优化技术
实现方式
效果
L1 ResidentA
GMM2 中 A 矩阵在 K 循环中保持不变，跳过 L1A 加载
节省 GM→L1 带宽
Workspace 4-stage
AIC 写 stage N，AIV 读 stage N-1
完全隐藏 AIC→AIV 数据传递延迟
UB 双缓冲
UB_STAGES=1，通过 event ID 切换
MTE2/V/MTE3 并行
INT4 打包
8 个 INT4 打包为 1 个 INT32
L1/L0 tile 大小减半，节省 50% 存储
Ping-Pong Buffer
magic % 2 选择 buffer
避免连续 kernel 间数据冲突
HFA Best-fit
双索引 (StartMap + SizeMap)
~99.7% HBM 利用率，~0.5µs 延迟


8.7 量化融合

传统方式：
  GMM(INT8) → INT32 → Cast FP32 → ×scale → SwiGLU → Quant → INT8
  ↑ 多步分离 ↑

融合方式：
  GMM(INT8) → INT32 → [Cast + Dequant + SwiGLU + Quant] → INT8
              ↑ 全部在 UB 中完成，零 GM 写回 ↑

W4A8 额外：
  INT32 → Cast FP32 → ×weight_scale → ×16 → +low → +assistMatrix → ×per_token → SwiGLU → Quant
          ↑ 高低 4-bit 分离与补偿矩阵融合 ↑

8.8 设计范式总结


设计决策
实现方式
优势
Bootstrap 网络
TCP Socket AllGather
与 NCCL 兼容，易于调试
设备内存共享
IPC 命名内存 + 全局地址表
零拷贝，低延迟
Host-Device 参数传递
CommArgs 结构体 + aclTensor 封装
单次拷贝，kernel 直接读取
版本同步
Magic 计数器
无锁，轻量级
多通信域
comm_id 索引的 g_comms[] 数组
支持 EP/TP 分离
算子调用
OpApi 动态加载 + OpCommand 异步提交
支持自定义算子覆盖
流水线并行
SHMEM 对称内存 + 融合算子
减少 bubble，提升吞吐

附录

A. 文件索引

A.1 EP Normal 算子

A.2 EP Low-Latency 算子

A.3 Fused Deep MoE

A.4 框架基础设施

A.5 共享工具

B. 关键常量表

C. Tiling Key 位定义

Bit 0 (0x01): EXEC_FLAG_DEEP_FUSE      - 多 expert 融合模式
Bit 1 (0x02): EXEC_FLAG_TENSOR_LIST    - 权重为 tensor list 格式
Bit 2 (0x04): EXEC_FLAG_X_ACTIVE_MASK  - 使用激活 mask
Bit 3 (0x08): EXEC_FLAG_SHARED_EXPERT  - 使用 shared expert
Bit 4 (0x10): EXEC_FLAG_SMOOTH_QUANT   - 使用 smooth scale
Bit 5 (0x20): EXEC_FLAG_ZERO_BUFFER    - 使用 SHMEM 零拷贝

D. 参考链接

• UMDK/CAM 仓库: 华为昇腾 NPU 上 MoE 通信加速库
• CANN 文档: Ascend Computing Language 官方文档
• Catlass: CUDA-like Template Library for Ascend (GEMM 模板库)
• HCCL: Huawei Collective Communication Library
• DeepSeek MoE 架构: Mixture-of-Experts 大规模语言模型设计
• ZeroBubble 流水线: 流水线并行气泡消除技术

E. 版本历史