mega_moe

Applicable Products

Ascend 950PR/Ascend 950DT: Supported
Atlas A3 training products/Atlas A3 inference products: Supported
Atlas A2 training products/Atlas A2 inference products: Supported
Atlas 200I/500 A2 inference products: Not supported
Atlas inference products: Not supported
Atlas training products: Not supported

Function

Description:

The Mega MoE operator fuses the complete computation flow of the expert FFN in the MoE layer and the preceding and following data communication (that is, Dispatch + Linear1 + SwiGLU + Linear2 + Combine) into a single operator, achieving overlap of communication and computation. This operator provides APIs such as mega_moe and get_symm_buffer_for_mega_moe, which must be used together.
- get_symm_buffer_for_mega_moe: Must be used together with mega_moe. It encapsulates input parameters and creates a SymmBuffer structure, generating the information required for the mega_moe operator to run, such as context, ep_world_size, and ccl_buffer_size.
Formula:
- Input:
  - $\mathbf{X} \in \mathbb{R}^{\text{total\_num\_tokens} \times \text{hidden}}$ : activation matrix, corresponding input parameter x. $\text{total\_num\_tokens}$ is the total number of global tokens, and $\text{hidden}$ is the hidden dimension.
  - $\mathbf{E} \in \mathbb{Z}^{\text{total\_num\_tokens} \times \text{num\_topk}}$ : expert index matrix selected by tokens, corresponding input parameter topk_ids. $\text{num\_topk}$ is the number of experts selected by each token.
  - $\mathbf{G} \in \mathbb{R}^{\text{total\_num\_tokens} \times \text{num\_topk}}$ : gating weight matrix of the experts selected by tokens, corresponding input parameter topk_weights.
  - $\mathbf{W}_1 \in \mathbb{R}^{\text{num\_experts} \times \text{hidden} \times (2 \text{intermediate\_hidden})}$ : weight matrix of Linear1, corresponding to input parameter l1_weights. $\text{num\_experts}$ is the number of experts, and $\text{intermediate\_hidden}$ is the intermediate dimension.
  - $\mathbf{W}_2 \in \mathbb{R}^{\text{num\_experts} \times \text{intermediate\_hidden} \times \text{hidden}}$ : weight matrix of Linear2, corresponding input parameter l2_weights.
- Output:
  - $\mathbf{Y} \in \mathbb{R}^{\text{total\_num\_tokens} \times \text{hidden}}$ : final output matrix, corresponding output parameter y.
- Conventions:
  - $⋅$ denotes matrix multiplication, and $⊙$ denotes element-wise multiplication.
  - $\left \lfloor z\right \rceil$ denotes rounding $z$ to the nearest integer, and $\left \lfloor z\right \rfloor$ denotes rounding $z$ down.
  - $|z|$ takes the absolute value, and $\max(z)$ takes the maximum value.
  - The set of all tokens is $\{ \text{token}_i \mid i \in \{0, 1, \dots, \text{total\_num\_tokens} - 1\} \}$ .
  - The token representation (i.e., hidden state vector) of $\text{token}_i$ is $\mathbf{x}_i \in \mathbb{R}^{1 \times \text{hidden}}$ , and $\mathbf{x}_i = \mathbf{X}[i,:]$ .
  - The expert index of $\text{token}_i$ is $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\}$ .
  - The activation function $\text{SiLU}(z) = z \cdot \sigma(z) = \frac{z}{1+e^{-z}}$ , where $\sigma$ is the Sigmoid function.
  - $\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 \}$ . Here, the superscript $\text{sym}$ of $\mathbb{Z}_8^{\text{sym}}$ denotes the symmetric quantization range: its range is symmetrically rounded about 0, with endpoints $-127$ and $127$ , which differs from the standard INT8 range of $[-128, 127]$ , hence the $\text{sym}$ superscript is used to distinguish it.
  - Tensor slicing uses Python-style start:stop:step notation. For example, $[0::2, :]$ selects even rows, and $[1::2, :]$ selects odd rows.
  - $\mathrm{bitcast}_{T}(\mathbf{Z})$ denotes a binary reinterpretation operation that reinterprets the underlying binary data of tensor $\mathbf{Z}$ as the target type $T$ .
Computation Process

The data types of the activation matrix A and weight matrix W supported by each product during Linear computation are as follows:

Use Case A W
A16W16 BFLOAT16 BFLOAT16
A8W8-INT INT8 INT8
A8W8-FP FLOAT8_E4M3FN, FLOAT8_E5M2 FLOAT8_E4M3FN, FLOAT8_E5M2
A8W4-INT INT8 INT4
- Atlas A2 training products/Atlas A2 inference products, Atlas A3 training products/Atlas A3 inference products: do not support the A8W8-FP use case in the preceding table.
- Ascend 950PR/Ascend 950DT: do not support the A16W16, A8W8-INT, and A8W4-INT use cases in the preceding table.
A16W16 non-quantized use case
- EP Dispatch
  
  In the Dispatch phase, each $\text{token}_i$ sends its token representation $\mathbf{x}_i$ to experts $e_{i,0}, e_{i,1}, \dots, e_{i,\text{num\_topk}-1}$ . That is, for each $k$ , expert $e_{i,k}$ receives a copy of $\mathbf{x}_i$ .
  
  Let $I_e = \{\, i \mid \exists k,\ \mathbf{E}[i,k]=e \,\}$ be the index set of all tokens dispatched to expert $e$ . The size of $I_e$ , $N_e = |I_e|$ , is the total number of tokens that expert $e$ must process. Then $\mathbf{X}_e \in \mathbb{R}^{N_e \times \text{hidden}}$ is the matrix formed by row-stacking in any fixed order the token representations $\mathbf{x}_i$ of all $i$ satisfying $i \in I_e$ . This matrix is all the token representations that expert $e$ receives after Dispatch.
  
  For each $\text{token}_i$ and its selected $k$ -th expert $e_{i,k}$ , there is a unique row index $\operatorname{row}(i,k) \in \{0,\dots,N_{e_{i,k}}-1\}$ such that $\mathbf{X}_{e_{i,k}}[\operatorname{row}(i,k), :] = \mathbf{x}_i$ . This mapping records the position of $\mathbf{x}_i$ in the input matrix $\mathbf{X}_{e_{i,k}}$ of expert $e_{i,k}$ .
- Expert Compute
  
  In an MoE layer, each expert is essentially an independent feedforward network (FFN) that uses the SwiGLU structure to improve expressive capacity. The entire computation process consists of the following three substeps.
  
  1. Linear1 Projection
  
  Linear1 projection is the first linear transformation of the expert network, producing the pre-activation values required by both the gate part and the up part. Its formula is
  $\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 Activation
  
  First split $\mathbf{H}_e$ along the column dimension into the gate part $\mathbf{H}_{\text{gate}}$ and the up part $\mathbf{H}_{\text{up}}$ . Then apply the SiLU activation function to the gate part and multiply it element-wise with the up part to obtain the intermediate activation representation $\mathbf{A}_e$ of expert $e$ .
  $\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]$ $\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 Projection
  
  Linear2 projection, as the second linear transformation, projects the intermediate activation representation $\mathbf{A}_e$ from the high-dimensional space back to the original hidden dimension $\text{hidden}$ , making the expert output compatible with subsequent operations such as residual connections.
  $\mathbf{Y}_e = \mathbf{A}_e \cdot \mathbf{W}_2[e] \quad\bigl(\in \mathbb{R}^{N_e \times \text{hidden}}\bigr)$
  After the above computation, each output row of expert $e$ corresponds to a token in its batch. For $\text{token}_i$ , its output row in expert $e_{i,k}$ is $\mathbf{Y}_{e_{i,k}}\bigl[\operatorname{row}(i,k),\,:\bigr]$ .
- Token Combine
  
  Combine collects the output vectors computed by all experts, performs weighted summation according to the expert weight originally assigned to each token, and finally generates a fused output for each token. Using the previously recorded position index $\operatorname{row}(i,k)$ , it retrieves the rows belonging to $\text{token}_i$ from the output matrix of expert $e_{i,k}$ , multiplies them by the gating weight, and sums them up:
  $\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)$
  where $w_k = \mathbf{G}[i,k]$ is the gating weight of $\text{token}_i$ for expert $e_{i,k}$ .
  
  The outputs of all tokens are stacked in input order as the final output $\mathbf{Y} \in \mathbb{R}^{\text{total\_num\_tokens} \times \text{hidden}}$ .

