MxMatmul Scenario

Background

Floating-point numbers are widely used in fields such as scientific computing, image processing, and neural networks. Take AI training as an example. The existing floating-point number formats or value ranges are insufficient, or the precision is not high enough, which affects the convergence speed and performance of models. Meeting both the value range and precision requirements will result in excessive memory usage, thereby increasing the costs of data storage and transmission. To address this issue, a new floating-point number format, namely, Microscaling (MX), is proposed in the industry. Floating-point numbers in MX format support AI training and inference with a lower bit width and occupy less memory. Data formats compliant with the MX standard deliver robust precision for AI training and inference models when using 8-bit or lower bit width.

All k elements P_i have the same bit width and data type and share a scale factor X. Each block containing k elements can be encoded using (w + k × d) bits. The data type and scale factor of each element can be selected independently.

The following figure shows the data structure of a floating-point number in MX format. S, E, and M represent the sign, exponent, and mantissa of the floating-point number, respectively. The shared scale factor X is a scale factor applied to the entire data block and determines the dynamic range of all elements in the data block. By introducing the shared scale factor, data in MX format can flexibly represent data of different ranges while maintaining a low bit width. The block size k refers to the size of low-bit data that form a data block (or group). The private element P_i refers to each low-bit data element in the data block. After being adjusted by the scale factor X, these elements together represent a high-precision floating-point number or integer.

Figure 1 MX format

There are multiple data types in MX format, such as MXFP8, MXFP4, MXFP16, and MXINT4. The following table lists the data types supported in the MxMatmul (short for "Microscaling Matmul") scenario.

Overview

MxMatmul involves matrix multiplication with quantization coefficients, where both the left and right matrices have corresponding quantization coefficient matrices (left quantization coefficient matrix scaleA and right quantization coefficient matrix scaleB). The left quantization coefficient matrix is multiplied by the left matrix, and the right quantization coefficient matrix is multiplied by the right matrix. The products of these multiplications are then subjected to matrix multiplication with each other.

MxMatmul compute formula: C = (scaleA ⊗ A) × (scaleB ⊗ B) + Bias, where "⊗" represents broadcasted multiplication. When the left/right matrix is multiplied by the left/right quantization coefficient matrix, every 32 elements in the K direction share a quantization factor, as shown in Figure 2.

• A, scaleA, B, and scaleB are source operands. A is a left matrix with the shape of [M, K]. scaleA is a left quantization coefficient matrix with the shape of [M, K/32]. B is a right matrix with the shape of [K, N]. scaleB is a right quantization coefficient matrix with the shape of [K/32, N].
• C is the destination operand, which is a matrix that stores the matrix multiplication result. Its shape is [M, N].
• Bias indicates the matrix multiplication bias, whose shape is [1, N]. Each row of the (scaleA ⊗ A) × (scaleB ⊗ B) result matrix is biased.

Figure 2 MxMatmul matrix multiplication

The following figures show the formats of matrices A, scaleA, B, and scaleB at different positions.

Figure 3 Formats of matrix A at different positions

Figure 4 Formats of matrix B at different positions

Figure 5 Formats of matrix scaleA at different positions

Figure 6 Formats of matrix scaleB at different positions

Application Scenarios

Scenario where matrices A and B need to be quantized before cube compute. The following table lists the data types supported by the Matmul input/output matrices in this scenario.

Implementation Procedure

Procedure for the host to automatically obtain tiling parameters:

1. Create a tiling object.

   1 2	auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); matmul_tiling::MatmulApiTiling cubeTiling(ascendcPlatform);

   You need to pass the hardware platform information, which can be obtained by calling GetPlatformInfo, to create a PlatformAscendC object. Then, create a tiling object.

