接口功能:根据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。
示例:
[object Object]
每个算子分为,必须先调用“aclnnGatherPaKvCacheGetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器,再调用“aclnnGatherPaKvCache”接口执行计算。
[object Object]
[object Object]
参数说明
[object Object]
[object Object][object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]:
- 输入keyCache、valueCache、keyRef、valueRef不支持FLOAT、UINT8、INT16、UINT16、INT32、UINT32、HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FN数据类型。
- 输入blockTables、seqLens、seqOffsetOptional不支持INT64数据类型。
返回值:
第一段接口完成入参校验,出现以下场景时报错:
[object Object]
- 确定性计算:
- aclnnGatherPaKvCache默认确定性实现。
[object Object]