@cannbot-plugin/cannbot-harness 0.0.0-stage → 0.1.5

package/assets/ops-direct-invoke/skills/ascendc-mc2-best-practice/references/foundations/blaze-shmem/all_to_all_matmul/include/kernel/all_to_all_comm_udma.h

@@ -0,0 +1,255 @@
+/**
+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+ * CANN Open Software License Agreement Version 2.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+// ============================================================================
+// MC2 通信层 —— AllToAll over UDMA（URMA 协议）
+// ----------------------------------------------------------------------------
+// 该头文件实现"所有 rank 互发 M 段数据"的 AllToAll 语义，承载通算融合中
+// "通"的部分。所有跨卡通信走 SHMEM/UDMA，**禁止使用 HCCL 高阶 API**
+// （Hccl::AllReduce / Hccl::AllGather / ... 禁止清单见 references/foundations/blaze-shmem/comm_shmem.md §5）。
+//
+// 改造场景：
+//   [MODIFY] N1  改通信原语（AllToAll → AllReduce/ReduceScatter/AllGather）：
+//                重写 PutToAllRanks / 新增 ReduceBuffer 等，参考 references/comm_shmem.md §4
+//   [MODIFY] C1  改 Put 策略（如只发给特定 rank）：改 PutToAllRanks 的 Block 分配
+//   [MODIFY] C2  改 SHMEM 空间布局：改 GetDataAddrGm / GetScaleAddrGm 的偏移计算
+//
+// [PITFALL] aclshmemx_udma_put_nbi 的 dst 是"本 rank 视角下本 rank 的 SHMEM 地址"，
+//           不是 remoteRank 的地址。SHMEM 内部做地址翻译。详见 comm_shmem.md §2。
+// [PITFALL] aclshmemx_udma_quiet(remoteRank) 必须在每次 Put 后调用，否则不保序。
+// [PITFALL] BarrierAll 用 aclshmem_barrier_all()，不能用 HCCL 同步原语。
+// [PITFALL] 每个 Block 负责发往对应的 remoteRank（不是单 Block 串行）——多 Block 并发
+//           才能让 UDMA 引擎并行下发，避免单 Block 成瓶颈。
+//
+// 进阶细节：references/comm_shmem.md
+// ============================================================================
+
+#pragma once
+
+#include "basic_api/kernel_basic_intf.h"
+#include "kernel_tiling/kernel_tiling.h"
+#include "blaze/gemm/utils/common_utils.h"
+#include "shmem.h"
+
+namespace AllToAllImpl {
+
+using namespace AscendC;
+using namespace Blaze::Gemm;
+
+/**
+ * @brief AllToAll 通信管理类，负责通过 UDMA 协议在多卡间交换数据和量化系数（Scale）
+ *
+ * @tparam XType 数据类型（通常为 fp8 或其他量化类型）
+ */
+template<typename XType>
+class AllToAllComm {
+public:
+  __aicore__ inline AllToAllComm() {};
+
+  /**
+   * @brief 初始化通信参数
+   *
+   * @param M 输入矩阵 A 的 M 轴总大小（单卡对应的行数）
+   * @param kPerRank 每个 Rank 负责的 K 轴长度
+   * @param commMSize 每次通信切分的 M 轴块大小（用于流水线）
+   * @param bufferSize 流水线 buffer 数（典型为 4）
+   * @param inputDataAddr 输入数据 A 的 GM 地址
+   * @param inputScaleAddr 输入量化系数 ScaleA 的 GM 地址
+   * @param shmemSpace 共享内存空间指针，用于跨卡 UDMA 访问
+   */
+  __aicore__ inline void InitParams(uint64_t M, uint64_t kPerRank, uint64_t commMSize,
+    uint64_t bufferSize,GM_ADDR inputDataAddr, GM_ADDR inputScaleAddr, __gm__ void* shmemSpace);
+
+  /**
+   * @brief 将当前 Rank 的数据段分发到所有其他 Rank
+   *
+   * @param srcMOffset 当前数据段在本地 M 轴的起始偏移
+   * @param mSize 当前传输的 M 轴块行数
+   * @param bufferId bufferId 缓存索引
+   */
+  __aicore__ inline void PutToAllRanks(uint64_t srcMOffset, uint64_t mSize, uint32_t bufferId);
+
+  /**
+   * @brief 将当前 Rank 的量化系数（Scale）分发到所有其他 Rank
+   *
+   * @param srcMOffset 起始偏移
+   * @param mSize 传输行数
+   */
+  __aicore__ inline void PutScaleToAllRanks(uint64_t srcMOffset, uint64_t mSize);
+
+  /**
+   * @brief 全局屏障同步，确保所有异步通信任务完成
+   */
+  __aicore__ inline void BarrierAll();
+
+  /**
+   * @brief 获取指定 bufferId 索引的数据 GM 地址（在共享内存中）
+   */
+  __aicore__ inline GM_ADDR GetDataAddrGm(uint32_t bufferId);
+
+  /**
+   * @brief 获取 Scale 数据在共享内存中的 GM 地址
+   */
+  __aicore__ inline GM_ADDR GetScaleAddrGm();
+
+private:
+  /**
+   * @brief 将特定数据段 Put 到远程指定的 Rank
+   */
+  __aicore__ inline void PutSegmentToRank(uint32_t remoteRank, uint64_t srcMOffset,
+    uint64_t mSize, uint32_t bufferId);
+
+  /**
+   * @brief 将特定 Scale 段 Put 到远程指定的 Rank
+   */
+  __aicore__ inline void PutScaleToRank(uint32_t remoteRank, uint64_t srcMOffset,
+    uint64_t mSize);
+
+  uint64_t rankSize_{0};    // 总卡数
+  uint64_t M_{0};       // 单卡负责的总行数
+  uint64_t commMSize_{0};   // 通信切块行数
+  uint32_t rankId_{0};    // 当前卡号
+  uint64_t bufferSize_;
+
+  uint64_t bytesPerMRow_{0};     // 数据每行的字节数
+  uint64_t scaleBytesPerMRow_{0};  // Scale 每行的字节数
+  uint64_t bufferBlockSize_{0};  // 单个 bufferId 块的总字节数（rankSize * commMSize * bytesPerMRow）
+
+  GM_ADDR inputDataAddr_{nullptr};
+  GM_ADDR inputScaleAddr_{nullptr};
+  __gm__ void* shmemContextGM_{nullptr};
+};
+
+template<typename XType>
+__aicore__ inline void AllToAllComm<XType>::InitParams(
+  uint64_t M, uint64_t kPerRank, uint64_t commMSize, uint64_t bufferSize,
+  GM_ADDR inputDataAddr, GM_ADDR inputScaleAddr, __gm__ void* shmemSpace)
+{
+  shmemContextGM_ = shmemSpace;
+  bufferSize_ = bufferSize;
+  inputDataAddr_ = inputDataAddr;
+  inputScaleAddr_ = inputScaleAddr;
+
+  rankSize_ = aclshmem_n_pes();
+  rankId_ = aclshmem_my_pe();
+  M_ = M;
+  commMSize_ = commMSize;
+
+  bytesPerMRow_ = kPerRank * sizeof(XType);
+  // 计算 MXFP8 格式下 Scale 的字节数（按块量化）
+  scaleBytesPerMRow_ = CeilDiv(kPerRank, Blaze::Gemm::MXFP_DIVISOR_SIZE) * Blaze::Gemm::MXFP_MULTI_BASE_SIZE;
+  // 单个 bufferId 缓存块承载所有卡发送过来的当前流水步的数据
+  bufferBlockSize_ = rankSize_ * commMSize_ * bytesPerMRow_;
+}
+
+template<typename XType>
+__aicore__ inline void AllToAllComm<XType>::PutToAllRanks(uint64_t srcMOffset, uint64_t mSize,
+  uint32_t bufferId)
+{
+  // [设计要点] 每个 Block 负责发送到对应的 remoteRank（Block i → remoteRank i）。
+  // 这样 UDMA 引擎能并行下发多条 Put，避免单 Block 串行 Put 所有 rank 成为瓶颈。
+  // [MODIFY C1] 若改为 ReduceScatter（只发给一个目标 rank），改为：
+  if ((AscendC::GetBlockIdx() < rankSize_) && (AscendC::GetBlockIdx() != rankId_)) {
+    PutSegmentToRank(AscendC::GetBlockIdx(), srcMOffset, mSize, bufferId);
+  }
+  // 等待所有卡完成 Put 任务
+  BarrierAll();
+}
+
+template<typename XType>
+__aicore__ inline void AllToAllComm<XType>::PutScaleToAllRanks(uint64_t srcMOffset, uint64_t mSize)
+{
+  if ((AscendC::GetBlockIdx() < rankSize_) && (AscendC::GetBlockIdx() != rankId_)) {
+    PutScaleToRank(AscendC::GetBlockIdx(), srcMOffset, mSize);
+  }
+  // 注意：PutScaleToAllRanks的同步由 PutToAllRanks 中的 aclshmemx_udma_quiet 保证
+}
+
+template<typename XType>
+__aicore__ inline void AllToAllComm<XType>::PutSegmentToRank(uint32_t remoteRank, uint64_t srcMOffset,
+  uint64_t mSize, uint32_t bufferId)
+{
+  // 远程 Rank 的目标窗口地址
+  // [地址语义] remoteWinAddr 是"本 rank 视角下本 rank 的 SHMEM 基地址"，对应 bufferId 块。
+  //            UDMA 引擎会把这个地址映射到 remoteRank 的 SHMEM；本 rank 不需要知道对端物理地址。
+  GM_ADDR remoteWinAddr = GetDataAddrGm(bufferId);
+
+  // [地址语义] dstDataOffset 用 rankId_（自己的 rank）：
+  //            对端按"我从哪个 rank 来"放到对应位置，避免多 rank 数据互相覆盖。
+  //            例：rank 2 发给 rank 0/1/3，对端都在 rankId_ * commMSize 偏移处收 rank 2 的数据。
+  uint64_t dstDataOffset = rankId_ * commMSize_ * bytesPerMRow_;
+  // [地址语义] srcDataOffset 用 remoteRank：
+  //            本 rank 的 A 矩阵逻辑上按 remoteRank 分段，第 i 段是要发给 remoteRank=i 的数据。
+  //            srcMOffset 是当前流水步的块内偏移。
+  uint64_t srcDataOffset = remoteRank * M_ * bytesPerMRow_ + srcMOffset * bytesPerMRow_;
+  uint64_t dataSize = mSize * bytesPerMRow_;
+
+  // 非阻塞 UDMA 发送
+  // [PITFALL] ubuf 参数（第三个）在 MC2 大块 Put 场景固定为 nullptr；
+  //            小消息聚合场景才需要 UB 中转，本工程不涉及。
+  aclshmemx_udma_put_nbi(
+    remoteWinAddr + dstDataOffset,
+    inputDataAddr_ + srcDataOffset,
+    (__ubuf__ uint8_t*)nullptr,
+    dataSize,
+    remoteRank
+  );
+  // 确保对该远程 Rank 的传输已完成（数据已到达对端 PE）
+  // [PITFALL] aclshmemx_udma_quiet 确保数据到达对端，不是"只下发不等对端"。
+  aclshmemx_udma_quiet(remoteRank);
+}
+
+template<typename XType>
+__aicore__ inline void AllToAllComm<XType>::PutScaleToRank(uint32_t remoteRank, uint64_t srcMOffset, uint64_t mSize)
+{
+  GM_ADDR remoteScaleAddr = GetScaleAddrGm();
+
+  // Scale 空间不使用 bufferId，通常直接一次性或同步流水发送
+  uint64_t dstScaleOffset = rankId_ * M_ * scaleBytesPerMRow_;
+  uint64_t srcScaleOffset = remoteRank * M_ * scaleBytesPerMRow_ + srcMOffset * scaleBytesPerMRow_;
+  uint64_t scaleSize = mSize * scaleBytesPerMRow_;
+
+  aclshmemx_udma_put_nbi(
+    remoteScaleAddr + dstScaleOffset,
+    inputScaleAddr_ + srcScaleOffset,
+    (__ubuf__ uint8_t*)nullptr,
+    scaleSize,
+    remoteRank
+  );
+}
+
+template<typename XType>
+__aicore__ inline void AllToAllComm<XType>::BarrierAll()
+{
+  // 全卡同步等待
+  aclshmem_barrier_all();
+}
+
+template<typename XType>
+__aicore__ inline GM_ADDR AllToAllComm<XType>::GetDataAddrGm(uint32_t bufferId)
+{
+  // 获取共享内存在本地卡对应的 GM 基地址
+  // [PITFALL] aclshmem_ptr 第二个参数是 rankId_（**本卡自己**），不是 remoteRank！
+  //           本卡视角下本卡的 SHMEM 基地址，就是其他卡 Put 过来的数据落点。
+  // [MODIFY C2] 若改 SHMEM 布局（如新增 bias 段），偏移计算同步改。
+  GM_ADDR baseAddr = (GM_ADDR)(aclshmem_ptr(shmemContextGM_, rankId_));
+  return baseAddr + bufferId * bufferBlockSize_;
+}
+
+template<typename XType>
+__aicore__ inline GM_ADDR AllToAllComm<XType>::GetScaleAddrGm()
+{
+  // Scale 存放在数据 bufferId 缓存之后
+  // [设计要点] Scale 不参与 M 轴流水（数据量小），所有 rank 一次性 Put 后存在这里。
+  //            布局：[Data buffers (bufferSize 个)][Scale buffers]
+  return (GM_ADDR)(aclshmem_ptr(shmemContextGM_, rankId_)) + bufferSize_ * bufferBlockSize_;
+}
+
+} // namespace AllToAllImpl