Use Case	A	W
A16W16	BFLOAT16	BFLOAT16
A8W8-INT	INT8	INT8
A8W8-FP	FLOAT8_E4M3FN, FLOAT8_E5M2	FLOAT8_E4M3FN, FLOAT8_E5M2
A8W4-INT	INT8	INT4

A8W8-INT Quantization Use Case

Input
- $\mathbf{S}^{W1} \in \mathbb{R}^{\text{num\_experts} \times (2\times \text{intermediate\_hidden})}$ : per-channel scaling factor of the Linear1 weight matrix, corresponding to the input parameter l1_weights_sf.
- $\mathbf{S}^{W2} \in \mathbb{R}^{\text{num\_experts} \times \text{hidden}}$ : per-channel scaling factor of the Linear2 weight matrix, corresponding to the input parameter l2_weights_sf.
EP Dispatch

Before Dispatch communication, the original BF16 activation matrix $\mathbf{X}$ is first quantized to INT8. For each $\text{token}_i$ , compute its per-token scaling factor:
$s^{X}_i = \frac{\max(|\mathbf{X}[i,:]|)}{127} \in \mathbb{R}$
Then quantize to obtain the INT8 representation:
$\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}}$
During the Dispatch communication phase, each $\text{token}_i$ sends its quantized vector $\mathbf{q}_i$ and scaling factor $s^{X}_i$ to experts $e_{i,0}, e_{i,1}, \dots, e_{i,\text{num\_topk}-1}$ ( $e_{i,k} = \mathbf{E}[i,k]$ ).

Let $I_e = \{\, i \mid \exists k,\ \mathbf{E}[i,k]=e \,\}$ be the index set of all tokens dispatched to expert $e$ , and let $N_e = |I_e|$ be the total number of tokens that expert $e$ must process. Then $\mathbf{Q}_e \in \left(\mathbb{Z}_8^{\text{sym}}\right)^{N_e \times \text{hidden}}$ is the matrix formed by row-stacking all $\mathbf{q}_i$ satisfying $i \in I_e$ in any fixed order. This matrix is the full set of token representations that expert $e$ receives after Dispatch. Similarly, the corresponding scaling factor vector received by expert $e$ is denoted as $\mathbf{s}^{X}_e \in \mathbb{R}^{N_e}$ , whose elements are formed by stacking all $s^{X}_i$ satisfying $i \in I_e$ in the same row order as $\mathbf{Q}_e$ .

For each $\text{token}_i$ and its selected $k$ -th expert $e_{i,k}$ , there exists a unique row index $\operatorname{row}(i,k) \in \{0,\dots,N_{e_{i,k}}-1\}$ such that $\mathbf{Q}_{e_{i,k}}[\operatorname{row}(i,k), :] = \mathbf{q}_i$ . This mapping records the position of $\mathbf{q}_i$ in the input matrix of expert $e_{i,k}$ .
Expert Compute

In a MoE layer, each expert is essentially an independent feedforward network (FFN) that uses the SwiGLU structure to improve expressive capacity. In the A8W8 use case, both linear layers perform matrix multiplication with INT8 inputs and INT8 weights, producing INT32 intermediate results that are then dequantized. This is divided into three sub-steps.

1. Linear1 projection (INT8 matrix multiplication + dequantization)

The Linear1 projection is the first linear transformation of the expert network, producing the pre-activation values required by both the gate part and the up part. INT8 matrix multiplication is performed to obtain the INT32 result:
$\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}}$
Then dequantize to the pre-activation value $\mathbf{H}_e$ :
$\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 activation

First split $\mathbf{H}_e$ along the column dimension into the gate part $\mathbf{H}_{\text{gate}}$ and the up part $\mathbf{H}_{\text{up}}$ . Then apply the SiLU activation function to the gate part and multiply it element-wise with the up part to obtain the intermediate activation representation $\mathbf{A}_e$ of expert $e$ .
$\mathbf{H}_{\text{gate}} = \mathbf{H}_e[:, :\text{intermediate\_hidden}], \qquad \mathbf{H}_{\text{up}} = \mathbf{H}_e[:, \text{intermediate\_hidden}:]$ $\mathbf{A}_e = \operatorname{SiLU}(\mathbf{H}_{\text{gate}}) \odot \mathbf{H}_{\text{up}}$
3. Linear2 projection (quantization + INT8 matrix multiplication + dequantization)

The Linear2 projection is the second linear transformation. It projects the intermediate activation representation $\mathbf{A}_e$ from the high-dimensional space back to the original hidden dimension $\text{hidden}$ , making the expert output compatible with subsequent operations such as residual connections. In the A8W8 use case, the activation values must be quantized to INT8. Therefore, first compute the scaling factor for each row (each token) of $\mathbf{A}_e$ :
$s^{A_e}_i = \frac{\max(|\mathbf{A}_e[i,:]|)}{127}, \quad i=0,\dots,N_e-1$
This yields the activation scaling factor $\mathbf{s}^{A_e}_e \in \mathbb{R}^{N_e}$ of expert $e$ in the Linear2 computation. Then quantize:
$\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}}$
Then perform INT8 matrix multiplication and dequantize:
$\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}}$ $\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}}$
After the above computation, each output row of expert $e$ corresponds to a token in its batch. For $\text{token}_i$ , its output row in expert $e_{i,k}$ is $\mathbf{Y}_{e_{i,k}}\bigl[\operatorname{row}(i,k),\,:\bigr]$ .
Token Combine

Combine collects the output vectors computed by all experts, performs weighted summation according to the expert weight originally assigned to each token, and finally generates a fused output for each token. Using the previously recorded position index $\operatorname{row}(i,k)$ , retrieve the rows belonging to $\text{token}_i$ from the output matrix of expert $e_{i,k}$ , multiply them by the gating weight, and sum them up:
$\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}}$
where $w_k = \mathbf{G}[i,k]$ is the gating weight of $\text{token}_i$ for expert $e_{i,k}$ .

The outputs of all tokens are stacked in input order as the final output $\mathbf{Y} \in \mathbb{R}^{\text{total\_num\_tokens} \times \text{hidden}}$ .

A8W4-INT Quantization Use Case

Input
- $\mathbf{S}^{W1} \in \mathbb{R}^{\text{num\_experts} \times (2\times \text{intermediate\_hidden})}$ : Per-channel scaling factor of the Linear1 weight matrix, corresponding to the input parameter l1_weights_sf.
- $\mathbf{S}^{W2} \in \mathbb{R}^{\text{num\_experts} \times \text{hidden}}$ : Per-channel scaling factor of the Linear2 weight matrix, corresponding to the input parameter l2_weights_sf.
- $\mathbf{B}_1 \in \mathbb{R}^{\text{num\_experts} \times (2\times \text{intermediate\_hidden})}$ : Bias matrix of Linear1, generated offline during INT4 quantization, corresponding to the input parameter l1_bias.
- $\mathbf{B}_2 \in \mathbb{R}^{\text{num\_experts} \times \text{hidden}}$ : Bias matrix of Linear2, generated offline during INT4 quantization, corresponding to the input parameter l2_bias.
EP Dispatch

Before the Dispatch communication, the original BF16 activation matrix $\mathbf{X}$ is first quantized to INT8. For each $\text{token}_i$ , its per-token scaling factor is computed as:
$s^{X}_i = \frac{\max(|\mathbf{X}[i,:]|)}{127} \in \mathbb{R}$
The INT8 representation is then obtained by quantization:
$\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}}$
During the Dispatch communication phase, each $\text{token}_i$ sends its quantized vector $\mathbf{q}_i$ and scaling factor $s^{X}_i$ to experts $e_{i,0}, e_{i,1}, \dots, e_{i,\text{num\_topk}-1}$

Let $I_e = \{\, i \mid \exists k,\ \mathbf{E}[i,k]=e \,\}$ be the index set of all tokens dispatched to expert $e$ , and let $N_e = |I_e|$ be the total number of tokens that expert $e$ must process. Then $\mathbf{Q}_e \in \left(\mathbb{Z}_8^{\text{sym}}\right)^{N_e \times \text{hidden}}$ is the matrix formed by row-stacking all $\mathbf{q}_i$ satisfying $i \in I_e$ in any fixed order. This matrix is the full set of token representations that expert $e$ receives after Dispatch. Similarly, the corresponding scaling factor vector received by expert $e$ is denoted as $\mathbf{s}^{X}_e \in \mathbb{R}^{N_e}$ , whose elements are all $s^{X}_i$ satisfying $i \in I_e$ , stacked in the same row order as $\mathbf{Q}_e$ .

