PagedCacheLoadOperation

You can connect PyTorch by calling the C API 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 (keyCache/valueCache in ND Format)	Dimension (keyCache/valueCache in NZ Format)	Data Type	Format	Default Value	Mandatory	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. elenum_aligned is 32 if this parameter is an int8 value, or 16 in other cases.
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. elenum_aligned is 32 if this parameter is an int8 value, or 16 in other cases.
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 block table. Only the ND format is supported.
kvCacheCfg	Scalar	-	-	int8_t	-	0	Yes	K_CACHE_V_CACHE_NZ = 0 (default): key_cache and value_cache in NZ format are passed. K_CACHE_V_CACHE_ND = 1: key_cache and value_cache in ND format are passed.
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	Sets SeqStart to provide the start position (similar to offset) of each batch in the block table. Only the ND format is supported.

Original API

See PagedCacheLoadOperation.

Parent topic: atb/acl.h