StridedBatchMatmulOperation

Description

Groups matrixes and specifies the step between each group of matrices to implement more flexible matrix multiplication.

Definition

struct StridedBatchMatmulParam {
    bool transposeA = false;
    bool transposeB = false;
    int32_t batch      = 1;
    int32_t headNum    = 1;
    std::vector<int32_t> m;
    std::vector<int32_t> n;
    std::vector<int32_t> k;
    std::vector<int32_t> lda;
    std::vector<int32_t> ldb;
    std::vector<int32_t> ldc;
    std::vector<int32_t> strideA;
    std::vector<int32_t> strideB;
    std::vector<int32_t> strideC;
    uint8_t rsv[8] = {0};
};

Parameters

Member	Type	Default Value	Description
transposeA	bool	false	Whether to transpose matrix A.
transposeB	bool	false	Whether to transpose matrix B.
batch	int32_t	1	Number of batches (batchSize).
headNum	int32_t	1	Number of heads in the multi-headed attention mechanism.
m	std::vector< int32_t >	-	Shape size of matrix A participating in a matrix multiplication instruction. The number of elements is batchSize.
n	std::vector< int32_t >	-	Shape size of matrix B participating in a matrix multiplication instruction. The number of elements is batchSize.
k	std::vector< int32_t >	-	Shape size of matrix C participating in a matrix multiplication instruction. The number of elements is batchSize.
lda	std::vector< int32_t >	-	Number of columns in matrix A. The number of elements is batchSize.
ldb	std::vector< int32_t >	-	Number of columns in matrix B. The number of elements is batchSize.
ldc	std::vector< int32_t >	-	Number of columns in matrix C. The number of elements is batchSize.
strideA	std::vector< int32_t >	-	Stride between two adjacent calculations of matrix A in the memory. The number of elements is batchSize.
strideB	std::vector< int32_t >	-	Stride between two adjacent calculations of matrix B in the memory. The number of elements is batchSize.
strideC	std::vector< int32_t >	-	Stride between two adjacent calculations of matrix C in the memory. The number of elements is batchSize.
rsv[8]	uint8_t	{0}	Reserved

Input

Parameter	Dimension	Data Type	Format	Description
A	bmm1, bmm1_grad2, bmm2_grad1: [nTokens, hiddenSize] bmm1_grad1, bmm2, bmm2_grad2: [nSquareTokens]	float16	ND	Input tensor
B	bmm1, bmm1_grad2, bmm2_grad1: [nTokens, hiddenSize] bmm1_grad1, bmm2, bmm2_grad2: [nSquareTokens]	float16	ND	Input tensor

Output

Parameter	Dimension	Data Type	Format	Description
output	[outdims]	float16	ND	Output tensor.

Restrictions

Currently, only the Atlas A2 inference products is supported.

Parent topic: atb/train_op_params.h