aclnnNsaCompressAttentionInfer
产品支持情况
| 产品 | 是否支持 |
|---|---|
| Ascend 950PR/Ascend 950DT | × |
| Atlas A3 训练系列产品/Atlas A3 推理系列产品 | √ |
| Atlas A2 训练系列产品/Atlas A2 推理系列产品 | √ |
| Atlas 200I/500 A2 推理产品 | × |
| Atlas 推理系列产品 | × |
| Atlas 训练系列产品 | × |
功能说明
- 接口功能：Native Sparse Attention推理过程中，Compress Attention的计算。
- 计算公式：
Pcmp=Softmax(scale∗query⋅keyT)attentionOut=Pcmp⋅valuePslc[j]=∑m=0l′/d−1∑n=0l/d−1Pcmp[l′/d∗j−m−n]Pslc′=∑g=1GPslcg,其中 G=GroupSize（分组大小），即：G=numHeadsnumKeyValueHeadstopkIndices=topk(Pslc′)P_{cmp}= Softmax(scale * query · key^T) \\ attentionOut = P_{cmp} · value\\ P_{slc}[j] = \sum\limits_{m=0}^{l'/d -1} \sum\limits_{n = 0}^{l/d -1} P_{cmp} [l'/d * j -m - n]\\ P_{slc'} = \sum\limits_{g=1}^{G} P_{slc} ^g,\quad \text{其中 } G = \text{GroupSize（分组大小），即：} G = \frac{\text{numHeads}}{\text{numKeyValueHeads}} \\ topkIndices = topk(P_{slc'})\\
函数原型
每个算子分为两段式接口，必须先调用“aclnnNsaCompressAttentionInferGetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器，再调用“aclnnNsaCompressAttentionInfer”接口执行计算。
aclnnStatus aclnnNsaCompressAttentionInferGetWorkspaceSize(
const aclTensor *query,
const aclTensor *key,
const aclTensor *value,
const aclTensor *attentionMaskOptional,
const aclTensor *blockTableOptional,
const aclIntArray *actualQSeqLenOptional,
const aclIntArray *actualCmpKvSeqLenOptional,
const aclIntArray *actualSelKvSeqLenOptional,
const aclTensor *topKMaskOptional,
int64_t numHeads,
int64_t numKeyValueHeads,
int64_t selectBlockSize,
int64_t selectBlockCount,
int64_t compressBlockSize,
int64_t compressBlockStride,
double scaleValue,
char *layoutOptional,
int64_t pageBlockSize,
int64_t sparseMode,
const aclTensor *output,
const aclTensor *topKOutput,
uint64_t *workspaceSize,
aclOpExecutor **executor
)
aclnnStatus aclnnNsaCompressAttentionInfer(
void *workspace,
uint64_t workspaceSize,
aclOpExecutor *executor,
aclrtStream stream
)
aclnnNsaCompressAttentionInferGetWorkspaceSize
-
参数说明
参数名 输入/输出 描述 使用说明 数据类型 数据格式 维度(shape) 非连续Tensor query 输入 attention结构的Query输入。 - query中的B是[1, 10000]区间内的整数，且与blockTableOptional中的B以及actualCmpKvSeqLenOptional数组的长度相等。
- query的S轴小于等于4。
- query中的N和numHeads值相等，且N轴必须是key/value的N轴（H/D）的整数倍，此外，query的N轴与key/value的N轴（H/D）的比值（即GQA中的group大小）小于等于128，且128是group的整数倍。
- query中的D和key的D(H/numKeyValueHeads)值相等，小于等于192且大于等于value的D轴。
- query，key，value输入的数据类型完全相同，为FLOAT16或BFLOAT16。
FLOAT16、BFLOAT16 ND [BSND]、[TND] x key 输入 attention结构的Key输入。 - key中的numBlocks和参数value中的numBlocks值相等。
- key中的blockSize和pageBlockSize值相等，且blockSize小于等于128，且是16的整数倍。
- key的S轴小于等于8192。
- key的N和numKeyValueHeads值相等。
- key的D轴小于等于192且大于等于value的D轴。
FLOAT16、BFLOAT16 ND [numBlocks, blockSize, numKeyValueHeads * headDimsQK] x value 输入 attention结构的Value输入。 - value的N和numKeyValueHeads值相等。
- value的D(H/numKeyValueHeads)和output的D值相等。
- value的D轴小于等于128。
- value中的blockSize和pageBlockSize值相等。
- value的S轴小于等于8192。
FLOAT16、BFLOAT16 ND [numBlocks, blockSize, numKeyValueHeads * headDimsV] x attentionMaskOptional 可选输入 attention掩码矩阵。 仅在Q_S大于1的情况下生效 - ND - - blockTableOptional 输入 paged attention中KV存储使用的block映射表。 当前只支持paged attention，因此该参数必须传入。 INT32 ND - x actualQSeqLenOptional 可选输入 query的S轴实际长度。 - INT64 ND [B] - actualCmpKvSeqLenOptional 可选输入 经过压缩后的key和value的S轴实际长度，也即该算子处理的key和value的S轴实际长度。 当前只支持paged attention，因此该参数必须传入。 INT64 ND - - actualSelKvSeqLenOptional 可选输入 压缩前的key和value的S轴实际长度。 - INT64 ND - - topKMaskOptional 可选输入 topK计算中的掩码矩阵。 预留参数暂未使用。 - ND - - numHeads 输入 head个数。 numHeads是numKeyValueHeads的倍数关系。 INT64 - - - numKeyValueHeads 输入 kvHead个数。 - INT64 - - - selectBlockSize 输入 select阶段的block大小，在计算importance score时使用。 仅支持selectBlockSize取值16、32、48、64、80、96、112、128，且selectBlockSize大于等于compressBlockSize，并且是compressBlockStride的整数倍。 INT64 - - - selectBlockCount 输入 topK阶段需要保留的block数量。 - INT64 - - - compressBlockSize 输入 压缩时的滑窗大小。 仅支持compressBlockSize取值16、32、48、64、80、96、112、128，且需要大于等于compressBlockStride。 INT64 - - - compressBlockStride 输入 两次压缩间的滑窗间隔大小。 仅支持compressBlockStride取值16、32、48、64。 INT64 - - - scaleValue 输入 缩放系数，作为计算流中Muls的scalar值。 - DOUBLE - - - layoutOptional 输入 输入query、key、value的数据排布格式。 当前支持取值“TND”和“BSND”。 STRING - - - pageBlockSize 输入 blockTableOptional中一个block的大小。 - INT64 - - - sparseMode 输入 sparse的模式，控制有attentionMaskOptional输入时的稀疏计算。 - INT64 - - - output 输出 attention的输出。 - FLOAT16、BFLOAT16 ND [BSND]、[TND] - topKOutput 输出 topK的输出。 - INT32 ND [T, N, selectBlockCount]、[B, S, N, selectBlockCount] - workspaceSize 输出 返回需要在Device侧申请的workspace大小。 - - - - - executor 输出 返回op执行器，包含了算子计算流程。 - - - - - -
返回值
aclnnStatus：返回状态码，具体参见aclnn返回码。
第一段接口完成入参校验，出现以下场景时报错：
返回值 错误码 描述 ACLNN_ERR_PARAM_NULLPTR 161001 传入参数是必选输入，输出或者必选属性，且是空指针。 ACLNN_ERR_PARAM_INVALID 161002 query、key、value、blockTableOptional、attentionOut、topKOutput的数据类型不在支持的范围内。 ACLNN_ERR_RUNTIME_ERROR 361001 API内存调用npu runtime的接口异常。
aclnnNsaCompressAttentionInfer
-
参数说明
参数名 输入/输出 描述 workspace 输入 在Device侧申请的workspace内存地址。 workspaceSize 输入 在Device侧申请的workspace大小，由第一段接口aclnnNsaCompressAttentionInfer获取。 executor 输入 op执行器，包含了算子计算流程。 stream 输入 指定执行任务的Stream。 -
返回值
aclnnStatus：返回状态码，具体参见aclnn返回码。
约束说明
- 确定性计算：
- aclnnNsaCompressAttentionInfer默认确定性实现。
-
参数query仅支持TND、BSND输入。T是B和S合轴紧密排列的数据（每个batch的actualQSeqLenOptional）、B（batch）表示输入样本批量大小、S（qSeqlen）表示输入样本序列长度、N（numHeads）表示多头数、D（headDimsQK）表示隐藏层最小的单元尺寸。
-
压缩前的kvSeqlen的上限可以表示为：actualSelKvSeqLenCeil=(actualCmpKvSeqLenOptional-1)*compressBlockStride+compressBlockSize，需要满足actualSelKvSeqLenCeil/selectBlockSize<=4096，且需要满足selectBlockCount<=actualSelKvSeqLenCeil/selectBlockSize。如果actualSelKvSeqLenOptional不满足actualCmpKvSeqLenOptional=(actualSelKvSeqLenOptional-compressBlockSize)/compressBlockStride+1，或者actualCmpKvSeqLenOptional的长度和blockTableOptional的batch维度不同，则会默认进入单token推理场景。
-
多token推理场景下，actualQSeqLenOptional参数必须传入，actualQSeqLenOptional的长度必须和blockTableOptional的batch维度相等，仅支持query的S轴最大等于4，并且此时要求每个batch单独的actualQSeqLenOptional<=actualSelKvSeqLenOptional。如果actualQSeqLenOptional的长度和blockTableOptional的batch维度不同，或者actualQSeqLenOptional的值小于1或者大于4，则会默认进入单token推理场景。
调用示例
示例代码如下，仅供参考，具体编译和执行过程请参考编译与运行样例。
#include <iostream>
#include <vector>
#include <cstring>
#include "acl/acl.h"
#include "aclnn/opdev/fp16_t.h"
#include "aclnnop/aclnn_nsa_compress_attention_infer.h"
using namespace std;
#define CHECK_RET(cond, return_expr) \
do { \
if (!(cond)) { \
return_expr; \
} \
} while (0)
#define LOG_PRINT(message, ...) \
do { \
printf(message, ##__VA_ARGS__); \
} while (0)
int64_t GetShapeSize(const std::vector<int64_t>& shape) {
int64_t shapeSize = 1;
for (auto i : shape) {
shapeSize *= i;
}
return shapeSize;
}
int Init(int32_t deviceId, aclrtStream* stream) {
// 固定写法，资源初始化
auto ret = aclInit(nullptr);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
ret = aclrtSetDevice(deviceId);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
ret = aclrtCreateStream(stream);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
return 0;
}
template <typename T>
int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
aclDataType dataType, aclTensor** tensor) {
auto size = GetShapeSize(shape) * sizeof(T);
// 调用aclrtMalloc申请device侧内存
auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
// 调用aclrtMemcpy将host侧数据拷贝到device侧内存上
ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
// 计算连续tensor的strides
std::vector<int64_t> strides(shape.size(), 1);
for (int64_t i = shape.size() - 2; i >= 0; i--) {
strides[i] = shape[i + 1] * strides[i + 1];
}
// 调用aclCreateTensor接口创建aclTensor
*tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
shape.data(), shape.size(), *deviceAddr);
return 0;
}
int main() {
// 1. （固定写法）device/stream初始化，参考AscendCL对外接口列表
// 根据自己的实际device填写deviceId
int32_t deviceId = 0;
aclrtStream stream;
auto ret = Init(deviceId, &stream);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
// 2. 构造输入与输出，需要根据API的接口自定义构造
int32_t batchSize = 20;
int32_t headDimsQK = 192;
int32_t blockNum = 640;
int32_t headDimsV = 128;
int32_t sequenceLengthK = 4096;
int32_t maxNumBlocksPerSeq = 32;
// attr
int64_t numHeads = 64;
int64_t numKeyValueHeads = 4;
int64_t selectBlockSize = 64;
int64_t selectBlockCount = 16;
int64_t compressBlockSize = 32;
int64_t compressStride = 16;
double scaleValue = 0.088388;
string sLayerOut = "TND";
char layOut[sLayerOut.length()+1];
strcpy(layOut, sLayerOut.c_str());
int64_t pageBlockSize = 128;
int64_t sparseMod = 0;
std::vector<int64_t> queryShape = {batchSize, numHeads, headDimsQK};
std::vector<int64_t> keyShape = {blockNum, pageBlockSize, numKeyValueHeads * headDimsQK};
std::vector<int64_t> valueShape = {blockNum, pageBlockSize, numKeyValueHeads * headDimsV};
std::vector<int64_t> blockTableOptionalShape = {batchSize, maxNumBlocksPerSeq};
std::vector<int64_t> outputShape = {batchSize, numHeads, headDimsV};
std::vector<int64_t> topkIndicesShape = {batchSize, numKeyValueHeads, selectBlockCount};
void *queryDeviceAddr = nullptr;
void *keyDeviceAddr = nullptr;
void *valueDeviceAddr = nullptr;
void *blockTableOptionalDeviceAddr = nullptr;
void *outputDeviceAddr = nullptr;
void *topkIndicesDeviceAddr = nullptr;
aclTensor *queryTensor = nullptr;
aclTensor *keyTensor = nullptr;
aclTensor *valueTensor = nullptr;
aclTensor *blockTableOptionalTensor = nullptr;
aclTensor *outputTensor = nullptr;
aclTensor *topkIndicesTensor = nullptr;
std::vector<op::fp16_t> queryHostData(batchSize * numHeads * headDimsQK, 1.0);
std::vector<op::fp16_t> keyHostData(blockNum * pageBlockSize * numKeyValueHeads * headDimsQK, 1.0);
std::vector<op::fp16_t> valueHostData(blockNum * pageBlockSize * numKeyValueHeads * headDimsV, 1.0);
std::vector<int32_t> blockTableOptionalHostData(batchSize * maxNumBlocksPerSeq, 1);
std::vector<op::fp16_t> outputHostData(batchSize * numHeads * headDimsV, 1.0);
std::vector<int32_t> topkIndicesHostData(batchSize * numKeyValueHeads * selectBlockCount, 1);
// 创建query aclTensor
ret = CreateAclTensor(queryHostData, queryShape, &queryDeviceAddr, aclDataType::ACL_FLOAT16, &queryTensor);
CHECK_RET(ret == ACL_SUCCESS, return ret);
// 创建key aclTensor
ret = CreateAclTensor(keyHostData, keyShape, &keyDeviceAddr, aclDataType::ACL_FLOAT16, &keyTensor);
CHECK_RET(ret == ACL_SUCCESS, return ret);
// 创建v aclTensor
ret = CreateAclTensor(valueHostData, valueShape, &valueDeviceAddr, aclDataType::ACL_FLOAT16, &valueTensor);
CHECK_RET(ret == ACL_SUCCESS, return ret);
// 创建blockTableOptional aclTensor
ret = CreateAclTensor(blockTableOptionalHostData, blockTableOptionalShape, &blockTableOptionalDeviceAddr, aclDataType::ACL_INT32, &blockTableOptionalTensor);
CHECK_RET(ret == ACL_SUCCESS, return ret);
// 创建output aclTensor
ret = CreateAclTensor(outputHostData, outputShape, &outputDeviceAddr, aclDataType::ACL_FLOAT16, &outputTensor);
CHECK_RET(ret == ACL_SUCCESS, return ret);
// 创建topkIndices aclTensor
ret = CreateAclTensor(topkIndicesHostData, topkIndicesShape, &topkIndicesDeviceAddr, aclDataType::ACL_INT32, &topkIndicesTensor);
CHECK_RET(ret == ACL_SUCCESS, return ret);
std::vector<int64_t> actualCmpKvSeqLenVector(batchSize, sequenceLengthK);
auto actualCmpKvSeqLen = aclCreateIntArray(actualCmpKvSeqLenVector.data(), actualCmpKvSeqLenVector.size());
// 3. 调用CANN算子库API
uint64_t workspaceSize = 0;
aclOpExecutor* executor;
// 调用第一段接口
ret = aclnnNsaCompressAttentionInferGetWorkspaceSize(queryTensor, keyTensor, valueTensor, nullptr, blockTableOptionalTensor, nullptr, actualCmpKvSeqLen,
nullptr, nullptr,
numHeads, numKeyValueHeads, selectBlockSize, selectBlockCount, compressBlockSize, compressStride,
scaleValue, layOut, pageBlockSize, sparseMod, outputTensor, topkIndicesTensor, &workspaceSize, &executor);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnNsaCompressAttentionInferGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
// 根据第一段接口计算出的workspaceSize申请device内存
void* workspaceAddr = nullptr;
if (workspaceSize > 0) {
ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
}
// 调用第二段接口
ret = aclnnNsaCompressAttentionInfer(workspaceAddr, workspaceSize, executor, stream);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnNsaCompressAttentionInfer failed. ERROR: %d\n", ret); return ret);
// 4. （固定写法）同步等待任务执行结束
ret = aclrtSynchronizeStream(stream);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
// 5. 获取输出的值，将device侧内存上的结果拷贝至host侧，需要根据具体API的接口定义修改
auto size = GetShapeSize(outputShape);
std::vector<op::fp16_t> resultData(size, 0);
ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outputDeviceAddr,
size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy [attn] result from device to host failed. ERROR: %d\n", ret); return ret);
uint64_t printNum = 10;
for (int64_t i = 0; i < printNum; i++) {
std::cout << "index: " << i << ": " << static_cast<float>(resultData[i]) << std::endl;
}
auto topksize = GetShapeSize(topkIndicesShape);
std::vector<op::fp16_t> topkresultData(topksize, 0);
ret = aclrtMemcpy(topkresultData.data(), topkresultData.size() * sizeof(topkresultData[0]), topkIndicesDeviceAddr,
topksize * sizeof(topkresultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy [top k] result from device to host failed. ERROR: %d\n", ret); return ret);
for (int64_t i = 0; i < printNum; i++) {
std::cout << "topk index: " << i << ": " << static_cast<int32_t>(topkresultData[i]) << std::endl;
}
// 6. 释放资源
aclDestroyTensor(queryTensor);
aclDestroyTensor(keyTensor);
aclDestroyTensor(valueTensor);
aclDestroyTensor(blockTableOptionalTensor);
aclDestroyIntArray(actualCmpKvSeqLen);
aclDestroyTensor(outputTensor);
aclDestroyTensor(topkIndicesTensor);
aclrtFree(queryDeviceAddr);
aclrtFree(keyDeviceAddr);
aclrtFree(valueDeviceAddr);
aclrtFree(blockTableOptionalDeviceAddr);
aclrtFree(outputDeviceAddr);
aclrtFree(topkIndicesDeviceAddr);
if (workspaceSize > 0) {
aclrtFree(workspaceAddr);
}
aclrtDestroyStream(stream);
aclrtResetDevice(deviceId);
aclFinalize();
return 0;
}