aclnnNsaCompressWithCache

Applicable Products

Function

• Description: Used for KV compression in the inference phase of Native-Sparse-Attention. Each batch generates a new token during each inference. When the number of tokens in a batch fills up one compress_block, this operator compresses the last compress_block tokens of that batch into one compress_token.

• Formula:

Function Prototype

Each operator uses a two-phase API. You must first call the "aclnnNsaCompressWithCacheGetWorkspaceSize" API to obtain the workspace size required for computation and the executor that contains the operator computation flow, and then call the "aclnnNsaCompressWithCache" API to perform the computation.

aclnnStatus aclnnNsaCompressWithCacheGetWorkspaceSize(
   const aclTensor   *input,
   const aclTensor   *weight,
   const aclTensor   *slotMapping,
   const aclIntArray *actSeqLenOptional,
   const aclTensor   *blockTableOptional,
   char              *layoutOptional,
   int64_t            compressBlockSize,
   int64_t            compressStride,
   int64_t            actSeqLenType,
   int64_t            pageBlockSize,
   aclTensor         *outputCache,
   uint64_t          *workspaceSize,
   aclOpExecutor    **executor)

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

aclnnNsaCompressWithCacheGetWorkspaceSize

• Parameters

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
  input	Input	Tensor to be compressed.	Does not support empty tensor. input and weight satisfy the broadcast relationship, and the size of the third dimension of input equals the size of the second dimension of weight. headDim is an integer multiple of 16, and headDim <= 256. headNum <= 64, and when headNum > 50, headNum % 2 = 0. N (Head-Num) indicates the number of multi-heads, and D (Head-Dim) indicates the smallest unit size of the hidden layer.	FLOAT16, BFLOAT16	ND	[blockNum, pageBlockSize, N, D], [TND]	✓
  weight	Input	Compression weight.	Does not support empty tensors. The data type must be consistent with input. N (Head-Num) indicates the number of multi-heads.	FLOAT16, BFLOAT16	ND	[compressBlockSize, N]	√
  slotMapping	Input	Index of the storage location for compressed data at the tail of each batch.	Does not support empty tensor. The values of slotMapping must be unique; otherwise, the computation result may be unstable.	INT32	ND	[B]	x
  actSeqLenOptional	optional input	Size of S for each batch.	Required in TND layout scenarios. Input nullptr in other scenarios. S (Seq-Length) indicates the sequence length of the input sample. The value of actSeqLenOptional must not exceed the maximum sequence length.	INT64	ND	[B]	-
  blockTableOptional	Optional input	Block mapping table used for KV storage in PageAttention.	Pass nullptr if this feature is not used. The value of blockTableOptional must not exceed blockNum; otherwise, an out-of-bounds access occurs.	INT32	ND	[batch, blockNumPerBatch]	-
  layoutOptional	Optional Input	Data layout format of the input.	Currently only "TND" is supported. This parameter is invalid when blockTableOptional is passed in; otherwise, it is a required parameter. Where T is the tightly packed data of B and S axes (actSeqLen of each batch), B (Batch) indicates the input sample batch size, S (Seq-Length) indicates the input sample sequence length, N (Head-Num) indicates the number of multi-heads, and D (Head-Dim) indicates the minimum unit size of the hidden layer.	STRING	-	-	-
  compressBlockSize	Input	Size of the compression sliding window.	Must be an integer multiple of 16, with compressBlockSize >= compressStride and compressBlockSize <= 64.	INT64	-	-	-
  compressStride	Input	Sliding window interval between two compressions.	Only supports compressStride values of 16, 32, 48, and 64.	INT64	-	-	-
  actSeqLenType	Input	Different expression forms of actSeqLenOptional.	Takes effect when actSeqLenOptional has input. The value can be 0 or 1. 0 indicates that the values in actSeqLenOptional are the cumulative sum of the sequence sizes of the preceding batches, and 1 indicates that the values in actSeqLenOptional are the sequence size in each batch. Currently, only 1 is supported.	INT64	-	-	-
  pageBlockSize	Input	Block size of the page in the page attention scenario.	Can only be 64 or 128.	INT64	-	-	-
  outputCache	Output	Compressed cache.	The data type is the same as that of input.	FLOAT16, BFLOAT16	ND	[result_len, N, D]	x
  workspaceSize	Output	Returns the workspace size to be applied on the Device side.	-	-	-	-	-
  executor	Output	Returns the op executor, which includes the operator computation flow.	-	-	-	-	-

