GatherPaKvCache

产品支持情况

产品	是否支持
Ascend 950PR/Ascend 950DT	√
Atlas A3 训练系列产品/Atlas A3 推理系列产品	√
Atlas A2 训练系列产品/Atlas A2 推理系列产品	√
Atlas 200I/500 A2 推理产品	×
Atlas 推理系列产品	×
Atlas 训练系列产品	×

功能说明

算子功能：根据blockTables中的blockId值、seqLens中key/value的seqLen从keyCache/valueCache中将内存不连续的token搬运、拼接成连续的key/value序列。
计算逻辑：
- keyRef/valueRef的第一个维度取决于seq_lens大小。
- 如果isSeqLensCumsum为true，则seqLens中最后一个值即为keyRef/valueRef的第一个维度大小： keyRef[dim0] = seqLens[-1]
- 如果isSeqLensCumsum为false，则seqLens中所有值的累加和即为keyRef/valueRef的第一个维度大小：keyRef[dim0] = sum(seqLens)
关于keyRef、valueRef的一些限制条件如下：
- 每个token大小控制在148k以内，例如，对于fp16/bf16类型，num_heads * head_size(keyRef/valueRef)取128*576。

示例：

  keyCache_shape: [128, 128, 16, 144]
  valueCache_shape: [128, 128, 16, 128]
  blockTables_shape: [16, 12]
  seqLens_shape: [16]
  keyRef_shape: [8931, 16, 144]
  valueRef_shape: [8931, 16, 128]
  seqOffset_shape: [16]
  out1_shape: [8931, 16, 144]  
  out2_shape: [8931, 16, 128]

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
keyCache	输入	当前层存储的key向量缓存	INT8, FLOAT16, BFLOAT16, FLOAT, UINT8, INT16, UINT16, INT32, UINT32, HIFLOAT8, FLOAT8_E5M2, FLOAT8_E4M3FN	ND
valueCache	输入	当前层存储的value向量缓存	INT8, FLOAT16, BFLOAT16, FLOAT, UINT8, INT16, UINT16, INT32, UINT32, HIFLOAT8, FLOAT8_E5M2, FLOAT8_E4M3FN	FRACTAL_NZ
blockTables	输入	每个batch中KV Cache的逻辑块到物理块的映射关系	INT32、INT64	ND
seqLens	输入	每个batch对应的序列长度	INT32、INT64	ND
keyRef	输入/输出	当前层的key向量	INT8, FLOAT16, BFLOAT16, FLOAT, UINT8, INT16, UINT16, INT32, UINT32, HIFLOAT8, FLOAT8_E5M2, FLOAT8_E4M3FN	ND
valueRef	输入/输出	当前层的value向量	INT8, FLOAT16, BFLOAT16, FLOAT, UINT8, INT16, UINT16, INT32, UINT32, HIFLOAT8, FLOAT8_E5M2, FLOAT8_E4M3FN	ND
seqOffset	输入	blockTables获取blockId时存在的首偏移	INT32、INT64	ND
cacheMode	输入	表示输入的数据排布格式，支持Norm、PA_NZ	String	ND
isSeqLensCumsum	输入	表示seqLens是否为累加和。false表示非累加和	BOOL	ND

约束说明

无

调用说明

调用方式	样例代码	说明
aclnn接口	test_aclnn_gather_pa_kv_cache	通过aclnnGatherPaKvCache调用GatherPaKvCache算子