PagedCacheLoadOperation

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

Definition

atb::Status AtbPagedCacheLoadGetWorkspaceSize(const aclTensor *keyCache, const aclTensor *valueCache,
                                              const aclTensor *blockTables, const aclTensor *contextLens,
                                              const aclTensor *key, const aclTensor *value, const aclTensor *seqStarts,
                                              int8_t kvCacheCfg, bool isSeqLensCumsumType, bool hasSeqStarts,
                                              uint64_t *workspaceSize, atb::Operation **op, atb::Context *context);
atb::Status AtbPagedCacheLoad(void *workspace, uint64_t workspaceSize, atb::Operation *op, atb::Context *context);

AtbPagedCacheLoadGetWorkspaceSize Members

Parameter	Scalar/Tensor	Dimension (ND Format for keyCache/valueCache)	Dimension (NZ Format for keyCache/valueCache)	Data Type	Format	Default Value	Mandatory or Not	Description
keyCache	Tensor	[num_blocks, block_size, num_heads, head_size_k]	[num_blocks, num_heads * head_size_k // elenum_aligned, block_size, elenum_aligned]	float16/bf16/int8	ND/NZ	-	Yes	Input tensor. If the data type is int8, elenum_aligned is 32. In other cases, the value is 16.
valueCache	Tensor	[num_blocks, block_size, num_heads, head_size_v]	[num_blocks, num_heads * head_size_v // elenum_aligned, block_size, elenum_aligned]	float16/bf16/int8	ND/NZ	-	Yes	Input tensor. If the data type is int8, elenum_aligned is 32. In other cases, the value is 16.
blockTables	Tensor	[batch, block_indices]	[len(contextLens), (max(contextLens) - 1) // block_size + 1]	int32	ND	-	Yes	Input tensor. len(contextLens) indicates the length of contextlens.
contextLens	Tensor	[batch]or[batch+1]	[len(contextLens)]	int32	ND	-	Yes	Input tensor. len(contextLens) indicates the length of contextlens.
key	Tensor	[num_tokens, num_heads, head_size_k]	[sum(contextLens), num_heads * head_size_k]	float16/bf16/int8	ND	-	Yes	Input/Output tensor. sum(contextLens) is the sum of the elements of contextlens.
value	Tensor	[num_tokens, num_heads, head_size_v]	[sum(contextLens), num_heads * head_size_v]	float16/bf16/int8	ND	-	Yes	Input/Output tensor. sum(contextLens) is the sum of the elements of contextlens.
seqStarts	Tensor	[batch]	-	int32	ND	-	No	(Optional) Start position of each batch in the blocktable. Only the ND format is supported.
kvCacheCfg	Scalar	-	-	int8_t	-	0	Yes	K_CACHE_V_CACHE_NZ = 0 (default value): key_cache and value_cache are passed, and the NZ format is used. K_CACHE_V_CACHE_ND = 1: key_cache and value_cache are passed, and the ND format is used.
isSeqLensCumsumType	Scalar	-	-	bool	-	false	No	The input seqLens can be in accumulated sum mode. That is, the number of elements to be extracted in batch n is the accumulated sum of the number of elements to be extracted in the first n batches. Only the ND format is supported.
hasSeqStarts	Scalar	-	-	bool	-	false	No	SeqStart is used to provide the initial position (similar to offset) of each batch in the blockTable. Only the ND format is supported.

Original APIs

For details, see PagedCacheLoadOperation.

Parent topic: atb/acl.h