2. Set the logical memory locations, formats, data types, and transpose information of matrices A, B, C, and Bias, and set the logical memory locations, formats, and transpose information of scaleA and scaleB.
   Call the SetScaleAType and SetScaleBType APIs to set the logical memory locations, formats, and transpose information of scaleA and scaleB.

   1 2 3 4 5 6

   cubeTiling.SetAType(AscendC::TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT8_E5M2, false); cubeTiling.SetBType(AscendC::TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT8_E5M2, true); cubeTiling.SetScaleAType(AscendC::TPosition::GM, CubeFormat::ND, false); cubeTiling.SetScaleBType(AscendC::TPosition::GM, CubeFormat::ND, true); cubeTiling.SetCType(AscendC::TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT); cubeTiling.SetBiasType(AscendC::TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT);

3. Enable the MxMatmul scenario.
   Call the SetMadType API to set the tiling compute logic to the MxMatmul scenario.

   1	cubetiling.SetMadType(MatrixMadType::MXMODE);

4. Set the matrix shape.

   1 2	cubeTiling.SetShape(M, N, K); cubeTiling.SetOrgShape(M, N, K); // Set the original complete shapes M, N, and K.

5. Set the size of the available space.
   Set the size of the available space on L1 Buffer, L0C Buffer, or Unified Buffer for Matmul compute. The value –1 indicates the size of the buffer corresponding to the AI processor.

   1	cubeTiling.SetBufferSpace(-1, -1, -1);

6. Set other parameters as required, for example, bias that will participate in the compute.

   1	cubeTiling.EnableBias(true);

7. Obtain tiling parameters.

   1 2 3 4

   MatmulCustomTilingData tiling; if (cubeTiling.GetTiling(tiling.cubeTilingData) == -1){ return ge::GRAPH_FAILED; }

8. Perform other operations such as serialization and saving of tiling parameters.

Procedure on the kernel side:

1. Create a Matmul object.

   1 2 3 4 5 6 7

   // In the MxMatmul scenario, define the parameter types of A, scaleA, B, and scaleB using MatmulTypeWithScale. typedef AscendC::MatmulTypeWithScale<AscendC::TPosition::GM, AscendC::TPosition::GM, CubeFormat::ND, fp8_e5m2_t, isTransposeA> aType; typedef AscendC::MatmulTypeWithScale<AscendC::TPosition::GM, AscendC::TPosition::GM, CubeFormat::ND, fp8_e5m2_t, isTransposeB> bType; typedef AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, float> cType; typedef AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, float> biasType; // When defining a Matmul object, input MatmulWithScalePolicy to enable the MxMatmul template strategy. AscendC::Matmul<aType, bType, cType, biasType, CFG_MDL, MatmulCallBackFunc<nullptr, nullptr, nullptr>, AscendC::Impl::Detail::MatmulWithScalePolicy> mm;

   During object creation, input the type information of parameters A, scaleA, B, scaleB, C, and Bias. The type information of A, scaleA, B, and scaleB is defined by MatmulTypeWithScale, and the type information of C and Bias is defined by MatmulType, including the logical memory location, data format, data type, and transpose information. In addition, input MatmulWithScalePolicy by using template parameter MatmulPolicy to enable the MxMatmul scenario.

   1 2 3 4 5 6 7

   template <TPosition POSITION, TPosition SCALE_POSITION, CubeFormat FORMAT, typename TYPE, bool ISTRANS = false, TPosition SRCPOS = TPosition::GM, CubeFormat SCALE_FORMAT = FORMAT, bool SCALE_ISTRANS = ISTRANS, TPosition SCALE_SRCPOS = SRCPOS> struct MatmulTypeWithScale: public MatmulType<POSITION, FORMAT, TYPE, ISTRANS> { constexpr static TPosition scalePosition = SCALE_POSITION; constexpr static CubeFormat scaleFormat = SCALE_FORMAT; constexpr static bool isScaleTrans = SCALE_ISTRANS; constexpr static TPosition srcScalePos = SCALE_SRCPOS; };

2. Perform initialization.

   1	REGIST_MATMUL_OBJ(&pipe, GetSysWorkSpacePtr(), mm, &tiling); // Initialization

