aclnnIncreFlashAttentionV2

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

Applicable Products

Function

• Description: Compatible with the aclnnIncreFlashAttention API, with the addition of the quantization feature.

  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 "aclnnIncreFlashAttentionV2GetWorkspaceSize" API to obtain the workspace size required for computation and the executor that encapsulates the operator computation process, and then call the "aclnnIncreFlashAttentionV2" API to perform the computation.

aclnnStatus aclnnIncreFlashAttentionV2GetWorkspaceSize(
    const aclTensor     *query,
    const aclTensorList *key,
    const aclTensorList *value,
    const aclTensor     *pseShift,
    const aclTensor     *attenMask,
    const aclIntArray   *actualSeqLengths,
    const aclTensor     *dequantScale1,
    const aclTensor     *quantScale1,
    const aclTensor     *dequantScale2,
    const aclTensor     *quantScale2,
    const aclTensor     *quantOffset2,
    int64_t              numHeads,
    double               scaleValue,
    char                *inputLayout,
    int64_t              numKeyValueHeads,
    const aclTensor     *attentionOut,
    uint64_t            *workspaceSize,
    aclOpExecutor       **executor)

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

aclnnIncreFlashAttentionV2GetWorkspaceSize

• 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, currently unused.	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	The actual length of the S axis of key and value.	For comprehensive constraints, see Constraints.	INT64	ND	(B)	-
  dequantScale1	Input	Quantization factor for dequantization after BMM1.	Supports empty tensor. Supports per-tensor (scalar).	UINT64, FLOAT32	ND	(1)	×
  quantScale1	Input	Quantization scale factor for the quantization before BMM2.	Supports empty tensors. Supports per-tensor (scalar).	FLOAT32, BFLOAT16	ND	-	×
  dequantScale2	Input	Quantization factor for dequantization after BMM2.	Supports empty tensor. Supports per-tensor (scalar).	UINT64, FLOAT32	ND	(1)	×
  quantScale2	Input	Quantization scale factor for output quantization.	Supports empty tensors. Supports per-tensor and per-channel.	FLOAT32, BFLOAT16	ND	(1) (D * N), where N is numKeyValueHeads.	×
  quantOffset2	Input	Quantization offset for output quantization.	Supports empty tensor. Supports per-tensor and per-channel.	FLOAT32, BFLOAT16	ND	(1) (D * N), where N is numKeyValueHeads	×
  numHeads	Input	Number of heads in 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. "BSH" is recommended when the user does not specify a format.	STRING	-	-	-
  numKeyValueHeads	Input	Number of heads in key and value.	Used to support GQA (Grouped-Query Attention) scenarios. Passing 0 indicates that the number of heads equals that of the query. For comprehensive constraints, see Constraints.	INT64	-	-	-
  attentionOut	Output	Output in the formula.	-	FLOAT16, BFLOAT16, INT8	ND	(B, N, S, D) (B, S, N, D) (B, S, H)	-
  workspaceSize	Output	Returns the workspace size that the user needs to apply for on the Device side.	-	-	-	-	-
  executor	Output	The returned 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	Returns 161001 if the input parameter is a mandatory input, output, or mandatory attribute and is a null pointer.
  ACLNN_ERR_PARAM_INVALID	161002	The data type or data format of query, key, value, actualSeqLengths, pseShift, attenMask, or attentionOut is not within the supported range.
  ACLNN_ERR_RUNTIME_ERROR	361001	An API memory call to the npu runtime interface is abnormal.

aclnnIncreFlashAttentionV2

• Parameters

  Parameter	Input/Output	Description
  workspace	Input	Memory address of the workspace applied for on the Device side.
  workspaceSize	Input	Size of the workspace applied for on the Device side, obtained through the first-phase API aclnnIncreFlashAttentionV2GetWorkspaceSize.
  executor	Input	Operator executor that contains the operator computation process.
  stream	Input	Specifies the stream for task execution.

• Return Value

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

Constraints

• Deterministic computation:

  • aclnnIncreFlashAttentionV2 defaults to a deterministic implementation.

• Atlas A2 training products/Atlas A2 inference products:

  • 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 dequantScale1 and dequantScale2 data types support UINT64 and FLOAT32.

  • The quantScale1, quantScale2, and quantOffset2 data types support FLOAT32.

  • The numKeyValueHeads data type supports INT64.

• Atlas inference accelerator cards:

  • Support 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.

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

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

  • dequantScale1, quantScale1, dequantScale2, quantScale2, and quantOffset2 only support nullptr.

  • 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 support query with S axis equal to 1.

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

• The scenario where query, key, and value inputs are all INT8 is not supported.

• In the scenario where query, key, and value inputs are FLOAT16 and the output is INT8: the input parameter quantScale2 is required, quantOffset2 is optional, and the dequantScale1, quantScale1, and dequantScale2 parameters cannot be passed (i.e., they must be nullptr).

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_v2.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 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);

  // 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. See the acl API manual.
  // Specify 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. Customize the construction 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 query aclTensor.
  ret = CreateAclTensor(queryHostData, queryShape, &queryDeviceAddr, aclDataType::ACL_FLOAT16, &queryTensor);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create 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 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 atten aclTensor.
  ret = CreateAclTensor(attenHostData, attenShape, &attenDeviceAddr, aclDataType::ACL_INT8, &attenTensor);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create 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 = aclnnIncreFlashAttentionV2GetWorkspaceSize(queryTensor, tensorKeyList, tensorValueList, nullptr, attenTensor, actualSeqLengths, nullptr, nullptr, nullptr, nullptr, nullptr,
                                                   numHeads, scaleValue, layerOut, numKeyValueHeads, outTensor, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnIncreFlashAttentionV2GetWorkspaceSize 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.
  ret = aclnnIncreFlashAttentionV2(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnIncreFlashAttentionV2 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 this 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;
}