PagedCacheLoadOperation

Applicable Products

Hardware Model	Supported or Not
Atlas 350 accelerator card	x
Atlas A3 inference products/Atlas A3 training products	√
Atlas A2 training products/Atlas A2 inference products	√
Atlas training products	x
Atlas inference products	x
Atlas 200I/500 A2 inference products	x

Description

Moves and concatenates tokens at the corresponding positions in the kvCache based on the blockId and kv request length stored in the blockTable to form a continuous key/value sequence.

Definition

struct PagedCacheLoadParam {
    enum KvCacheCfg : int8_t {
        K_CACHE_V_CACHE_NZ = 0, 
        K_CACHE_V_CACHE_ND,    
    };
    KvCacheCfg kvCacheCfg = K_CACHE_V_CACHE_NZ;
    bool isSeqLensCumsumMode = false;
    bool hasSeqStarts = false;
    uint8_t rsv[61] = {0};
};

Parameters

Member	Type	Default Value	Value Range	Mandatory or Not	Description
kvCacheCfg	KvCacheCfg	0	0, 1	No	KvCache configuration. K_CACHE_V_CACHE_NZ: default value. key_cache and value_cache are passed, and the NZ format is used. K_CACHE_V_CACHE_ND: key_cache and value_cache are passed, and the ND format is used.
isSeqLensCumsumMode	bool	false	true,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. This function does not support the NZ format.
hasSeqStarts	bool	false	true,false	No	SeqStart is used to provide the initial position (similar to offset) of each batch in the blockTable. This function does not support the NZ format.
rsv[61]	uint8_t	{0}	[0]	No	Reserved

Input and Output List

Input when KvCacheCfg is set to K_CACHE_V_CACHE_NZ:

Parameter	Dimension	Data Type	Format	Description
keyCache	[num_blocks, num_heads * head_size_k // elenum_aligned, block_size, elenum_aligned]	float16/bf16/int8	NZ	When the data type is int8: elenum_aligned = 32. In other cases, the value is 16. The data type of valueCache can be float16 or bf16.
valueCache	[num_blocks, num_heads * head_size_v// elenum_aligned, block_size, elenum_aligned]	float16/bf16/int8	NZ	When the data type is int8, the value of elenum_aligned is 32. In other cases, the value of elenum_aligned is 16.
blocktable	[len(contextLens), (max(contextLens)-1)//block_size +1]	int32	ND	len(contextLens) indicates the length of input contextlens.
contextlens	[len(contextLens)]	int32	ND	len(contextLens) indicates the length of input contextlens.
key	[sum(contextLens), num_heads * head_size_k]	float16/bf16/int8	ND	sum(contextLens) is the sum of the elements of the input contextlens. The data type is the same as that of keyCache.
value	[sum(contextLens), num_heads * head_size_v]	float16/bf16/int8	ND	sum(contextLens) is the sum of the elements of the input contextlens. The data type is the same as that of valueCache.

Output when KvCacheCfg is set to K_CACHE_V_CACHE_NZ:

Parameter	Dimension	Data Type	Format	Description
key	[sum(contextLens), num_heads * head_size_k]	float16/bf16/int8	ND	The address is the same as that of the input key. The data type is the same as that of keyCache.
value	[sum(contextLens), num_heads * head_size_v]	float16/bf16/int8	ND	The address is the same as that of the input value. The data type is the same as that of valueCache.

Input when KvCacheCfg is set to K_CACHE_V_CACHE_ND:

Parameter	Dimension	Data Type	Format	Description
keyCache	[num_blocks, block_size, num_heads, head_size_k]	float16/bf16/int8	ND	num_heads × head_size_k must be 32-byte aligned.
valueCache	[num_blocks, block_size, num_heads, head_size_v]	float16/bf16/int8	ND	num_heads × head_size_v must be 32-byte aligned.
blocktable	[len(contextLens), (max(contextLens)-1)//block_size +1]	int32	ND	len(contextLens) indicates the length of input contextlens.
contextlens	isSeqLensCumsumMode = false: [len(contextLens)] isSeqLensCumsumMode = true: [len(contextLens)+1]	int32	ND	isSeqLensCumsumMode = false: len(contextLens) indicates the length of input contextlens. isSeqLensCumsumMode = true: sum of each len(contextLens) sequence length.
key	[num_tokens, num_heads, head_size_k]	float16/bf16/int8	ND	sum(contextLens) is the sum of the elements of the input contextlens.
value	[num_tokens, num_heads, head_size_v]	float16/bf16/int8	ND	sum(contextLens) is the sum of the elements of the input contextlens.
seq_starts	[len(contextLens)]	int32	ND	Start position of each batch in the blocktable. This tensor is used only when hasSeqStarts is set to true.

Output when KvCacheCfg is set to K_CACHE_V_CACHE_ND:

Parameter	Dimension	Data Type	Format	Description
key	[num_tokens, num_heads, head_size_k]	float16/bf16/int8	ND	The address is the same as that of the input key.
value	[num_tokens, num_heads, head_size_v]	float16/bf16/int8	ND	The address is the same as that of the input value.

Constraints

NZ: num_heads × head_size_v is less than 128 × 576. For example, for the float16 type, the size of each token is less than 144 KB.
block_size must be a multiple of 16.

Parent topic: atb/infer_op_params.h