aclnnGatherPaKvCache
产品支持情况
| 产品 | 是否支持 |
|---|---|
| [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object] | √ |
| [object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object] | √ |
| [object Object]Atlas 200I/500 A2 推理产品[object Object] | √ |
| [object Object]Atlas 推理系列产品 [object Object] | √ |
| [object Object]Atlas 训练系列产品[object Object] | √ |
功能说明
算子功能：根据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以内，例如，对于float16/bf16类型， num_heads * head_size(keyRef/valueRef)取128*576。
示例：
[object Object]
函数原型
每个算子分为，必须先调用“aclnnGatherNdGetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器，再调用“aclnnGatherNd”接口执行计算。
aclnnStatus aclnnGatherPaKvCacheGetWorkspaceSize(const aclTensor *keyCache, const aclTensor *valueCache, const aclTensor *blockTables, const aclTensor *seqLens, const aclTensor *keyRef, const aclTensor *valueRef, const aclTensor *seqOffsetOptional, char* cacheMode, const bool isSeqLensCumsum, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnGatherPaKvCache(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)
aclnnGatherNdGetWorkspaceSize
参数说明：
keyCache（aclTensor*，计算输入） shape为[num_blocks, block_size, num_heads, head_size_k]。Device侧的aclTensor，表示在当前层存储的key向量缓存。数据类型支持INT8、FLOAT16、BFLOAT16，支持ND。
valueCache(aclTensor *,计算输入)：shape为[num_blocks, block_size, num_heads, head_size_v]。Device侧的aclTensor，表示在当前层存储的value向量缓存。数据类型支持INT8、FLOAT16、BFLOAT16，支持ND。
blockTables（aclTensor*，计算输入）：shape为[batch, block_indices]。数据类型支持int32_t，储存向量的blockId。支持ND。
seqLens（aclTensor*，计算输入）：shape为[batch]or[batch+1]。数据类型支持int32_t，表示缓存向量的seqLen。支持ND。
keyRef（aclTensor*，计算输入/输出）：shape为[num_tokens, num_heads, head_size_k]。表示key向量。数据类型支持INT8、FLOAT16、BFLOAT16，支持ND。
valueRef（aclTensor*，计算输入/输出）：shape为[num_tokens, num_heads, head_size_v]。表示value向量。数据类型支持INT8、FLOAT16、BFLOAT16，支持ND。
seqOffsetOptional（aclTensor*，可选计算输入）：shape为[batch]。如果传入，表示需要进行数据偏移计算，不传入表示不需要进行数据偏移。数据类型支持int32_t，支持ND。
cacheMode（char*，计算输入）：支持["Norm", "PA_NZ"]两种模式，分别表示ND、NZ的数据格式。表示输入keyCache/valueCache/keyRef/valueRef 的shape输入模式。
isSeqLensCumsum（bool，计算输入）：表示是否对seqLens做累加。
workspaceSize(uint64_t*, 出参)：返回需要在Device侧申请的workspace大小。
executor(aclOpExecutor**, 出参)：返回op执行器，包含了算子计算流程。
返回值：
[object Object]
aclnnGatherPaKvCache
参数说明：
workspace(void*, 入参)：在Device侧申请的workspace内存地址。
workspaceSize(uint64_t, 入参)：在Device侧申请的workspace大小，由第一段接口aclnnGatherNdGetWorkspaceSize获取。
executor(aclOpExecutor*, 入参)：op执行器，包含了算子计算流程。
stream(aclrtStream, 入参)：指定执行任务的Stream。
返回值：
约束说明
无。