For each $\text{token}_i$ and its selected $k$ -th expert $e_{i,k}$ , there exists a unique row index $\operatorname{row}(i,k) \in \{0,\dots,N_{e_{i,k}}-1\}$ such that $\mathbf{Q}_{e_{i,k}}[\operatorname{row}(i,k), :] = \mathbf{q}_i$ . This mapping records the position of $\mathbf{q}_i$ in the input matrix of expert $e_{i,k}$ .
Expert Compute

In an MoE layer, each expert is essentially an independent feedforward network (FFN) that uses the SwiGLU structure to improve expressive capacity. In the A8W4-INT use case, both linear layers use the MSD (Mixed-precision Split-activation Decomposition) scheme for matrix multiplication. By splitting an INT8 value into two signed INT4 values, the high 4 bits and the low 4 bits, an INT8×INT4 matrix multiplication can be decomposed into two INT4×INT4 matrix multiplications, thereby leveraging the hardware INT4 matrix multiplication acceleration. For the mathematical principle and implementation logic of this scheme, see GroupedMatmul W4A8 Quantization and MSD Scheme.

1. Generate a precision-compensated bias matrix (generated offline, completed outside the operator, and passed as an operator input)

The core of the MSD scheme is to split the quantized INT8 activation value into two INT4 components through binary reinterpretation. Because the low-bit INT4 component is defined as $(\mathbf{X}^{\text{int8}} \mathbin{\&} 0x0F) - 8$ , its value range is mapped to $[-8, 7]$ , which is equivalent to subtracting 8 from the original unsigned low 4 bits (0 to 15). When the high-bit and low-bit INT4 values are separately multiplied by the weight matrix and then merged, this offset introduces a constant term that must be compensated in the final result.

Let the original INT8 activation matrix be $\mathbf{X}^{\text{int8}}$ , the split high-bit and low-bit INT4 activation matrices be $\mathbf{X}_1^{\text{int4}}$ and $\mathbf{X}_2^{\text{int4}}$ , and the weight matrix be $\mathbf{W}$ . The recovery relationship is:
$\mathbf{X}^{\text{int8}} = 16 \times \mathbf{X}_1^{\text{int4}} + (\mathbf{X}_2^{\text{int4}} + 8 \cdot \mathbf{1}_{\text{mat}})$
Here $\mathbf{1}_{\text{mat}}$ is an all-ones matrix with the same shape as $\mathbf{X}_2^{\text{int4}}$ , used to add 8 element-wise. Let $\mathbf{1}$ be an all-ones column vector whose shape matches the input feature dimension. The matrix multiplication expands to:
$\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})$
Here $\mathbf{1}^\top \cdot \mathbf{W}$ denotes the sum over each column of the weight matrix $\mathbf{W}$ , and the result is a row vector (with the same shape as the output dimension). In the implementation, this row vector is broadcast to the batch dimension and added to the output of all tokens. As shown, if only the first two terms are computed, a positive term is missing. To recover the result exactly, the compensation bias must be precomputed:
$\text{bias} = 8 \cdot (\mathbf{1}^\top \cdot \mathbf{W}),$
$\mathbf{B}_1$ and $\mathbf{B}_2$ are handled the same way. Each uses its corresponding all-ones column vector (with dimensions $\text{hidden}$ and $\text{intermediate\_hidden}$ , respectively) multiplied by its own weight matrix to obtain a bias row vector with matching shape. These biases are computed offline and then passed into the operator. They directly participate in the addition when merging the high-bit and low-bit results, thereby eliminating the error introduced by the offset.

2. Linear1 Projection (INT4 × INT4 Matrix Multiplication + Dequantization)

The Linear1 projection is the first linear transformation of the expert network, producing the pre-activation values required by both the gate part and the up part.

(1) Reinterpreting the Activation as INT4

Reinterpret the INT8 activation tensor $\mathbf{Q}_e^{\text{int8}} \in \left(\mathbb{Z}_8^{\text{sym}}\right)^{N_e \times \text{hidden}}$ as a binary view of two concatenated INT4 values:
$\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}}$
The high-bit and low-bit components are directly given by the even rows and odd rows of $\mathbf{Q}_e^{\text{int4}}$ :
$\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}}$
They can be numerically obtained from the original INT8 as follows:
$\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$
The recovery relation is $\mathbf{Q}_e^{\text{int8}} = 16\mathbf{Q}_e^{\text{high}} + (\mathbf{Q}_e^{\text{low}} + 8)$ . Since $\mathrm{bitcast}$ only changes the type view, $\mathbf{Q}_e^{\text{int4}}$ and $\mathbf{Q}_e^{\text{int8}}$ share the same underlying physical memory, requiring no data rearrangement or copy.

(2) INT4 × INT4 matrix multiplication and weight dequantization

Perform matrix multiplication between the reinterpreted INT4 activation view $\mathbf{Q}_e^{\text{int4}} \in \mathbb{Z}_4^{2N_e \times \text{hidden}}$ and the weight matrix $\mathbf{W}_1[e] \in \mathbb{R}^{\text{hidden} \times 2\cdot\text{intermediate\_hidden}}$ , and apply the weight scaling factor $\mathbf{s}^{W1}_e$ for dequantization to obtain the result:
$\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}}$
Record the results by even rows and odd rows as the following matrix views, which correspond to the computation results of $\mathbf{Q}_e^{\text{high}}$ and $\mathbf{Q}_e^{\text{low}}$ :
$\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) Precision compensation and activation dequantization

Use the bias $\mathbf{B}_1[e]$ and the activation scaling factor $\mathbf{s}^{X}_e$ (row vector) for precision compensation and activation dequantization respectively. The final pre-activation value is:
$\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}}$
Here $\odot \mathbf{s}^{X}_e$ means multiplying each row of the matrix by the corresponding scalar in $\mathbf{s}^{X}_e$ .

3. SwiGLU activation

First split $\mathbf{H}_e$ along the column dimension into the gate part $\mathbf{H}_{\text{gate}}$ and the up part $\mathbf{H}_{\text{up}}$ . Then apply the SiLU activation function to the gate part and multiply it element-wise with the up part to obtain the intermediate activation representation $\mathbf{A}_e$ of expert $e$ .
$\mathbf{H}_{\text{gate}} = \mathbf{H}_e[:,\; :\!\text{intermediate\_hidden}], \qquad \mathbf{H}_{\text{up}} = \mathbf{H}_e[:,\; \text{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 projection (quantization + INT4 × INT4 matrix multiplication + dequantization)

The Linear2 projection serves as the second linear layer, projecting the intermediate activation representation $\mathbf{A}_e$ from the high-dimensional space back to the original hidden dimension $\text{hidden}$ , so that the expert output is compatible with subsequent operations such as residual connections. In the A8W4 use case, the activation values are first quantized to INT8, then binary-reinterpreted as a concatenated view of two INT4 values, completing the computation with a single matrix multiplication. The process is as follows.

(1) Activation quantization

Compute the scaling factor for each row of $\mathbf{A}_e$ and quantize it to INT8:
$s^{A_e}_i = \frac{\max(|\mathbf{A}_e[i,:]|)}{127}, \quad i=0,\dots,N_e-1$
Obtain the activation scaling factor $\mathbf{s}^{A_e}_e \in \mathbb{R}^{N_e}$ for expert $e$ in the Linear2 computation, and compute the quantization result:
$\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) Activation reinterpretation as INT4

Binary-reinterpret the INT8 activation tensor $\mathbf{A}_e^{\text{int8}} \in \left(\mathbb{Z}_8^{\text{sym}}\right)^{N_e \times \text{intermediate\_hidden}}$ as a view of two alternately concatenated INT4 values:
$\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}}$
The high and low components are directly given by the even rows and odd rows of $\mathbf{A}_e^{\text{int4}}$ :
$\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}}$
Numerically, they can be obtained from the original INT8 as follows:
$\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$
The recovery relationship is $\mathbf{A}_e^{\text{int8}} = 16\mathbf{A}_e^{\text{high}} + (\mathbf{A}_e^{\text{low}} + 8)$ . Since $\mathrm{bitcast}$ only changes the type view, $\mathbf{A}_e^{\text{int4}}$ and $\mathbf{A}_e^{\text{int8}}$ share the underlying physical memory, without any data rearrangement or copy.

(3) INT4 × INT4 matrix multiplication and weight dequantization