3. Set the left matrix A, right matrix B, left quantization coefficient matrix scaleA, right quantization coefficient matrix scaleB, and bias.
   Set the left quantization coefficient matrix scaleA and right quantization coefficient matrix scaleB by using SetTensorScaleA and SetTensorScaleB.

   1 2 3 4 5

   mm.SetTensorA(gm_a, isTransposeA); // Set the left matrix A. mm.SetTensorB(gm_b, isTransposeB); // Set the right matrix B. mm.SetTensorScaleA(gm_scaleA, isTransposeScaleA); // Set the left quantization coefficient matrix scaleA. mm.SetTensorScaleB(gm_scaleB, isTransposeScaleB); // Set the right quantization coefficient matrix scaleB. mm.SetBias(gm_bias); // Set the bias.

4. Execute the matrix multiplication.
   • Call Iterate to complete a single iteration computation, and use a while loop to compute the full data on a single core. The Iterate method allows for flexible control over the number of iterations required to compute the desired amount of data.

     1 2 3

     while (mm.Iterate()) { mm.GetTensorC(gm_c); }

   • Call IterateAll to compute all data on a single core. The IterateAll method does not require cyclic iterations and is relatively simple to use.

     1	mm.IterateAll(gm_c);

5. End the matrix multiplication.

   1

   mm.End();

For more complete operator examples, see MxMatmul sample with the K direction size of Scale being even, MxMatmul sample with the K direction size of Scale being odd, mx_ub_tscm_nz sample, and matmul_mx_typepara sample.

Parameters

Constraints

• In the MxMatmul scenario, only the Norm template and MDL template are supported.

• In the MxMatmul scenario, if both matrices A and B are located in the GM, there is no special restriction on singleKIn. In this case, if the K direction size of scaleA and scaleB (that is, Ceil(singleKIn, 32)) is odd, you must manually pad zeros in the K direction until it becomes even. For example, when singleKIn is 30 and Ceil(singleKIn, 32) is 1, you must manually pad zeros in the K direction of scaleA and scaleB until the K direction size becomes even. For other combinations (matrices A and B not both in the GM), the value of singleKIn, after being aligned upward to 32 elements, must be an even multiple of 32.
• In the MxMatmul scenario, if the input data type is fp4x2_e2m1_t or fp4x2_e1m2_t, the inner axis must be even.
• In the MxMatmul scenario, the GEMV mode can be enabled by setting the data format of matrices A and scaleA to VECTOR. In this mode, matrices A and scaleA only support GM as their logical memory location, and neither of them can be transposed.
• When the data of matrices A and B is directly read from the UB, the inner axis of the matrices must be 32-byte aligned. For example, if the shape of matrix A is (M, K), K must be 32-byte aligned. If the shape of matrix A is (K, M), M must be 32-byte aligned.
• When the data of matrices scaleA and scaleB is directly read from the UB, the inner axis of the matrices must be 32-byte aligned. For example, if the shape of matrix scaleA is (M, K/32), K/32 must be 32-byte aligned. If the shape of matrix scaleA is (K/32, M), M must be 32-byte aligned.
• When matrices scaleA and scaleB are input in ND format, the high-level API needs to occupy the UB temporary space during internal format conversion. You need to use the SetLocalWorkspace API to configure the temporary space. The formula for computing the size of the temporary space (in bytes) is as follows:

  int32_t scaleATmpBuf = 0;
  int32_t scaleBTmpBuf = 0;
  if constexpr (A_TYPE::scalePosition == TPosition::VECOUT) {
      if (A_TYPE::isScaleTrans) {
          scaleATmpBuf = CeilAlign(SingleCoreM, 32) * scaleK;
      } else {
          scaleATmpBuf = CeilAlign(scaleK, 32) * SingleCoreM;
      }
  }
  if constexpr (B_TYPE::scalePosition == TPosition::VECOUT) {
      if (B_TYPE::isScaleTrans) {
          scaleBTmpBuf = SingleCoreN * CeilAlign(scaleK, 32);
      } else {
          scaleBTmpBuf = scaleK * CeilAlign(SingleCoreN, 32);
      }
  }
  int32_t totalTmpBuf = scaleATmpBuf + scaleBTmpBuf;