aclnnLightningIndexerGrad

Applicable Products

Function

• Description: In training scenarios, this operator implements the backward propagation of LightningIndexer. The inputs include Query, Key, Weights, Dy, and Indices. The backward propagation primarily uses the Indices from forward computation to extract the TopK sequence from Key, thereby reducing the Matmul computation workload.

• Formula:

  The backward propagation formula of LightningIndexer is as follows:

  $$S = Relu(Matmul(Query, Gather(Key, Indices)))$$ $$Y = Dy*Weights$$ $$dW = Reduce(S * dy)$$ $$dQ = Matmul(ReluGrad(Y, S), Gather(Key, Indices))$$ $$dK = ScatterAdd(Matmul(ReluGrad(Y, S), Q), Indices)$$

Function Prototype

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

aclnnStatus aclnnLightningIndexerGradGetWorkspaceSize(
  const aclTensor   *query,
  const aclTensor   *key,
  const aclTensor   *dy,
  const aclTensor   *spareIndices,
  const aclTensor   *weights,
  const aclTensor   *actualSeqLengthsQuery,
  const aclTensor   *actualSeqLengthsKey,
  int64_t           headNum,
  char              *layout,
  int64_t           sparseMode,
  int64_t           preTokens,
  int64_t           nextTokens,
  bool              deterministic,
  const aclTensor   *dQuery,
  const aclTensor   *dKey,
  const aclTensor   *dWeights,
  uint64_t         *workspaceSize,
  aclOpExecutor    **executor)

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

aclnnLightningIndexerGradGetWorkspaceSize

• Parameters

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (shape)	Non-contiguous Tensor
  query	Input	query in the formula.	Shape supports [B, S1, N1, D]/[T1, N1, D].	BFLOAT16, FLOAT16	ND	3-4	√
  key	Input	key in the Formula.	Supports shape [B, S2, N2, D]/[T2, N2, D].	BFLOAT16, FLOAT16	ND	3-4	√
  dy	Input	Value in the formula.	Shape supports [B, S1, N1, D]/[T1, N1, D].	BFLOAT16, FLOAT16	ND	3-4	√
  sparseIndices	Input	sparseIndices in the formula.	The shape supports [B, S1, K]/[T1, K].	INT64	ND	2-3	√
  weights	Input	weights in the Formula.	Supports shapes [B, S1, N1]/[T1, N1].	BFLOAT16, FLOAT16	ND	2-3	√
  actualSeqLengthsQuery	Input	Indicates the cumulative sum length of each batch S in the query.	Required when the layout is TND; input nullptr in other scenarios.	INT64	ND	1	√
  actualSeqLengthsKey	Input	Indicates the cumulative sum length of each Batch S of the key.	Required for TND layout; input nullptr in other scenarios.	INT64	ND	1	√
  headNum	Input	Indicates the number of heads.	-	INT64	-	-	-
  layout	Input	Indicates the data layout format of query/key.	Currently supports TND/BSND.	String	-	-	-
  sparseMode	Input	Indicates the sparse mode.	Supports values 0 or 3.	INT64	-	-	-
  preTokens	Input	Used for sparse computation. Indicates the left boundary of the slides window.	-	INT64	-	-	-
  nextTokens	Input	Indicates the right boundary of the slides window, used for sparse computation.	-	INT64	-	-	-
  deterministic	Input	Indicates whether deterministic computation is supported.	-	BOOL	-	-	-
  dQuery	Output	dQuery gradient.	The data type is consistent with query.	BFLOAT16, FLOAT16	ND	3-4	√
  dKey	Output	Gradient of dKey.	The data type is consistent with key.	BFLOAT16, FLOAT16	ND	3-4	√
  dWeights	Output	Gradient of dWeights.	The data type is consistent with weights.	BFLOAT16, FLOAT16	ND	2-3	√
  workspaceSize	Output	Returns the workspace size to be applied by the Device side.	-	-	-	-	-
  executor	Output	Returns the op executor, which contains the operator computation process.	-	-	-	-	-

