PrefixEncoder

The PyTorch can be directly connected through the C interface to replace affinity operators on the entire network.

Definition

atb::Status AtbSelfAttentionPrefixEncoderGetWorkspaceSize(const aclTensor *query, const aclTensor *key,
                                                          const aclTensor *value, const aclTensor *blockTables,
                                                          const aclTensor *mask, const aclTensor *seqLen,
                                                          const aclTensor *kvSeqLen, const aclTensor *slopes,
                                                          int maskType, int32_t headNum, int32_t kvHeadNum,
                                                          float qkScale, aclTensor *attnOut, uint64_t *workspaceSize,
                                                          atb::Operation **op, atb::Context *context);
atb::Status AtbSelfAttentionPrefixEncoder(void *workspace, uint64_t workspaceSize, atb::Operation *op,
                                          atb::Context *context);

AtbSelfAttentionPrefixEncoderGetWorkspaceSize Members

Parameter	Scalar/Tensor	Dimension	Data Type	Format	Default Value	Mandatory or Not	Description
query	Tensor	[nTokens, qHiddenSize] [nTokens, headNum, headSize]	float16/bf16	ND	-	Yes	Query matrix.
key	Tensor	[numBlocks, blockSize, kvHiddenSize] [numBlocks, blockSize, headNum, headSize]	float16/bf16	ND	-	Yes	Key matrix.
value	Tensor	[numBlocks, blockSize, kvHiddenSize] [numBlocks, blockSize, headNum, headSize]	float16/bf16	ND	-	Yes	Value matrix.
blockTables	Tensor	[batch, maxBlockNum]	int32	ND	-	Yes	Block table of kvcache of each query. The first dimension is the token index, and the second dimension indicates the block index.
mask	Tensor	[256, 256] [headNum, seqLen, 128]	float16/bf16	ND	-	No	alibiMask. If maskType is not MASK_TYPE_ALIBI_COMPRESS or MASK_TYPE_ALIBI_COMPRESS_SQRT, set this parameter to nullptr.
seqLen	Tensor	[batch]	int32/uint32	ND	-	Yes	Sequence length of each batch corresponding to query. Host side.
kvSeqLen	Tensor	[batch]	int32/uint32	ND	-	Yes	Sequence length of each batch corresponding to key and value. Tensor on the host side.
slopes	Tensor	[head]	fp32	ND	-	No	ALiBi coefficient. When alibiMask is used and a tensor is passed, the value is the coefficient of each head of alibiMask. When maskType is set to MASK_TYPE_CAUSAL_MASK, set this parameter to nullptr.
maskType	Scalar	-	int	-	0	Yes	Mask type. Support list: 3: MASK_TYPE_ALIBI_COMPRESS, ALiBi compression mask. 4: MASK_TYPE_ALIBI_COMPRESS_SQRT, square compression of ALiBi mask. 5: MASK_TYPE_ALIBI_COMPRESS_LEFT_ALIGN, left-aligned ALiBi compression mask. It is supported only by the Atlas A2 training products/Atlas A2 inference products and Atlas A3 inference products/Atlas A3 training products. 9: MASK_TYPE_CAUSAL_MASK, internally generated mask.
headNum	Scalar	-	int32_t	-	0	Yes	Query head size, which must be greater than 0.
kvHeadNum	Scalar	-	int32_t	-	0	Yes	Number of KV heads. The value is determined based on the actual model.
qkScale	Scalar	-	float	-	1	Yes	Scaling coefficient post-multiplied by Q*K^T.
attnOut	Tensor	[nTokens, qHiddenSize] [nTokens, headNum, headSize]	float16/bf16	ND	-	Yes	Attention output.

Original APIs

For details, see Input and Output When calcType Is PREFIX_ENCODER.

Parent topic: SelfAttentionOperation