PrefixEncoder

You can connect PyTorch by calling the C API 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	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. Set it to a nullptr if maskType is not MASK_TYPE_ALIBI_COMPRESS or MASK_TYPE_ALIBI_COMPRESS_SQRT.
seqLen	Tensor	[batch]	int32/uint32	ND	-	Yes	Sequence length of each batch corresponding to query. Used on the host.
kvSeqLen	Tensor	[batch]	int32/uint32	ND	-	Yes	Sequence length of each batch corresponding to key and value. Tensor on the host.
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. Set it to a nullptr if maskType is set to MASK_TYPE_CAUSAL_MASK.
maskType	Scalar	-	int	-	0	Yes	Mask type. Options: 3: MASK_TYPE_ALIBI_COMPRESS, ALiBi-compressed mask. 4: MASK_TYPE_ALIBI_COMPRESS_SQRT, ALiBi-compressed mask with square root. 5: MASK_TYPE_ALIBI_COMPRESS_LEFT_ALIGN, ALiBi-compressed left-aligned mask. Supported only for 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, depending on the model in use.
qkScale	Scalar	-	float	-	1	Yes	Scale post-multiplied by Q*K^T.
attnOut	Tensor	[nTokens, qHiddenSize] [nTokens, headNum, headSize]	float16/bf16	ND	-	Yes	Attention output.

Original API

See Input and Output When calcType Is PREFIX_ENCODER.

Parent topic: SelfAttentionOperation