aclnnIncreFlashAttention

Notice: This API will be deprecated in a later version. Use the latest API aclnnIncreFlashAttentionV4 instead.

Applicable Products

Function

• Description: For auto-regressive language models, as new tokens are generated, the inference input length keeps increasing. Incremental inference is implemented on top of the original full inference. The S axis of the query is fixed to 1. The key and value, after KV Cache, superimpose the previously inferred state information. The actual length of the S axis for each batch may differ, and the input data is fixed-length data after padding.

  Compared with the FlashAttention operator in the full-attention scenario (PromptFlashAttention), the incremental inference process is not completely equivalent to normal full-attention inference, but the precision of incremental inference does not exhibit significant degradation.

  NOTE

  KV Cache is a common technique for optimizing the inference performance of large models. During sampling, the Transformer model uses the given prompt/context as the initial input for inference (which can be processed in parallel), and then generates additional tokens one by one to further refine the generated sequence (reflecting the autoregressive nature of the model). During the sampling process, the Transformer performs self-attention operations, which requires extracting key-value (KV) vectors for each item in the current sequence (whether prompt/context or generated tokens). These vectors are stored in a matrix, commonly referred to as the KV cache.

• Formula:

  Self-attention leverages the relationships within the input sample itself to construct an attention model. The principle is as follows: assuming an input sample sequence $x$ of length $n$, where each element of $x$ is a $d$-dimensional vector, each $d$-dimensional vector can be regarded as a token embedding. This sequence is transformed through three weight matrices to obtain three matrices of dimension $n*d$.

  The calculation formula of self-attention is generally defined as follows, where $Q$, $K$, and $V$ are important attribute elements of the input sample, obtained through spatial transformation of the input sample, and can be unified into a single feature space.

  $$Attention(Q,K,V)=Score(Q,K)V$$

  In this operator, the Score function uses the Softmax function, and the self-attention formula is:

  $$Attention(Q,K,V)=Softmax(\frac{QK^T}{\sqrt{d}})V$$

  Where the product of $Q$ and $K^T$ represents the attention of input $x$. To prevent this value from becoming too large, it is typically divided by the square root of $d$ for scaling, and each row is normalized by softmax. After multiplying with $V$, an $n*d$ matrix is obtained.

Function Prototype

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

aclnnStatus aclnnIncreFlashAttentionGetWorkspaceSize(
    const aclTensor     *query,
    const aclTensorList *key,
    const aclTensorList *value,
    const aclTensor     *pseShift,
    const aclTensor     *attenMask,
    const aclIntArray   *actualSeqLengths,
    int64_t              numHeads,
    double               scaleValue,
    char                *inputLayout,
    int64_t              numKeyValueHeads,
    const aclTensor     *attentionOut,
    uint64_t            *workspaceSize,
    aclOpExecutor       **executor)

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

aclnnIncreFlashAttentionGetWorkspaceSize

• Parameters

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (shape)	Non-contiguous Tensor
  query	Input	Input Q in the formula.	The shapes of query and attentionOut must be exactly the same.	FLOAT16, BFLOAT16	ND	(B, N, S, D) (B, S, N, D) (B, S, H)	×
  key	Input	Input K in the formula.	The shapes of the corresponding tensors in key and value must be exactly the same.	FLOAT16, BFLOAT16	ND	(B, N, S, D) (B, S, N, D) (B, S, H)	×
  value	Input	Input V in the formula.	The shapes of the corresponding tensors in key and value must be exactly the same.	FLOAT16, BFLOAT16	ND	(B, N, S, D) (B, S, N, D) (B, S, H)	×
  pseShift	Input.	Positional encoding.	Reserved parameter, not used currently.	FLOAT16, BFLOAT16	ND	-	-
  attenMask	Input	Attention mask matrix.	Supports empty tensors. When the data type of attenMask is INT8 or UINT8, the values in its tensor must be 0 or 1.	BOOL, INT8, UINT8	ND	(B, N, 1, S) (1, N, 1, S)	×
  actualSeqLengths	Input	Actual length of the S axis of key and value.	For comprehensive constraints, see Constraints.	INT64	ND	(B)	-
  numHeads	Input	Number of heads for query.	numHeads is a multiple of numKeyValueHeads.	INT64	-	-	-
  scaleValue	Input	The reciprocal of the square root of d in the formula.	-	DOUBLE	-	-	-
  inputLayout	Input	Specifies the data layout format of the input query, key, and value.	Currently supports BSH, BNSD, and BSND. If not explicitly specified by the user, "BSH" is recommended.	STRING	-	-	-
  numKeyValueHeads	Input	Number of heads in key and value.	Used to support GQA (Grouped-Query Attention) scenarios. Passing 0 means the number of heads equals that of the query. For comprehensive constraints, see Constraints.	INT64	-	-	-
  attentionOut	Output	Output in the formula.	-	FLOAT16, BFLOAT16	ND	(B, N, S, D) (B, S, N, D) (B, S, H)	-
  workspaceSize	Output	Size of the workspace that the user needs to apply for 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 Codes.

  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	If the input parameter is a required input, output, or required attribute and is a null pointer, 161001 is returned.
  ACLNN_ERR_PARAM_INVALID	161002	The data type or data format of query, key, value, pseShift, attenMask, or attentionOut is not within the supported range.
  ACLNN_ERR_RUNTIME_ERROR	361001	API memory call to the NPU runtime API exception.