Perform matrix multiplication on the reinterpreted INT4 activation view $\mathbf{A}_e^{\text{int4}} \in \mathbb{Z}_4^{2N_e \times \text{intermediate\_hidden}}$ and the weight matrix $\mathbf{W}_2[e] \in \mathbb{R}^{\text{intermediate\_hidden} \times \text{hidden}}$ , and apply the weight scaling factor $\mathbf{s}^{W2}_e$ for dequantization to obtain the result:
$\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}}$
Record the result as the following matrix views by even rows and odd rows, which correspond to the computation results of $\mathbf{A}_e^{\text{high}}$ and $\mathbf{A}_e^{\text{low}}$ :
$\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) Precision compensation and activation dequantization

Use the bias $\mathbf{B}_2[e]$ and the activation scaling factor $\mathbf{s}^{A_e}_e$ (row vector) to perform precision compensation and activation dequantization respectively. The final output is:
$\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}}$
Here $\odot \mathbf{s}^{A_e}_e$ means multiplying each row of the matrix by the corresponding scalar in $\mathbf{s}^{A_e}_e$ .

After the preceding computation, each output row of expert $e$ corresponds to a token in its batch. For $\text{token}_i$ , its output row in expert $e_{i,k}$ is $\mathbf{Y}_{e_{i,k}}\bigl[\operatorname{row}(i,k),\,:\bigr]$ .
Token Combine

Combine collects the output vectors computed by all experts, performs a weighted sum according to the expert weight originally assigned to each token, and finally generates a fused output for each token. Using the previously recorded position index $\operatorname{row}(i,k)$ , it retrieves the rows belonging to $\text{token}_i$ from the output matrix of expert $e_{i,k}$ , multiplies them by the gating weight, and sums them up:
$\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)$
where $w_k = \mathbf{G}[i,k]$ is the gating weight of $\text{token}_i$ for expert $e_{i,k}$ .

The outputs of all tokens are stacked in input order to form the final output $\mathbf{Y} \in \mathbb{R}^{\text{total\_num\_tokens} \times \text{hidden}}$ .

A8W8-FP Quantization Use Case

In the first phase, the input tokens are grouped and collected by expert, then quantized with MXFP8 to generate the quantized input and scaling factor of each expert:

\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

Note: Tokens are sorted and gathered by expert according to topkIds. $\mathcal{T}_e$ is the index set of tokens assigned to expert $e$ . $E_{local}$ indicates the maximum number of tokens received by the current expert, which may differ for each expert. $X[\mathcal{T}_e]$ is the corresponding submatrix. $\mathrm{Q}_{\text{MX}}$ indicates MX per-group quantization (group size = 32). For each group of 32 elements, the shared exponent is extracted and then quantized to the FP8 target type (FLOAT8_E5M2 or FLOAT8_E4M3FN), while outputting the FLOAT8_E8M0 scaling factor. The quantized data serves as the input to GMM1.

In the second phase, GMM1 matrix multiplication (splitting $W_1$ into two halves along the column direction and computing them separately), SwiGLU activation, and MX quantization are performed for each expert:

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)})

U_e = Z_e^{(x)} \odot \sigma\!\left(Z_e^{(x)}\right) \odot Z_e^{(y)}

\hat{U}_e,\ S_{U,e} = \mathrm{Q}_{\text{MX}}(U_e)

Note: The first $N/2$ columns $W_{1,e}^{(x)}$ and the last $N/2$ columns $W_{1,e}^{(y)}$ of $W_1$ are respectively multiplied with the MX-dequantized input to obtain the Swish branch $Z_e^{(x)}$ and the gating branch $Z_e^{(y)}$ . The SwiGLU activation performs an element-wise product $x \cdot \sigma(x) \cdot y$ on the two branches, where $\sigma$ is the Sigmoid function, reducing the intermediate dimension from $N$ to $N/2$ . The SwiGLU output is then MX-quantized to obtain the quantized input $\hat{U}_e$ for GMM2.

In the third phase, GMM2 matrix multiplication is performed for each expert, and the results are distributed to the target ranks:

O_e = \mathrm{DQ}_{\text{MX}}(\hat{U}_e, S_{U,e}) \cdot \mathrm{DQ}_{\text{MX}}(W_{2,e}, S_{2,e})

Note: The quantized SwiGLU output is multiplied with the second weight group $W_2$ after MX dequantization, mapping the $N/2$ -dimensional intermediate representation back to the $H$ -dimensional hidden space to obtain each expert's output $O_e$ . After the computation, the results are written to the remote expert offset address of the target rank through RDMA peermem, achieving cross-rank aggregation.

In the fourth stage, all tokens are weighted and summed by routing weights to restore the output to the same shape as the input:

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

Description: For each token $i$ , the corresponding row is fetched from the aggregated expert results according to the sorted routing index $\pi(i,k)$ , and element-wise weighted accumulation is performed using the weights in topkWeights to obtain the final output $Y$ .

Here, $X$ denotes the parameter x, $W$ denotes the parameter topkWeights, $W_1$ denotes the parameter weight1, $W_2$ denotes the parameter weight2, $Y$ denotes the parameter y, and $E_{\text{local}}$ denotes the attribute moeExpertNum / epWorldSize (the number of experts per rank). $K$ denotes the second dimension of topkIds (the top-K value, which is 6 or 8).

Local variable description:

$\mathcal{T}_e$ : The index set of tokens routed to expert $e$ , determined after sorting topkIds.
$\hat{X}_e,\ S_{X,e}$ : The quantized input of expert $e$ and its MX scaling factor, an intermediate result of the first stage.
$W_{1,e}^{(x)}$ and $W_{1,e}^{(y)}$ : The first $N/2$ columns and the last $N/2$ columns of the submatrix of $W_1$ corresponding to expert $e$ , derived by splitting weight1 according to SwiGLU.
$S_{1,e}^{(x)}$ and $S_{1,e}^{(y)}$ : The MX scaling factors corresponding to $W_{1,e}^{(x)}$ and $W_{1,e}^{(y)}$ , truncated from weightScales1 by dimension.
$S_{2,e}$ : The MX scaling factor corresponding to $W_{2,e}$ , from the parameter weightScales2.
$Z_e^{(x)},\ Z_e^{(y)}$ : Output of the two matrix multiplications in GMM1 (Swish branch and gating branch), an intermediate result.
$U_e$ : SwiGLU activation output, with dimension $m_e \times N/2$ , an intermediate result.
$\hat{U}_e,\ S_{U,e}$ : Quantized SwiGLU output and its MX scaling factor, an intermediate result.
$O_e$ : Expert-level output of GMM2, with dimension $m_e \times H$ , an intermediate result.
$\pi(i, k)$ : Row index of the $k$ -th top-k expert of token $i$ after expansion and sorting, determined by routing order.
$\mathrm{Q}_{\text{MX}}(\cdot)$ : MX per-group quantization operation, block size = 32, outputting FP8 data and E8M0 scaling factors.
$\mathrm{DQ}_{\text{MX}}(\cdot)$ : MX per-group dequantization operation, implicitly performed inside matmul.

Prototype

get_symm_buffer_for_mega_moe:

get_symm_buffer_for_mega_moe(group, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, *, max_recv_token_num=0, dispatch_quant_mode=0, dispatch_quant_out_dtype=None, combine_quant_mode=0, comm_alg="") -> SymmBuffer

mega_moe:

mega_moe(x, topk_ids, topk_weights, l1_weights, l2_weights, sym_buffer, *, l1_weights_sf=None, l2_weights_sf=None, l1_bias=None, l2_bias=None, x_active_mask=None, activation="swiglu", activation_clamp=None, weight1_type=None, weight2_type=None) -> (Tensor, Tensor)

Parameters

get_symm_buffer_for_mega_moe

Parameter	Type	Optional/Required	Description
group	str	Required	EP communicator name (expert parallel communicator).
num_experts	int	Required	Total number of experts in the MoE model.
num_max_tokens_per_rank	int	Required	Number of tokens on each device. When num_tokens differs across ranks, this is the maximum num_tokens.
num_topk	int	Required	Number of experts each token is sent to.
hidden	int	Required	Size of each token.
intermediate_hidden	int	Required	Projection dimension of the intermediate layer.
max_recv_token_num	int	Optional	Maximum number of tokens that each rank can receive. The default value 0 indicates automatic calculation. The default value is 0.
dispatch_quant_mode	int	Optional	Quantization mode for dispatch communication. 0 indicates non-quantization (A16W16 use case), 2 indicates INT8 quantization (A8W8-INT and A8W4-INT use cases), and 4 indicates MXFP quantization (A8W8-FP use case). For supported values of each product, see Constraints. The default value is 0.
dispatch_quant_out_dtype	torch.dtype	Optional	Data type of the output after dispatch quantization. Supports torch.int8, torch.float8_e5m2, and torch.float8_e4m3fn. For the supported values of each product, see Constraints. The default value is None.
combine_quant_mode	int	Optional	Quantization mode for combine communication. 0 indicates non-quantization, 3 indicates MXFP float8_e5m2, and 4 indicates MXFP float8_e4m3. For the supported values of each product, see Constraints. The default value is 0.
comm_alg	str	Optional	This parameter is not supported yet. Use the default value. The default value is "".