• 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	Indicates that a required input, output, or required attribute is a null pointer.
  ACLNN_ERR_PARAM_INVALID	161002	The data types and data formats of query, key, dy, sparseIndices, and weights are not within the supported range.
  		The input type of input_layout is not within the supported range.

aclnnLightningIndexerGrad

• Parameters

  Parameter	Input/Output	Description
  workspace	Input	Workspace memory address applied by the Device side.
  workspaceSize	Input	Workspace size applied by the Device side, obtained from the first-phase API aclnnLightningIndexerGradGetWorkspaceSize.
  executor	Input	Operator executor, which contains 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:

  • aclnnLightningIndexerGrad defaults to a non-deterministic implementation and does not support enabling determinism through aclrtCtxSetSysParamOpt.

• When this API is used with PyTorch, ensure that the versions of CANN-related packages and PyTorch-related packages match.

• inputLayout supports TND/BSND.

• Constraints on data shape, using the BSND layout as an example:

  • B (Batchsize): Value range is 1 to 1024.

  • N (Head-Num): Value is 64.

  • G (Group): Value is 64.

  • S1 (Seq-LengthQ): Value range is 1 to 128K.

  • S2 (Seq-LengthK): Value Range is topK to 128K.

  • D (Head-Dim): Value is 128.

  • TopK: Value is 2048.

Example

The following provides sample code for calling the aclnn single operator for reference only. For details about compilation and execution, see Compile and Run Samples.

#include <iostream>
#include <cstdio>
#include <string>
#include <vector>
#include <fstream>
#include <sys/stat.h>
#include "acl/acl.h"
#include "aclnnop/aclnn_lightning_indexer_grad.h"

#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;
}

template <typename T>
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 < 10; i++) {
        LOG_PRINT("mean result[%ld] is: %f\n", i, resultData[i]);
    }
}

int Init(int32_t deviceId, aclrtContext *context, aclrtStream *stream)
{
    // Boilerplate: 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); aclFinalize(); return ret);
    ret = aclrtCreateContext(context, deviceId);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateContext failed. ERROR: %d\n", ret); aclrtResetDevice(deviceId);
        aclFinalize(); return ret);
    ret = aclrtSetCurrentContext(*context);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetCurrentContext failed. ERROR: %d\n", ret);
        aclrtDestroyContext(context); aclrtResetDevice(deviceId); aclFinalize(); return ret);
    ret = aclrtCreateStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret);
        aclrtDestroyContext(context); aclrtResetDevice(deviceId); aclFinalize(); 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 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 = static_cast<int64_t>(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;
}


void FreeResource(aclTensor *q, aclTensor *k, aclTensor *dy, aclTensor *sparseIndices, aclTensor *weights,
                  aclTensor *dQuery, aclTensor *dKey, aclTensor *dWeights, void *qDeviceAddr, void *kDeviceAddr,
                  void *dyDeviceAddr, void *sparseIndicesDeviceAddr, void *weightsDeviceAddr, void *dQueryAddr,
                  void *dKeyAddr, void *dWeightsAddr, uint64_t workspaceSize, void *workspaceAddr, int32_t deviceId,
                  aclrtContext *context, aclrtStream *stream)
{
    // Release aclTensor and aclScalar. Modify based on the specific API definition.
    if (q != nullptr) {
        aclDestroyTensor(q);
    }
    if (k != nullptr) {
        aclDestroyTensor(k);
    }
    if (dy != nullptr) {
        aclDestroyTensor(dy);
    }
    if (sparseIndices != nullptr) {
        aclDestroyTensor(sparseIndices);
    }
    if (weights != nullptr) {
        aclDestroyTensor(weights);
    }
    if (dQuery != nullptr) {
        aclDestroyTensor(dQuery);
    }
    if (dKey != nullptr) {
        aclDestroyTensor(dKey);
    }
    if (dWeights != nullptr) {
        aclDestroyTensor(dWeights);
    }

    // Release device resources.
    if (qDeviceAddr != nullptr) {
        aclrtFree(qDeviceAddr);
    }
    if (kDeviceAddr != nullptr) {
        aclrtFree(kDeviceAddr);
    }
    if (dyDeviceAddr != nullptr) {
        aclrtFree(dyDeviceAddr);
    }
    if (sparseIndicesDeviceAddr != nullptr) {
        aclrtFree(sparseIndicesDeviceAddr);
    }
    if (weightsDeviceAddr != nullptr) {
        aclrtFree(weightsDeviceAddr);
    }
    if (dQueryAddr != nullptr) {
        aclrtFree(dQueryAddr);
    }
    if (dKeyAddr != nullptr) {
        aclrtFree(dKeyAddr);
    }
    if (dWeightsAddr != nullptr) {
        aclrtFree(dWeightsAddr);
    }
    if (workspaceSize > 0) {
        aclrtFree(workspaceAddr);
    }
    if (stream != nullptr) {
        aclrtDestroyStream(stream);
    }
    if (context != nullptr) {
        aclrtDestroyContext(context);
    }
    aclrtResetDevice(deviceId);
    aclFinalize();
}

