aclnnNsaCompressGrad

Applicable Products

Product	Supported
Ascend 950PR/Ascend 950DT	✕
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

Operator function: backward computation of the aclnnNsaCompress operator.
Formula:

The forward computation formula for selective attention is as follows:
$\text{dw} = \text{dk\_cmp} \cdot K^\top$ $\text{dk} = W^\top \cdot \text{dk\_cmp}$

Function Prototype

Each operator is divided into a two-phase API. The "aclnnNsaCompressGradGetWorkspaceSize" API must be called first to obtain the workspace size required for computation and the executor that encapsulates the operator computation flow. Then, the "aclnnNsaCompressGrad" API is called to perform the computation.

aclnnStatus aclnnNsaCompressGradGetWorkspaceSize(
  const aclTensor   *outputGrad,
  const aclTensor   *input,
  const aclTensor   *weight,
  const aclIntArray *actSeqLenOptionalOptional,
  int64_t            compressBlockSize,
  int64_t            compressStride,
  int64_t            actSeqLenType,
  char              *layoutOptionalOptional,
  const aclTensor   *inputGradOut,
  const aclTensor   *weightGradOut,
  uint64_t          *workspaceSize,
  aclOpExecutor    **executor)

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

aclnnNsaCompressGradGetWorkspaceSize

Parameters

Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
outputGrad	Input	Backward gradient output by the forward operator.	-	BFLOAT16, FLOAT16	ND	3	√
input	Input	Tensor to be compressed.	-	BFLOAT16, FLOAT16	ND	3	√
weight	Input	Compressed weight, which satisfies the broadcast relationship with the shape of input.	The data type is the same as that of input.	BFLOAT16, FLOAT16	ND	2	√
actSeqLenOptionalOptional	Input	Describes the S size corresponding to each batch. Required when batches have different lengths.	-	INT64	ND	1
compressBlockSize	Input	Specifies the compressed sliding window size.	-	INT64	-	-
compressStride	Input	The stride between two compression sliding windows.	-	INT64	-	-
actSeqLenType	Input	Takes the value 0 or 1. Currently only 0 is supported.	-	INT64	-	-
layoutOptionalOptional	Input	Represents the data layout format of the input, supporting TND.	-	string	-	-
inputGradOut	Output	Gradient of input, with the same shape as input.	The data type is the same as that of input.	BFLOAT16, FLOAT16	ND	3	√
weightGradOut	Output	Gradient of weight, with the same shape as weight.	The data type is the same as that of weight.	BFLOAT16, FLOAT16	ND	2	√
workspaceSize	Output	Returns the workspace size to be requested on the Device side.	-	-	-	-
executor	Output	Returns the op executor, which contains 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 input, weight, outputGrad, inputGrad, or weightGrad is a null pointer.
ACLNN_ERR_PARAM_INVALID	161002	The data types of input and weight are not within the supported range.
		The shapes of input and weight cannot be broadcast.
		layoutOptional is invalid.

aclnnNsaCompressGrad

Parameters

Parameter	Input/Output	Description
workspace	Input	The workspace memory address applied for on the Device side.
workspaceSize	Input	The workspace size applied for on the Device side, obtained by the first-phase API aclnnNsaCompressGradGetWorkspaceSize.
executor	Input	The operator executor, which contains the operator computation flow.
stream	Input	Specifies the stream for task execution.

Return Value

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

Constraints

Deterministic computation:
- aclnnNsaCompressGrad defaults to a deterministic implementation.
compressBlockSize and compressStride must be integer multiples of 16, and compressBlockSize > compressStride.

Example

The following is a sample code for calling the aclnn single-operator API (using Atlas A2 training products as an example). This is for reference only. For the specific compilation and execution process, see Compile and Run Samples.

#include <algorithm>
#include <cstdint>
#include <iostream>
#include <vector>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <fstream>
#include <fcntl.h>

#include "acl/acl.h"
#include "aclnnop/aclnn_nsa_compress_grad.h"

#define SUCCESS 0
#define FAILED 1

#define INFO_LOG(fmt, args...) fprintf(stdout, "[INFO]  " fmt "\n", ##args)
#define WARN_LOG(fmt, args...) fprintf(stdout, "[WARN]  " fmt "\n", ##args)
#define ERROR_LOG(fmt, args...) fprintf(stderr, "[ERROR]  " fmt "\n", ##args)

#define CHECK_RET(cond, return_expr)     \
    do {                                 \
        if (!(cond)) {                   \
            return_expr;                 \
        }                                \
    } while (0)