aclnnIncreFlashAttention

• Parameters

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

• Return Value

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

Constraints

• Deterministic computation:

  • aclnnIncreFlashAttention defaults to a deterministic implementation.

• Atlas A2 training products/Atlas A2 inference products, Ascend 950PR/Ascend 950DT:

  • Supports B-axis less than or equal to 65536, N-axis less than or equal to 256, and D-axis less than or equal to 512.

  • The query data type supports FLOAT16 and BFLOAT16, and the attentionOut, key, and value data types support FLOAT16 and BFLOAT16.

  • The numKeyValueHeads data type supports INT64.

• Atlas inference accelerator cards:

  • Supports B-axis less than or equal to 256, N-axis less than or equal to 256, and D-axis less than or equal to 512.

  • Supports key and value S-axis less than or equal to 65536.

  • query, key, value, and attentionOut: Only support FLOAT16 data type.

  • numKeyValueHeads only supports the value 0.

• In non-contiguous scenarios, the number of tensors in the tensorlist of the key and value parameters is equal to B of query (due to tensorlist limitations, B must be less than or equal to 256 in non-contiguous scenarios). The shape must be fully consistent except for S, and batch can only be 1.

• The N in the query parameter is equal to the numHeads value, the N of key and value is equal to the numKeyValueHeads value, and numHeads is a multiple of numKeyValueHeads.

• Only the S axis of query equal to 1 is supported.

• When the data type of attenMask is INT8 or UINT8, the values in its tensor must be 0 or 1.

Example

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

