aclnnNsaCompressAttention

Supported Products

Product	Supported
Atlas A3 training products/Atlas A3 inference products	√
Atlas A2 training products/Atlas A2 inference products	√
Atlas 200I/500 A2 inference products	×
Atlas inference products	×
Atlas training products	×

Function

Description: Computes the compressed attention and selected top-k indexes in Native Sparse Attention (NSA). For details, see the paper (https://arxiv.org/pdf/2502.11089).
Computation formulas: compression block size: $l$ ; selection block size: $l'$ ; and compression stride: $d$

P_{cmp} = Softmax(query*key^T) \\

attentionOut = Softmax(atten\_mask(scale*query*key^T, atten\_mask)))*value

P_{slc}[j] = \sum_{m=0}^{l'/d-1}\sum_{n=0}^{l/d-1}P_{cmp} [l'/d*j-m-n],

P_{slc'} = \sum_{h=1}^{H}P_{slc}^{h}

P_{slc'} = topk\_mask(P_{slc'})

topkIndices = topk(P_{slc'})

NsaCompressAttention input tensors (query, key, and value) support flexible data layouts interpretable across multiple dimensions. Use inputLayout to specify the desired format (only TND is supported).

B (Batch): input batch size
T: total length of all input sample sequences
S (Seq-Length): sequence length of input samples
H (Head-Size): hidden-layer size
N (Head-Num): number of heads
D (Head-Dim): minimum unit size of the hidden layer (D = H/N)

Prototype

Each operator has two-phase API calls. First, aclnnNsaCompressAttentionGetWorkspaceSize is called to obtain the workspace size required for computation and the executor that contains the operator computation flow. Then, aclnnNsaCompressAttention is called to perform computation.

aclnnStatus aclnnNsaCompressAttentionGetWorkspaceSize(
  const aclTensor   *query,
  const aclTensor   *key,
  const aclTensor   *value,
  const aclTensor   *attenMaskOptional,
  const aclTensor   *topkMaskOptional,
  const aclIntArray *actualSeqQLenOptional,
  const aclIntArray *actualCmpSeqKvLenOptional,
  const aclIntArray *actualSelSeqKvLenOptional,
  double             scaleValue,
  int64_t            headNum,
  char              *inputLayout,
  int64_t            sparseMode,
  int64_t            compressBlockSize,
  int64_t            compressStride,
  int64_t            selectBlockSize,
  int64_t            selectBlockCount,
  const aclTensor   *softmaxMaxOut,
  const aclTensor   *softmaxSumOut,
  const aclTensor   *attentionOutOut,
  const aclTensor   *topkIndicesOut,
  uint64_t          *workspaceSize,
  aclOpExecutor    **executor)

aclnnStatus aclnnNsaCompressAttention(
  void             *workspace,
  uint64_t          workspaceSize,
  aclOpExecutor    *executor,
  const aclrtStream stream)

aclnnNsaCompressAttentionGetWorkspaceSize

Parameters:

Name	Input/Output	Description	Usage Notes	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
query	Input	`query` in the formulas.	-	FLOAT16, BFLOAT16	ND	3–4	√
key	Input	`key` in the formulas.	-	FLOAT16, BFLOAT16	ND	3–4	√
value	Input	`value` in the formula.	-	FLOAT16, BFLOAT16	ND	3–4	√
attenMaskOptional	Input	`atten_mask` in the formula.	The input shape must be `[S, S]`. In the TND scenario, only the `[S, S]` format is supported, which indicates `max(Sq)` and `max(CmqSkv)`, respectively.	BOOL	ND	2	√
actualSeqQLenOptional	Input	`S` of `query` (`Sq`) corresponding to each batch.	-	INT64	-	1	-
actualCmpSeqKvLenOptional	Input	`S` of `key`/`value` (`CmpSkv`) corresponding to each batch of the compressed attention.	-	INT64	-	1	-
actualSelSeqKvLenOptional	Input	`S` of `key`/`value` (`SelSkv`) corresponding to each batch after the importance score computation and compression.	-	INT64	-	1	-
topkMaskOptional	Input	`topk_mask` in the formula.	The input shape must be `[S, S]`. In the TND scenario, only the `[S, S]` format is supported, which indicates `max(Sq)` and `max(SelSkv)`, respectively. If this parameter is not used, a null pointer can be passed.	BOOL	ND	2	√
scaleValue	Input	`scale` in the formula, indicating the scaling coefficient.	Generally, this parameter is set to `D^–0.5`.	DOUBLE	-	-	-
headNum	Input	Number of heads in `query`.	-	INT64	-	-	-
inputLayout	Input	Layout of the input `query`, `key`, and `value`.	Currently, only `TND` is supported.	String	-	-	-
sparseMode	Input	Sparse mode.	Only `0` and `1` are supported.	INT64	-	-	-
compressBlockSize	Input	Size of the compression sliding window, corresponding to `l` in the formula.	A multiple of 16, up to `128`.	INT64	-	-	-
compressStride	Input	Sliding window interval between two compressions, corresponding to `d` in the formula.	A multiple of 16, up to `64`.	INT64	-	-	-
selectBlockSize	Input	Size of the selected block, corresponding to `l'` in the formula.	A multiple of 16, up to `128`.	INT64	-	-	-
selectBlockCount	Input	Number of selected blocks, corresponding to the k value in the formula.	The value range is [1, 32] and `selectBlockCount` must be less than or equal to `min(SelSkv)`.	INT64	-	-	-
softmaxMaxOut	Output	Intermediate result of the Max operation in Softmax.	Used for backward propagation.	FLOAT	ND	3	√
softmaxSumOut	Output	Intermediate result of the Sum operation in Softmax.	Used for backward propagation.	FLOAT	ND	3	√
attentionOut	Output	`attentionOut` in the formula.	The data type and shape must be the same as those of `query`.	FLOAT16, BFLOAT16	ND	3–4	√
topkIndicesOut	Output	`topkIndices` in the formula.	-	INT32	-	3	√
workspaceSize	Output	Size of the workspace required to be allocated on the device.	-	-	-	-	-
executor	Output	Operator executor, containing the operator computation process.	-	-	-	-	-

Returns:

aclnnStatus: status code. For details, see aclnn Return Codes.

The first-phase API implements input parameter verification. The following errors may be thrown.

Return	Error Code	Description
ACLNN_ERR_PARAM_NULLPTR	161001	The input `query`, `key`, or `value` is a null pointer.
ACLNN_ERR_PARAM_INVALID	161002	The data type of `query`, `key`, or `value` is not supported.
		`inputLayout` is invalid.
		`sparseMode` is invalid.

aclnnNsaCompressAttention

Parameters:

Name	Input/Output	Description
workspace	Input	Address of the workspace to be allocated on the device.
workspaceSize	Input	Size of the workspace to be allocated on the device, which is obtained by the first-phase API `aclnnNsaCompressAttentionGetWorkspaceSize`.
executor	Input	Operator executor, containing the operator computation process.
stream	Input	Stream for executing the task.

Returns:

aclnnStatus: status code. For details, see aclnn Return Codes.

Constraints

Deterministic computation:
- aclnnNsaCompressAttention defaults to deterministic implementation.
When this API is used together with PyTorch, ensure that the CANN package versions match the PyTorch package versions.
The values of compressBlockSize, compressStride, and selectBlockSize must be integer multiples of 16 and meet the following requirements: compressBlockSize>=compressStride && selectBlockSize>=compressBlockSize && selectBlockSize%compressStride==0.
actualSeqQLenOptional, actualCmpSeqKvLenOptional, and actualSelSeqKvLenOptional must use the cumulative sum mode and must be passed in TND format.
Due to the UB restriction, CmpSkv must be less than or equal to 14000.
SelSkv = CeilDiv(CmpSkv, selectBlockSize // compressStride)
The constraints for the input query, key, and value are as follows:
- Their data types must be the same.
- Their batch sizes must be the same.
- headDim must meet the following requirement: qD == kD && kD >= vD
- inputLayout must be consistent.
If headNum of input query is N1 and headNum of input key and value is N2, then N1 >= N2 && N1% N2 == 0.
Assume G = N1/N2. G must meet the following constraint: G < 128 && 128 % G == 0.
The usage of attenMask and topkMask must comply with the description in the paper.

Example

The following example is for reference only (using Atlas A2 training products/Atlas A2 inference products as an example). For details, see Compilation and Running Sample.

#include <iostream>
#include <vector>
#include <cstring>
#include "acl/acl.h"
#include "aclnn/opdev/fp16_t.h"
#include "aclnnop/aclnn_nsa_compress_attention.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) {
    // (Boilerplate) Initialize resources.
    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;
}

void PrintOutResult(std::vector<int64_t> &shape, void** deviceAddr) {
    auto size = GetShapeSize(shape);
    std::vector<float> resultData(size, 0);
    auto ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]),
                           *deviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return);
    for (int64_t i = 0; i < size; i++) {
        LOG_PRINT("mean result[%ld] is: %f\n", i, resultData[i]);
    }
}

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);
    // Call aclrtMalloc to allocate memory on the 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);
    // Call aclrtMemcpy to copy the data on the host to the memory on the 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);
    
    // Compute the strides of the contiguous tensor.
    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];
    }

    // Call aclCreateTensor to create an 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. (Boilerplate) Initialize the device and stream. For details, see the AscendCL API manual.
    // Set the device ID (deviceId) based on the actual device.
    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. Construct the inputs and outputs based on the API definition.
    int64_t T1 = 1024;
    int64_t T2 = 64;
    int64_t N1 = 16;
    int64_t N2 = 4;
    int64_t D1 = 192;
    int64_t D2 = 128;
    int64_t selectBlockSize = 64;
    int64_t selectBlockCount = 16;
    int64_t compressBlockSize = 32;
    int64_t compressStride = 16;
    std::vector<int64_t> qShape = {T1, N1, D1};
    std::vector<int64_t> kShape = {T2, N2, D1};
    std::vector<int64_t> vShape = {T2, N2, D2};
    std::vector<int64_t> attenmaskShape = {T1, T2};                        //[maxS1, maxS2]
    std::vector<int64_t> topkmaskShape = {T1, T1 / selectBlockSize};       //[maxS1, maxSelS2]
    std::vector<int64_t> softmaxMaxShape = {T1, N1, 8};
    std::vector<int64_t> softmaxSumShape = {T1, N1, 8};
    std::vector<int64_t> attenOutShape = {T1, N1, D2};                     //[T1, N1, D2]
    std::vector<int64_t> topkIndicesOutShape = {T1, N2, selectBlockCount}; //[T1, N2, selectBlockCount]

    void* qDeviceAddr = nullptr;
    void* kDeviceAddr = nullptr;
    void* vDeviceAddr = nullptr;
    void* attenmaskDeviceAddr = nullptr;
    void* topkmaskDeviceAddr = nullptr;
    void* softmaxMaxDeviceAddr = nullptr;
    void* softmaxSumDeviceAddr = nullptr;
    void* attentionOutDeviceAddr = nullptr;
    void* topkIndicesOutDeviceAddr = nullptr;

    aclTensor* q = nullptr;
    aclTensor* k = nullptr;
    aclTensor* v = nullptr;
    aclTensor* attenmask = nullptr;
    aclTensor* topkmask = nullptr;
    aclTensor* softmaxMax = nullptr;
    aclTensor* softmaxSum = nullptr;
    aclTensor* attentionOut = nullptr;
    aclTensor* topkIndicesOut = nullptr;

    std::vector<op::fp16_t> qHostData(T1 * N1 * D1, 1.0);
    std::vector<op::fp16_t> kHostData(T2 * N2 * D1, 1.0);
    std::vector<op::fp16_t> vHostData(T2 * N2 * D2, 1.0);
    std::vector<uint8_t> attenmaskHostData(T1 * T2, 0);
    std::vector<uint8_t> topkmaskHostData(T1 * (T1 / selectBlockSize), 0);
    std::vector<float> softmaxMaxHostData(N1 * T1 * 8, 1.0);
    std::vector<float> softmaxSumHostData(N1 * T1 * 8, 1.0);
    std::vector<op::fp16_t> attenOutHostData(T1 * N1 * D2, 1.0);
    std::vector<int32_t> topkIndicesHostData(T1 * N2 * selectBlockCount, 1);

    ret = CreateAclTensor(qHostData, qShape, &qDeviceAddr, aclDataType::ACL_FLOAT16, &q);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(kHostData, kShape, &kDeviceAddr, aclDataType::ACL_FLOAT16, &k);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(vHostData, vShape, &vDeviceAddr, aclDataType::ACL_FLOAT16, &v);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(attenmaskHostData, attenmaskShape, &attenmaskDeviceAddr, aclDataType::ACL_UINT8, &attenmask);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(topkmaskHostData, topkmaskShape, &topkmaskDeviceAddr, aclDataType::ACL_UINT8, &topkmask);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(softmaxMaxHostData, softmaxMaxShape, &softmaxMaxDeviceAddr, aclDataType::ACL_FLOAT, &softmaxMax);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(softmaxSumHostData, softmaxSumShape, &softmaxSumDeviceAddr, aclDataType::ACL_FLOAT, &softmaxSum);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(attenOutHostData, attenOutShape, &attentionOutDeviceAddr, aclDataType::ACL_FLOAT16, &attentionOut);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(topkIndicesHostData, topkIndicesOutShape, &topkIndicesOutDeviceAddr, aclDataType::ACL_INT32, &topkIndicesOut);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    std::vector<int64_t> actualSeqQLenVec(1, T1);
    auto actualSeqQLen = aclCreateIntArray(actualSeqQLenVec.data(), actualSeqQLenVec.size());
    std::vector<int64_t> actualCmpKvSeqVec(1, T2);
    auto actualCmpKvSeqLen = aclCreateIntArray(actualCmpKvSeqVec.data(), actualCmpKvSeqVec.size());
    std::vector<int64_t> actualSelKvSeqVec(1, T1 / selectBlockSize);
    auto actualSelKvSeqLen = aclCreateIntArray(actualSelKvSeqVec.data(), actualSelKvSeqVec.size());

    double scale = 1.0;
    int64_t headNum = N1;
    char inputLayout[5] = {'T', 'N', 'D', 0};
    int64_t sparseMode = 1;

    // 3. Call the CANN operator library API.
    uint64_t workspaceSize = 0;
    aclOpExecutor* executor;

    // Call the first-phase API.
    ret = aclnnNsaCompressAttentionGetWorkspaceSize(q, k, v, attenmask, topkmask, actualSeqQLen, actualCmpKvSeqLen,
        actualSelKvSeqLen, scale, headNum, inputLayout, sparseMode, compressBlockSize, compressStride, selectBlockSize, selectBlockCount, 
        softmaxMax, softmaxSum, attentionOut, topkIndicesOut, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnNsaCompressAttentionGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);

    // Allocate device memory based on workspaceSize computed by the first-phase API.
    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);
    }

    // Call the second-phase API.
    ret = aclnnNsaCompressAttention(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnNsaCompressAttention failed. ERROR: %d\n", ret); return ret);

    // 4. (Boilerplate) Wait until the task execution is complete.
    ret = aclrtSynchronizeStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);

    // 5. Obtain the output value and copy the result from the device memory to the host. Modify the code based on the API definition.
    PrintOutResult(attenOutShape, &attentionOutDeviceAddr);
    PrintOutResult(softmaxMaxShape, &softmaxMaxDeviceAddr);
    PrintOutResult(softmaxSumShape, &softmaxSumDeviceAddr);
    PrintOutResult(topkIndicesOutShape, &topkIndicesOutDeviceAddr);

    // 6. Release resources.
    aclDestroyTensor(q);
    aclDestroyTensor(k);
    aclDestroyTensor(v);
    aclDestroyTensor(attenmask);
    aclDestroyTensor(topkmask);
    aclDestroyTensor(softmaxMax);
    aclDestroyTensor(softmaxSum);
    aclDestroyTensor(attentionOut);
    aclDestroyTensor(topkIndicesOut);
    aclrtFree(qDeviceAddr);
    aclrtFree(kDeviceAddr);
    aclrtFree(vDeviceAddr);
    aclrtFree(attenmaskDeviceAddr);
    aclrtFree(topkmaskDeviceAddr);
    aclrtFree(softmaxMaxDeviceAddr);
    aclrtFree(softmaxSumDeviceAddr);
    aclrtFree(attentionOutDeviceAddr);
    aclrtFree(topkIndicesOutDeviceAddr);
    if (workspaceSize > 0) {
        aclrtFree(workspaceAddr);
    }
    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    aclFinalize();
    return 0;
}