#define LOG_PRINT(message, ...)         \
    do {                                \
        printf(message, ##__VA_ARGS__); \
    } while (0)

bool ReadFile(const std::string &filePath, size_t &fileSize, void *buffer, size_t bufferSize)
{
    struct stat sBuf;
    int fileStatus = stat(filePath.data(), &sBuf);
    if (fileStatus == -1) {
        ERROR_LOG("failed to get file %s", filePath.c_str());
        return false;
    }
    if (S_ISREG(sBuf.st_mode) == 0) {
        ERROR_LOG("%s is not a file, please enter a file", filePath.c_str());
        return false;
    }

    std::ifstream file;
    file.open(filePath, std::ios::binary);
    if (!file.is_open()) {
        ERROR_LOG("Open file failed. path = %s", filePath.c_str());
        return false;
    }

    std::filebuf *buf = file.rdbuf();
    size_t size = buf->pubseekoff(0, std::ios::end, std::ios::in);
    if (size == 0) {
        ERROR_LOG("file size is 0");
        file.close();
        return false;
    }
    if (size > bufferSize) {
        ERROR_LOG("file size is larger than buffer size");
        file.close();
        return false;
    }
    buf->pubseekpos(0, std::ios::in);
    buf->sgetn(static_cast<char *>(buffer), size);
    fileSize = size;
    file.close();
    return true;
}

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, aclrtContext* context, aclrtStream* stream) {
    // Boilerplate: ACL 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 = aclrtCreateContext(context, deviceId);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateContext failed. ERROR: %d\n", ret); return ret);
    ret = aclrtSetCurrentContext(*context);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetCurrentContext 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** xOrResult) {
    auto size = GetShapeSize(shape) * sizeof(T);
    // Call aclrtMalloc to allocate 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 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);

    // Compute the strides of the contiguous xOrResult.
    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.
    *xOrResult = 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 device/context/stream. Refer to the acl external API list.
    // Fill in the deviceId based on your actual device.
    int32_t deviceId = 0;
    aclrtContext context;
    aclrtStream stream;
    auto ret = Init(deviceId, &context, &stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);

    // 2. Construct inputs and outputs. Custom construction is required based on the API.
    int64_t headNum = 64;
    int64_t headDim = 128;
    int64_t blockSize = 32;
    int64_t blockStride = 16;
    int64_t blockNum = 15;
    int64_t seqLensSum = 272;
    int64_t seqLen = 3;
    std::vector<int64_t> outputGradShape = {blockNum, headNum, headDim};
    std::vector<int64_t> inputKVShape = {seqLensSum, headNum, headDim};
    std::vector<int64_t> weightShape = {blockSize, headNum};
    std::vector<int64_t> inputGradOutShape = {seqLensSum, headNum, headDim};
    std::vector<int64_t> weightGradOutShape = {blockSize, headNum};
    int64_t SeqLenType = 0;
    char layOut[] = "TND";

    void* outputGradDeviceAddr = nullptr;
    void* inputKVDeviceAddr = nullptr;
    void* weightDeviceAddr = nullptr;
    void* inputGradOutDeviceAddr = nullptr;
    void* weightGradOutDeviceAddr = nullptr;

    aclTensor* outputGrad = nullptr;
    aclTensor* inputKV = nullptr;
    aclTensor* weight = nullptr;
    aclTensor* inputGradOut = nullptr;
    aclTensor* weightGradOut = nullptr;

    std::vector<float> inputGradOutHostData(seqLensSum * headNum * headDim, 0.0);
    std::vector<float> weightGradOutHostData(blockSize * headNum, 0.0);

    std::vector<float> outputGradHostData(blockNum * headNum * headDim, 1.0);
    std::vector<float> inputKVHostData(seqLensSum * headNum * headDim, 1.0);
    std::vector<float> weightHostData(blockSize * headNum, 1.0);
    std::vector<int64_t> actSeqLenOptionalHostData = {0, 128, 272};

    aclIntArray *actSeqLenOptional = aclCreateIntArray(actSeqLenOptionalHostData.data(), actSeqLenOptionalHostData.size());

    // Create the dy aclTensor.
    ret = CreateAclTensor(outputGradHostData, outputGradShape, &outputGradDeviceAddr, aclDataType::ACL_FLOAT16,
                          &outputGrad);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the x aclTensor.
    ret = CreateAclTensor(inputKVHostData, inputKVShape, &inputKVDeviceAddr, aclDataType::ACL_FLOAT16, &inputKV);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the gelu aclTensor.
    ret = CreateAclTensor(weightHostData, weightShape, &weightDeviceAddr, aclDataType::ACL_FLOAT16, &weight);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    ret = CreateAclTensor(inputGradOutHostData, inputGradOutShape, &inputGradOutDeviceAddr, aclDataType::ACL_FLOAT16, &inputGradOut);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    ret = CreateAclTensor(weightGradOutHostData, weightGradOutShape, &weightGradOutDeviceAddr, aclDataType::ACL_FLOAT16, &weightGradOut);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    // 3. Call the CANN operator library API. Replace with the specific API name.
    uint64_t workspaceSize = 0;
    aclOpExecutor* executor;
    // Call the first-phase API of aclnnNsaCompressGrad.
    ret = aclnnNsaCompressGradGetWorkspaceSize(
        outputGrad, inputKV, weight, actSeqLenOptional, blockSize, blockStride, SeqLenType, layOut,
        inputGradOut, weightGradOut, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnGeGluGradV2GetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
    // Allocate device memory based on the 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 of aclnnNsaCompressGrad.
    ret = aclnnNsaCompressGrad(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnGeGluGradV2 failed. ERROR: %d\n", ret); return ret);

    // 4. (Boilerplate) Synchronously 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 memory to the host memory. Modify based on the specific API definition.
    auto size = GetShapeSize(inputGradOutShape);
    std::vector<float> resultData(size, 0);
    ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), inputGradOutDeviceAddr,
                        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 ret);
    for (int64_t i = 0; i < size; i++) {
        LOG_PRINT("result[%ld] is: %f\n", i, resultData[i]);
    }

    // 6. Release aclTensor and aclScalar. Modify based on the specific API definition.
    aclDestroyTensor(outputGrad);
    aclDestroyTensor(inputKV);
    aclDestroyTensor(weight);
    aclDestroyTensor(inputGradOut);
    aclDestroyTensor(weightGradOut);

    // 7. Release device resources. Modify based on the specific API definition.
    aclrtFree(outputGradDeviceAddr);
    aclrtFree(inputKVDeviceAddr);
    aclrtFree(weightDeviceAddr);
    aclrtFree(inputGradOutDeviceAddr);
    aclrtFree(weightGradOutDeviceAddr);
    if (workspaceSize > 0) {
        aclrtFree(workspaceAddr);
    }
    aclrtDestroyStream(stream);
    aclrtDestroyContext(context);
    aclrtResetDevice(deviceId);
    aclFinalize();
    return 0;
}