mega_moe

The superscript ¹ indicates that Atlas A2 training products/Atlas A2 inference products and Atlas A3 training products/Atlas A3 inference products do not support it. The superscript ² indicates that Ascend 950PR/Ascend 950DT do not support it. For unsupported parameters, use the default value.

Parameter	Parameter Type	Optional/Required	Description	Data Type	Data Format	Dimension (shape)
x	Tensor	Required	Token hidden states input to the MoE layer.	BFLOAT16	ND	(num_tokens, hidden)
topk_ids	Tensor	Required	Expert index matrix, indicating the num_topk experts selected for each token. Element values range within `[0, num_experts)`, and the num_topk experts selected for the same token cannot be duplicated.	INT32	ND	(num_tokens, num_topk)
topk_weights	Tensor	Mandatory	Gating weight coefficients corresponding to the num_topk experts selected by the expert gating network of the MoE model for the current input token.	BFLOAT16², FLOAT32	ND	(num_tokens, num_topk)
l1_weights	list[Tensor]	Required	Weight matrix of the first linear layer of the expert network (including gate and up-projection), used to map the input to the intermediate dimension, with the output fed to the activation function.	BFLOAT16²	ND	(hidden, 2 × intermediate_hidden)
				INT8²	FRACTAL_NZ	(hidden, 2 × intermediate_hidden)
				INT4(INT32)²	FRACTAL_NZ	(hidden, 2 × intermediate_hidden // 8)
				FLOAT8_E5M2¹	ND	(num_experts_per_rank, 2 × intermediate_hidden, hidden)
				FLOAT8_E4M3FN¹	ND	(num_experts_per_rank, 2 × intermediate_hidden, hidden)
l2_weights	list[Tensor]	Required	Weight matrix of the second linear layer of the expert network, which projects the activated intermediate features back to the hidden dimension. The data type is the same as that of l1_weights.	BFLOAT16²	ND	(intermediate_hidden, hidden)
				INT8²	FRACTAL_NZ	(intermediate_hidden, hidden)
				INT4²	FRACTAL_NZ	(intermediate_hidden, hidden)
				FLOAT8_E5M2¹	ND	(num_experts_per_rank, hidden, intermediate_hidden)
				FLOAT8_E4M3FN¹	ND	(num_experts_per_rank, hidden, intermediate_hidden)
sym_buffer	SymmBuffer	Mandatory	Structure created by the get_symm_buffer_for_mega_moe API.	SymmBuffer	SymmBuffer	SymmBuffer
l1_weights_sf	list[Tensor]	Optional	Quantization scaling factor of the weight matrix of the first linear layer of the expert network.	UINT64²	ND	(2 × intermediate_hidden, )
l1_weights_sf	list[Tensor]	Optional		FLOAT8_E8M0¹	ND	(num_experts_per_rank, 2 × intermediate_hidden, CeilDiv(hidden, 64), 2)
l2_weights_sf	list[Tensor]	Optional	Quantization scaling factor of the weight matrix of the second linear layer of the expert network.	UINT64²	ND	(hidden, )
l2_weights_sf	list[Tensor]	Optional		FLOAT8_E8M0¹	ND	(num_experts_per_rank, hidden, CeilDiv(intermediate_hidden, 64), 2)
l1_bias²	list[Tensor]	Optional	Bias of the first linear layer of the expert network. Required only in the A8W4-INT quantization use case for precision compensation.	FLOAT32	ND	(2 × intermediate_hidden, )
l2_bias²	list[Tensor]	Optional	Bias of the second linear layer of the expert network. Required only in the A8W4-INT quantization use case for precision compensation.	FLOAT32	ND	(hidden, )
x_active_mask²	Tensor	Optional	Indicates whether the token participates in communication.	INT8	ND	(num_tokens, )
activation	str	Optional	Activation function type. Defaults to "swiglu". Currently only "swiglu" is supported.	str	Not applicable	Not applicable
activation_clamp	float	Optional	Symmetric clamp threshold for the activation function input. Limits the input tensor to the range [-activation_clamp, activation_clamp]. None means no clamping. The value must be ≥ 0.	float	Not applicable	Not applicable
weight1_type	Reserved parameter	Optional	Reserved parameter. Use the default value.	Reserved parameter	Not involved	Not involved
weight2_type	Reserved parameter	Optional	Reserved parameter. Use the default value.	Reserved parameter	Not applicable	Not applicable
In the table, `CeilDiv(x, y) = ⌈x / y⌉ = ⌊(x + y - 1) / y⌋`
In the table, T1(T2) indicates that data type T1 must be reinterpreted as data type T2 before being passed in. For example, INT4(INT32) indicates that the actual INT4 data must be reinterpreted as INT32 before being passed in, and its shape is the shape after reinterpretation.

Return Value

get_symm_buffer_for_mega_moe

Parameter	Type	Optional/Required	Description	Data Type	Dimension (shape)
sym_buffer	SymmBuffer	Required	Encapsulates the input parameters and generates `context`, `ep_world_size`, and `ccl_buffer_size`.	SymmBuffer	-

mega_moe

Parameter	Parameter Type	Optional/Required	Description	Data Type	Shape
y	Tensor	Required	Token data received by this device, corresponding to Y in the formula. The data type is the same as the input `x`. It must be a 2D tensor in ND format. Non-contiguous tensors are supported.	bfloat16	(num_tokens, hidden)
expert_token_nums	Tensor	Required	Number of tokens actually received by each expert on this device. It must be a 1D tensor in ND format. Non-contiguous tensors are supported.	int32	(num_experts_per_rank,)

Constraints

The constraints on the list[Tensor] length, whether to transpose, and whether non-contiguous tensors are supported for each tensor parameter are as follows:

Parameter	list[Tensor] Length	Transposed	Non-contiguous Tensor Supported
x	Not applicable	No	Yes
topk_ids	Not applicable	No	Supported
topk_weights	Not applicable	No	Supported
l1_weights	num_experts_per_rank (BFLOAT16/INT8/INT4 use case) or 1 (FLOAT8_E5M2/FLOAT8_E4M3FN use case)	No (BFLOAT16/INT8/INT4 use case)/Yes (FLOAT8_E5M2/FLOAT8_E4M3FN use case)	Not supported
l2_weights	num_experts_per_rank (BFLOAT16/INT8/INT4 use case) or 1 (FLOAT8_E5M2/FLOAT8_E4M3FN use case)	No (BFLOAT16/INT8/INT4 use case)/Yes (FLOAT8_E5M2/FLOAT8_E4M3FN use case)	Not supported
l1_weights_sf	num_experts_per_rank (UINT64 use case) or 1 (FLOAT8_E8M0 use case)	No	Not supported
l2_weights_sf	num_experts_per_rank (UINT64 use case) or 1 (FLOAT8_E8M0 use case)	No	Not supported
l1_bias	num_experts_per_rank	No	Not supported
l2_bias	num_experts_per_rank	No	Not supported
x_active_mask	Not involved	No	Supported

Parameter consistency constraints:
- All input parameters of the mega_moe API and their corresponding tensor dimensions must be consistent with the same-named parameters of get_symm_buffer_for_mega_moe (for example, num_experts, hidden, intermediate_hidden, etc.).
- The values of parameters such as num_experts, max_recv_token_num, dispatch_quant_mode, dispatch_quant_out_dtype, and num_max_tokens_per_rank used during the API call must be consistent across all devices and across different layers in the network.
Communicator and networking constraints:
- The values of ep_world_size and ccl_buffer_size must be consistent across all devices.
- The communicator buffer size must be consistent across all devices. ccl_buffer_size is the total size (in bytes) of the CCL communication buffer allocated on HBM, containing two equal-sized spaces: windowIn and windowOut. Validation is based on a single space of ccl_buffer_size / 2, which must satisfy:
  
  $ccl\_buffer\_size\ /\ 2 \ge \mathrm{offsetTokenPerExpert} + \mathrm{offsetTensor} + \mathrm{offsetFlag} + 10\,\mathrm{MB}$
  
  Atlas A2 training products/Atlas A2 inference products:
```
offsetTokenPerExpert = ep_world_size × CeilAlign(ep_world_size × maxExpertPerRank + 1, 128) × 4B

// winIn
offsetAAfterDispatch = max_recv_token_num × (quant ? hidden + 512 : hidden × 2)
offsetD              = num_max_tokens_per_rank × num_topk × hidden × 2B
winInTensorSize      = offsetAAfterDispatch + offsetD

// winOut
offsetA              = num_max_tokens_per_rank × num_topk × (quant ? hidden + 512 : hidden × 2)
offsetC              = max_recv_token_num × hidden × 2B
winOutTensorSize     = offsetA + offsetC

offsetTensor         = max(winInTensorSize, winOutTensorSize)
                    + (quant ? max_recv_token_num × 4B : 0)

// sync flags
offsetFlag           = ep_world_size × 512B
                    + ep_world_size × maxExpertPerRank × 64B
                    + ep_world_size × 64B
```
  Atlas A3 training products/Atlas A3 inference products:
```
offsetTokenPerExpert = ep_world_size × CeilAlign(ep_world_size × maxExpertPerRank + 1, 128) × 4B

// winIn (A3 uses only winIn)
offsetAAfterDispatch = num_max_tokens_per_rank × num_topk × (quant ? hidden + 512 : hidden × 2)
offsetD              = num_max_tokens_per_rank × num_topk × hidden × 2B
winInTensorSize      = offsetAAfterDispatch + offsetD

offsetTensor         = winInTensorSize
                    + (quant ? num_max_tokens_per_rank × num_topk × 4B : 0)

// sync flags
syncStateReservedSize = 512KB
offsetFlag            = max(ep_world_size × 512B, syncStateReservedSize)
```
  Here, ep_world_size is the communicator size, and $maxExpertPerRank$ indicates the maximum possible number of experts on each device. $\mathrm{quant}$ indicates whether dispatch quantization is enabled (dispatch_quant_mode = 2). The reserved 10 MB space is for internal metadata alignment and safety margin.
- The driver versions of all nodes in the communicator must be the same.
- Atlas A2 training products/Atlas A2 inference products: Multi-device communicators require switch-based networking. Direct connection between two devices is not supported.
- Atlas A3 training products/Atlas A3 inference products: Multi-device communicators must be within one SuperPoD. Direct connection between two devices and cross-SuperPoD networking are not supported.

Parameter constraints:

Common constraints:
- num_experts: The value range is $[\text{ep_world_size},\ 1024], and num_experts % ep_world_size == 0.

Atlas A2 training products/Atlas A2 inference products, Atlas A3 training products/Atlas A3 inference products:

num_tokens must be consistent across devices, and num_tokens (that is, x.dim0) on each device must be less than or equal to num_max_tokens_per_rank.
ep_world_size: only 2, 4, 8, 16, 32, and 64 are supported.
num_experts_per_rank: value range $[1,\ 128], and num_experts_per_rank = num_experts / ep_world_size.
num_max_tokens_per_rank: value range $[1,\ 4096].
max_recv_token_num must be greater than 0. Input 0 indicates automatic calculation, with the formula num_tokens × ep_world_size × min(num_topk, num_experts_per_rank).
num_topk: value range $[1,\ 16].
hidden: value range $[1024,\ 8192], and hidden % 512 == 0.
intermediate_hidden: value range $[1024,\ 3072], and intermediate_hidden % 512 == 0
dispatch_quant_mode: only 0 (non-quantized) and 2 (INT8 quantization) are supported.
dispatch_quant_out_dtype: only torch.int8 (corresponding to integer encoding 1) is supported.

Use case matrix for A2/A3 products:

Use case	x	l1_weights	l2_weights	l1_weights_sf	l2_weights_sf	l1_bias	l2_bias	y	dispatch_quant_mode	dispatch_quant_out_dtype
A16W16	BFLOAT16	BFLOAT16	BFLOAT16	–	–	–	–	BFLOAT16	0	–
A8W8-INT	BFLOAT16	INT8	INT8	UINT64	UINT64	–	–	BFLOAT16	2	torch.int8
A8W4-INT	BFLOAT16	INT4(INT32)	INT4(INT32)	UINT64	UINT64	FLOAT32	FLOAT32	BFLOAT16	2	torch.int8

Ascend 950PR/Ascend 950DT:
- num_tokens: value range $[1,\ 512].
- hidden: only 4096, 5120, and 7168 are supported.
- num_topk: only 6 and 8 are supported.
- num_experts_per_rank: value range $[1,\ 16], and num_experts_per_rank = num_experts / ep_world_size.
- intermediate_hidden: Only 1024, 2048, 3072, 4096, and 7168 are supported.
- ep_world_size: Value range $[2,\ 768]$ .
- num_max_tokens_per_rank: 0 or satisfies num_tokens == num_max_tokens_per_rank.
- dispatch_quant_out_dtype: Only torch.float8_e5m2 (corresponding integer encoding 23) or torch.float8_e4m3fn (corresponding integer encoding 24) is supported.
- dispatch_quant_mode: Only 4 (MXFP quantization mode) is supported. The dispatch phase uses MX per-group quantization (group size = 32), and the quantization scaling factor type is FLOAT8_E8M0.
- combine_quant_mode: Must be 0.
- comm_alg: Reserved parameter. Must be an empty string "".
- x_active_mask and scales: In the current version, a null pointer must be passed. Non-null input is not supported.
- The data type of y is the same as that of x.
- The dim1 (intermediate_hidden) of weight1 must be twice the dim2 of weight2. The shape of weight1 is (num_experts_per_rank, intermediate_hidden, hidden), and the shape of weight2 is (num_experts_per_rank, hidden, intermediate_hidden / 2). In the MXFP use case, when dispatch_quant_out_dtype is torch.float8_e5m2, weight1 and weight2 must be FLOAT8_E5M2. When it is torch.float8_e4m3fn, weight1 and weight2 must be FLOAT8_E4M3FN.
- weight_scales1, weight_scales2: must not be null pointers. The shape of weight_scales1 is (num_experts_per_rank, intermediate_hidden, CeilDiv(hidden, 64), 2), and the shape of weight_scales2 is (num_experts_per_rank, hidden, CeilDiv(intermediate_hidden / 2, 64), 2). The dim3 of both must be 2. Where $CeilDiv(hidden,\ 64) = \lceil hidden\ /\ 64 \rceil = \lfloor (hidden + 63)\ /\ 64 \rfloor$ , $CeilDiv(intermediate\_hidden / 2,\ 64) = \lceil (intermediate\_hidden / 2)\ /\ 64 \rceil = \lfloor (intermediate\_hidden / 2 + 63)\ /\ 64 \rfloor$ .

Deterministic Computation

Defaults to a deterministic computation.

Example

Single-operator mode call:

The following example chains the two APIs in call order: first initialize the communicator, then construct sym_buffer with get_symm_buffer_for_mega_moe, and finally call mega_moe to run the operator.

Ascend 950PR/Ascend 950DT:

import os
import torch
import torch_npu
from torch.multiprocessing import Process, Manager
import torch.distributed as dist
from torch.distributed import ReduceOp
import torch.multiprocessing as mp
from cann_ops_transformer.ops import get_symm_buffer_for_mega_moe, mega_moe
import torchair

E = 4
num_tokens = 256
H = 4096
N = 1024
topK = 6
num_experts = 8

server_num = 1
rank_per_dev = 2
world_size = server_num * rank_per_dev
ep_ranks_list = [list(range(tp_id, world_size, 1)) for tp_id in range(1)]
server_index = 0


def ceil(a, b):
    return (a + b - 1) // b

def set_device(rank):
    torch_npu.npu.set_device(rank % rank_per_dev)
    print(f"current device set: {torch_npu.npu.current_device()}")

def init_hccl_comm(rank):
    # Create the HCCL communication link and initialize the EP domain.
    print(f'[INFO] device_{rank} creates the HCCL communication link')
    master_ip = '127.0.0.1'
    dist.init_process_group(backend="hccl", rank=rank, world_size=world_size, init_method=f'tcp://{master_ip}:50001')
    print(f"device_{rank} init_process_group success")

    print(f"device {rank} initializes the EP domain")
    for ep_ranks in ep_ranks_list:
        tmp_group = dist.new_group(backend="hccl", ranks=ep_ranks)
        if rank in ep_ranks:
            ep_group = tmp_group

    ep_hcomm_info = ep_group._get_backend(torch.device("npu")).get_hccl_comm_name(rank)

    return ep_hcomm_info, ep_group

def get_megamoe_kwargs(
    x, expert_ids, weights1, weights_scales1, weights2, weights_scales2, expert_scales
):
    x = x.to(torch.bfloat16).npu()
    expert_ids = expert_ids.to(torch.int32).npu()
    weights1 = weights1.to(torch.float8_e5m2).npu()
    weights_scales1 = weights_scales1.to(torch.float8_e8m0fnu).npu()
    weights2 = weights2.to(torch.float8_e5m2).npu()
    weights_scales2 = weights_scales2.to(torch.float8_e8m0fnu).npu()
    expert_scales = expert_scales.to(torch.bfloat16).npu()

    return {
        'x': x,
        'topk_ids': expert_ids,
        'topk_weights': expert_scales,
        'l1_weights': [weights1],
        'l1_weights_sf': [weights_scales1],
        'l2_weights': [weights2],
        'l2_weights_sf': [weights_scales2],
    }

def run_megamoe_npu(
    queue, rank, x, expert_ids, weights1, weights_scales1, weights2, weights_scales2, expert_scales
):
    print(f"{os.getpid()=}{rank=}")
    set_device(rank)
    print(f'[INFO] device_{rank} constructs the megamoe operator input data')
    megamoe_kwargs = get_megamoe_kwargs(
        x=x,
        expert_ids=expert_ids,
        weights1=weights1,
        weights_scales1=weights_scales1,
        weights2=weights2,
        weights_scales2=weights_scales2,
        expert_scales=expert_scales,
    )
    # Step 1: Initialize the communicator.
    ep_hcomm_info, ep_group = init_hccl_comm(rank)
    # Step 2: Construct distribute_buffer (SymmBuffer struct).
    distribute_buffer = get_symm_buffer_for_mega_moe(
        ep_group, num_experts=num_experts,
        num_max_tokens_per_rank=num_tokens, num_topk=topK,
        hidden=H, intermediate_hidden=0,
        dispatch_quant_mode=4, dispatch_quant_out_dtype=torch.float8_e5m2
    )
    # Step 3: Run mega_moe and pass in the sym_buffer constructed in the previous step.
    y, expert_token_nums = mega_moe(**megamoe_kwargs, sym_buffer=distribute_buffer)

    torch.npu.synchronize()
    print(f"[INFO] device_{rank} finish\n")
    dist.destroy_process_group()
    print(f'rank {rank} epid {rank} npu finished! \n')

    queue.put([
        rank,
        [
            y.cpu(), expert_token_nums.cpu()
        ]
    ])

def gen_npu(target_func, **server_kwargs):
    def parse_rank_input(target_func, result_queue, rank, server_kwargs):

        ep_id = rank // 1

        if target_func == run_megamoe_npu:
            return {
                "queue": result_queue,
                "rank": rank,
                "x": server_kwargs["x_list"][ep_id],
                "expert_ids": server_kwargs["expert_ids_list"][ep_id],
                "weights1": server_kwargs["weights1_list"][ep_id],
                "weights_scales1": server_kwargs["weights_scales1_list"][ep_id],
                "weights2": server_kwargs["weights2_list"][ep_id],
                "weights_scales2": server_kwargs["weights_scales2_list"][ep_id],
                "expert_scales": server_kwargs["expert_scales_list"][ep_id]
            }


    print("single_server scene!!!!!")
    rank_list = list(range(world_size))
    print(f"rank list is: {rank_list}")

    proc_list = []
    manager = Manager()
    result_queue = manager.Queue()
    mp.set_start_method("forkserver", force=True)
    for rank in rank_list:
        rank_kwargs = parse_rank_input(target_func, result_queue, rank, server_kwargs)
        proc = Process(target=target_func, kwargs=rank_kwargs)
        proc.start()
        proc_list.append(proc)


    rank_outputs = [None] * rank_per_dev
    for proc in proc_list:
        rank_id, rank_output = result_queue.get()
        local_rank_id = rank_id - server_index * rank_per_dev
        rank_outputs[local_rank_id] = rank_output


    for proc in proc_list:
        proc.join()

    if None in rank_outputs:
        print("[ERROR] Task failed! Please check the detailed error logs printed by the subprocesses.")
        exit(1)

    # Put the outputs of each category into the same list. category_outputs stores the list of outputs of each category.
    category_outputs = []
    category_num = len(rank_outputs[0])
    for category_id in range(category_num):
        specific_category_output = [rank_output[category_id] for rank_output in rank_outputs]
        category_outputs.append(specific_category_output)

    return category_outputs

if __name__ == "__main__":
    x_shape = [num_tokens, H]
    expert_idx_shape = [num_tokens, topK]
    weight_shape = [E, N, H]
    weight_scale_shape = [E, N, ceil(H, 64), 2]
    output_shape = [num_tokens, N//2]
    weight2_shape = [E, H, N//2]
    weight2_scale_shape = [E, H, ceil(N//2, 64), 2]
    expert_scales_shape = [num_tokens, topK]
    # Construct the input.
    x = torch.randn(x_shape, dtype=torch.bfloat16)
    expert_scales = torch.randn(expert_scales_shape, dtype=torch.bfloat16)
    expert_ids = torch.stack(
        [torch.randperm(num_experts)[:topK] for _ in range(num_tokens)]
    ).to(torch.int32)
    weight1 = torch.randn(weight_shape, dtype=torch.float32).to(torch.float8_e5m2)
    weight_scales1 = torch.randint(125, 130, weight_scale_shape, dtype=torch.uint8).view(torch.float8_e8m0fnu)
    weight2 = torch.randn(weight2_shape, dtype=torch.float32).to(torch.float8_e5m2)
    weight_scales2 = torch.randint(125, 130, weight2_scale_shape, dtype=torch.uint8).view(torch.float8_e8m0fnu)

    golden_x_list = [x.clone() for _ in range(rank_per_dev)]
    golden_expert_ids_list = [expert_ids.clone() for _ in range(rank_per_dev)]
    golden_weights1_list = [weight1.clone() for _ in range(rank_per_dev)]
    golden_weights_scales1_list = [weight_scales1.clone() for _ in range(rank_per_dev)]
    golden_weights2_list = [weight2.clone() for _ in range(rank_per_dev)]
    golden_weights_scales2_list = [weight_scales2.clone() for _ in range(rank_per_dev)]
    golden_expert_scales_list = [expert_scales.clone() for _ in range(rank_per_dev)]

    [y, expert_token_nums] = gen_npu(
        run_megamoe_npu,
        x_list=golden_x_list,
        expert_ids_list=golden_expert_ids_list,
        weights1_list=golden_weights1_list,
        weights_scales1_list=golden_weights_scales1_list,
        weights2_list=golden_weights2_list,
        weights_scales2_list=golden_weights_scales2_list,
        expert_scales_list=golden_expert_scales_list,
    )

Atlas A2 training products/Atlas A2 inference products, Atlas A3 training products/Atlas A3 inference products:

import torch
import torch_npu
import torch.multiprocessing as mp
import torch.distributed as dist
from cann_ops_transformer.ops import get_symm_buffer_for_mega_moe, mega_moe

num_tokens = 64
hidden = 4096
intermediate_hidden = 1024
num_topk = 6
num_experts = 16
num_max_tokens_per_rank = 256
num_servers = 1
num_ranks_per_server = 2
server_index = 0  # The number of the current machine in the cluster.
master_addr = '127.0.0.1'
master_port = 50001
scene = 'A16W16'

world_size = num_servers * num_ranks_per_server
num_experts_per_rank = num_experts // world_size

torch_npu.npu.config.allow_internal_format = True


def init_hccl_comm(rank):
    print(f'device_{rank} start init_process_group')
    options = torch_npu._C._distributed_c10d.ProcessGroupHCCL.Options()
    # When using it, configure the value of hccl_buffer_size according to the constraint on the communicator buffer size given in the operator API document.
    options.hccl_config = {'hccl_buffer_size': 200}
    dist.init_process_group(
        backend='hccl',
        rank=server_index * num_ranks_per_server + rank,
        world_size=world_size,
        init_method=f'tcp://{master_addr}:{master_port}',
        pg_options=options,
    )
    print(f'device_{rank} init_process_group success')
    ep_group = dist.new_group(backend='hccl', ranks=list(range(world_size)))
    _ = ep_group._get_backend(torch.device('npu')).get_hccl_comm_name(rank)
    return ep_group


def bf16_mega_moe(rank):
    torch_npu.npu.set_device(rank % num_ranks_per_server)

    x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='npu')
    scores = torch.randn((num_tokens, num_experts), dtype=torch.float, device='npu')
    topk_weights, topk_ids = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False)
    topk_ids = topk_ids.to(torch.int32)
    l1_weights = [
        torch.randn((hidden, intermediate_hidden * 2), dtype=torch.bfloat16, device='npu')
        for _ in range(num_experts_per_rank)
    ]
    l2_weights = [
        torch.randn((intermediate_hidden, hidden), dtype=torch.bfloat16, device='npu')
        for _ in range(num_experts_per_rank)
    ]

    ep_group = init_hccl_comm(rank)

    sym_buffer = get_symm_buffer_for_mega_moe(
        ep_group,
        num_experts=num_experts,
        num_max_tokens_per_rank=num_max_tokens_per_rank,
        num_topk=num_topk,
        hidden=hidden,
        intermediate_hidden=intermediate_hidden,
        dispatch_quant_mode=0,
        dispatch_quant_out_dtype=None,
    )
    y, expert_token_nums = mega_moe(x, topk_ids, topk_weights, l1_weights, l2_weights, sym_buffer)

    dist.barrier()
    dist.destroy_process_group()


def int8_int8_mega_moe(rank):
    torch_npu.npu.set_device(rank % num_ranks_per_server)

    x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='npu')
    scores = torch.randn((num_tokens, num_experts), dtype=torch.float, device='npu')
    topk_weights, topk_ids = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False)
    topk_ids = topk_ids.to(torch.int32)
    l1_weights_bf16 = [
        torch.randn((hidden, intermediate_hidden * 2), dtype=torch.bfloat16, device='npu')
        for _ in range(num_experts_per_rank)
    ]
    l2_weights_bf16 = [
        torch.randn((intermediate_hidden, hidden), dtype=torch.bfloat16, device='npu')
        for _ in range(num_experts_per_rank)
    ]

    def per_channel_cast_to_int8(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        max_abs = torch.amax(torch.abs(w), dim=0, keepdim=True)
        sf = torch.clamp(max_abs / 127.0, min=1e-8)
        q_weight = torch.round(w / sf).clamp(-127, 127).to(torch.int8)
        return q_weight, sf.squeeze().to(torch.float32)

    l1_weights_int8, l1_weights_sf_float = map(list, zip(*[per_channel_cast_to_int8(w) for w in l1_weights_bf16]))
    l2_weights_int8, l2_weights_sf_float = map(list, zip(*[per_channel_cast_to_int8(w) for w in l2_weights_bf16]))

    # Cast weights to NZ format for better performance.
    l1_weights = [torch_npu.npu_format_cast(w, torch_npu.Format.FRACTAL_NZ) for w in l1_weights_int8]
    l2_weights = [torch_npu.npu_format_cast(w, torch_npu.Format.FRACTAL_NZ) for w in l2_weights_int8]

    # Pack float32 scaling factors into uint64, which is required for the process of fixpipe dequantization.
    l1_weights_sf = [sf.view(torch.int32).to(torch.int64).view(torch.uint64) for sf in l1_weights_sf_float]
    l2_weights_sf = [sf.view(torch.int32).to(torch.int64).view(torch.uint64) for sf in l2_weights_sf_float]

    ep_group = init_hccl_comm(rank)

    sym_buffer = get_symm_buffer_for_mega_moe(
        ep_group,
        num_experts=num_experts,
        num_max_tokens_per_rank=num_max_tokens_per_rank,
        num_topk=num_topk,
        hidden=hidden,
        intermediate_hidden=intermediate_hidden,
        dispatch_quant_mode=2,
        dispatch_quant_out_dtype=torch.int8,
    )
    y, expert_token_nums = mega_moe(
        x,
        topk_ids,
        topk_weights,
        l1_weights,
        l2_weights,
        sym_buffer,
        l1_weights_sf=l1_weights_sf,
        l2_weights_sf=l2_weights_sf,
    )

    dist.barrier()
    dist.destroy_process_group()


def int8_int4_mega_moe(rank):
    torch_npu.npu.set_device(rank % num_ranks_per_server)

    x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='npu')
    scores = torch.randn((num_tokens, num_experts), dtype=torch.float, device='npu')
    topk_weights, topk_ids = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False)
    topk_ids = topk_ids.to(torch.int32)
    l1_weights_bf16 = [
        torch.randn((hidden, intermediate_hidden * 2), dtype=torch.bfloat16, device='npu')
        for _ in range(num_experts_per_rank)
    ]
    l2_weights_bf16 = [
        torch.randn((intermediate_hidden, hidden), dtype=torch.bfloat16, device='npu')
        for _ in range(num_experts_per_rank)
    ]

    def per_channel_cast_to_int4(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        max_abs = torch.amax(torch.abs(w), dim=0, keepdim=True)
        sf = torch.clamp(max_abs / 7.0, min=1e-8)
        q_weight = torch.round(w / sf).clamp(-7, 7).to(torch.int8)
        return q_weight, sf.squeeze().to(torch.float)

    def pack_int4_to_int8(x: torch.Tensor) -> torch.Tensor:
        x = (x + 8).to(torch.uint8).reshape(x.shape[0], -1, 2)
        x = ((x[..., 1] << 4) | x[..., 0]).view(torch.int8)
        return x

    l1_weights_int8, l1_weights_sf_float = map(list, zip(*[per_channel_cast_to_int4(w) for w in l1_weights_bf16]))
    l2_weights_int8, l2_weights_sf_float = map(list, zip(*[per_channel_cast_to_int4(w) for w in l2_weights_bf16]))

    l1_weights_int4 = [pack_int4_to_int8(w) for w in l1_weights_int8]
    l2_weights_int4 = [pack_int4_to_int8(w) for w in l2_weights_int8]

    # Cast weights to NZ format for better performance.
    l1_weights = [torch_npu.npu_format_cast(w, torch_npu.Format.FRACTAL_NZ).view(torch.int32) for w in l1_weights_int4]
    l2_weights = [torch_npu.npu_format_cast(w, torch_npu.Format.FRACTAL_NZ).view(torch.int32) for w in l2_weights_int4]

    # Pack float32 scaling factors into uint64, which is required for the process of fixpipe dequantization.
    l1_weights_sf = [sf.view(torch.int32).to(torch.int64).view(torch.uint64) for sf in l1_weights_sf_float]
    l2_weights_sf = [sf.view(torch.int32).to(torch.int64).view(torch.uint64) for sf in l2_weights_sf_float]

    l1_bias = [(w.float() * sf.unsqueeze(0)).sum(dim=0) * 8.0 for w, sf in zip(l1_weights_int8, l1_weights_sf_float)]
    l2_bias = [(w.float() * sf.unsqueeze(0)).sum(dim=0) * 8.0 for w, sf in zip(l2_weights_int8, l2_weights_sf_float)]

    ep_group = init_hccl_comm(rank)

    sym_buffer = get_symm_buffer_for_mega_moe(
        ep_group,
        num_experts=num_experts,
        num_max_tokens_per_rank=num_max_tokens_per_rank,
        num_topk=num_topk,
        hidden=hidden,
        intermediate_hidden=intermediate_hidden,
        dispatch_quant_mode=2,
        dispatch_quant_out_dtype=torch.int8,
    )
    y, expert_token_nums = mega_moe(
        x,
        topk_ids,
        topk_weights,
        l1_weights,
        l2_weights,
        sym_buffer,
        l1_weights_sf=l1_weights_sf,
        l2_weights_sf=l2_weights_sf,
        l1_bias=l1_bias,
        l2_bias=l2_bias,
    )

    dist.barrier()
    dist.destroy_process_group()


if __name__ == '__main__':
    if scene == 'A16W16':
        torch.multiprocessing.spawn(bf16_mega_moe, nprocs=num_ranks_per_server)
    elif scene == 'A8W8-INT':
        torch.multiprocessing.spawn(int8_int8_mega_moe, nprocs=num_ranks_per_server)
    elif scene == 'A8W4-INT':
        torch.multiprocessing.spawn(int8_int4_mega_moe, nprocs=num_ranks_per_server)
    else:
        raise ValueError(f"Unsupported scene: {scene}, please choose from ['A16W16', 'A8W8-INT', 'A8W4-INT']")