PagedCacheLoadOperation（代码开放）

产品支持情况

硬件型号	是否支持
Atlas A3 推理系列产品/Atlas A3 训练系列产品	√
Atlas A2 训练系列产品/Atlas 800I A2 推理产品	√
Atlas 训练系列产品	x
Atlas 推理系列产品	x
Atlas 200I/500 A2 推理产品	x

功能说明

根据blockTable上存储的blockId和kv请求长度在kvCache上将对应位置的token搬运拼接成1条连续的key/value序列。

定义

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 cuSeqLens = false;
    bool hasSeqStarts = false;
    uint8_t rsv[61] = {0};
};

参数列表

成员名称	类型	默认值	取值范围	是否必选	描述
kvCacheCfg	KvCacheCfg	0	0，1	否	KvCache配置。 K_CACHE_V_CACHE_NZ：默认值，传入key_cache和value_cache，且为NZ格式。 K_CACHE_V_CACHE_ND：传入key_cache和value_cache，且为ND格式。
isSeqLensCumsumMode	bool	false	true,false	否	支持输入SeqLens为累加和模式，即第n个batch中所需提取的元素数量为前n个batch的所需提取的元素数量的累加和。该功能不支持NZ格式。
hasSeqStarts	bool	false	true,false	否	提供SeqStart用于为每个batch提供在blockTable中的初始位置（类似于offset）。该功能不支持NZ格式。
rsv[61]	uint8_t	{0}	[0]	否	预留参数。

输入输出列表

“KvCacheCfg”为“K_CACHE_V_CACHE_NZ”时输入：

参数	维度	数据类型	格式	描述
keyCache	[num_blocks, num_heads * head_size_k // elenum_aligned, block_size, elenum_aligned]	float16/bf16/int8	NZ	int8时：elenum_aligned=32，其他情况为16。
valueCache	[num_blocks, num_heads * head_size_v// elenum_aligned, block_size, elenum_aligned]	float16/bf16/int8	NZ	int8时：elenum_aligned=32，其他情况为16。
blocktable	[len(contextLens), ((max(contextLens)-1)//block_size +1]	int32	ND	len(contextLens)为输入contextlens的长度。
contextlens	[len(contextLens)]	int32	ND	len(contextLens)为输入contextlens的长度。
key	[sum(contextLens), num_heads * head_size_k]	float16/bf16/int8	ND	sum(contextLens)为输入contextlens的各元素求和。
value	[sum(contextLens), num_heads * head_size_v]	float16/bf16/int8	ND	sum(contextLens)为输入contextlens的各元素求和。

“KvCacheCfg”为“K_CACHE_V_CACHE_NZ”时输出：

参数	维度	数据类型	格式	描述
key	[sum(contextLens), num_heads * head_size_k]	float16/bf16/int8	ND	和输入key同地址。
value	[sum(contextLens), num_heads * head_size_v]	float16/bf16/int8	ND	和输入value同地址。

“KvCacheCfg”为“K_CACHE_V_CACHE_ND”时输入：

参数	维度	数据类型	格式	描述
keyCache	[num_blocks, block_size, num_heads, head_size_k]	float16/bf16/int8	ND	num_heads*head_size_k 需要是32B对齐。
valueCache	[num_blocks, block_size, num_heads, head_size_v]	float16/bf16/int8	ND	num_heads*head_size_v 需要是32B对齐。
blocktable	[len(contextLens), ((max(contextLens)-1)//block_size +1]	int32	ND	len(contextLens)为输入contextlens的长度。
contextlens	isSeqLensCumsumMode为false：[len(contextLens)] isSeqLensCumsumMode为true：[len(contextLens)+1]	int32	ND	isSeqLensCumsumMode为false：len(contextLens)为输入contextlens的长度。 isSeqLensCumsumMode为true：每个len(contextLens)序列长度的累加和。
key	[num_tokens, num_heads, head_size_k]	float16/bf16/int8	ND	sum(contextLens)为输入contextlens的各元素求和。
value	[num_tokens, num_heads, head_size_v]	float16/bf16/int8	ND	sum(contextLens)为输入contextlens的各元素求和。
seq_starts	[len(contextLens)]	int32	ND	每个batch在blocktable中对应的起始位置。该tensor仅在hasSeqStarts为true时使用。

“KvCacheCfg”为“K_CACHE_V_CACHE_ND”时输出：

参数	维度	数据类型	格式	描述
key	[num_tokens, num_heads, head_size_k]	float16/bf16/int8	ND	和输入key同地址。
value	[num_tokens, num_heads, head_size_v]	float16/bf16/int8	ND	和输入value同地址。

约束说明

NZ：num_heads * head_size_v控制在128*576以内，例如：对于float16类型，将每个token大小控制在144kB以内。
block_size为16的倍数。

父主题： atb/infer_op_params.h