昇腾社区首页
中文
注册
开发者
下载
支持
积分兑换 NEW

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以内，例如，对于fp16/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*，计算输入)：Device侧的aclTensor，表示在当前层存储的key向量缓存。当cacheMode为"Norm"时，shape为[num_blocks, block_size, num_heads, head_size_k]。当cacheMode为"PA_NZ"时，shape为[num_blocks, num_heads * head_size_k // elenum_aligned, block_size, elenum_aligned](b8场景 ：elenum_aligned=32，b16场景为16，b32场景为8。b8表示每个数据元素位宽是8bit，如INT8；b16表示每个数据元素位宽是16bit，如INT16；b32表示每个数据元素位宽是32bit，如INT32)。

      • [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]：数据类型支持INT8、FLOAT16、BFLOAT16，支持ND。

    • valueCache(aclTensor *,计算输入)：Device侧的aclTensor，表示在当前层存储的value向量缓存。当cacheMode为"Norm"时，shape为[num_blocks, block_size, num_heads, head_size_v]。当cacheMode为"PA_NZ"时，shape为[num_blocks, num_heads * head_size_v // elenum_aligned, block_size, elenum_aligned](b8场景 ：elenum_aligned=32，b16场景为16，b32场景为8)。

      • [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]：数据类型支持INT8、FLOAT16、BFLOAT16，支持ND。

    • blockTables(aclTensor*，计算输入)：Device侧的aclTensor，存储了每个batch中KV Cache的逻辑块到物理块的映射关系。支持ND，shape为[batch, block_indices]。

      • [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]：数据类型支持INT32。

    • seqLens(aclTensor*，计算输入)：Device侧的aclTensor，表示每个batch对应的序列长度。支持ND，shape为[batch]或[batch + 1]。

      • [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]：数据类型支持INT32。

    • keyRef(aclTensor*，计算输入/输出)：Device侧的aclTensor，表示key向量。支持ND。当cacheMode为"Norm"时，shape为[num_tokens, num_heads, head_size_k]。当cacheMode为"PA_NZ"时，shape为[num_tokens, num_heads * head_size_k]。

      • [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]：数据类型支持INT8、FLOAT16、BFLOAT16。

    • valueRef(aclTensor*，计算输入/输出)：Device侧的aclTensor，表示value向量。支持ND。当cacheMode为"Norm"时，shape为[num_tokens, num_heads, head_size_v]。当cacheMode为"PA_NZ"时，shape为[num_tokens, num_heads * head_size_v]。

      • [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]：数据类型支持INT8、FLOAT16、BFLOAT16。

    • seqOffsetOptional(aclTensor*，可选计算输入)：Device侧的aclTensor，支持ND，shape为[batch]。如果传入，表示在从blockTables获取blockId时存在首偏移（偏移量为seqOffsetOptional[i] / block_size，i表示某一个batch）；不传入表示不需要偏移。

      • [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]：数据类型支持INT32。

    • cacheMode(char*，计算输入)：支持["Norm", "PA_NZ"]两种模式，分别表示ND、NZ的数据格式。表示输入keyCache/valueCache/keyRef/valueRef 的shape输入模式。

    • isSeqLensCumsum(bool，计算输入)：表示seqLens是否为累加和。false表示非累加和，例如seqLens为[1, 3, 5, 3, 7]。true表示累加和，例如seqLens为[0, 1, 4, 9, 12, 19]，此时第0个元素必定是0。累加和的seqlens[i + 1] - seqlens[i]等于非累加的seqlens[i]。

    • workspaceSize(uint64_t*, 出参)：返回需要在Device侧申请的workspace大小。

    • executor(aclOpExecutor**, 出参)：返回op执行器，包含了算子计算流程。

  • 返回值：

    aclnnStatus：返回状态码，具体参见。

    [object Object]

aclnnGatherPaKvCache

  • 参数说明：

    • workspace(void*, 入参)：在Device侧申请的workspace内存地址。

    • workspaceSize(uint64_t, 入参)：在Device侧申请的workspace大小，由第一段接口aclnnGatherNdGetWorkspaceSize获取。

    • executor(aclOpExecutor*, 入参)：op执行器，包含了算子计算流程。

    • stream(aclrtStream, 入参)：指定执行任务的Stream。

  • 返回值：

    aclnnStatus：返回状态码，具体参见。

约束说明

无。

调用示例

示例代码如下，仅供参考，具体编译和执行过程请参考。

[object Object]