昇腾社区首页
开发者
下载
支持
积分兑换 NEW
文档中心
mega_moe
[object Object][object Object]
  • [object Object]Ascend 950PR/Ascend 950DT[object Object]：支持
  • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：支持
  • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：支持
  • [object Object]Atlas 200I/500 A2 推理产品[object Object]：不支持
  • [object Object]Atlas 推理系列产品[object Object]：不支持
  • [object Object]Atlas 训练系列产品[object Object]：不支持
[object Object]

  • 接口功能：

    Mega MoE算子 将 MoE 层的专家 FFN 的完整计算流程及前后数据通信（即 Dispatch + Linear1 + SwiGLU + Linear2 + Combine）融合为单个算子，实现了通信和计算的掩盖。 该算子提供了 mega_moe 与 get_symm_buffer_for_mega_moe等接口，这些接口需配套使用。

    • get_symm_buffer_for_mega_moe：需与mega_moe配套使用，用于封装输入参数并创建SymmBuffer结构体，生成[object Object]、[object Object]和[object Object]等mega_moe算子运行所需信息。

  • 计算公式：

    • 输入：

      • X∈Rtotal_num_tokens×hidden\mathbf{X} \in \mathbb{R}^{\text{total\_num\_tokens} \times \text{hidden}}：激活矩阵，对应入参 [object Object]。total_num_tokens\text{total\_num\_tokens} 是全局总 token 数，hidden\text{hidden} 是隐藏层维度。
      • E∈Ztotal_num_tokens×num_topk\mathbf{E} \in \mathbb{Z}^{\text{total\_num\_tokens} \times \text{num\_topk}}：token 选择的专家编号矩阵，对应入参 [object Object]。num_topk\text{num\_topk} 是每个 token 选择的专家数量。
      • G∈Rtotal_num_tokens×num_topk\mathbf{G} \in \mathbb{R}^{\text{total\_num\_tokens} \times \text{num\_topk}}：token 选择的专家的门控权重矩阵，对应入参 [object Object]。
      • W1∈Rnum_experts×hidden×(2intermediate_hidden)\mathbf{W}_1 \in \mathbb{R}^{\text{num\_experts} \times \text{hidden} \times (2 \text{intermediate\_hidden})}：Linear1 的权重矩阵，对应入参 [object Object]。num_experts\text{num\_experts} 是专家数量，intermediate_hidden\text{intermediate\_hidden} 是中间层维度。
      • W2∈Rnum_experts×intermediate_hidden×hidden\mathbf{W}_2 \in \mathbb{R}^{\text{num\_experts} \times \text{intermediate\_hidden} \times \text{hidden}}：Linear2 的权重矩阵，对应入参 [object Object]。

    • 输出：

      • Y∈Rtotal_num_tokens×hidden\mathbf{Y} \in \mathbb{R}^{\text{total\_num\_tokens} \times \text{hidden}}：最终输出矩阵，对应出参 [object Object]。

    • 约定：

      • ⋅⋅ 表示矩阵乘法，⊙⊙ 表示逐元素乘法。
      • ⌊z⌉\left \lfloor z\right \rceil 表示将 zz 四舍五入到最近的整数，⌊z⌋\left \lfloor z\right \rfloor 表示将 zz 向下取整。
      • ∣z∣|z| 表示取绝对值，max⁡(z)\max(z) 表示取最大值。
      • 全体 token 的集合为 {tokeni∣i∈{0,1,…,total_num_tokens−1}}\{ \text{token}_i \mid i \in \{0, 1, \dots, \text{total\_num\_tokens} - 1\} \}。
      • tokeni\text{token}_i 的 token 表示（即隐藏状态向量）为 xi∈R1×hidden\mathbf{x}_i \in \mathbb{R}^{1 \times \text{hidden}}，且 xi=X[i,:]\mathbf{x}_i = \mathbf{X}[i,:]。
      • tokeni\text{token}_i 的专家索引为 ei,k=E[i,k],k∈{0,…,num_topk−1},ei,k∈{0,…,num_experts−1}e_{i,k} = \mathbf{E}[i,k],\quad k \in \{0,\dots,\text{num\_topk} - 1\},\quad e_{i,k} \in \{0,\dots,\text{num\_experts} - 1\}。
      • 激活函数 SiLU(z)=z⋅σ(z)=z1+e−z\text{SiLU}(z) = z \cdot \sigma(z) = \frac{z}{1+e^{-z}}，其中 σ\sigma 为 Sigmoid 函数。
      • Z4={x∈Z∣−8≤x≤7},Z8sym={x∈Z∣−127≤x≤127},Z32={x∈Z∣−231≤x≤231−1}\mathbb{Z}_4 = \{ x \in \mathbb{Z} \mid -8 \le x \le 7 \}, \quad \mathbb{Z}_8^{\text{sym}} = \{ x \in \mathbb{Z} \mid -127 \le x \le 127 \}, \quad \mathbb{Z}_{32} = \{ x \in \mathbb{Z} \mid -2^{31} \le x \le 2^{31}-1 \}。其中 Z8sym\mathbb{Z}_8^{\text{sym}} 的上标 sym\text{sym} 表示对称量化值域区间：其值域关于 −127-127 与 127127 对称取整，与标准 INT8 的 [−128,127][-128, 127] 值域不同，故以 sym\text{sym} 上标区分。
      • 张量切片操作采用Python风格的 [object Object] 表示法，例如[0::2,:][0::2, :] 代表取偶数行、[1::2,:][1::2, :] 代表取奇数行。
      • bitcastT(Z)\mathrm{bitcast}_{T}(\mathbf{Z}) 表示二进制重解释操作，将张量Z\mathbf{Z}的底层二进制数据按目标类型 TT 重新解释。

  • 计算过程

    各产品支持的Linear计算时的激活矩阵A和权重矩阵W的数据类型如下：

    [object Object]undefined
    • [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：不支持上表中A8W8-FP场景。
    • [object Object]Ascend 950PR/Ascend 950DT[object Object]：不支持上表中A16W16 、A8W8-INT、A8W4-INT场景。
    [object Object]

    • EP Dispatch

      在 Dispatch 阶段，每个 tokeni\text{token}_i 将其token表示 xi\mathbf{x}_i 发送给专家 ei,0,ei,1,…,ei,num_topk−1e_{i,0}, e_{i,1}, \dots, e_{i,\text{num\_topk}-1}。即对于每个 kk，专家 ei,ke_{i,k} 接收一份 xi\mathbf{x}_i。

      记 Ie={ i∣∃k, E[i,k]=e }I_e = \{\, i \mid \exists k,\ \mathbf{E}[i,k]=e \,\} 为所有被分派给专家 ee 的 token 索引集合，集合 IeI_e 的大小 Ne=∣Ie∣N_e = |I_e| 即为专家 ee 需要处理的 token 总数，则有 Xe∈RNe×hidden\mathbf{X}_e \in \mathbb{R}^{N_e \times \text{hidden}} 是由所有满足 i∈Iei \in I_e 的token表示 xi\mathbf{x}_i 按任意固定顺序行堆叠而成的矩阵，该矩阵即为专家 ee 经过 Dispatch 后收到的全部 token 表示。

      对于每个 tokeni\text{token}_i 及其选中的第 kk 个专家 ei,ke_{i,k}，存在唯一的行索引 row⁡(i,k)∈{0,…,Nei,k−1}\operatorname{row}(i,k) \in \{0,\dots,N_{e_{i,k}}-1\}，使得 Xei,k[row⁡(i,k),:]=xi\mathbf{X}_{e_{i,k}}[\operatorname{row}(i,k), :] = \mathbf{x}_i。该映射记录了 xi\mathbf{x}_i 在专家 ei,ke_{i,k} 的输入矩阵 Xei,k\mathbf{X}_{e_{i,k}} 中的位置。

    • Expert Compute

      在 MoE 层中，每个专家本质上是一个独立的前馈网络（FFN），采用 SwiGLU 结构以提升表达能力。整个计算过程分为如下三个子步骤。

      1. Linear1 投影

      Linear1 投影是专家网络的第一层线性变换，同时产生 gate 部分 和 up 部分 所需的预激活值。其计算公式为

      He=Xe⋅W1[e](∈RNe×2⋅intermediate_hidden)\mathbf{H}_e = \mathbf{X}_e \cdot \mathbf{W}_1[e] \quad\bigl(\in \mathbb{R}^{N_e \times 2\cdot\text{intermediate\_hidden}}\bigr)

      2. SwiGLU 激活

      首先将 He\mathbf{H}_e 沿列维度拆分为 gate 部分 Hgate\mathbf{H}_{\text{gate}} 和 up 部分 Hup\mathbf{H}_{\text{up}}，然后对 gate 部分应用 SiLU 激活函数，再与 up 部分逐元素相乘，得到专家 ee 的中间激活表示 Ae\mathbf{A}_e。

      Hgate=He[:,  :intermediate_hidden],Hup=He[:,  intermediate_hidden:]\mathbf{H}_{\text{gate}} = \mathbf{H}_e\bigl[:,\;\text{:intermediate\_hidden}\bigr], \qquad \mathbf{H}_{\text{up}} = \mathbf{H}_e\bigl[:,\;\text{intermediate\_hidden:}\bigr] Ae=SiLU(Hgate)  ⊙  Hup(∈RNe×intermediate_hidden)\mathbf{A}_e = \text{SiLU}\bigl(\mathbf{H}_{\text{gate}}\bigr) \;\odot\; \mathbf{H}_{\text{up}} \quad\bigl(\in \mathbb{R}^{N_e \times \text{intermediate\_hidden}}\bigr)

      3. Linear2 投影

      Linear2 投影作为第二层线性变换，将中间激活表示 Ae\mathbf{A}_e 从高维空间投影回原始的隐藏维度 hidden\text{hidden}，使专家输出能够与残差连接等后续操作兼容。

      Ye=Ae⋅W2[e](∈RNe×hidden)\mathbf{Y}_e = \mathbf{A}_e \cdot \mathbf{W}_2[e] \quad\bigl(\in \mathbb{R}^{N_e \times \text{hidden}}\bigr)

      经过以上计算，专家 ee 的每一行输出对应其批次中的一个 token。对于 tokeni\text{token}_i，它在专家 ei,ke_{i,k} 中的输出行即为 Yei,k[row⁡(i,k), :]\mathbf{Y}_{e_{i,k}}\bigl[\operatorname{row}(i,k),\,:\bigr]。

    • Token Combine

      Combine 负责收集所有专家计算出的输出向量，按照每个 token 原先分配到的专家权重进行加权求和，最终为每个 token 生成一个融合后的输出。利用之前记录的位置索引 row⁡(i,k)\operatorname{row}(i,k)，从专家 ei,ke_{i,k} 的输出矩阵中收回属于 tokeni\text{token}_i 的行，并与门控权重相乘后求和：

      yi=∑k=0num_topk−1wk  ⋅  Yei,k ⁣[ row⁡(i,k), : ](∈R1×hidden)\mathbf{y}_i = \sum_{k=0}^{\text{num\_topk} - 1} w_k \;\cdot\; \mathbf{Y}_{e_{i,k}}\!\bigl[\,\operatorname{row}(i,k),\,:\,\bigr] \qquad\bigl(\in \mathbb{R}^{1 \times \text{hidden}}\bigr)

      其中 wk=G[i,k]w_k = \mathbf{G}[i,k] 为 tokeni\text{token}_i 对专家 ei,ke_{i,k} 的门控权重。

      所有 token 的输出按输入顺序堆叠为最终输出 Y∈Rtotal_num_tokens×hidden\mathbf{Y} \in \mathbb{R}^{\text{total\_num\_tokens} \times \text{hidden}}。

    [object Object][object Object]

    • 输入

      • SW1∈Rnum_experts×(2×intermediate_hidden)\mathbf{S}^{W1} \in \mathbb{R}^{\text{num\_experts} \times (2\times \text{intermediate\_hidden})}：Linear1 权重矩阵的逐通道缩放因子，对应入参 [object Object]。
      • SW2∈Rnum_experts×hidden\mathbf{S}^{W2} \in \mathbb{R}^{\text{num\_experts} \times \text{hidden}}：Linear2 权重矩阵的逐通道缩放因子，对应入参 [object Object]。

    • EP Dispatch

      在 Dispatch 通信之前，首先将原始 BF16 激活矩阵 X\mathbf{X} 量化为 INT8。对每个 tokeni\text{token}_i，计算其逐 token 缩放因子：

      siX=max⁡(∣X[i,:]∣)127∈Rs^{X}_i = \frac{\max(|\mathbf{X}[i,:]|)}{127} \in \mathbb{R}

      然后量化得到 INT8 表示：

      qi=⌊X[i,:]siX⌉∈(Z8sym)1×hidden\mathbf{q}_i = \left\lfloor \frac{\mathbf{X}[i,:]}{s^{X}_i} \right\rceil \quad \in \left(\mathbb{Z}_8^{\text{sym}}\right)^{1 \times \text{hidden}}

      在 Dispatch 通信阶段，每个 tokeni\text{token}_i 将其量化后的向量 qi\mathbf{q}_i 和缩放因子 siXs^{X}_i 发送给专家 ei,0,ei,1,…,ei,num_topk−1e_{i,0}, e_{i,1}, \dots, e_{i,\text{num\_topk}-1}（ei,k=E[i,k]e_{i,k} = \mathbf{E}[i,k]）。

      记 Ie={ i∣∃k, E[i,k]=e }I_e = \{\, i \mid \exists k,\ \mathbf{E}[i,k]=e \,\} 为所有被分派给专家 ee 的 token 索引集合，集合 IeI_e 的大小 Ne=∣Ie∣N_e = |I_e| 即为专家 ee 需要处理的 token 总数。则有 Qe∈(Z8sym)Ne×hidden\mathbf{Q}_e \in \left(\mathbb{Z}_8^{\text{sym}}\right)^{N_e \times \text{hidden}} 是由所有满足 i∈Iei \in I_e 的 qi\mathbf{q}_i 按任意固定顺序行堆叠而成的矩阵，该矩阵即为专家 ee 经过 Dispatch 后收到的全部 token 表示；同理，对应的专家 ee 收到的缩放因子向量记为 seX∈RNe\mathbf{s}^{X}_e \in \mathbb{R}^{N_e}，其元素由所有满足 i∈Iei \in I_e 的 siXs^{X}_i 按与 Qe\mathbf{Q}_e 相同的行顺序堆叠而成。

      对于每个 tokeni\text{token}_i 及其选中的第 kk 个专家 ei,ke_{i,k}，存在唯一的行索引 row⁡(i,k)∈{0,…,Nei,k−1}\operatorname{row}(i,k) \in \{0,\dots,N_{e_{i,k}}-1\}，使得 Qei,k[row⁡(i,k),:]=qi\mathbf{Q}_{e_{i,k}}[\operatorname{row}(i,k), :] = \mathbf{q}_i。该映射记录了 qi\mathbf{q}_i 在专家 ei,ke_{i,k} 的输入矩阵中的位置。

    • Expert Compute

      在 MoE 层中，每个专家本质上是一个独立的前馈网络（FFN），采用 SwiGLU 结构以提升表达能力。在A8W8场景下，两个线性层都使用 INT8 输入和 INT8 权重进行矩阵乘，得到 INT32 中间结果并反量化。具体分为三个子步骤。

      1. Linear1 投影（INT8 矩阵乘 + 反量化）

      Linear1 投影是专家网络的第一层线性变换，同时产生 gate 部分 和 up 部分 所需的预激活值。计算时执行 INT8 矩阵乘法，得到 INT32 计算结果：

      Ceint32=Qe⋅W1[e]int8∈Z32Ne×2⋅intermediate_hidden\mathbf{C}_e^{\text{int32}} = \mathbf{Q}_e \cdot \mathbf{W}_1[e]^{\text{int8}} \quad \in \mathbb{Z}_{32}^{N_e \times 2\cdot\text{intermediate\_hidden}}

      然后反量化为预激活值 He\mathbf{H}_e：

      He=(Ceint32⊙seW1)⊙seX∈RNe×2⋅intermediate_hidden\mathbf{H}_e = \left( \mathbf{C}_e^{\text{int32}} \odot \mathbf{s}^{W1}_e \right) \odot \mathbf{s}^{X}_e \quad \in \mathbb{R}^{N_e \times 2\cdot\text{intermediate\_hidden}}

      2. SwiGLU 激活

      首先将 He\mathbf{H}_e 沿列维度拆分为 gate 部分 Hgate\mathbf{H}_{\text{gate}} 和 up 部分 Hup\mathbf{H}_{\text{up}}，然后对 gate 部分应用 SiLU 激活函数，再与 up 部分逐元素相乘，得到专家 ee 的中间激活表示 Ae\mathbf{A}_e。

      Hgate=He[:,:intermediate_hidden],Hup=He[:,intermediate_hidden:]\mathbf{H}_{\text{gate}} = \mathbf{H}_e[:, :\text{intermediate\_hidden}], \qquad \mathbf{H}_{\text{up}} = \mathbf{H}_e[:, \text{intermediate\_hidden}:] Ae=SiLU⁡(Hgate)⊙Hup\mathbf{A}_e = \operatorname{SiLU}(\mathbf{H}_{\text{gate}}) \odot \mathbf{H}_{\text{up}}

      3. Linear2 投影（量化 + INT8 矩阵乘 + 反量化）

      Linear2 投影作为第二层线性变换，将中间激活表示 Ae\mathbf{A}_e 从高维空间投影回原始的隐藏维度 hidden\text{hidden}，使专家输出能够与残差连接等后续操作兼容。 在 A8W8 场景下，需要将激活值量化为 INT8，因此先对 Ae\mathbf{A}_e 的每一行（每个 token）计算缩放因子：

      siAe=max⁡(∣Ae[i,:]∣)127,i=0,…,Ne−1s^{A_e}_i = \frac{\max(|\mathbf{A}_e[i,:]|)}{127}, \quad i=0,\dots,N_e-1

      得到专家 ee 在 Linear2 计算时的激活缩放因子 seAe∈RNe\mathbf{s}^{A_e}_e \in \mathbb{R}^{N_e}。然后量化：

      Aeint8[i,:]=⌊Ae[i,:]siAe⌉∈(Z8sym)Ne×intermediate_hidden\mathbf{A}_e^{\text{int8}}[i,:] = \left\lfloor \frac{\mathbf{A}_e[i,:]}{s^{A_e}_i} \right\rceil \quad \in \left(\mathbb{Z}_8^{\text{sym}}\right)^{N_e \times \text{intermediate\_hidden}}

      再执行 INT8 矩阵乘法并反量化：

      Deint32=Aeint8⋅W2[e]int8∈Z32Ne×hidden\mathbf{D}_e^{\text{int32}} = \mathbf{A}_e^{\text{int8}} \cdot \mathbf{W}_2[e]^{\text{int8}} \quad \in \mathbb{Z}_{32}^{N_e \times \text{hidden}} Ye=(Deint32⊙seW2)⊙seAe∈RNe×hidden\mathbf{Y}_e = \left( \mathbf{D}_e^{\text{int32}} \odot \mathbf{s}^{W2}_e \right) \odot \mathbf{s}^{A_e}_e \quad \in \mathbb{R}^{N_e \times \text{hidden}}

      经过以上计算，专家 ee 的每一行输出对应其批次中的一个 token。对于 tokeni\text{token}_i，它在专家 ei,ke_{i,k} 中的输出行即为 Yei,k[row⁡(i,k), :]\mathbf{Y}_{e_{i,k}}\bigl[\operatorname{row}(i,k),\,:\bigr]。

    • Token Combine

      Combine 负责收集所有专家计算出的输出向量，按照每个 token 原先分配到的专家权重进行加权求和，最终为每个 token 生成一个融合后的输出。利用之前记录的位置索引 row⁡(i,k)\operatorname{row}(i,k)，从专家 ei,ke_{i,k} 的输出矩阵中收回属于 tokeni\text{token}_i 的行，并与门控权重相乘后求和：

      yi=∑k=0num_topk−1wk  ⋅  Yei,k ⁣[ row⁡(i,k), : ]∈R1×hidden\mathbf{y}_i = \sum_{k=0}^{\text{num\_topk} - 1} w_k \;\cdot\; \mathbf{Y}_{e_{i,k}}\!\bigl[\,\operatorname{row}(i,k),\,:\,\bigr] \quad \in \mathbb{R}^{1 \times \text{hidden}}

      其中 wk=G[i,k]w_k = \mathbf{G}[i,k] 为 tokeni\text{token}_i 对专家 ei,ke_{i,k} 的门控权重。

      所有 token 的输出按输入顺序堆叠为最终输出 Y∈Rtotal_num_tokens×hidden\mathbf{Y} \in \mathbb{R}^{\text{total\_num\_tokens} \times \text{hidden}}。

    [object Object][object Object]

    • 输入

      • SW1∈Rnum_experts×(2×intermediate_hidden)\mathbf{S}^{W1} \in \mathbb{R}^{\text{num\_experts} \times (2\times \text{intermediate\_hidden})}：Linear1 权重矩阵的逐通道缩放因子，对应入参 [object Object]。
      • SW2∈Rnum_experts×hidden\mathbf{S}^{W2} \in \mathbb{R}^{\text{num\_experts} \times \text{hidden}}：Linear2 权重矩阵的逐通道缩放因子，对应入参 [object Object]。
      • B1∈Rnum_experts×(2×intermediate_hidden)\mathbf{B}_1 \in \mathbb{R}^{\text{num\_experts} \times (2\times \text{intermediate\_hidden})}：Linear1 的偏置矩阵，由 INT4 量化过程离线生成，对应入参 [object Object]。
      • B2∈Rnum_experts×hidden\mathbf{B}_2 \in \mathbb{R}^{\text{num\_experts} \times \text{hidden}}：Linear2 的偏置矩阵，由 INT4 量化过程离线生成，对应入参 [object Object]。

    • EP Dispatch

      在 Dispatch 通信之前，首先将原始 BF16 激活矩阵 X\mathbf{X} 量化为 INT8。对每个 tokeni\text{token}_i，计算其逐 token 缩放因子：

      siX=max⁡(∣X[i,:]∣)127∈Rs^{X}_i = \frac{\max(|\mathbf{X}[i,:]|)}{127} \in \mathbb{R}

      然后量化得到 INT8 表示：

      qi=⌊X[i,:]siX⌉∈(Z8sym)1×hidden\mathbf{q}_i = \left\lfloor \frac{\mathbf{X}[i,:]}{s^{X}_i} \right\rceil \quad \in \left(\mathbb{Z}_8^{\text{sym}}\right)^{1 \times \text{hidden}}

      在 Dispatch 通信阶段，每个 tokeni\text{token}_i 将其量化后的向量 qi\mathbf{q}_i 和缩放因子 siXs^{X}_i 发送给专家 ei,0,ei,1,…,ei,num_topk−1e_{i,0}, e_{i,1}, \dots, e_{i,\text{num\_topk}-1}

      记 Ie={ i∣∃k, E[i,k]=e }I_e = \{\, i \mid \exists k,\ \mathbf{E}[i,k]=e \,\} 为所有被分派给专家 ee 的 token 索引集合，集合 IeI_e 的大小 Ne=∣Ie∣N_e = |I_e| 即为专家 ee 需要处理的 token 总数。则有 Qe∈(Z8sym)Ne×hidden\mathbf{Q}_e \in \left(\mathbb{Z}_8^{\text{sym}}\right)^{N_e \times \text{hidden}} 是由所有满足 i∈Iei \in I_e 的 qi\mathbf{q}_i 按任意固定顺序行堆叠而成的矩阵，该矩阵即为专家 ee 经过 Dispatch 后收到的全部 token 表示；同理，对应的专家 ee 收到的缩放因子向量记为 seX∈RNe\mathbf{s}^{X}_e \in \mathbb{R}^{N_e}，其元素由所有满足 i∈Iei \in I_e 的 siXs^{X}_i 按与 Qe\mathbf{Q}_e 相同的行顺序堆叠而成。

      对于每个 tokeni\text{token}_i 及其选中的第 kk 个专家 ei,ke_{i,k}，存在唯一的行索引 row⁡(i,k)∈{0,…,Nei,k−1}\operatorname{row}(i,k) \in \{0,\dots,N_{e_{i,k}}-1\}，使得 Qei,k[row⁡(i,k),:]=qi\mathbf{Q}_{e_{i,k}}[\operatorname{row}(i,k), :] = \mathbf{q}_i。该映射记录了 qi\mathbf{q}_i 在专家 ei,ke_{i,k} 的输入矩阵中的位置。

    • Expert Compute

      在 MoE 层中，每个专家本质上是一个独立的前馈网络（FFN），采用 SwiGLU 结构以提升表达能力。在A8W4-INT场景下，两个线性层都采用MSD（Mixed-precision Split-activation Decomposition，混合精度激活拆分分解）方案进行矩阵乘，通过将 INT8 值拆为高4位和低4位两个有符号 INT4，使得 INT8×INT4 的矩阵乘可分解为两个 INT4×INT4 的矩阵乘，从而利用硬件的 INT4 矩阵乘加速。该方案的数学原理和实现逻辑可参阅。

      1. 生成精度补偿的偏置矩阵（离线生成，在算子外完成，并作为算子输入）

      MSD 方案的核心是将量化后的 INT8 激活值二进制重解释地拆分为两个 INT4 分量。由于低位 INT4 分量按 (Xint8&0x0F)−8(\mathbf{X}^{\text{int8}} \mathbin{\&} 0x0F) - 8 定义，其数值范围被映射到 [−8,7][-8, 7]，这相当于在原始无符号低 4 位（0~15）的基础上减去了 8。当将高、低位的 INT4 分别与权重矩阵做矩阵乘并合并时，该偏移会引入一个常数项，需要在最终结果中予以补偿。

      记原始 INT8 激活矩阵为 Xint8\mathbf{X}^{\text{int8}}，拆分后的高位和低位 INT4 激活矩阵分别为 X1int4\mathbf{X}_1^{\text{int4}} 和 X2int4\mathbf{X}_2^{\text{int4}}，权重矩阵为 W\mathbf{W}。恢复关系为：

      Xint8=16×X1int4+(X2int4+8⋅1mat)\mathbf{X}^{\text{int8}} = 16 \times \mathbf{X}_1^{\text{int4}} + (\mathbf{X}_2^{\text{int4}} + 8 \cdot \mathbf{1}_{\text{mat}})

      其中 1mat\mathbf{1}_{\text{mat}} 为与 X2int4\mathbf{X}_2^{\text{int4}} 形状相同的全 1 矩阵，用于逐元素加 8。令 1\mathbf{1} 为形状与输入特征维度一致的全 1 列向量，则矩阵乘展开为：

      Xint8⋅W=16⋅(X1int4⋅W)+(X2int4⋅W)+8⋅(1⊤⋅W)\mathbf{X}^{\text{int8}} \cdot \mathbf{W} = 16 \cdot (\mathbf{X}_1^{\text{int4}} \cdot \mathbf{W}) + (\mathbf{X}_2^{\text{int4}} \cdot \mathbf{W}) + 8 \cdot (\mathbf{1}^\top \cdot \mathbf{W})

      这里 1⊤⋅W\mathbf{1}^\top \cdot \mathbf{W} 表示对权重矩阵 W\mathbf{W} 的每一列求和，其结果为一个行向量（形状与输出维度相同）。在具体实现中，该行向量会被广播到批次维度，加到所有 token 的输出上。可见，若只计算前两项，会缺失一个正数项。为了精确恢复结果，需预先计算补偿偏置：

      bias=8⋅(1⊤⋅W),\text{bias} = 8 \cdot (\mathbf{1}^\top \cdot \mathbf{W}),

      B1\mathbf{B}_1 和 B2\mathbf{B}_2 均按此法，分别使用对应的全 1 列向量（维度分别为 hidden\text{hidden} 和 intermediate_hidden\text{intermediate\_hidden}）与各自的权重矩阵相乘，得到形状匹配的偏置行向量。这些偏置在离线阶段计算完成后传入算子，在合并高低位结果时直接参与加法，从而消除偏移引入的误差。

      2. Linear1 投影（INT4 × INT4 矩阵乘 + 反量化）

      Linear1 投影是专家网络的第一层线性变换，同时产生 gate 部分 和 up 部分 所需的预激活值。

      (1) 激活重解释为INT4

      将 INT8 激活张量 Qeint8∈(Z8sym)Ne×hidden\mathbf{Q}_e^{\text{int8}} \in \left(\mathbb{Z}_8^{\text{sym}}\right)^{N_e \times \text{hidden}} 二进制重解释为两个 INT4 交替拼接的视图：

      Qeint4=bitcastZ42Ne×hidden(Qeint8)∈Z42Ne×hidden\mathbf{Q}_e^{\text{int4}} = \mathrm{bitcast}_{\mathbb{Z}_4^{2N_e \times \text{hidden}}} \left( \mathbf{Q}_e^{\text{int8}} \right) \in \mathbb{Z}_4^{2N_e \times \text{hidden}}

      其中高位和低位分量直接由 Qeint4\mathbf{Q}_e^{\text{int4}} 的偶数行和奇数行给出：

      Qehigh=Qeint4[0::2,:]∈Z4Ne×hidden,Qelow=Qeint4[1::2,:]∈Z4Ne×hidden\mathbf{Q}_e^{\text{high}} = \mathbf{Q}_e^{\text{int4}}[0::2, :] \in \mathbb{Z}_4^{N_e \times \text{hidden}}, \quad \mathbf{Q}_e^{\text{low}} = \mathbf{Q}_e^{\text{int4}}[1::2, :] \in \mathbb{Z}_4^{N_e \times \text{hidden}}

      它们在数值上可以由原 INT8 经如下计算得到：

      Qehigh=⌊Qeint816⌋,Qelow=(Qeint8&0x0F)−8\mathbf{Q}_e^{\text{high}} = \left\lfloor \frac{\mathbf{Q}_e^{\text{int8}}}{16} \right\rfloor, \quad \mathbf{Q}_e^{\text{low}} = (\mathbf{Q}_e^{\text{int8}} \mathbin{\&} 0x0F) - 8

      恢复关系为 Qeint8=16Qehigh+(Qelow+8)\mathbf{Q}_e^{\text{int8}} = 16\mathbf{Q}_e^{\text{high}} + (\mathbf{Q}_e^{\text{low}} + 8)。由于 bitcast\mathrm{bitcast} 仅改变类型视图，Qeint4\mathbf{Q}_e^{\text{int4}} 与 Qeint8\mathbf{Q}_e^{\text{int8}} 共享底层物理内存，无需任何数据重排或拷贝。

      (2) INT4 × INT4 矩阵乘与权重反量化

      将重解释后的 INT4 激活视图 Qeint4∈Z42Ne×hidden\mathbf{Q}_e^{\text{int4}} \in \mathbb{Z}_4^{2N_e \times \text{hidden}} 与权重矩阵 W1[e]∈Rhidden×2⋅intermediate_hidden\mathbf{W}_1[e] \in \mathbb{R}^{\text{hidden} \times 2\cdot\text{intermediate\_hidden}} 执行矩阵乘，并应用权重缩放因子 seW1\mathbf{s}^{W1}_e 进行反量化，得到结果：

      Ce=(Qeint4⋅W1[e])⊙seW1  ∈  R2Ne×2⋅intermediate_hidden\mathbf{C}_e = \bigl( \mathbf{Q}_e^{\text{int4}} \cdot \mathbf{W}_1[e] \bigr) \odot \mathbf{s}^{W1}_e \;\in\; \mathbb{R}^{2N_e \times 2\cdot\text{intermediate\_hidden}}

      将结果按偶数行和奇数行分别记为如下的矩阵视图，它们即对应于Qehigh\mathbf{Q}_e^{\text{high}}和Qelow\mathbf{Q}_e^{\text{low}}的计算结果：

      Cehigh=Ce[0::2,:]∈RNe×2⋅intermediate_hidden,Celow=Ce[1::2,:]∈RNe×2⋅intermediate_hidden\mathbf{C}_e^{\text{high}} = \mathbf{C}_e[0::2, :] \in \mathbb{R}^{N_e \times 2\cdot\text{intermediate\_hidden}}, \qquad \mathbf{C}_e^{\text{low}} = \mathbf{C}_e[1::2, :] \in \mathbb{R}^{N_e \times 2\cdot\text{intermediate\_hidden}}

      (3) 精度补偿与激活反量化

      利用偏置 B1[e]\mathbf{B}_1[e] 和激活缩放因子 seX\mathbf{s}^{X}_e（行向量）分别进行精度补偿和激活反量化，最终预激活值为：

      He=(16⋅Cehigh+Celow+B1[e])⊙seX∈RNe×2⋅intermediate_hidden\mathbf{H}_e = \Bigl( 16 \cdot \mathbf{C}_e^{\text{high}} + \mathbf{C}_e^{\text{low}} + \mathbf{B}_1[e] \Bigr) \odot \mathbf{s}^{X}_e \quad\in \mathbb{R}^{N_e \times 2\cdot\text{intermediate\_hidden}}

      其中 ⊙seX\odot \mathbf{s}^{X}_e 表示将矩阵的每一行乘以 seX\mathbf{s}^{X}_e 中对应的标量。

      3. SwiGLU 激活

      首先将 He\mathbf{H}_e 沿列维度拆分为 gate 部分 Hgate\mathbf{H}_{\text{gate}} 和 up 部分 Hup\mathbf{H}_{\text{up}}，然后对 gate 部分应用 SiLU 激活函数，再与 up 部分逐元素相乘，得到专家 ee 的中间激活表示 Ae\mathbf{A}_e。

      Hgate=He[:,  : ⁣intermediate_hidden],Hup=He[:,  intermediate_hidden:]\mathbf{H}_{\text{gate}} = \mathbf{H}_e[:,\; :\!\text{intermediate\_hidden}], \qquad \mathbf{H}_{\text{up}} = \mathbf{H}_e[:,\; \text{intermediate\_hidden}:] Ae=SiLU⁡(Hgate)⊙Hup∈RNe×intermediate_hidden\mathbf{A}_e = \operatorname{SiLU}(\mathbf{H}_{\text{gate}}) \odot \mathbf{H}_{\text{up}} \quad\in \mathbb{R}^{N_e \times \text{intermediate\_hidden}}

      4. Linear2 投影（量化 + INT4 × INT4 矩阵乘 + 反量化）

      Linear2 投影作为第二层线性变换，将中间激活表示 Ae\mathbf{A}_e 从高维空间投影回原始的隐藏维度 hidden\text{hidden}，使专家输出能够与残差连接等后续操作兼容。 在 A8W4 场景下，首先将激活值量化为 INT8，再通过二进制重解释为两个 INT4 的拼接视图，以一次矩阵乘完成计算，过程如下。

      (1) 激活量化

      对 Ae\mathbf{A}_e 的每一行计算缩放因子并量化至 INT8：

      siAe=max⁡(∣Ae[i,:]∣)127,i=0,…,Ne−1s^{A_e}_i = \frac{\max(|\mathbf{A}_e[i,:]|)}{127}, \quad i=0,\dots,N_e-1

      得到专家 ee 在 Linear2 计算时的激活缩放因子 seAe∈RNe\mathbf{s}^{A_e}_e \in \mathbb{R}^{N_e}，并计算量化结果：

      Aeint8[i,:]=⌊Ae[i,:]siAe⌉∈(Z8sym)Ne×intermediate_hidden\mathbf{A}_e^{\text{int8}}[i,:] = \left\lfloor \frac{\mathbf{A}_e[i,:]}{s^{A_e}_i} \right\rceil \quad \in \left(\mathbb{Z}_8^{\text{sym}}\right)^{N_e \times \text{intermediate\_hidden}}

      (2) 激活重解释为 INT4

      将 INT8 激活张量 Aeint8∈(Z8sym)Ne×intermediate_hidden\mathbf{A}_e^{\text{int8}} \in \left(\mathbb{Z}_8^{\text{sym}}\right)^{N_e \times \text{intermediate\_hidden}} 二进制重解释为两个 INT4 交替拼接的视图：

      Aeint4=bitcastZ42Ne×intermediate_hidden(Aeint8)∈Z42Ne×intermediate_hidden\mathbf{A}_e^{\text{int4}} = \mathrm{bitcast}_{\mathbb{Z}_4^{2N_e \times \text{intermediate\_hidden}}} \left( \mathbf{A}_e^{\text{int8}} \right) \in \mathbb{Z}_4^{2N_e \times \text{intermediate\_hidden}}

      其中高位和低位分量直接由 Aeint4\mathbf{A}_e^{\text{int4}} 的偶数行和奇数行给出：

      Aehigh=Aeint4[0::2,:]∈Z4Ne×intermediate_hidden,Aelow=Aeint4[1::2,:]∈Z4Ne×intermediate_hidden\mathbf{A}_e^{\text{high}} = \mathbf{A}_e^{\text{int4}}[0::2, :] \in \mathbb{Z}_4^{N_e \times \text{intermediate\_hidden}}, \quad \mathbf{A}_e^{\text{low}} = \mathbf{A}_e^{\text{int4}}[1::2, :] \in \mathbb{Z}_4^{N_e \times \text{intermediate\_hidden}}

      它们在数值上可以由原 INT8 经如下计算得到：

      Aehigh=⌊Aeint816⌋,Aelow=(Aeint8&0x0F)−8\mathbf{A}_e^{\text{high}} = \left\lfloor \frac{\mathbf{A}_e^{\text{int8}}}{16} \right\rfloor, \quad \mathbf{A}_e^{\text{low}} = (\mathbf{A}_e^{\text{int8}} \mathbin{\&} 0x0F) - 8

      恢复关系为 Aeint8=16Aehigh+(Aelow+8)\mathbf{A}_e^{\text{int8}} = 16\mathbf{A}_e^{\text{high}} + (\mathbf{A}_e^{\text{low}} + 8)。由于 bitcast\mathrm{bitcast} 仅改变类型视图，Aeint4\mathbf{A}_e^{\text{int4}} 与 Aeint8\mathbf{A}_e^{\text{int8}} 共享底层物理内存，无需任何数据重排或拷贝。

      (3) INT4 × INT4 矩阵乘与权重反量化

      将重解释后的 INT4 激活视图 Aeint4∈Z42Ne×intermediate_hidden\mathbf{A}_e^{\text{int4}} \in \mathbb{Z}_4^{2N_e \times \text{intermediate\_hidden}} 与权重矩阵 W2[e]∈Rintermediate_hidden×hidden\mathbf{W}_2[e] \in \mathbb{R}^{\text{intermediate\_hidden} \times \text{hidden}} 执行矩阵乘，并应用权重缩放因子 seW2\mathbf{s}^{W2}_e 进行反量化，得到结果：

      De=(Aeint4⋅W2[e])⊙seW2  ∈  R2Ne×hidden\mathbf{D}_e = \bigl( \mathbf{A}_e^{\text{int4}} \cdot \mathbf{W}_2[e] \bigr) \odot \mathbf{s}^{W2}_e \;\in\; \mathbb{R}^{2N_e \times \text{hidden}}

      将结果按偶数行和奇数行分别记为如下的矩阵视图，它们即对应于 Aehigh\mathbf{A}_e^{\text{high}} 和 Aelow\mathbf{A}_e^{\text{low}} 的计算结果：

      Dehigh=De[0::2,:]∈RNe×hidden,Delow=De[1::2,:]∈RNe×hidden\mathbf{D}_e^{\text{high}} = \mathbf{D}_e[0::2, :] \in \mathbb{R}^{N_e \times \text{hidden}}, \qquad \mathbf{D}_e^{\text{low}} = \mathbf{D}_e[1::2, :] \in \mathbb{R}^{N_e \times \text{hidden}}

      (4) 精度补偿与激活反量化

      利用偏置 B2[e]\mathbf{B}_2[e] 和激活缩放因子 seAe\mathbf{s}^{A_e}_e（行向量）分别进行精度补偿和激活反量化，最终输出为：

      Ye=(16⋅Dehigh+Delow+B2[e])⊙seAe∈RNe×hidden\mathbf{Y}_e = \Bigl( 16 \cdot \mathbf{D}_e^{\text{high}} + \mathbf{D}_e^{\text{low}} + \mathbf{B}_2[e] \Bigr) \odot \mathbf{s}^{A_e}_e \quad\in \mathbb{R}^{N_e \times \text{hidden}}

      其中 ⊙seAe\odot \mathbf{s}^{A_e}_e 表示将矩阵的每一行乘以 seAe\mathbf{s}^{A_e}_e 中对应的标量。

      经过以上计算，专家 ee 的每一行输出对应其批次中的一个 token。对于 tokeni\text{token}_i，它在专家 ei,ke_{i,k} 中的输出行即为 Yei,k[row⁡(i,k), :]\mathbf{Y}_{e_{i,k}}\bigl[\operatorname{row}(i,k),\,:\bigr]。

    • Token Combine

      Combine 负责收集所有专家计算出的输出向量，按照每个 token 原先分配到的专家权重进行加权求和，最终为每个 token 生成一个融合后的输出。利用之前记录的位置索引 row⁡(i,k)\operatorname{row}(i,k)，从专家 ei,ke_{i,k} 的输出矩阵中收回属于 tokeni\text{token}_i 的行，并与门控权重相乘后求和：

      yi=∑k=0num_topk−1wk  ⋅  Yei,k ⁣[ row⁡(i,k), : ](∈R1×hidden)\mathbf{y}_i = \sum_{k=0}^{\text{num\_topk} - 1} w_k \;\cdot\; \mathbf{Y}_{e_{i,k}}\!\bigl[\,\operatorname{row}(i,k),\,:\,\bigr] \qquad\bigl(\in \mathbb{R}^{1 \times \text{hidden}}\bigr)

      其中 wk=G[i,k]w_k = \mathbf{G}[i,k] 为 tokeni\text{token}_i 对专家 ei,ke_{i,k} 的门控权重。

      所有 token 的输出按输入顺序堆叠为最终输出 Y∈Rtotal_num_tokens×hidden\mathbf{Y} \in \mathbb{R}^{\text{total\_num\_tokens} \times \text{hidden}}。

    [object Object][object Object]

    第一阶段对输入 Token 按专家分组收集后做 MXFP8 量化，生成各专家的量化输入与缩放因子：

    X^e, SX,e=QMX ⁣(X[Te]),e=0,1,…,Elocal−1\hat{X}_e,\ S_{X,e} = \mathrm{Q}_{\text{MX}}\!\left(X[\mathcal{T}_e]\right), \quad e = 0, 1, \ldots, E_{\text{local}}-1

    说明：根据 [object Object] 将 Token 按专家排序收集，Te\mathcal{T}_e 为分配到专家 ee 的 Token 索引集合，ElocalE_{local}表示当前专家收到的最大token数，每个专家数值可能不同，X[Te]X[\mathcal{T}_e] 为对应的子矩阵。QMX\mathrm{Q}_{\text{MX}} 表示 MX 逐组量化（group size = 32），对每组 32 个元素提取共享指数后量化为 FP8 目标类型（FLOAT8_E5M2 或 FLOAT8_E4M3FN），同时输出 FLOAT8_E8M0 缩放因子。量化后的数据将作为 GMM1 的输入。

    第二阶段对每个专家执行 GMM1 矩阵乘法（将 W1W_1 沿列方向分为两半分别计算）、SwiGLU 激活和 MX 量化：

    Ze(x)=DQMX(X^e,SX,e)⋅DQMX(W1,e(x),S1,e(x)),Ze(y)=DQMX(X^e,SX,e)⋅DQMX(W1,e(y),S1,e(y))Z_e^{(x)} = \mathrm{DQ}_{\text{MX}}(\hat{X}_e, S_{X,e}) \cdot \mathrm{DQ}_{\text{MX}}(W_{1,e}^{(x)}, S_{1,e}^{(x)}), \quad Z_e^{(y)} = \mathrm{DQ}_{\text{MX}}(\hat{X}_e, S_{X,e}) \cdot \mathrm{DQ}_{\text{MX}}(W_{1,e}^{(y)}, S_{1,e}^{(y)}) Ue=Ze(x)⊙σ ⁣(Ze(x))⊙Ze(y)U_e = Z_e^{(x)} \odot \sigma\!\left(Z_e^{(x)}\right) \odot Z_e^{(y)} U^e, SU,e=QMX(Ue)\hat{U}_e,\ S_{U,e} = \mathrm{Q}_{\text{MX}}(U_e)

    说明：将 W1W_1 的前 N/2N/2 列 W1,e(x)W_{1,e}^{(x)} 和后 N/2N/2 列 W1,e(y)W_{1,e}^{(y)} 分别与 MX 反量化后的输入做矩阵乘法，得到 Swish 分支 Ze(x)Z_e^{(x)} 和门控分支 Ze(y)Z_e^{(y)}。SwiGLU 激活对两个分支做逐元素乘积 x⋅σ(x)⋅yx \cdot \sigma(x) \cdot y，其中 σ\sigma 为 Sigmoid 函数，将中间维度从 NN 减半为 N/2N/2。随后对 SwiGLU 输出做 MX 量化，得到 GMM2 的量化输入 U^e\hat{U}_e。

    第三阶段对每个专家执行 GMM2 矩阵乘法，并将结果按目标 Rank 分发：

    Oe=DQMX(U^e,SU,e)⋅DQMX(W2,e,S2,e)O_e = \mathrm{DQ}_{\text{MX}}(\hat{U}_e, S_{U,e}) \cdot \mathrm{DQ}_{\text{MX}}(W_{2,e}, S_{2,e})

    说明：将量化后的 SwiGLU 输出与第二组权重 W2W_2 做 MX 反量化后的矩阵乘法，将 N/2N/2 维中间表示映射回 HH 维隐藏空间，得到每个专家的输出 OeO_e。计算完成后通过 RDMA peermem 将结果按目标 Rank 的专家偏移地址写入远端，实现跨 Rank 聚合。

    第四阶段对所有 Token 按路由权重加权求和，恢复为与输入相同形状的输出：

    Y[i]=∑k=0K−1W[i, k]⋅O[π(i, k)]Y[i] = \sum_{k=0}^{K-1} W[i,\, k] \cdot O[\pi(i,\, k)]

    说明：对每个 Token ii，根据排序后的路由索引 π(i,k)\pi(i,k) 从聚合后的专家结果中取出对应行，按 [object Object] 中的权重逐元素加权累加，得到最终输出 YY。

    其中，XX 表示参数 [object Object]，WW 表示参数 [object Object]，W1W_1 表示参数 [object Object]，W2W_2 表示参数 [object Object]，YY 表示参数 [object Object]，ElocalE_{\text{local}} 表示属性 [object Object]（每个 Rank 的专家数），KK 表示 [object Object] 的第二维度（top-K 值，取值 6 或 8）。

    局部变量说明：

    • Te\mathcal{T}_e：被路由到专家 ee 的 Token 索引集合，由 [object Object] 排序后确定。
    • X^e, SX,e\hat{X}_e,\ S_{X,e}：专家 ee 的量化输入及其 MX 缩放因子，第一阶段中间结果。
    • W1,e(x)W_{1,e}^{(x)}、W1,e(y)W_{1,e}^{(y)}：W1W_1 对应专家 ee 的前 N/2N/2 列和后 N/2N/2 列子矩阵，由 [object Object] 按 SwiGLU 拆分推导。
    • S1,e(x)S_{1,e}^{(x)}、S1,e(y)S_{1,e}^{(y)}：W1,e(x)W_{1,e}^{(x)} 和 W1,e(y)W_{1,e}^{(y)} 对应的 MX 缩放因子，从 [object Object] 按维度截取。
    • S2,eS_{2,e}：W2,eW_{2,e} 对应的 MX 缩放因子，来自参数 [object Object]。
    • Ze(x), Ze(y)Z_e^{(x)},\ Z_e^{(y)}：GMM1 的两路矩阵乘法输出（Swish 分支和门控分支），中间结果。
    • UeU_e：SwiGLU 激活输出，维度 me×N/2m_e \times N/2，中间结果。
    • U^e, SU,e\hat{U}_e,\ S_{U,e}：量化后的 SwiGLU 输出及其 MX 缩放因子，中间结果。
    • OeO_e：GMM2 的专家级输出，维度 me×Hm_e \times H，中间结果。
    • π(i,k)\pi(i, k)：Token ii 的第 kk 个 top-k 专家在展开排序后的行索引，由路由排序确定。
    • QMX(⋅)\mathrm{Q}_{\text{MX}}(\cdot)：MX 逐组量化操作，block size = 32，输出 FP8 数据和 E8M0 缩放因子。
    • DQMX(⋅)\mathrm{DQ}_{\text{MX}}(\cdot)：MX 逐组反量化操作，在 matmul 内部隐式执行。[object Object]
[object Object]
  • get_symm_buffer_for_mega_moe：
[object Object]
  • mega_moe：
[object Object]
[object Object][object Object][object Object][object Object]

上角标[object Object]1[object Object]表示Atlas A2 训练系列产品/Atlas A2 推理系列产品、Atlas A3 训练系列产品/Atlas A3 推理系列产品不支持，上角标[object Object]2[object Object]表示Ascend 950PR/Ascend 950DT不支持，产品不支持的参数使用默认值即可。

[object Object][object Object][object Object][object Object][object Object][object Object][object Object]

  • 各张量参数的list[Tensor]长度、是否转置、是否支持非连续Tensor约束如下：

    [object Object]

  • 参数一致性约束：

    • mega_moe 接口的所有输入参数及其对应的张量维度，必须与 get_symm_buffer_for_mega_moe 的同名参数（例如 [object Object]、[object Object]、[object Object] 等）保持一致。
    • 调用算子过程中使用的[object Object]、[object Object]、[object Object]、[object Object]、[object Object]等参数取值，所有卡需保持一致，网络中不同层中也需保持一致。

  • 通信域和组网约束：

    • 所有卡的[object Object]、[object Object]参数取值需保持一致。

    • 各卡的通信域缓存区大小应当一致。[object Object] 为 HBM 上分配的 CCL 通信缓冲区总大小（Bytes），包含等大小的 windowIn 和 windowOut 两块空间，校验时以单个空间 [object Object] 为准，需满足：

      ccl_buffer_size / 2≥offsetTokenPerExpert+offsetTensor+offsetFlag+10 MBccl\_buffer\_size\ /\ 2 \ge \mathrm{offsetTokenPerExpert} + \mathrm{offsetTensor} + \mathrm{offsetFlag} + 10\,\mathrm{MB}

      Atlas A2 训练系列产品/Atlas A2 推理系列产品：

      [object Object]

      Atlas A3 训练系列产品/Atlas A3 推理系列产品：

      [object Object]

      其中 [object Object] 即通信域大小，maxExpertPerRankmaxExpertPerRank 表示每张卡上可能专家数的最大值，quant\mathrm{quant} 表示是否开启 dispatch 量化（[object Object]）。预留空间 10 MB 为内部元数据对齐与安全余量。

    • 通信域各节点的驱动版本应当相同。

    • Atlas A2 训练系列产品/Atlas A2 推理系列产品：多机通信域要求交换机组网，不支持双机直连组网。

    • Atlas A3 训练系列产品/Atlas A3 推理系列产品：多机通信域要求在一个超节点内，不支持双机直连组网和跨超节点组网。

  • 参数约束：

    • Atlas A2 训练系列产品/Atlas A2 推理系列产品、Atlas A3 训练系列产品/Atlas A3 推理系列产品：
      • 各卡 [object Object] 需保持一致。
      • [object Object]：取值为 [object Object]、[object Object]、[object Object]、[object Object]、[object Object]。
      • [object Object]：取值范围为[object Object]，且 [object Object]。
      • [object Object]：取值范围为 [object Object]，且 [object Object]。
      • [object Object]：取值范围为 [object Object]。
      • [object Object] 需大于0，输入0表示自动计算，公式为 [object Object]。
      • [object Object]：取值范围为 [object Object]。
      • [object Object]：取值范围为 [object Object]，且 [object Object]。
      • [object Object]：取值范围为 [object Object]，且 [object Object]。
      • [object Object]：取值范围为 [object Object]（非量化）、[object Object]（pertoken量化）。
      • [object Object]：取值为 [object Object]。
      • 支持三种计算场景（A16W16、A8W8-INT、A8W4-INT），不同场景下可选入参（缩放因子、偏置等）的必需性及数据类型有严格配套要求。调用时必须根据所选场景完整提供对应参数，不可混用或遗漏，配套关系见下表。
      [object Object]
    • Ascend 950PR/Ascend 950DT：

      • num_tokens（x.dim0）范围 [1, maxBs], maxBs = (8 (totalUbSize - 48 1024)) / (num_topk * (64 + ep_world_size)), 其中totalUbSize为UB总大小，950系列产品该值为256K。

      • hidden（x.dim1）仅支持1024、2048、3072、4096、5120、6144、7168、8192。

      • num_topk（topk_ids.dim1）支持[1, 16]。

      • num_experts_per_rank（weight1.dim0）范围 [1, 1024]。

      • intermediate_hidden（weight1.dim1）仅支持1024、2048、3072、4096、7168。

      • ep_world_size范围 [2, 1024]。

      • num_experts范围 [ep_world_size, 2048]，且num_experts % ep_world_size == 0。

      • max_recv_token_num范围 [0, num_tokens × ep_world_size × min(num_topk, num_experts_per_rank)]。

      • dispatch_quant_out_dtype仅支持torch.float8_e5m2或torch.float8_e4m3fn或torch.float4_e2m1。

      • 当前版本仅支持MXFP量化模式（dispatch_quant_mode = 4），dispatch阶段使用MX逐组量化（group size = 32），量化缩放因子类型为FLOAT8_E8M0。

      • x_active_mask和scales参数当前版本必须传入None，不支持非空输入。

      • combine_quant_mode当前支持0（非量化），3（MX模式float8_e5m2类型），4（MX模式float8_e4m3类型）。

      • comm_alg预留参数，必须为空字符串""。

      • y的数据类型与x相同。

      • weight1的dim1（intermediate_hidden）必须等于weight2的dim2的二倍，这是因为SwiGLU激活需要将中间维度从intermediate_hidden减半为intermediate_hidden/2。

      • weight_scales1和weight_scales2不可为空指针。

      • num_experts_per_rank = num_experts / ep_world_size，必须为整数且在 [1, 1024] 范围内。

      • weight_scales1和weight_scales2不可为空指针。

      • MXFP量化场景约束：

        • weight1 shape为(num_experts_per_rank, intermediate_hidden, hidden)，weight2 shape为(num_experts_per_rank, hidden, intermediate_hidden / 2)。
        • weightScales1 shape为(num_experts_per_rank, intermediate_hidden, CeilDiv(hidden, 64), 2)，其中 CeilDiv(hidden, 64) = ⌈hidden / 64⌉ = ⌊(hidden + 63) / 64⌋。
        • weightScales2 shape为(num_experts_per_rank, hidden, CeilDiv(intermediate_hidden / 2, 64), 2)，其中 CeilDiv(intermediate_hidden / 2, 64) = ⌈(intermediate_hidden / 2) / 64⌉ = ⌊(intermediate_hidden / 2 + 63) / 64⌋。
        • weightScales1的dim3和weightScales2的dim3必须等于2。
        • MXFP场景下，dispatch_quant_out_dtype=torch.float8_e5m2时weight1和weight2必须为FLOAT8_E5M2，dispatch_quant_out_dtype=torch.float8_e4m3fn时必须为FLOAT8_E4M3FN，dispatch_quant_out_dtype=torch.float4_e2m1时必须为FLOAT4_E2M1。
        • x_active_mask和scales必须为None。

      • 支持三种计算场景（A8W8-FP、A8W4-FP、A4W4-FP），不同场景下可选入参（缩放因子、偏置等）的必需性及数据类型有严格配套要求。调用时必须根据所选场景完整提供对应参数，不可混用或遗漏，配套关系见下表。

        [object Object]
[object Object]

默认支持确定性计算。

[object Object]

  • 单算子模式调用：

    下面示例将两个接口按调用顺序串联：先初始化通信域，再用 get_symm_buffer_for_mega_moe 构造 sym_buffer，最后调用 mega_moe 运行算子。

    • Ascend 950PR/Ascend 950DT：

      [object Object]

    • Atlas A2 训练系列产品/Atlas A2 推理系列产品、Atlas A3 训练系列产品/Atlas A3 推理系列产品：

      [object Object]
我们使用cookie来确保您的高速浏览体验。继续浏览本站，即表示您同意我们使用cookie。查看详情