mhc_post

Applicable Products

• Ascend 950PR/Ascend 950DT: Not 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:

  Implements the forward computation of the MHC Post component for the post-processing stage of multi-layer residual connections in Transformer models. This operator fuses the residual matrix transformation and the output state projection into a single computation, avoiding the overhead of multiple independent operator calls.

• Formula:

  • The core formula of the MHC Post operator is:

    $$x_{l+1} = (H_{l}^{res})^{T} \cdot x_{l} + h_{l}^{out} \cdot H_{t}^{post}$$

    Where:

    • $(H_{l}^{res})^{T} \cdot x_{l}$ denotes the transpose matrix multiplication of the residual matrix on the input $x_{l}$. For the $i$-th row of the output (corresponding to the $i$-th head), the computation is:

      $$x_{l+1}[i] = \sum_{j=0}^{n-1} H_{l}^{res}[j, i] \cdot x_{l}[j]$$

      That is, the $H_{l}^{res}$ matrix is multiplied with $x_{l}$ in transposed form. $H_{l}^{res}[j, i]$ is a scalar. Scalar multiplication is performed on the $j$-th row of $x_{l}$, and the result is accumulated into the $i$-th row of the output.

    • $h_{l}^{out} \cdot H_{t}^{post}$ denotes the element-wise multiplication and broadcast of the output state $h_{l}^{out}$ with the post-processing weight $H_{t}^{post}$. For the $i$-th head:

      $$x_{l+1}[i] += H_{t}^{post}[i] \cdot h_{l}^{out}$$

      That is, $H_{t}^{post}[i]$ is a scalar. Scalar multiplication is performed on the entire row of $h_{l}^{out}$, and the result is added to the $i$-th row of the output.

  • The complete computation is:

    $$x_{l+1}[i, :] = H_{t}^{post}[i] \cdot h_{l}^{out}[:] + \sum_{j=0}^{n-1} H_{l}^{res}[j, i] \cdot x_{l}[j, :]$$

    Here, $x_{l}$ denotes the parameter x, $H_{l}^{res}$ denotes the parameter h_res, $h_{l}^{out}$ denotes the parameter h_out, $H_{t}^{post}$ denotes the parameter h_post, and $x_{l+1}$ denotes the output y.

  Note

  The input supports two dimension formats: BSND (4D) and TND (3D). B (Batch) denotes the batch size, S (Seq-Length) denotes the sequence length, T denotes the sum of sequence lengths across all batches ($T = B \times S$), n denotes the number of heads, and D denotes the hidden dimension size of each head (headdim).

Prototype

mhc_post(x, h_res, h_out, h_post) -> Tensor

Parameters

Return Value

Constraints

• This API supports use in training and inference use cases.

• This API supports single-operator mode calls.

• Data type constraints:

  • The data types of x and h_out must be the same.
  • The data type of the output y remains consistent with x.

• Dimension constraints:

  • The dimension of h_res must match the dimension format of x: (B, S, n, n) for 4D, and (T, n, n) for 3D.

• Shape consistency constraints:

  • In the 4D (BSND) format:
    • The (B, S) dimension of h_res must be consistent with the (B, S) dimension of x. The last two dimensions of h_res are (n, n), where n is consistent with the third dimension of x.
    • The (B, S) dimension of h_out must be consistent with the (B, S) dimension of x. The D dimension of h_out must be consistent with the D dimension of x.
    • The (B, S) dimension of h_post must be consistent with the (B, S) dimension of x. The n dimension of h_post must be consistent with the n dimension of x.
  • In the 3D (TND) format:
    • The T dimension of h_res must be consistent with the T dimension of x. The last two dimensions are (n, n).
    • The T dimension of h_out must be consistent with the T dimension of x. The D dimension must be consistent with the D dimension of x.
    • The T dimension of h_post must be consistent with the T dimension of x. The n dimension must be consistent with the n dimension of x.

• The value of each dimension in the shape of all input tensors must be a positive number (greater than 0).

Deterministic Computation

Defaults to a deterministic computation.

Example

• Single-operator mode call:

  import torch
  import torch_npu
  from cann_ops_transformer.ops import mhc_post
  
  B = 2
  S = 8
  n = 4
  D = 128
  
  x = torch.randn(B, S, n, D, dtype=torch.bfloat16).npu()
  h_res = torch.randn(B, S, n, n, dtype=torch.float32).npu()
  h_out = torch.randn(B, S, D, dtype=torch.bfloat16).npu()
  h_post = torch.randn(B, S, n, dtype=torch.float32).npu()
  
  y = mhc_post(x, h_res, h_out, h_post)
  print(f"output shape: {y.shape}")