741 lines
25 KiB
Markdown
741 lines
25 KiB
Markdown
# aclnnLightningIndexer
|
||
|
||
[📄 查看源码](https://gitcode.com/cann/ops-transformer/tree/master/attention/lightning_indexer)
|
||
|
||
## 产品支持情况
|
||
|
||
|产品 | 是否支持 |
|
||
|:----------------------------|:-----------:|
|
||
|<term>Ascend 950PR/Ascend 950DT</term>| × |
|
||
|<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>| √ |
|
||
|<term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>| √ |
|
||
|<term>Atlas 200I/500 A2 推理产品</term>| × |
|
||
|<term>Atlas 推理系列产品</term>| × |
|
||
|<term>Atlas 训练系列产品</term>| × |
|
||
|
||
## 功能说明
|
||
|
||
- 接口功能：`lightning_indexer`基于一系列操作得到每一个token对应的Top-$k$个位置。
|
||
|
||
- 计算公式：
|
||
|
||
$$
|
||
Indices=\text{Top-}k\left\{[1]_{1\times g}@\left[(W@[1]_{1\times S_{k}})\odot\text{ReLU}\left(Q_{index}@K_{index}^T\right)\right]\right\}
|
||
$$
|
||
|
||
对于某个token对应的Index Query $Q_{index}\in\R^{g\times d}$，给定上下文Index Key $K_{index}\in\R^{S_{k}\times d},W\in\R^{g\times 1}$，其中$g$为GQA对应的group size，$d$为每一个头的维度，$S_{k}$是上下文的长度。
|
||
|
||
## 函数原型
|
||
|
||
每个算子分为[两段式接口](../../../docs/zh/context/两段式接口.md)，必须先调用“aclnnLightningIndexerGetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器，再调用“aclnnLightningIndexer”接口执行计算。
|
||
|
||
```Cpp
|
||
aclnnStatus aclnnLightningIndexerGetWorkspaceSize(
|
||
const aclTensor *query,
|
||
const aclTensor *key,
|
||
const aclTensor *weights,
|
||
const aclTensor *actualSeqLengthsQueryOptional,
|
||
const aclTensor *actualSeqLengthsKeyOptional,
|
||
const aclTensor *blockTableOptional,
|
||
char *layoutQueryOptional,
|
||
char *layoutKeyOptional,
|
||
int64_t sparseCount,
|
||
int64_t sparseMode,
|
||
int64_t preTokens,
|
||
int64_t nextTokens,
|
||
bool returnValues,
|
||
const aclTensor *sparseIndicesOut,
|
||
const aclTensor *sparseValuesOut,
|
||
uint64_t *workspaceSize,
|
||
aclOpExecutor **executor)
|
||
```
|
||
|
||
```Cpp
|
||
aclnnStatus aclnnLightningIndexer(
|
||
void *workspace,
|
||
uint64_t workspaceSize,
|
||
aclOpExecutor *executor,
|
||
const aclrtStream stream)
|
||
```
|
||
|
||
## aclnnLightningIndexerGetWorkspaceSize
|
||
|
||
- **参数说明：**
|
||
|
||
> [!NOTE]
|
||
>
|
||
> - query、key、weights参数维度含义：B（Batch Size）表示输入样本批量大小、S（Sequence Length）表示输入样本序列长度、H（Head Size）表示hidden层的大小、N（Head Num）表示多头数、D（Head Dim）表示hidden层最小的单元尺寸，且满足D=H/N、T表示所有Batch输入样本序列长度的累加和。
|
||
> - S1表示query shape中的S，S2表示key shape中的S，T1表示query shape中的T，T2表示key shape中的T，N1表示query shape中的N，N2表示key shape中的N。
|
||
|
||
<table style="undefined;table-layout: fixed; width: 1601px"><colgroup>
|
||
<col style="width: 264px">
|
||
<col style="width: 132px">
|
||
<col style="width: 232px">
|
||
<col style="width: 330px">
|
||
<col style="width: 164px">
|
||
<col style="width: 119px">
|
||
<col style="width: 215px">
|
||
<col style="width: 145px">
|
||
</colgroup>
|
||
<thead>
|
||
<tr>
|
||
<th>参数名</th>
|
||
<th>输入/输出</th>
|
||
<th>描述</th>
|
||
<th>使用说明</th>
|
||
<th>数据类型</th>
|
||
<th>数据格式</th>
|
||
<th>维度(shape)</th>
|
||
<th>非连续Tensor</th>
|
||
</tr></thead>
|
||
<tbody>
|
||
<tr>
|
||
<td>query</td>
|
||
<td>输入</td>
|
||
<td>公式中的输入Q。</td>
|
||
<td>不支持空tensor。</td>
|
||
<td>FLOAT16、BFLOAT16</td>
|
||
<td>ND</td>
|
||
<td>
|
||
<ul>
|
||
<li>layout_query为BSND时，shape为(B,S1,N1,D)。</li>
|
||
<li>layout_query为TND时，shape为(T1,N1,D)。</li>
|
||
</ul>
|
||
</td>
|
||
<td>x</td>
|
||
</tr>
|
||
<tr>
|
||
<td>key</td>
|
||
<td>输入</td>
|
||
<td>公式中的输入K。</td>
|
||
<td>
|
||
<ul>
|
||
<li>不支持空tensor。</li>
|
||
<li>block_num为PageAttention时block总数，block_size为一个block的token数。</li>
|
||
</ul>
|
||
</td>
|
||
<td>FLOAT16、BFLOAT16</td>
|
||
<td>ND</td>
|
||
<td>
|
||
<ul>
|
||
<li>layout_key为PA_BSND时，shape为(block_num, block_size, N2, D)。</li>
|
||
<li>layout_kv为BSND时，shape为(B, S2, N2, D)。</li>
|
||
<li>layout_kv为TND时，shape为(T2, N2, D)。</li>
|
||
</ul>
|
||
</td>
|
||
<td>x</td>
|
||
</tr>
|
||
<tr>
|
||
<td>weights</td>
|
||
<td>输入</td>
|
||
<td>公式中的输入W。</td>
|
||
<td>不支持空tensor。</td>
|
||
<td>FLOAT16、BFLOAT16、FLOAT</td>
|
||
<td>ND</td>
|
||
<td>
|
||
<ul>
|
||
<li>layout_query为BSND时，shape为(B,S1,N1)。</li>
|
||
<li>layout_query为TND时，shape为(T1,N1)。</li>
|
||
</ul>
|
||
</td>
|
||
<td>x</td>
|
||
</tr>
|
||
<tr>
|
||
<td>actualSeqLengthsQueryOptional</td>
|
||
<td>输入</td>
|
||
<td>每个Batch中，Query的有效token数。</td>
|
||
<td>
|
||
<ul>
|
||
<li>不支持空tensor。</li>
|
||
<li>如果不指定seqlen可传入None，表示和`query`的shape的S长度相同。</li>
|
||
<li>该入参中每个Batch的有效token数不超过`query`中的维度S大小且不小于0，支持长度为B的一维tensor。</li>
|
||
<li>当`layout_query`为TND时，该入参必须传入，且以该入参元素的数量作为B值，该入参中每个元素的值表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须大于等于前一个元素的值。</li>
|
||
</ul>
|
||
</td>
|
||
<td>INT32</td>
|
||
<td>ND</td>
|
||
<td>(B,)</td>
|
||
<td>x</td>
|
||
</tr>
|
||
<tr>
|
||
<td>actualSeqLengthsKeyOptional</td>
|
||
<td>输入</td>
|
||
<td>每个Batch中，Key的有效token数。</td>
|
||
<td>
|
||
<ul>
|
||
<li>不支持空tensor。</li>
|
||
<li>如果不指定seqlen可传入None，表示和key的shape的S长度相同。</li>
|
||
<li> 该参数中每个Batch的有效token数不超过`key/value`中的维度S大小且不小于0，支持长度为B的一维tensor。</li>
|
||
<li>当`layout_key`为TND或PA_BSND时，该入参必须传入，`layout_key`为TND，该参数中每个元素的值表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须大于等于前一个元素的值。</li>
|
||
</ul>
|
||
</td>
|
||
<td>INT32</td>
|
||
<td>ND</td>
|
||
<td>(B,)</td>
|
||
<td>x</td>
|
||
</tr>
|
||
<tr>
|
||
<td>blockTableOptional</td>
|
||
<td>输入</td>
|
||
<td>表示PageAttention中KV存储使用的block映射表。</td>
|
||
<td>
|
||
<ul>
|
||
<li>不支持空tensor。</li>
|
||
<li>PageAttention场景下，block\_table必须为二维，第一维长度需要等于B，第二维长度不能小于maxBlockNumPerSeq（maxBlockNumPerSeq为每个batch中最大actual\_seq\_lengths\_key对应的block数量）</li>
|
||
</ul>
|
||
</td>
|
||
<td>INT32</td>
|
||
<td>ND</td>
|
||
<td>shape支持(B,S2/block_size)</td>
|
||
<td>x</td>
|
||
</tr>
|
||
<tr>
|
||
<td>layoutQueryOptional</td>
|
||
<td>输入</td>
|
||
<td>用于标识输入Query的数据排布格式。</td>
|
||
<td>
|
||
<ul>
|
||
<li>用户不特意指定时可传入默认值"BSND"。</li>
|
||
<li>当前支持BSND、TND。</li>
|
||
</ul>
|
||
</td>
|
||
<td>STRING</td>
|
||
<td>-</td>
|
||
<td>-</td>
|
||
<td>-</td>
|
||
</tr>
|
||
<tr>
|
||
<td>layoutKeyOptional</td>
|
||
<td>输入</td>
|
||
<td>用于标识输入Key的数据排布格式。</td>
|
||
<td>
|
||
<ul>
|
||
<li>用户不特意指定时可传入默认值"BSND"。</li>
|
||
<li>当前支持PA_BSND、BSND、TND。</li>
|
||
</ul>
|
||
</td>
|
||
<td>STRING</td>
|
||
<td>-</td>
|
||
<td>-</td>
|
||
<td>-</td>
|
||
</tr>
|
||
<tr>
|
||
<td>sparseCount</td>
|
||
<td>输入</td>
|
||
<td>topK阶段需要保留的block数量。</td>
|
||
<td>支持[1, 2048]，以及3072、4096、5120、6144、7168、8192</td>
|
||
<td>INT32</td>
|
||
<td>-</td>
|
||
<td>-</td>
|
||
<td>-</td>
|
||
</tr>
|
||
<tr>
|
||
<td>sparseMode</td>
|
||
<td>输入</td>
|
||
<td>表示sparse的模式。</td>
|
||
<td>
|
||
<ul>
|
||
<li>sparse_mode为0时，代表defaultMask模式。</li>
|
||
<li>sparse_mode为3时，代表rightDownCausal模式的mask，对应以右顶点为划分的下三角场景。</li>
|
||
</ul>
|
||
</td>
|
||
<td>INT32</td>
|
||
<td>-</td>
|
||
<td>-</td>
|
||
<td>-</td>
|
||
</tr>
|
||
<tr>
|
||
<td>preTokens</td>
|
||
<td>输入</td>
|
||
<td>用于稀疏计算，表示attention需要和前几个Token计算关联。</td>
|
||
<td>仅支持默认值2^63-1。</td>
|
||
<td>INT64</td>
|
||
<td>-</td>
|
||
<td>-</td>
|
||
<td>-</td>
|
||
</tr>
|
||
<tr>
|
||
<td>nextTokens</td>
|
||
<td>输入</td>
|
||
<td>用于稀疏计算，表示attention需要和后几个Token计算关联。</td>
|
||
<td>仅支持默认值2^63-1。</td>
|
||
<td>INT64</td>
|
||
<td>-</td>
|
||
<td>-</td>
|
||
<td>-</td>
|
||
</tr>
|
||
<tr>
|
||
<td>returnValues</td>
|
||
<td>输入</td>
|
||
<td>表示是否输出sparseValuesOut。</td>
|
||
<td>
|
||
<ul>
|
||
<li>True表示输出，但图模式下不支持，False表示不输出；默认值为False</li>
|
||
<li>仅在训练且layout_key不为PA_BSND场景支持</li>
|
||
</ul>
|
||
</td>
|
||
<td>BOOL</td>
|
||
<td>-</td>
|
||
<td>-</td>
|
||
<td>-</td>
|
||
</tr>
|
||
<tr>
|
||
<td>sparseIndicesOut</td>
|
||
<td>输出</td>
|
||
<td>公式中的Indices输出。</td>
|
||
<td>不支持空tensor。</td>
|
||
<td>INT32</td>
|
||
<td>-</td>
|
||
<td>
|
||
<ul>
|
||
<li>layout_query为"BSND"时输出shape为[B, S1, N2, sparseCount]。</li>
|
||
<li>layout_query为"TND"时输出shape为[T1, N2, sparseCount]。</li>
|
||
</ul>
|
||
</td>
|
||
<td>x</td>
|
||
</tr>
|
||
<tr>
|
||
<td>sparseValuesOut</td>
|
||
<td>输出</td>
|
||
<td>公式中的Indices输出对应的value值。</td>
|
||
<td>不支持空tensor。</td>
|
||
<td>FLOAT16、BFLOAT16</td>
|
||
<td>ND</td>
|
||
<td>shape与sparseIndicesOut保持一致</td>
|
||
<td>x</td>
|
||
</tr>
|
||
<tr>
|
||
<td>workspaceSize</td>
|
||
<td>输出</td>
|
||
<td>返回需要在Device侧申请的workspace大小。</td>
|
||
<td>-</td>
|
||
<td>-</td>
|
||
<td>-</td>
|
||
<td>-</td>
|
||
<td>-</td>
|
||
</tr>
|
||
<tr>
|
||
<td>executor</td>
|
||
<td>输出</td>
|
||
<td>返回op执行器，包含了算子计算流程。</td>
|
||
<td>-</td>
|
||
<td>-</td>
|
||
<td>-</td>
|
||
<td>-</td>
|
||
<td>-</td>
|
||
</tr>
|
||
</tbody>
|
||
</table>
|
||
|
||
- **返回值：**
|
||
|
||
aclnnStatus：返回状态码，具体参见[aclnn返回码](../../../docs/zh/context/aclnn返回码.md)。
|
||
|
||
第一段接口会完成入参校验，出现以下场景时报错：
|
||
|
||
<table style="undefined;table-layout: fixed;width: 1155px"><colgroup>
|
||
<col style="width: 319px">
|
||
<col style="width: 144px">
|
||
<col style="width: 671px">
|
||
</colgroup>
|
||
<thead>
|
||
<th>返回值</th>
|
||
<th>错误码</th>
|
||
<th>描述</th>
|
||
</thead>
|
||
<tbody>
|
||
<tr>
|
||
<td>ACLNN_ERR_PARAM_NULLPTR</td>
|
||
<td>161001</td>
|
||
<td>如果传入参数是必选输入，输出或者必选属性，且是空指针，则返回161001。</td>
|
||
</tr>
|
||
<tr>
|
||
<td>ACLNN_ERR_PARAM_INVALID</td>
|
||
<td>161002</td>
|
||
<td>query、key、weights、actualSeqLengthsQueryOptional、actualSeqLengthsKeyOptional、layoutQueryOptional、layoutKeyOptional、sparseCount、sparseMode、returnValues、sparseIndicesOut、sparseValuesOut的数据类型和数据格式不在支持的范围内。</td>
|
||
</tr>
|
||
</tbody>
|
||
</table>
|
||
|
||
## aclnnLightningIndexer
|
||
|
||
- **参数说明：**
|
||
|
||
<table style="undefined;table-layout: fixed; width: 1151px"><colgroup>
|
||
<col style="width: 184px">
|
||
<col style="width: 134px">
|
||
<col style="width: 833px">
|
||
</colgroup>
|
||
<thead>
|
||
<tr>
|
||
<th>参数名</th>
|
||
<th>输入/输出</th>
|
||
<th>描述</th>
|
||
</tr></thead>
|
||
<tbody>
|
||
<tr>
|
||
<td>workspace</td>
|
||
<td>输入</td>
|
||
<td>在Device侧申请的workspace内存地址。</td>
|
||
</tr>
|
||
<tr>
|
||
<td>workspaceSize</td>
|
||
<td>输入</td>
|
||
<td>在Device侧申请的workspace大小，由第一段接口aclnnLightningIndexerGetWorkspaceSize获取。</td>
|
||
</tr>
|
||
<tr>
|
||
<td>executor</td>
|
||
<td>输入</td>
|
||
<td>op执行器，包含了算子计算流程。</td>
|
||
</tr>
|
||
<tr>
|
||
<td>stream</td>
|
||
<td>输入</td>
|
||
<td>指定执行任务的Stream。</td>
|
||
</tr>
|
||
</tbody>
|
||
</table>
|
||
|
||
- **返回值：**
|
||
|
||
aclnnStatus：返回状态码，具体参见[aclnn返回码](../../../docs/zh/context/aclnn返回码.md)。
|
||
|
||
## 约束说明
|
||
|
||
- 参数query中的N支持小于等于64，key的N支持1。
|
||
- headdim支持128。
|
||
- block_size取值为16的倍数，最大支持1024。
|
||
- 参数query、key的数据类型应保持一致。
|
||
- 参数weights不为`float32`时，参数query、key、weights的数据类型应保持一致。
|
||
|
||
## 调用示例
|
||
|
||
示例代码如下，仅供参考，具体编译和执行过程请参考[编译与运行样例](../../../docs/zh/context/编译与运行样例.md)。
|
||
|
||
```Cpp
|
||
/**
|
||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||
* See LICENSE in the root of the software repository for the full text of the License.
|
||
*/
|
||
|
||
/*!
|
||
* \file test_incre_flash_attention_v4.cpp
|
||
* \brief
|
||
*/
|
||
//testci
|
||
#include <iostream>
|
||
#include <vector>
|
||
#include <cmath>
|
||
#include <cstring>
|
||
#include "securec.h"
|
||
#include "acl/acl.h"
|
||
#include "aclnnop/aclnn_lightning_indexer.h"
|
||
|
||
using namespace std;
|
||
|
||
namespace {
|
||
|
||
#define CHECK_RET(cond) ((cond) ? true :(false))
|
||
|
||
#define LOG_PRINT(message, ...) \
|
||
do { \
|
||
(void)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);
|
||
if (!CHECK_RET(ret == ACL_SUCCESS)) {
|
||
LOG_PRINT("aclInit failed. ERROR: %d\n", ret);
|
||
return ret;
|
||
}
|
||
ret = aclrtSetDevice(deviceId);
|
||
if (!CHECK_RET(ret == ACL_SUCCESS)) {
|
||
LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret);
|
||
return ret;
|
||
}
|
||
ret = aclrtCreateStream(stream);
|
||
if (!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);
|
||
auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
|
||
if (!CHECK_RET(ret == ACL_SUCCESS)) {
|
||
LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret);
|
||
return ret;
|
||
}
|
||
|
||
ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
|
||
if (!CHECK_RET(ret == ACL_SUCCESS)) {
|
||
LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret);
|
||
return ret;
|
||
}
|
||
|
||
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];
|
||
}
|
||
|
||
*tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
|
||
shape.data(), shape.size(), *deviceAddr);
|
||
return 0;
|
||
}
|
||
|
||
struct TensorResources {
|
||
void* queryDeviceAddr = nullptr;
|
||
void* keyDeviceAddr = nullptr;
|
||
void* weightsDeviceAddr = nullptr;
|
||
void* sparseIndicesDeviceAddr = nullptr;
|
||
void* sparseValuesDeviceAddr = nullptr;
|
||
|
||
aclTensor* queryTensor = nullptr;
|
||
aclTensor* keyTensor = nullptr;
|
||
aclTensor* weightsTensor = nullptr;
|
||
aclTensor* sparseIndicesTensor = nullptr;
|
||
aclTensor* sparseValuesTensor = nullptr;
|
||
};
|
||
|
||
int InitializeTensors(TensorResources& resources) {
|
||
std::vector<int64_t> queryShape = {1, 2, 1, 128};
|
||
std::vector<int64_t> keyShape = {1, 2, 1, 128};
|
||
std::vector<int64_t> weightsShape = {1, 2, 1};
|
||
std::vector<int64_t> sparseIndicesShape = {1, 2, 1, 2048};
|
||
std::vector<int64_t> sparseValuesShape = {1, 2, 1, 2048};
|
||
|
||
int64_t queryShapeSize = GetShapeSize(queryShape);
|
||
int64_t keyShapeSize = GetShapeSize(keyShape);
|
||
int64_t weightsShapeSize = GetShapeSize(weightsShape);
|
||
int64_t sparseIndicesShapeSize = GetShapeSize(sparseIndicesShape);
|
||
int64_t sparseValuesShapeSize = GetShapeSize(sparseValuesShape);
|
||
|
||
std::vector<float> queryHostData(queryShapeSize, 1);
|
||
std::vector<float> keyHostData(keyShapeSize, 1);
|
||
std::vector<float> weightsHostData(weightsShapeSize, 1);
|
||
std::vector<int32_t> sparseIndicesHostData(sparseIndicesShapeSize, 1);
|
||
std::vector<float> sparseValuesHostData(sparseValuesShapeSize, 1);
|
||
|
||
int ret = CreateAclTensor(queryHostData, queryShape, &resources.queryDeviceAddr,
|
||
aclDataType::ACL_FLOAT16, &resources.queryTensor);
|
||
if (!CHECK_RET(ret == ACL_SUCCESS)) {
|
||
return ret;
|
||
}
|
||
|
||
ret = CreateAclTensor(keyHostData, keyShape, &resources.keyDeviceAddr,
|
||
aclDataType::ACL_FLOAT16, &resources.keyTensor);
|
||
if (!CHECK_RET(ret == ACL_SUCCESS)) {
|
||
return ret;
|
||
}
|
||
|
||
ret = CreateAclTensor(weightsHostData, weightsShape, &resources.weightsDeviceAddr,
|
||
aclDataType::ACL_FLOAT16, &resources.weightsTensor);
|
||
if (!CHECK_RET(ret == ACL_SUCCESS)) {
|
||
return ret;
|
||
}
|
||
|
||
ret = CreateAclTensor(sparseIndicesHostData, sparseIndicesShape, &resources.sparseIndicesDeviceAddr,
|
||
aclDataType::ACL_INT32, &resources.sparseIndicesTensor);
|
||
if (!CHECK_RET(ret == ACL_SUCCESS)) {
|
||
return ret;
|
||
}
|
||
|
||
ret = CreateAclTensor(sparseValuesHostData, sparseValuesShape, &resources.sparseValuesDeviceAddr,
|
||
aclDataType::ACL_FLOAT16, &resources.sparseValuesTensor);
|
||
if (!CHECK_RET(ret == ACL_SUCCESS)) {
|
||
return ret;
|
||
}
|
||
return ACL_SUCCESS;
|
||
}
|
||
|
||
int ExecuteLightningIndexer(TensorResources& resources, aclrtStream stream,
|
||
void** workspaceAddr, uint64_t* workspaceSize) {
|
||
int64_t sparseCount = 2048;
|
||
int64_t sparseMode = 3;
|
||
int64_t preTokens = 9223372036854775807;
|
||
int64_t nextTokens = 9223372036854775807;
|
||
bool returnValue = true;
|
||
constexpr const char layerOutStr[] = "BSND";
|
||
constexpr size_t layerOutLen = sizeof(layerOutStr);
|
||
char layoutQuery[layerOutLen];
|
||
char layoutKey[layerOutLen];
|
||
errno_t memcpyRet = memcpy_s(layoutQuery, sizeof(layoutQuery), layerOutStr, layerOutLen);
|
||
if (!CHECK_RET(memcpyRet == 0)) {
|
||
LOG_PRINT("memcpy_s layoutQuery failed. ERROR: %d\n", memcpyRet);
|
||
return -1;
|
||
}
|
||
memcpyRet = memcpy_s(layoutKey, sizeof(layoutKey), layerOutStr, layerOutLen);
|
||
if (!CHECK_RET(memcpyRet == 0)) {
|
||
LOG_PRINT("memcpy_s layoutKey failed. ERROR: %d\n", memcpyRet);
|
||
return -1;
|
||
}
|
||
aclOpExecutor* executor;
|
||
|
||
int ret = aclnnLightningIndexerGetWorkspaceSize(resources.queryTensor, resources.keyTensor, resources.weightsTensor, nullptr, nullptr, nullptr,
|
||
layoutQuery, layoutKey, sparseCount, sparseMode, preTokens, nextTokens,returnValue,
|
||
resources.sparseIndicesTensor, resources.sparseValuesTensor, workspaceSize, &executor);
|
||
|
||
if (!CHECK_RET(ret == ACL_SUCCESS)) {
|
||
LOG_PRINT("aclnnLightningIndexerGetWorkspaceSize failed. ERROR: %d\n", ret);
|
||
return ret;
|
||
}
|
||
|
||
if (*workspaceSize > 0ULL) {
|
||
ret = aclrtMalloc(workspaceAddr, *workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
|
||
if (!CHECK_RET(ret == ACL_SUCCESS)) {
|
||
LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret);
|
||
return ret;
|
||
}
|
||
}
|
||
|
||
ret = aclnnLightningIndexer(*workspaceAddr, *workspaceSize, executor, stream);
|
||
if (!CHECK_RET(ret == ACL_SUCCESS)) {
|
||
LOG_PRINT("aclnnLightningIndexer failed. ERROR: %d\n", ret);
|
||
return ret;
|
||
}
|
||
|
||
return ACL_SUCCESS;
|
||
}
|
||
|
||
int PrintValueOutResult(std::vector<int64_t> &shape, void** deviceAddr) {
|
||
auto size = GetShapeSize(shape);
|
||
std::vector<aclFloat16> resultData(size, 0);
|
||
auto ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]),
|
||
*deviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
|
||
if (!CHECK_RET(ret == ACL_SUCCESS)) {
|
||
LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret);
|
||
return ret;
|
||
}
|
||
for (int64_t i = 0; i < size; i++) {
|
||
LOG_PRINT("mean result[%ld] is: %f\n", i, aclFloat16ToFloat(resultData[i]));
|
||
}
|
||
return ACL_SUCCESS;
|
||
}
|
||
|
||
int PrintIndicesOutResult(std::vector<int64_t> &shape, void** deviceAddr) {
|
||
auto size = GetShapeSize(shape);
|
||
std::vector<int32_t> resultData(size, 0);
|
||
auto ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]),
|
||
*deviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
|
||
if (!CHECK_RET(ret == ACL_SUCCESS)) {
|
||
LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret);
|
||
return ret;
|
||
}
|
||
for (int64_t i = 0; i < size; i++) {
|
||
LOG_PRINT("mean result[%ld] is: %d\n", i, resultData[i]);
|
||
}
|
||
return ACL_SUCCESS;
|
||
}
|
||
|
||
void CleanupResources(TensorResources& resources, void* workspaceAddr,
|
||
aclrtStream stream, int32_t deviceId) {
|
||
if (resources.queryTensor) {
|
||
aclDestroyTensor(resources.queryTensor);
|
||
}
|
||
if (resources.keyTensor) {
|
||
aclDestroyTensor(resources.keyTensor);
|
||
}
|
||
if (resources.weightsTensor) {
|
||
aclDestroyTensor(resources.weightsTensor);
|
||
}
|
||
if (resources.sparseIndicesTensor) {
|
||
aclDestroyTensor(resources.sparseIndicesTensor);
|
||
}
|
||
if (resources.sparseValuesTensor) {
|
||
aclDestroyTensor(resources.sparseValuesTensor);
|
||
}
|
||
|
||
if (resources.queryDeviceAddr) {
|
||
aclrtFree(resources.queryDeviceAddr);
|
||
}
|
||
if (resources.keyDeviceAddr) {
|
||
aclrtFree(resources.keyDeviceAddr);
|
||
}
|
||
if (resources.weightsDeviceAddr) {
|
||
aclrtFree(resources.weightsDeviceAddr);
|
||
}
|
||
if (resources.sparseIndicesDeviceAddr) {
|
||
aclrtFree(resources.sparseIndicesDeviceAddr);
|
||
}
|
||
if (resources.sparseValuesDeviceAddr) {
|
||
aclrtFree(resources.sparseValuesDeviceAddr);
|
||
}
|
||
|
||
if (workspaceAddr) {
|
||
aclrtFree(workspaceAddr);
|
||
}
|
||
if (stream) {
|
||
aclrtDestroyStream(stream);
|
||
}
|
||
aclrtResetDevice(deviceId);
|
||
aclFinalize();
|
||
}
|
||
|
||
} // namespace
|
||
|
||
int main() {
|
||
int32_t deviceId = 0;
|
||
aclrtStream stream = nullptr;
|
||
TensorResources resources = {};
|
||
void* workspaceAddr = nullptr;
|
||
uint64_t workspaceSize = 0;
|
||
std::vector<int64_t> sparseIndicesShape = {1, 2, 1, 2048};
|
||
std::vector<int64_t> sparseValuesShape = {1, 2, 1, 2048};
|
||
int ret = ACL_SUCCESS;
|
||
|
||
// 1. Initialize device and stream
|
||
ret = Init(deviceId, &stream);
|
||
if (!CHECK_RET(ret == ACL_SUCCESS)) {
|
||
LOG_PRINT("Init acl failed. ERROR: %d\n", ret);
|
||
return ret;
|
||
}
|
||
|
||
// 2. Initialize tensors
|
||
ret = InitializeTensors(resources);
|
||
if (!CHECK_RET(ret == ACL_SUCCESS)) {
|
||
CleanupResources(resources, workspaceAddr, stream, deviceId);
|
||
return ret;
|
||
}
|
||
|
||
// 3. Execute the operation
|
||
ret = ExecuteLightningIndexer(resources, stream, &workspaceAddr, &workspaceSize);
|
||
if (!CHECK_RET(ret == ACL_SUCCESS)) {
|
||
CleanupResources(resources, workspaceAddr, stream, deviceId);
|
||
return ret;
|
||
}
|
||
|
||
// 4. Synchronize stream
|
||
ret = aclrtSynchronizeStream(stream);
|
||
if (!CHECK_RET(ret == ACL_SUCCESS)) {
|
||
LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret);
|
||
CleanupResources(resources, workspaceAddr, stream, deviceId);
|
||
return ret;
|
||
}
|
||
|
||
// 5. Process results
|
||
PrintIndicesOutResult(sparseIndicesShape, &resources.sparseIndicesDeviceAddr);
|
||
PrintValueOutResult(sparseValuesShape, &resources.sparseValuesDeviceAddr);
|
||
|
||
// 6. Cleanup resources
|
||
CleanupResources(resources, workspaceAddr, stream, deviceId);
|
||
return 0;
|
||
}
|
||
```
|