#include <iostream>
#include <vector>
#include <math.h>
#include <cstring>
#include "acl/acl.h"
#include "aclnn/opdev/fp16_t.h"
#include "aclnnop/aclnn_incre_flash_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: 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 allocate device 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 data to the device 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 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() {
  // 1. (Boilerplate) Initialize the device/stream. Refer to the ACL API manual.
  // Fill in the deviceId based on your 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. Custom construction is required based on the API.
  int32_t batchSize = 1;
  int32_t numHeads = 2;
  int32_t headDims = 16;
  int32_t keyNumHeads = 2;
  int32_t sequenceLengthKV = 16;
  std::vector<int64_t> queryShape = {batchSize, numHeads, 1, headDims}; // BNSD
  std::vector<int64_t> keyShape = {batchSize, keyNumHeads, sequenceLengthKV, headDims}; // BNSD
  std::vector<int64_t> valueShape = {batchSize, keyNumHeads, sequenceLengthKV, headDims}; // BNSD
  std::vector<int64_t> attenShape = {batchSize, 1, 1, sequenceLengthKV}; // B11S
  std::vector<int64_t> outShape = {batchSize, numHeads, 1, headDims}; // BNSD
  void *queryDeviceAddr = nullptr;
  void *keyDeviceAddr = nullptr;
  void *valueDeviceAddr = nullptr;
  void *attenDeviceAddr = nullptr;
  void *outDeviceAddr = nullptr;
  aclTensor *queryTensor = nullptr;
  aclTensor *keyTensor = nullptr;
  aclTensor *valueTensor = nullptr;
  aclTensor *attenTensor = nullptr;
  aclTensor *outTensor = nullptr;
  std::vector<float> queryHostData(batchSize * numHeads * headDims, 1.0f);
  std::vector<float> keyHostData(batchSize * keyNumHeads * sequenceLengthKV * headDims, 1.0f);
  std::vector<float> valueHostData(batchSize * keyNumHeads * sequenceLengthKV * headDims, 1.0f);
  std::vector<int8_t> attenHostData(batchSize * sequenceLengthKV, 0);
  std::vector<float> outHostData(batchSize * numHeads * headDims, 1.0f);

  // Create the query aclTensor.
  ret = CreateAclTensor(queryHostData, queryShape, &queryDeviceAddr, aclDataType::ACL_FLOAT16, &queryTensor);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the key aclTensor.
  ret = CreateAclTensor(keyHostData, keyShape, &keyDeviceAddr, aclDataType::ACL_FLOAT16, &keyTensor);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  int kvTensorNum = 1;
  aclTensor *tensorsOfKey[kvTensorNum];
  tensorsOfKey[0] = keyTensor;
  auto tensorKeyList = aclCreateTensorList(tensorsOfKey, kvTensorNum);
  // Create the value aclTensor.
  ret = CreateAclTensor(valueHostData, valueShape, &valueDeviceAddr, aclDataType::ACL_FLOAT16, &valueTensor);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  aclTensor *tensorsOfValue[kvTensorNum];
  tensorsOfValue[0] = valueTensor;
  auto tensorValueList = aclCreateTensorList(tensorsOfValue, kvTensorNum);
  // Create the atten aclTensor.
  ret = CreateAclTensor(attenHostData, attenShape, &attenDeviceAddr, aclDataType::ACL_INT8, &attenTensor);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the out aclTensor.
  ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT16, &outTensor);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  std::vector<int64_t> actualSeqlenVector = {sequenceLengthKV};
  auto actualSeqLengths = aclCreateIntArray(actualSeqlenVector.data(), actualSeqlenVector.size());

  int64_t numKeyValueHeads = numHeads;
  double scaleValue = 1 / sqrt(headDims); // 1/sqrt(d)
  string sLayerOut = "BNSD";
  char layerOut[sLayerOut.length()+1];
  strcpy(layerOut, sLayerOut.c_str());
  // 3. Call the CANN operator library API.
  uint64_t workspaceSize = 0;
  aclOpExecutor* executor;
  // Call the first-phase API.
  ret = aclnnIncreFlashAttentionGetWorkspaceSize(queryTensor, tensorKeyList, tensorValueList, nullptr, attenTensor, actualSeqLengths, numHeads, scaleValue, layerOut, numKeyValueHeads, outTensor,
                                                 &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnIncreFlashAttentionGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
  // Apply for 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.
  ret = aclnnIncreFlashAttention(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnIncreFlashAttention failed. ERROR: %d\n", ret); return ret);

  // 4. (Boilerplate) Synchronously wait for task execution to finish.
  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 based on the specific API definition.
  auto size = GetShapeSize(outShape);
  std::vector<op::fp16_t> resultData(size, 0);
  ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr,
                    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++) {
      std::cout << "index: " << i << ": " << static_cast<float>(resultData[i]) << std::endl;
  }

  // 6. Release resources.
  aclDestroyTensor(queryTensor);
  aclDestroyTensor(keyTensor);
  aclDestroyTensor(valueTensor);
  aclDestroyTensor(attenTensor);
  aclDestroyTensor(outTensor);
  aclDestroyIntArray(actualSeqLengths);
  aclrtFree(queryDeviceAddr);
  aclrtFree(keyDeviceAddr);
  aclrtFree(valueDeviceAddr);
  aclrtFree(attenDeviceAddr);
  aclrtFree(outDeviceAddr);
  if (workspaceSize > 0) {
      aclrtFree(workspaceAddr);
  }
  aclrtDestroyStream(stream);
  aclrtResetDevice(deviceId);
  aclFinalize();
  return 0;
}