Input and Output

input

[m, k]/[batch, m, k]

• Floating point: float16/bf16
• Quantization:
  • When quantType is per_group: float16/bf16
  • Other scenarios: int8

ND

Matrix A for matrix multiplication.

• When type is ALL_GATHER_LINEAR_REDUCE_SCATTER, quantization is not supported.
• When type is LINEAR_REDUCE_SCATTER:
  • When batch is greater than 1, batch must be exactly divided by rackSize.
  • In other cases, m must be exactly divided by rankSize.
• When mc2 is used, only two dimensions are supported, and the value range of k is [256, 65535).

weight

ND: [k, n]

NZ:

• Non-quantization: [1, n/16, k, 16] or [k, n]
• Quantization: [1, n/32, k, 32] or [k, n]

mc2:

• ND: [k, n]
• NZ: [k, n]

• Floating point: float16/bf16
• Quantization: int8

ND/NZ

Weight, matrix B for matrix multiplication.

• If backend is set to lcoc, type is LINEAR_ALL_REDUCE.
  • In non-quantization scenarios, float16/bf16 supports NZ.
  • When NZ is supported for quantization, the data type of input must be int8, and bias and deqScale must be W8A8.
• If backend is set to hccl or lccl, only float16 supports NZ.
• In NZ format:
  • k must be a multiple of 16.
  • hasResidual cannot be set to true.
  • quantType cannot be set to per_group.
  • keepIntermediate cannot be set to true.
  • Transposition is not supported. That is, transWeight cannot be set to true.
• Quantization is not supported when ALL_GATHER_LINEAR_REDUCE_SCATTER is used.
• When mc2 is used, only two dimensions are supported, and the value range of k is [256, 65535).

bias

• quantType = per_tensor: [1]
• quantType = per_channel: [1, n]/[n]
• quantType = per_group: []
• quantType = per_token: [1, n]/[n]

• W8A16 quantization: float16/bf16
• W8A8 quantization: int32

ND

Added bias matrix.

Not supported when ALL_GATHER_LINEAR_REDUCE_SCATTER is used.

deqScale

• quantType = per_tensor: [1]
• quantType = per_channel: [1, n]/[n]
• quantType = per_group: [k/quantGroupSize, n]
• quantType = per_token: [1, n]/[n]

• W8A16 quantization: float16/bf16
• W8A8 quantization:
  • If the output is of type float16, deqScale is of type int64.
  • If the output is of type bf16, both int64 and float are supported.

ND

Scale of dequantization.

Required only for quantization.

Not supported when type is ALL_GATHER_LINEAR_REDUCE_SCATTER.

allgatherV2 does not support quantization. allgatherV2 refers to the scenario where type is ALL_GATHER_LINEAR and keepIntermediate is true.

residual

[n]

float16/bf16

ND

Residual, added to the final output result.

Not supported when type is ALL_GATHER_LINEAR_REDUCE_SCATTER.

Not supported for non-quantization.

perTokenScale

[m]

float

ND

Per-token dequantization scale, input during quantization and required only when quantType is QUANT_TYPE_PER_TOKEN.

output

• When type is set to linear_all_reduce/pure_linear: [m, n]/[batch, m, n]
• When type is linear_reduce_scatter: [m/rankSize, n]/[batch/rankSize, m, n]
• When type is set to all_gather_linear: [m*rankSize, n]/[batch*rankSize, m, n]
• When type is set to ALL_GATHER_LINEAR_REDUCE_SCATTER: [m*agDim/rsDim, n]

float16/bf16

ND

Output tensor. The number of dimensions is the same as that of x.

intermediateOutput

[m*rankSize, n]/[batch*rankSize, m, n]

float16/bf16

ND

Output tensor. The number of dimensions is the same as that of x.

Valid only when keepIntermediate is set to true and ParallelType is set to ALL_GATHER_LINEAR.