• Return Value

  aclnnStatus: return code. For details, see aclnn Return Code.

  The first-phase API performs input parameter validation and returns an error in the following scenarios:

  Return Value	Error Code	Description
  ACLNN_ERR_PARAM_NULLPTR	161001	The compute input and required compute output are null pointers.
  ACLNN_ERR_PARAM_INVALID	161002	The data types or formats of the input and output are not within the supported range.
  		input, weight, or outputCache is an empty tensor.
  ACLNN_ERR_INNER_TILING_ERROR	561002	input and weight do not satisfy the broadcast relationship, that is, the size of the third dimension of input is not equal to the size of the second dimension of weight.
  		The values of activeNum, expertNum, and expertCapacity are less than 0.
  		compress_block_size or compress_stride is not an integer multiple of 16.
  		seq_lens_type is not equal to 1, or the value of layout is not one of BSH, SBH, BSND, BNSD, or TND.
  		The value of page_block_size is not 64 or 128.
  		headDim is not aligned to 16.

aclnnNsaCompressWithCache

• Parameters

  Parameter	Input/Output	Description
  workspace	Input	Memory address of the workspace applied on the Device side.
  workspaceSize	Input	Size of the workspace applied on the Device side, obtained from the first-phase API aclnnNsaCompressWithCacheGetWorkspaceSize.
  executor	Input	Operator executor that includes the operator computation flow.
  stream	Input	Specifies the stream for executing the task.

• Return Value

  aclnnStatus: return code. For details, see aclnn Return Code.

Constraints

• Deterministic computation:

  • aclnnNsaCompressWithCache defaults to a deterministic implementation.

• The N and D of outputCache must be consistent with those of input, and result_len must satisfy result_len > (blockNum * pageBlockSize - compressBlockSize) / compressStride.

• In the page attention scenario, the shape of input supports [blockNum, pageBlockSize, N, D]; in other scenarios, the shape of input supports [T, N, D].

Example

The following provides example code for reference only. For details about compilation and running, see Compile and Run Samples.

