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算子