API function:
The
[object Object]API extends the functions of the[object Object]API. In addition to supporting the x1 and x2 input types of FLOAT16/BFLOAT16, the following functions are added:Ascend 950PR/Ascend 950DT:
The low-precision data type FLOAT8_E4M3FN/FLOAT8_E5M2/HIFLOAT8 is supported. It supports pertensor, perblock, and mx .
[object Object]Atlas A3 training products/Atlas A3 inference products[object Object], [object Object]Atlas A2 training products/Atlas A2 inference products[object Object]:
The low-precision data types INT8 and INT4 are supported. It supports per-token/per-channel .
Formula:
Case 1: If the data types of x1 and x2 are FLOAT16 or BFLOAT16, MatMul is performed on x1 and x2 after AllGather is performed on x1.
Case 2: If the data types of x1 and x2 are FLOAT8_E4M3FN/FLOAT8_E5M2/HIFLOAT8 in the pertensor scenario, or INT8/INT4 in the perchannel or pertoken scenario, and amaxOut is not output, MatMul is performed on x1 and x2 after AllGather is performed on x1, and then dequant is performed.
Case 3: If the data types of x1 and x2 are FLOAT8_E4M3FN/FLOAT8_E5M2/HIFLOAT8 in the perblock scenario, and amaxOut is not output, when x1 is (m, k) and x2 is (k, n), x1Scale is (ceilDiv(m, 128), ceilDiv(k, 128)) and x2Scale is (ceilDiv(k, 128), ceilDiv(n, 128)). After AllGather is performed on x1 and x1Scale, perblock quantized MatMul is performed on x1 and x2, and then dequant is performed.
Case 4: If the data types of x1 and x2 are FLOAT8_E4M3FN/FLOAT8_E5M2 in the mx quantization scenario, x1 is (a0, a1), x2 is (b0, b1), x1Scale is (a0, ceilDiv(a1, 64), 2), and x2Scale is (b0, ceilDiv(b1, 64), 2), after AllGather is performed on x1 and x1Scale, MatMul is performed on x1 and x2, and then dequant is performed.
Each operator has calls. First, [object Object] is called to obtain the workspace size required for computation and the executor that contains the operator computation flow. Then, [object Object] is called to perform computation.
Parameters
[object Object]- [object Object]Atlas A2 training products/Atlas A2 inference products[object Object] and [object Object]Atlas A3 training products/Atlas A3 inference products[object Object]:
- x1 and x2: The data type can be FLOAT16, BFLOAT16, INT8, or INT4.
- bias: When commMode is set to aiv, the current version supports only nullptr as the input.
- x1Scale: The data type can be FLOAT. When the data type of x1 and x2 is FLOAT16/BFLOAT16, only
[object Object]is supported. In per-token scenarios, the shape is (m, 1). - x2Scale: The data type can be FLOAT or INT64. The INT64 data type is supported only when the data type of x1 and x2 is INT4 or the data type of output is FLOAT16. When the data type of x1 and x2 is FLOAT16/BFLOAT16, only
[object Object]is supported. In per-channel scenarios, the shape is (1, n). - groupSize: In the current version, only 0 is supported.
- commMode: Currently, only the aiv mode is supported. In the
[object Object]mode, the AI Vector Core is used to complete communication tasks. In the current version, only "aiv" is supported. - output: The data type can be FLOAT16 or BFLOAT16. If the data type of x1 is FLOAT16 or BFLOAT16, the data type of output is the same as that of x1.
- gatherOut: The data type can be FLOAT16, BFLOAT16, INT8, or INT4.
- Ascend 950PR/Ascend 950DT:
- x1 and x2: The data type can be FLOAT16, BFLOAT16, FLOAT8_E4M3FN, FLOAT8_E5M2 or HIFLOAT8.
- bias: If the data type of x1 is FLOAT16 or BFLOAT16, the data type of bias must be FLOAT16 or BFLOAT16. If the data type of x1 is FLOAT8_E4M3FN, FLOAT8_E5M2 or HIFLOAT8, the data type of bias must be FLOAT in the pertensor and mx quantization scenarios. In the perblock scenario, only nullptr can be input.
- x1Scale: When the data type of x1 and x2 is FLOAT16 or BFLOAT16, only nullptr can be input. In the pertensor scenario, the shape is [1]. In perblock mode, the shape is [ceilDiv(m, 128), ceilDiv(k, 128)]. In pertensor and perblock modes, the data type is FLOAT. In the mx quantization scenario, the data type is FLOAT8_E8M0 and the shape is (m, ceilDiv(k, 64), 2).
- x2Scale: When the data types of x1 and x2 are FLOAT16 or BFLOAT16, only nullptr is supported. In pertensor mode, the shape is [1]. In perblock mode, the shape is [ceilDiv(k, 128), ceilDiv(n, 128)]. In pertensor and perblock modes, the data type is FLOAT. In the mx scenario, the data type is FLOAT8_E8M0 and the shape is (ceilDiv(k, 64), n, 2). Only the transpose scenario is supported.
- groupSize: When both x1Scale and x2Scale are 2D inputs and their data types are FLOAT, the value of groupSize must be 549764202624, which corresponds to the value of [groupSizeM, groupSizeN, groupSizeK] being [128, 128, 128]. When both x1Scale and x2Scale are 3D inputs and their data types are FLOAT8_E8M0, the value of groupSize must be 4295032864, which corresponds to the value of [groupSizeM, groupSizeN, groupSizeK] being [1, 1, 32]. In other scenarios, only 0 is supported.
- commMode: Only "ccu" is supported in the current version.
- output: If the type of x1 is FLOAT16 or BFLOAT16, the type of output is the same as that of x1. If the type of x1 is FLOAT8_E4M3FN, FLOAT8_E5M2 or HIFLOAT8, the data type can be FLOAT16, BFLOAT16, or FLOAT.
- gatherOut: The data type can be FLOAT16, BFLOAT16, FLOAT8_E4M3FN, FLOAT8_E5M2 or HIFLOAT8.
- Restrictions on groupSize:
- The groupSize value is valid only when x1Scale and x2Scale are 2D or higher-dimensional inputs. In other scenarios, the input must be 0.
- The input groupSize is decomposed into groupSizeM, groupSizeN, and groupSizeK according to the following formulas. If one or more of them are 0, groupSizeM, groupSizeN, and groupSizeK are reset based on the input shape of x1/x2/x1Scale/x2Scale for calculation. Principle: If groupSizeM is 0, the quantization group size in the m direction is inferred by the API. The inference formula is groupSizeM = m/scaleM (m must be exactly divided by scaleM). m is the same as that in the x1 shape, and scaleM is the same as that in the x1Scale shape.
- [object Object]Atlas A2 training products/Atlas A2 inference products[object Object] and [object Object]Atlas A3 training products/Atlas A3 inference products[object Object]:
Returns
[object Object]: status code. For details, see .The first-phase API implements input parameter validation. The following error codes may be returned.
[object Object]
Deterministic computation:
- The default deterministic implementation of
[object Object]is used.
- The default deterministic implementation of
Ascend 950PR/Ascend 950DT:
- The input x1 is 2-dimensional, and its dimension is (m, k). x2 must be 2-dimensional, with the shape of (k, n). The axis meets the input parameter requirements of the mm operator, the k axis is equal, and the value range of the k axis is [256, 65535).
- Input
[object Object]must be 1D (n,). - The output is 2-dimensional, with the shape of (m x rank_size, n), where rank_size indicates the number of devices.
- The gatherout is 2-dimensional, with the shape of (m x rank_size, k), where rank_size indicates the number of devices.
- When the data type of x1 and x2 is FLOAT16 or BFLOAT16, the output data type is the same as that of x1 and x2.
- When the data type of x1 and x2 is FLOAT8_E4M3FN/FLOAT8_E5M2/HIFLOAT8, the output data type can be FLOAT16, BFLOAT16, or FLOAT.
- When the data type of x1 and x2 is FLOAT16, BFLOAT16, or HIFLOAT8, the data types of x1 and x2 must be the same.
- When the data type of x1 and x2 is FLOAT8_E4M3FN/FLOAT8_E5M2, the data type of x1 and x2 can be either of them.
- When the data type of x1 and x2 is FLOAT16/BFLOAT16/HIFLOAT8/FLOAT8_E4M3FN/FLOAT8_E5M2, x2 can be transposed or not, and x1 can only be not transposed.
- When groupSize is set to 549764202624, bias must be empty.
- 2, 4, 8, 16, 32, or 64 devices are supported.
- The total size of the AllGather (x1) collective communication data cannot exceed 16 x 256 MB. The total size of the collective communication data is calculated as follows: m x k x sizeof(x1_dtype) x Number of devices. The internal implementation of the operator may vary according to the shape. Therefore, the actual supported total communication volume may be slightly less than this value.
[object Object]Atlas A3 training products/Atlas A3 inference products[object Object], [object Object]Atlas A2 training products/Atlas A2 inference products[object Object]:
- The
[object Object]matrix cannot be transposed. The[object Object]matrix can be transposed or not transposed. - Input
[object Object]is 2D (m, k). - Input
[object Object]must be 2D (k, n). The axes must meet the input parameter requirements of the MatMul operator. The k axes of[object Object]and[object Object]must be equal and fall within the range of [256, 65535). - The bias supports only nullptr as the input.
- The output is 2D (m*rank_size, n). rank_size indicates the number of devices.
- Empty tensors are not supported.
- The data types of
[object Object]and[object Object]must be the same. - When the data types of x1 and x2 are INT4, k and n must be even numbers.
- Two, four, and eight devices are supported.
- The
Note: This sample code calls some HCCL collective communication library APIs, including HcclGetCommName, HcclCommInitAll, and HcclCommDestroy. For details, see .
The following example is for reference only. For details, see .
[object Object]Atlas A2 training products/Atlas A2 inference products[object Object]
[object Object]Ascend 950PR/Ascend 950DT:
[object Object]