package/assets/ops-direct-invoke/skills/ascendc-mc2-best-practice/references/foundations/blaze-shmem/all_to_all_matmul/include/kernel/all_to_all_matmul_impl.h

@@ -0,0 +1,370 @@
+/**
+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+ * CANN Open Software License Agreement Version 2.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+// ============================================================================
+// MC2 通算融合主类 —— AllToAll（通信）+ Blaze Matmul（计算）流水编排
+// ----------------------------------------------------------------------------
+// 该头文件是 MC2 算子的"大脑"，决定 AIV/AIC 如何协同：
+//   - AIV（Vector 核）：跑 AllToAllProcess()，下发 UDMA Put，把 M 段发给所有 rank；
+//   - AIC（Cube 核）：跑 MatmulProcess()，在通信 buffer 上做完整 Blaze Matmul
+//                   （遍历所有 rank，在 L0C 上累加部分和）；
+//   - 通过 CrossCoreSetFlag/WaitFlag 在 AIV↔AIC 间传递 buffer 就绪/释放信号；
+//   - 通过 4-buffer 流水（bufferSize=4）掩盖通信延迟。
+//
+// 改造场景：
+//   [MODIFY] N1  改通算融合类型（AllToAll+Matmul → AllReduce+Matmul 等）：
+//                改 AllToAllProcess 内部 + 替换 allToAllComm_ 成员类型
+//   [MODIFY] C1  改流水深度（bufferSize 4 → 8）：
+//                改 host 侧 tilingData.commTilingData.bufferSize + 同步调 SHMEM 空间
+//   [MODIFY] A1  改 AIC/AIV 协同顺序：改 Process() 中的 ASCEND_IS_AIV/AIC 分支
+//
+// [PITFALL] CrossCoreSetFlag 和 CrossCoreWaitFlag 必须按 idx 一一配对（1:1）。
+//           多次 Wait 同一 flagId 会死锁；Set 了不 Wait 会让 AIV 写覆盖未读数据。
+// [PITFALL] PIPE_MTE3（AIV 输出）↔ PIPE_MTE2（AIC 输入）是 AIV→AIC 的固定方向；
+//           反向（AIC→AIV）用 PIPE_FIX → PIPE_MTE3。modeId 0x2 是 AIV↔AIC 跨核同步模式。
+//
+// 进阶细节：references/mc2_architecture.md（AIV/AIC 分工） + references/matmul_blaze.md
+// ============================================================================
+
+#pragma once
+
+#include "basic_api/kernel_basic_intf.h"
+#include "kernel_tiling/kernel_tiling.h"
+#include "kernel_operator.h"
+#include "all_to_all_comm_udma.h"
+#include "shmem.h"
+
+#include "blaze/gemm/utils/common_utils.h"
+#include "blaze/gemm/block/block_mmad_qbmm_mx.h"
+#include "blaze/gemm/policy/dispatch_policy.h"
+#include "qbmm_mx_kernel.h"
+#include "../tiling/quant_matmul_tiling_data.h"
+#include "../tiling/all_to_all_matmul_tiling_data.h"
+
+namespace AllToAllQuantMatmulImpl {
+
+using namespace AscendC;
+using namespace AllToAllImpl;
+using namespace Blaze::Gemm;
+
+// 定义问题形状：[M, N, K, Batch]
+using ProblemShape = AscendC::Te::Shape<int64_t, int64_t, int64_t, int64_t>;
+
+/**
+ * @brief All2All-Matmul 核心实现类
+ * 实现思路：
+ * 1. AIC 负责本地计算 (LocalMatmul) 和远程数据计算 (RemoteMatmul)。
+ * 2. AIV 负责背景通信 (AllToAllProcess)，通过 UDMA 将数据分发到各卡。
+ * 3. 通过流水线 (4 buffer) 掩盖通信开销。
+ */
+template<typename AType, typename BType, typename CType, bool TransA, bool TransB>
+class AllToAllQuantMatmulImpl {
+public:
+  __aicore__ inline AllToAllQuantMatmulImpl() {};
+
+  /**
+   * @brief 初始化算子状态和参数
+   */
+  __aicore__ inline void Init(__gm__ void* shmemSpace,
+                 GM_ADDR aGM, GM_ADDR scaleAGM,
+                 GM_ADDR bGM, GM_ADDR scaleBGM,
+                 GM_ADDR cGM,
+                 const allToAllMatmulTilingData *tilingData);
+  /**
+   * @brief 执行算子逻辑（包含 AIC/AIV 分离逻辑）
+   */
+  __aicore__ inline void Process();
+
+  using TypeA = AType;
+  using TypeB = BType;
+  using TypeC = CType;
+  using TypeScaleA = ::fp8_e8m0_t;
+  using TypeScaleB = ::fp8_e8m0_t;
+  static constexpr int32_t SCALE_C0 = 2;
+
+  // Layout 定义
+  using LayoutA = typename AscendC::Std::conditional_t<TransA, AscendC::Te::DNExtLayoutPtn, AscendC::Te::NDExtLayoutPtn>;
+  using LayoutB = typename AscendC::Std::conditional_t<TransB, AscendC::Te::DNExtLayoutPtn, AscendC::Te::NDExtLayoutPtn>;
+  using LayoutC = AscendC::Te::NDExtLayoutPtn;
+  using LayoutScaleA = typename AscendC::Te::FrameLayoutFormat<
+    AscendC::Std::conditional_t<TransA, AscendC::Te::ScaleADNLayoutPtn, AscendC::Te::ScaleANDLayoutPtn>,
+    AscendC::Std::Int<SCALE_C0>>;
+  using LayoutScaleB = typename AscendC::Te::FrameLayoutFormat<
+    AscendC::Std::conditional_t<TransB, AscendC::Te::ScaleBDNLayoutPtn, AscendC::Te::ScaleBNDLayoutPtn>,
+    AscendC::Std::Int<SCALE_C0>>;
+
+  // 组件定义
+  using BlockScheduler = Blaze::Gemm::Block::BlockSchedulerQuantBatchMatmulV3<ProblemShape,
+    NONE_FULL_LOAD_MODE, LayoutA, LayoutB, TypeA>;
+  using DispatchPolicy = Blaze::Gemm::MatmulWithScaleMx<NONE_FULL_LOAD_MODE, false>;
+  using BiasType = float;
+  using LayoutBias = AscendC::Te::NDExtLayoutPtn;
+  using BlockMmad = Blaze::Gemm::Block::BlockMmad<
+    DispatchPolicy, TypeA, LayoutA, TypeB, LayoutB, TypeC, LayoutC, BiasType, LayoutBias>;
+  using QuantMatmulKernelImpl = Kernel::QuantMatmulMxKernelSwat<ProblemShape, BlockMmad, BlockScheduler>;
+
+  // 参数类型
+  using Params = typename QuantMatmulKernelImpl::Params;
+  using BlockMmadParams = typename BlockMmad::Params;
+  using L1Params = typename QuantMatmulKernelImpl::L1Params;
+  using LocalParams = typename QuantMatmulKernelImpl::LocalParams;
+  using BlockSchedulerParams = typename QuantMatmulKernelImpl::BlockSchedulerParams;
+  using QBMMTiling = typename QuantMatmulKernelImpl::QBMMTiling;
+
+  QuantMatmulKernelImpl quantMatmulKernelImpl_;
+
+private:
+  __aicore__ inline void InitBaseParams(const allToAllMatmulTilingData *tilingData);
+  __aicore__ inline void SetupParams(const QuantMatmulTilingData* mmTile, Params& out);
+  __aicore__ inline void AllToAllProcess();  // AIV 通信任务
+  __aicore__ inline void MatmulProcess();    // AIC 通算计算任务
+  __aicore__ inline uint64_t CalcMOffset(uint32_t loopIdx) const;
+  __aicore__ inline uint64_t GetCurMSize(uint32_t loopIdx) const;
+
+  AllToAllComm<AType> allToAllComm_;
+
+  GM_ADDR aGm_;
+  GM_ADDR scaleAGm_;
+  GM_ADDR bGm_;
+  GM_ADDR scaleBGm_;
+  GM_ADDR cGm_;
+
+  uint32_t rankId_{0};
+  uint64_t axisM_{0};     // M 轴总大小
+  uint64_t axisKa_{0};    // K 轴大小
+  uint32_t axisN_{0};     // N 轴大小
+  uint32_t rankSize_{0};    // 卡数
+  uint64_t commMSize_{0};   // 每次通信的 M 大小
+  int32_t  commTurn_{0};     // 总流水步数
+  uint64_t bufferSize_{0};   // 缓冲区大小
+
+  uint64_t headMSize_{0};   // 标准块 M 大小
+  uint32_t tileCnt_{0};     // 块数量
+  uint64_t scaleKaSize_{0};   // Scale 的 K 轴字节长度
+
+  const allToAllMatmulTilingData* tilingData_{nullptr};
+};
+
+template<typename AType, typename BType, typename CType, bool TransA, bool TransB>
+__aicore__ inline void AllToAllQuantMatmulImpl<AType, BType, CType, TransA, TransB>::Init(
+  __gm__ void* shmemSpace, GM_ADDR aGM, GM_ADDR scaleAGM, GM_ADDR bGM, GM_ADDR scaleBGM, GM_ADDR cGM,
+  const allToAllMatmulTilingData *tilingData)
+{
+  tilingData_ = tilingData;
+  InitBaseParams(tilingData);
+  aGm_ = aGM;
+  scaleAGm_ = scaleAGM;
+  bGm_ = bGM;
+  scaleBGm_ = scaleBGM;
+  cGm_ = cGM;
+
+  // 初始化通信参数
+  allToAllComm_.InitParams(axisM_, axisKa_, commMSize_, bufferSize_, aGM, scaleAGM, shmemSpace);
+}
+
+template<typename AType, typename BType, typename CType, bool TransA, bool TransB>
+__aicore__ inline void AllToAllQuantMatmulImpl<AType, BType, CType, TransA, TransB>::InitBaseParams(
+  const allToAllMatmulTilingData *tilingData)
+{
+  rankSize_ = aclshmem_n_pes();
+  rankId_ = aclshmem_my_pe();
+
+  headMSize_ = tilingData->tileQbmmTilingData.m;
+  tileCnt_ = tilingData->commTilingData.tileCnt;
+
+  axisM_ = tileCnt_ * headMSize_;
+  axisKa_ = tilingData->tileQbmmTilingData.k;
+  axisN_ = tilingData->tileQbmmTilingData.n;
+  // 计算 MXFP8 格式下 Scale 每一行的字节数
+  scaleKaSize_ = CeilDiv(axisKa_, Blaze::Gemm::MXFP_DIVISOR_SIZE) * Blaze::Gemm::MXFP_MULTI_BASE_SIZE;
+
+  commTurn_ = tileCnt_;
+  commMSize_ = headMSize_;
+  bufferSize_ = tilingData->commTilingData.bufferSize;
+}
+
+/**
+ * @brief 设置量化矩阵乘内核参数
+ */
+template<typename AType, typename BType, typename CType, bool TransA, bool TransB>
+__aicore__ inline void AllToAllQuantMatmulImpl<AType, BType, CType, TransA, TransB>::SetupParams(
+  const QuantMatmulTilingData* mmTile, Params& out)
+{
+  ProblemShape problemShape{mmTile->m, mmTile->n, mmTile->k, 1UL};
+  BlockMmadParams mmadParams;
+  // splitKNum = rankSize：遍历所有 rank 在 L0C 上累加（参与 rank 数 = rankSize）
+  LocalParams localParams{rankId_, rankSize_, axisM_, aGm_, scaleAGm_, rankSize_};
+  L1Params l1Params{static_cast<uint64_t>(mmTile->stepK) * mmTile->baseK, mmTile->scaleKL1,
+    mmTile->nBufferNum};
+
+  mmadParams.bGmAddr = bGm_;
+  mmadParams.scaleBGmAddr = scaleBGm_;
+
+  // 调度器参数
+  BlockSchedulerParams schedulerParams{
+    mmTile->baseM, mmTile->baseN, mmTile->mTailTile, mmTile->nTailTile,
+    mmTile->mBaseTailSplitCnt, mmTile->nBaseTailSplitCnt, mmTile->mTailMain, mmTile->nTailMain};
+  // 基础 Tiling 参数
+  QBMMTiling qbmmParams{mmTile->baseM, mmTile->baseN, mmTile->baseK, mmTile->dbL0c, false};
+
+  out = {problemShape, mmadParams, l1Params, schedulerParams, qbmmParams, localParams};
+}
+
+/**
+ * @brief AIC 负责的通算计算主流程
+ *
+ * wait_flag 始终在 kernel 外部、每次 kernel 调用前发出。
+ * 等待 AIV 把当前 bufferId 的数据 Put 完 → 跑 Blaze Matmul（遍历所有 rank 累加）→ 释放 buffer。
+ */
+template<typename AType, typename BType, typename CType, bool TransA, bool TransB>
+__aicore__ inline void AllToAllQuantMatmulImpl<AType, BType, CType, TransA, TransB>::MatmulProcess()
+{
+  Params params;
+  SetupParams(&tilingData_->tileQbmmTilingData, params);
+  for (int32_t mLoopIdx = 0; mLoopIdx < tileCnt_; ++mLoopIdx) {
+    // [设计要点] bufferId 用位掩码（mLoopIdx & (bufferSize_ - 1)）做环形索引。
+    //            bufferSize 必须是 2 的幂（4/8/16），否则位掩码失效。
+    uint32_t bufferId = mLoopIdx & (bufferSize_ - 1);
+    uint64_t mOffset = CalcMOffset(mLoopIdx);
+
+    // 更新当前流水步的地址偏移
+    params.mmadParams.aGmAddr = allToAllComm_.GetDataAddrGm(bufferId);
+    params.mmadParams.scaleAGmAddr = allToAllComm_.GetScaleAddrGm() + mOffset * scaleKaSize_;
+    params.mmadParams.cGmAddr = cGm_ + mOffset * axisN_ * sizeof(CType);
+    params.localParams.localAGmAddr = aGm_ + mOffset * axisKa_;
+    params.localParams.localScaleAGmAddr = scaleAGm_ + mOffset * scaleKaSize_;
+
+    // 等 AIV 把当前 bufferId 的数据 Put 完
+    // [PITFALL] CrossCoreWaitFlag<0x2, PIPE_MTE2> 与 AIV 侧 CrossCoreSetFlag<0x2, PIPE_MTE3>
+    //           的 idx 必须一致（都是 mLoopIdx）；不一致会死锁或拿到错位数据。
+    CrossCoreWaitFlag<0x2, PIPE_MTE2>(mLoopIdx);
+    quantMatmulKernelImpl_(params);
+
+    // 通知通信层：当前计算已完成，空间可释放（针对流水线控制）
+    // [设计要点] 流水深度控制：AIC 算完第 i 轮后，AIV 才能写第 (i + bufferSize) 轮的同一 bufferId。
+    //            commTurn_ > bufferSize_ 时才有"释放"概念（前 bufferSize 轮 AIV 不需要等）。
+    if ((commTurn_ > bufferSize_) && (mLoopIdx < (commTurn_ - bufferSize_))) {
+      CrossCoreSetFlag<0x2, PIPE_FIX>(mLoopIdx);
+    }
+  }
+}
+
+/**
+ * @brief AIV 负责的异步通信流程
+ */
+template<typename AType, typename BType, typename CType, bool TransA, bool TransB>
+__aicore__ inline void AllToAllQuantMatmulImpl<AType, BType, CType, TransA, TransB>::AllToAllProcess()
+{
+  // [设计要点] 先行发送全局 Scale：Scale 数据量小（M * CeilDiv(K, 64) * 2 字节），
+  //            不值得做 M 轴流水，AIV 启动时一次性 Put 给所有 rank。
+  allToAllComm_.PutScaleToAllRanks(0, axisM_);
+  for (uint32_t mLoopIdx = 0; mLoopIdx < commTurn_; ++mLoopIdx) {
+    uint32_t bufferId = mLoopIdx & (bufferSize_ - 1);
+    uint64_t curMSize = GetCurMSize(mLoopIdx);
+    uint64_t mOffset = CalcMOffset(mLoopIdx);
+
+    // 流水线深度控制，防止覆盖尚未被 AIC 读取的旧数据
+    // [PITFALL] 必须先 WaitFlag（等 AIC 算完）再 BarrierAll（等其他卡释放），
+    //           顺序反了会让本卡 AIV 覆盖未读数据。
+    if (mLoopIdx >= bufferSize_) {
+      CrossCoreWaitFlag<0x2, PIPE_MTE3>(mLoopIdx - bufferSize_);
+      // 确保所有卡释放了当前的buffer，才能开始写入下一个数据块
+      allToAllComm_.BarrierAll();
+    }
+
+    // 分发当前数据块
+    allToAllComm_.PutToAllRanks(mOffset, curMSize, bufferId);
+
+    // 设置 Flag，通知 AIC 数据已就绪
+    // [PITFALL] CrossCoreSetFlag<0x2, PIPE_MTE3>(mLoopIdx) 与 AIC 侧
+    //           CrossCoreWaitFlag<0x2, PIPE_MTE2>(mLoopIdx) 的 idx 必须一致。
+    CrossCoreSetFlag<0x2, PIPE_MTE3>(mLoopIdx);
+  }
+}
+
+/**
+ * @brief 算子执行入口
+ */
+template<typename AType, typename BType, typename CType, bool TransA, bool TransB>
+__aicore__ inline void AllToAllQuantMatmulImpl<AType, BType, CType, TransA, TransB>::Process()
+{
+  // [设计要点] 同一份 kernel 二进制同时跑在 AIV 和 AIC 上，靠编译期分支隔离。
+  //            ASCEND_IS_AIV / ASCEND_IS_AIC 是 bisheng 编译器提供的宏。
+  //            AIV 跑通信，AIC 跑通算计算，二者通过 CrossCore Flag 同步。
+  if ASCEND_IS_AIV {
+    AllToAllProcess(); // AIV 通信
+  }
+
+  if ASCEND_IS_AIC {
+    MatmulProcess();    // AIC 通算计算（带同步等待）
+  }
+}
+
+template<typename AType, typename BType, typename CType, bool TransA, bool TransB>
+__aicore__ inline uint64_t AllToAllQuantMatmulImpl<AType, BType, CType, TransA, TransB>::CalcMOffset(uint32_t loopIdx) const
+{
+  return loopIdx * headMSize_;
+}
+
+template<typename AType, typename BType, typename CType, bool TransA, bool TransB>
+__aicore__ inline uint64_t AllToAllQuantMatmulImpl<AType, BType, CType, TransA, TransB>::GetCurMSize(uint32_t loopIdx) const
+{
+  return headMSize_;
+}
+
+} // namespace AllToAllQuantMatmulImpl
+
+[[bisheng::core_ratio(1, 1)]]
+__global__ __aicore__ __schedmode__(1) void  AllToAllQuantMatmulKernelE4M3E4M3(__gm__ void* shmemSpace,
+                           GM_ADDR aGM, GM_ADDR scaleAGM,
+                           GM_ADDR bGM, GM_ADDR scaleBGM,
+                           GM_ADDR cGM,
+                           allToAllMatmulTilingData tilingData)
+{
+  AllToAllQuantMatmulImpl::AllToAllQuantMatmulImpl<fp8_e4m3fn_t, fp8_e4m3fn_t, bfloat16_t, false, true> impl;
+  impl.Init(shmemSpace, aGM, scaleAGM, bGM, scaleBGM, cGM, &tilingData);
+  impl.Process();
+}
+
+[[bisheng::core_ratio(1, 1)]]
+__global__ __aicore__ __schedmode__(1) void  AllToAllQuantMatmulKernelE5M2E5M2(__gm__ void* shmemSpace,
+                           GM_ADDR aGM, GM_ADDR scaleAGM,
+                           GM_ADDR bGM, GM_ADDR scaleBGM,
+                           GM_ADDR cGM,
+                           allToAllMatmulTilingData tilingData)
+{
+  AllToAllQuantMatmulImpl::AllToAllQuantMatmulImpl<fp8_e5m2_t, fp8_e5m2_t, bfloat16_t, false, true> impl;
+  impl.Init(shmemSpace, aGM, scaleAGM, bGM, scaleBGM, cGM, &tilingData);
+  impl.Process();
+}
+
+[[bisheng::core_ratio(1, 1)]]
+__global__ __aicore__ __schedmode__(1) void  AllToAllQuantMatmulKernelE4M3E5M2(__gm__ void* shmemSpace,
+                           GM_ADDR aGM, GM_ADDR scaleAGM,
+                           GM_ADDR bGM, GM_ADDR scaleBGM,
+                           GM_ADDR cGM,
+                           allToAllMatmulTilingData tilingData)
+{
+  AllToAllQuantMatmulImpl::AllToAllQuantMatmulImpl<fp8_e4m3fn_t, fp8_e5m2_t, bfloat16_t, false, true> impl;
+  impl.Init(shmemSpace, aGM, scaleAGM, bGM, scaleBGM, cGM, &tilingData);
+  impl.Process();
+}
+
+[[bisheng::core_ratio(1, 1)]]
+__global__ __aicore__ __schedmode__(1) void  AllToAllQuantMatmulKernelE5M2E4M3(__gm__ void* shmemSpace,
+                           GM_ADDR aGM, GM_ADDR scaleAGM,
+                           GM_ADDR bGM, GM_ADDR scaleBGM,
+                           GM_ADDR cGM,
+                           allToAllMatmulTilingData tilingData)
+{
+  AllToAllQuantMatmulImpl::AllToAllQuantMatmulImpl<fp8_e5m2_t, fp8_e4m3fn_t, bfloat16_t, false, true> impl;
+  impl.Init(shmemSpace, aGM, scaleAGM, bGM, scaleBGM, cGM, &tilingData);
+  impl.Process();
+}