int main()
{
    // 1. (Boilerplate) Initialize device/context/stream. See the AscendCL API list.
    // Fill in 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. Customize the construction based on the API.
    // Set the shape values for query, key, dy, sparseIndices, and weights, regenerate data, and then execute.
    int64_t batch = 2;
    int64_t s1 = 3;
    int64_t s2 = 2048;
    int64_t d = 128;
    int64_t g = 64;
    int64_t n2 = 1;
    int64_t topK = 2048;

    std::vector<int64_t> qShape = {batch, s1, n2 * g, d};
    std::vector<int64_t> kShape = {batch, s2, n2, d};
    std::vector<int64_t> dyShape = {batch, s1, n2 * g, d};
    std::vector<int64_t> sparseIndicesShape = {batch, s1, topK};
    std::vector<int64_t> weightsShape = {batch, s1, n2 * g};
    std::vector<int64_t> dQueryShape = {batch, s1, n2 * g, d};
    std::vector<int64_t> dKeyShape = {batch, s2, n2, d};
    std::vector<int64_t> dWeightsShape = {batch, s1, n2 * g};

    int64_t headNum = 64;
    int64_t sparseMode = 3;
    char layoutStr[] = "BSND";
    bool deteminstic = true;
    int64_t preToken = 65536;
    int64_t nextToken = 65536;

    void *qDeviceAddr = nullptr;
    void *kDeviceAddr = nullptr;
    void *dyDeviceAddr = nullptr;
    void *sparseIndicesDeviceAddr = nullptr;
    void *weightsDeviceAddr = nullptr;
    void *dQueryAddr = nullptr;
    void *dKeyAddr = nullptr;
    void *dWeightsAddr = nullptr;

    aclTensor *q = nullptr;
    aclTensor *k = nullptr;
    aclTensor *dy = nullptr;
    aclTensor *sparseIndices = nullptr;
    aclTensor *weights = nullptr;
    aclTensor *dQuery = nullptr;
    aclTensor *dKey = nullptr;
    aclTensor *dWeights = nullptr;

    std::vector<aclFloat16> qHostData(GetShapeSize(qShape), 1.0);
    std::vector<aclFloat16> kHostData(GetShapeSize(kShape), 1.0);
    std::vector<aclFloat16> dyHostData(GetShapeSize(dyShape), 1.0);
    std::vector<int32_t> sparseIndicesHostData(GetShapeSize(sparseIndicesShape), 1);
    std::vector<aclFloat16> weightsHostData(GetShapeSize(weightsShape), 1.0);
    std::vector<aclFloat16> dQueryHostData(GetShapeSize(dQueryShape), 1.0);
    std::vector<aclFloat16> dKeyHostData(GetShapeSize(dKeyShape), 1.0);
    std::vector<aclFloat16> dWeightsHostData(GetShapeSize(dWeightsShape), 1.0);

    uint64_t workspaceSize = 0;
    void *workspaceAddr = nullptr;

    // Create an ACL tensor.
    ret = CreateAclTensor(qHostData, qShape, &qDeviceAddr, aclDataType::ACL_FLOAT16, &q);
    CHECK_RET(ret == ACL_SUCCESS,
              FreeResource(q, k, dy, sparseIndices, weights, dQuery, dKey, dWeights, qDeviceAddr, kDeviceAddr, dyDeviceAddr,
                sparseIndicesDeviceAddr, weightsDeviceAddr, dQueryAddr, dKeyAddr, dWeightsAddr, workspaceSize, workspaceAddr,
                  deviceId, &context, &stream);
              return ret);
    ret = CreateAclTensor(kHostData, kShape, &kDeviceAddr, aclDataType::ACL_FLOAT16, &k);
    CHECK_RET(ret == ACL_SUCCESS,
              FreeResource(q, k, dy, sparseIndices, weights, dQuery, dKey, dWeights, qDeviceAddr, kDeviceAddr, dyDeviceAddr,
                sparseIndicesDeviceAddr, weightsDeviceAddr, dQueryAddr, dKeyAddr, dWeightsAddr, workspaceSize, workspaceAddr,
                  deviceId, &context, &stream);
              return ret);
    ret = CreateAclTensor(dyHostData, dyShape, &dyDeviceAddr, aclDataType::ACL_FLOAT16, &dy);
    CHECK_RET(ret == ACL_SUCCESS,
              FreeResource(q, k, dy, sparseIndices, weights, dQuery, dKey, dWeights, qDeviceAddr, kDeviceAddr, dyDeviceAddr,
                sparseIndicesDeviceAddr, weightsDeviceAddr, dQueryAddr, dKeyAddr, dWeightsAddr, workspaceSize, workspaceAddr,
                  deviceId, &context, &stream);
              return ret);
    ret = CreateAclTensor(sparseIndicesHostData, sparseIndicesShape, &sparseIndicesDeviceAddr, aclDataType::ACL_INT32, &sparseIndices);
    CHECK_RET(ret == ACL_SUCCESS,
              FreeResource(q, k, dy, sparseIndices, weights, dQuery, dKey, dWeights, qDeviceAddr, kDeviceAddr, dyDeviceAddr,
                sparseIndicesDeviceAddr, weightsDeviceAddr, dQueryAddr, dKeyAddr, dWeightsAddr, workspaceSize, workspaceAddr,
                  deviceId, &context, &stream);
              return ret);
    ret = CreateAclTensor(weightsHostData, weightsShape, &weightsDeviceAddr, aclDataType::ACL_FLOAT16, &weights);
    CHECK_RET(ret == ACL_SUCCESS,
              FreeResource(q, k, dy, sparseIndices, weights, dQuery, dKey, dWeights, qDeviceAddr, kDeviceAddr, dyDeviceAddr,
                sparseIndicesDeviceAddr, weightsDeviceAddr, dQueryAddr, dKeyAddr, dWeightsAddr, workspaceSize, workspaceAddr,
                  deviceId, &context, &stream);
              return ret);

    ret = CreateAclTensor(dQueryHostData, dQueryShape, &dQueryAddr, aclDataType::ACL_FLOAT16, &dQuery);
    CHECK_RET(ret == ACL_SUCCESS,
              FreeResource(q, k, dy, sparseIndices, weights, dQuery, dKey, dWeights, qDeviceAddr, kDeviceAddr, dyDeviceAddr,
                sparseIndicesDeviceAddr, weightsDeviceAddr, dQueryAddr, dKeyAddr, dWeightsAddr, workspaceSize, workspaceAddr,
                  deviceId, &context, &stream);
              return ret);
    ret = CreateAclTensor(dKeyHostData, dKeyShape, &dKeyAddr, aclDataType::ACL_FLOAT16, &dKey);
    CHECK_RET(ret == ACL_SUCCESS,
              FreeResource(q, k, dy, sparseIndices, weights, dQuery, dKey, dWeights, qDeviceAddr, kDeviceAddr, dyDeviceAddr,
                sparseIndicesDeviceAddr, weightsDeviceAddr, dQueryAddr, dKeyAddr, dWeightsAddr, workspaceSize, workspaceAddr,
                  deviceId, &context, &stream);
              return ret);
    ret = CreateAclTensor(dWeightsHostData, dWeightsShape, &dWeightsAddr, aclDataType::ACL_FLOAT16, &dWeights);
    CHECK_RET(ret == ACL_SUCCESS,
              FreeResource(q, k, dy, sparseIndices, weights, dQuery, dKey, dWeights, qDeviceAddr, kDeviceAddr, dyDeviceAddr,
                sparseIndicesDeviceAddr, weightsDeviceAddr, dQueryAddr, dKeyAddr, dWeightsAddr, workspaceSize, workspaceAddr,
                  deviceId, &context, &stream);
              return ret);

    // 3. Call the CANN operator library API. Replace with the specific API name.
    aclOpExecutor *executor;

    // Call the first-phase API of aclnnLightningIndexerGrad.
    ret = aclnnLightningIndexerGradGetWorkspaceSize(
        q, k, dy, sparseIndices, weights, nullptr, nullptr, headNum,
        layoutStr, sparseMode, preToken, nextToken, deteminstic, dQuery, dKey, dWeights, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnLightningIndexerGradGetWorkspaceSize failed. ERROR: %d\n", ret);
              FreeResource(q, k, dy, sparseIndices, weights, dQuery, dKey, dWeights, qDeviceAddr, kDeviceAddr, dyDeviceAddr,
                sparseIndicesDeviceAddr, weightsDeviceAddr, dQueryAddr, dKeyAddr, dWeightsAddr, workspaceSize, workspaceAddr,
                  deviceId, &context, &stream);
              return ret);

    // Allocate device memory based on the workspaceSize calculated by the first-phase API.
    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);
            FreeResource(q, k, dy, sparseIndices, weights, dQuery, dKey, dWeights, qDeviceAddr, kDeviceAddr, dyDeviceAddr,
                sparseIndicesDeviceAddr, weightsDeviceAddr, dQueryAddr, dKeyAddr, dWeightsAddr, workspaceSize, workspaceAddr,
                  deviceId, &context, &stream);
            return ret);
    }

    // Call the second-phase API of aclnnLightningIndexerGrad.
    ret = aclnnLightningIndexerGrad(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnLightningIndexerGrad failed. ERROR: %d\n", ret);
              FreeResource(q, k, dy, sparseIndices, weights, dQuery, dKey, dWeights, qDeviceAddr, kDeviceAddr, dyDeviceAddr,
                sparseIndicesDeviceAddr, weightsDeviceAddr, dQueryAddr, dKeyAddr, dWeightsAddr, workspaceSize, workspaceAddr,
                  deviceId, &context, &stream);
              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);
              FreeResource(q, k, dy, sparseIndices, weights, dQuery, dKey, dWeights, qDeviceAddr, kDeviceAddr, dyDeviceAddr,
                sparseIndicesDeviceAddr, weightsDeviceAddr, dQueryAddr, dKeyAddr, dWeightsAddr, workspaceSize, workspaceAddr,
                  deviceId, &context, &stream);
              return ret);

    // 5. Obtain the output values and copy the results from the device-side memory to the host side. Modify based on the specific API definition.
    PrintOutResult<aclFloat16>(dQueryShape, &dQueryAddr);
    PrintOutResult<aclFloat16>(dKeyShape, &dKeyAddr);
    PrintOutResult<aclFloat16>(dWeightsShape, &dWeightsAddr);

    // 6. Release aclTensor and aclScalar. Modify based on the specific API definition; release device resources.
    FreeResource(q, k, dy, sparseIndices, weights, dQuery, dKey, dWeights, qDeviceAddr, kDeviceAddr, dyDeviceAddr,
                  sparseIndicesDeviceAddr, weightsDeviceAddr, dQueryAddr, dKeyAddr, dWeightsAddr, workspaceSize, workspaceAddr,
                  deviceId, &context, &stream);

    return 0;
}