#include "acl/acl.h"
#include "aclnnop/aclnn_nsa_compress_with_cache.h"
#include <iostream>
#include <vector>
#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 shape_size = 1;
    for (auto i : shape) {
        shape_size *= i;
    }
    return shape_size;
}
int Init(int32_t deviceId, aclrtStream* stream) {
    // Boilerplate for resource initialization.
    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);
    // Call aclrtMalloc to apply for device-side memory.
    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 host-side data to the device-side memory.
    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);
    // Calculate the strides of a 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 the aclCreateTensor API 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() {
    // Set input shape-related parameters.
    constexpr int64_t compress_block_size = 32;
    constexpr int64_t compress_stride = 16;
    constexpr int64_t heads_num = 24;
    constexpr int64_t heads_dim = 192;
    constexpr int64_t batch_size = 4;
    constexpr int64_t page_block_size = 128;
    constexpr int64_t max_seq_len = 512;
    constexpr int64_t result_len = 512;
    constexpr int64_t block_num_per_batch = max_seq_len / page_block_size;
    constexpr int64_t blocks_num = block_num_per_batch * batch_size;
    // 1. Boilerplate: device/stream initialization. See the ACL external API list.
    // Fill in deviceId based on your actual device.
    int32_t deviceId = 0;
    aclrtStream stream;
    auto ret = Init(deviceId, &stream);
    // Handle the check based on your needs.
    CHECK_RET(ret == 0, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
    // 2. Construct inputs and outputs based on the API definition.
    std::vector<int64_t> inputShape = {blocks_num, page_block_size, heads_num, heads_dim};
    std::vector<int64_t> weightShape = {compress_block_size, heads_num};
    std::vector<int64_t> slotMappingShape = {batch_size};
    std::vector<int64_t> outputCacheRefShape = {result_len, heads_num, heads_dim};
    std::vector<int64_t> actSeqLenShape = {batch_size};
    std::vector<int64_t> blockTableShape = {batch_size, block_num_per_batch};

    void *inputDeviceAddr = nullptr;
    void *weightDeviceAddr = nullptr;
    void *slotMappingDeviceAddr = nullptr;
    void *outputCacheRefDeviceAddr = nullptr;
    void *actSeqLenDeviceAddr = nullptr;
    void *blockTableDeviceAddr = nullptr;

    aclTensor *input = nullptr;
    aclTensor *weight = nullptr;
    aclTensor *slotMapping = nullptr;
    aclTensor *outputCacheRef = nullptr;
    aclIntArray *actSeqLen = nullptr;
    aclTensor *blockTable = nullptr;

    std::vector<aclFloat16> inputHostData(inputShape[0] * inputShape[1] * inputShape[2] * inputShape[3],
                                          aclFloatToFloat16(1.0));
    std::vector<aclFloat16> weightHostData(weightShape[0] * weightShape[1], aclFloatToFloat16(1.0));
    std::vector<int32_t> slotMappingHostData(slotMappingShape[0], 0);
    std::vector<aclFloat16> outputCacheRefHostData(outputCacheRefShape[0] * outputCacheRefShape[1] *
                                                   outputCacheRefShape[2], aclFloatToFloat16(1.0));
    std::vector<int64_t> actSeqLenHostData(actSeqLenShape[0], 0);
    std::vector<int32_t> blockTableHostData(blockTableShape[0] * blockTableShape[1]);
    actSeqLenHostData[0]=32;
    // Create the self aclTensor.
    ret = CreateAclTensor(inputHostData, inputShape, &inputDeviceAddr, aclDataType::ACL_FLOAT16, &input);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(weightHostData, weightShape, &weightDeviceAddr, aclDataType::ACL_FLOAT16, &weight);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(slotMappingHostData, slotMappingShape, &slotMappingDeviceAddr, aclDataType::ACL_INT32,
                          &slotMapping);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(outputCacheRefHostData, outputCacheRefShape, &outputCacheRefDeviceAddr,
                          aclDataType::ACL_FLOAT16, &outputCacheRef);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    actSeqLen = aclCreateIntArray(actSeqLenHostData.data(), actSeqLenHostData.size());
    ret = CreateAclTensor(blockTableHostData, blockTableShape, &blockTableDeviceAddr, aclDataType::ACL_INT32,
                          &blockTable);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    char layout[4] = "TND";
    int64_t actSeqLenType = 1;
    // 3. Call the CANN operator library API. Replace with the specific API.
    uint64_t workspaceSize = 0;
    aclOpExecutor* executor;
    // Call the first-phase API of aclnnNsaCompressWithCache.
    ret = aclnnNsaCompressWithCacheGetWorkspaceSize(input, weight, slotMapping, actSeqLen, blockTable, layout,
                                                    compress_block_size, compress_stride, actSeqLenType,
                                                    page_block_size, outputCacheRef, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnNsaCompressWithCacheGetWorkspaceSize failed. ERROR: %d\n", ret);
              return ret);
    // Apply for device memory based on the workspaceSize calculated 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 of aclnnNsaCompressWithCache.
    ret = aclnnNsaCompressWithCache(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnNsaCompressWithCache failed. ERROR: %d\n", ret); return ret);
    // 4. Boilerplate: synchronize and wait for task execution to 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 side memory to the host side. Modify based on the specific API definition.
    auto size = GetShapeSize(outputCacheRefShape);
    std::vector<aclFloat16> resultData(size, 0);
    ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(aclFloat16), outputCacheRefDeviceAddr,
                      size * sizeof(aclFloat16), ACL_MEMCPY_DEVICE_TO_HOST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);
    for (int64_t i = heads_dim * heads_num - 16; i < heads_dim * heads_num + 16; i++) {
        printf("outputCache[%ld]:%f\n", i, aclFloat16ToFloat(resultData[i]));
    }
    // 6. Release aclTensor. Modify based on the specific API definition.
    aclDestroyTensor(input);
    aclDestroyTensor(weight);
    aclDestroyTensor(slotMapping);
    aclDestroyTensor(outputCacheRef);
    aclDestroyIntArray(actSeqLen);
    aclDestroyTensor(blockTable);

    // 7. Release device resources. Modify based on the specific API definition.
    aclrtFree(inputDeviceAddr);
    aclrtFree(weightDeviceAddr);
    aclrtFree(slotMappingDeviceAddr);
    aclrtFree(outputCacheRefDeviceAddr);
    aclrtFree(blockTableDeviceAddr);
    if (workspaceSize > 0) {
        aclrtFree(workspaceAddr);
    }
    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    aclFinalize();
    return 0;
}