aclnnFlashAttentionVarLenScoreV3

Applicable Products

Product	Supported or Not
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 Description

Uses the FlashAttention algorithm to perform self-attention computation in training scenarios. Compared with aclnnFlashAttentionVarLenScoreV2, this API supports multiple query/key inputs, that is, query, queryRope, key, and keyRope are used as inputs. In non-multi-input scenarios, use aclnnFlashAttentionVarLenScoreV2 or other APIs.
Formula

The forward computation formula for attention is as follows:
$attention\_out=Dropout(Softmax(Mask(scale*(query*key^T + queryRope*keyRope^T) + pse),atten\_mask),keep\_prob)*value$

Prototype

Each operator has two-phase API calls. First, aclnnFlashAttentionVarLenScoreV3GetWorkspaceSize is called to obtain the workspace size required for computation and the executor that contains the operator computation process. Then, aclnnFlashAttentionVarLenScoreV3 is called to perform computation.

aclnnStatus aclnnFlashAttentionVarLenScoreV3GetWorkspaceSize(
  const aclTensor   *query,
  const aclTensor   *queryRope,
  const aclTensor   *key,
  const aclTensor   *keyRope,
  const aclTensor   *value,
  const aclTensor   *realShiftOptional,
  const aclTensor   *dropMaskOptional,
  const aclTensor   *paddingMaskOptional,
  const aclTensor   *attenMaskOptional,
  const aclIntArray *prefixOptional,
  const aclIntArray *actualSeqQLenOptional,
  const aclIntArray *actualSeqKvLenOptional,
  const aclIntArray *qStartIdxOptional,
  const aclIntArray *kvStartIdxOptional,
  double             scaleValue,
  double             keepProb,
  int64_t            preTokens,
  int64_t            nextTokens,
  int64_t            headNum,
  char              *inputLayout,
  int64_t            innerPrecise,
  int64_t            sparseMode,
  int64_t            pseType,
  const aclTensor   *softmaxMaxOut,
  const aclTensor   *softmaxSumOut,
  const aclTensor   *softmaxOutOut,
  const aclTensor   *attentionOutOut,
  uint64_t          *workspaceSize,
  aclOpExecutor    **executor)

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

aclnnFlashAttentionVarLenScoreV3GetWorkspaceSize

Parameters

Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
query	Input	query in the formula.	The data type must be the same as that of key and value.	BFLOAT16	ND	[TND]	√
queryRope	Input	queryRope in the formula.	The data type must be the same as that of key and value.	BFLOAT16	ND	[TND]	√
key	Input	key in the formula.	The data type must be the same as that of query and value.	BFLOAT16	ND	[TND]	√
keyRope	Input	keyRope in the formula.	The data type must be the same as that of query and value.	BFLOAT16	ND	[TND]	√
value	Input	value in the formula.	The data type must be the same as that of query and key.	BFLOAT16	ND	[TND]	√
realShiftOptional	Input	pse in the formula.	Must be nullptr.	-	-	-	-
dropMaskOptional	Input	Dropout in the formula.	Must be nullptr.	-	-	-	-
attenMaskOptional	Input	atten_mask in the formula.	A value of 1 indicates the bit is excluded from computation, while 0 indicates it is included.	BOOL, UINT8	ND	[B,N,Sq,Skv], [B,1,Sq,Skv], [1,1,Sq,Skv], [Sq,Skv]	√
prefixOptional	Input	N of each batch in the prefix sparse computation scenario.	-	INT64	ND	0, 1	-
actualSeqQLenOptional	Input	Sequence length of query corresponding to each batch.	-	INT64	ND	0, 1	-
actualSeqKvLenOptional	Input	Sequence length of key/value corresponding to each batch.	-	INT64	ND	0, 1	-
qStartIdxOptional	Input	Global start index of the query sequence for the current chunk in an outer splitting scenario.	-	INT64	ND	0, 1	-
kvStartIdxOptional	Input	Global start index of the query sequence for the current chunk in an outer splitting scenario.	-	INT64	ND	0, 1	-
scaleValue	Input	scale in the formula, indicating the scale factor.	-	DOUBLE	-	-	-
keepProb	Input	Proportion of 1s in dropMaskOptional.	-	DOUBLE	-	-	-
preTokens	Input	Left boundary of the sliding window, used for sparse computation.	-	INT64	-	-	-
nextTokens	Input	Right boundary of the sliding window, used for sparse computation.	-	INT64	-	-	-
headNum	Input	Number of heads on a single device, that is, the length of the N axis of the input query.	-	INT64	-	-	-
inputLayout	Input	Layout of the input query, key, and value.	TND is supported.	String	-	-	-
innerPrecise	Input	Used to improve precision.	Reserved.	INT64	-	-	-
sparseMode	Input	Sparse mode.	The value can be 0, 1, 2, 3, 4, 7, or 8.	INT64	-	-	-
pseType	Input	Calculation sequence of multiplication and addition. The value can only be 1.	-	INT64	-	-	-
softmaxMaxOut	Output	Intermediate result of the Max operation in softmax, used for backward computation.	-	FLOAT	ND	[N,T,8]	√
softmaxSumOut	Output	Intermediate result of the Sum operation in softmax, used for backward computation.	-	FLOAT	ND	[N,T,8]	√
attentionOutOut	Output	Final output of the formula.	The data type and shape must be the same as those of query.	BFLOAT16	ND	[TND]	√
workspaceSize	Output	Size of the workspace required to be allocated on the device.	-	-	-	-	-
executor	Output	Operator executor, containing the operator computation process.	-	-	-	-	-

Returns

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

The first-phase API implements input parameter verification. The following errors may be thrown.

Return	Error Code	Description
ACLNN_ERR_PARAM_NULLPTR	161001	The required input, output, or attribute is passed as a null pointer.
ACLNN_ERR_PARAM_INVALID	161002	The data type of query, queryRope, key, keyRope, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, softmaxMaxOut, softmaxSumOut, softmaxOutOut, or attentionOutOut is not supported.
ACLNN_ERR_PARAM_INVALID	161002	The data format of query, queryRope, key, keyRope, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, softmaxMaxOut, softmaxSumOut, softmaxOutOut, or attentionOutOut is not supported.

aclnnFlashAttentionVarLenScoreV3

Parameters

Parameter	Input/Output	Description
workspace	Input	Address of the workspace to be allocated on the device.
workspaceSize	Input	Size of the workspace to be allocated on the device, which is obtained by calling the first-phase API aclnnFlashAttentionVarLenScoreV3GetWorkspaceSize.
executor	Input	Operator executor, containing the operator computation process.
stream	Input	Stream for executing the task.

Returns

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

Constraints

Deterministic computation
- aclnnFlashAttentionVarLenScoreV3 defaults to a deterministic implementation.
When this API is used together with PyTorch, ensure that the CANN package versions match the PyTorch package versions.
B (batch size) of the input query, queryRope, key, keyRope, and value must be the same.
D (Head-Dim) of the input query, key, and value must satisfy (qD == kD && kD >= vD). D must be an integer multiple of 8.
D (Head-Dim) of the input queryRope and keyRope must satisfy (qRopeD == kRopeD). D must be an integer multiple of 8 and less than or equal to the D of query, key, and value.
The inputLayout of the input query, key, and value must be TND.
Constraints on the data shape:
- T: The value ranges from 1 to 1M.
- N: The value ranges from 1 to 256.
- D: The value ranges from 1 to 768.
- The data shape must be TND.
- KeepProb must be 1.
The data format of query, key, and value can only be TND. T indicates the data closely arranged on the B and S axes (SeqLenQ and SeqLenKV of each batch). B (Batch) indicates the batch size of the input sample, and S (Seq-Length) indicates the length of the input sample sequence. H (Head-Size) indicates the size of the hidden layer. N (Head-Num) indicates the number of heads. D (Head-Dim) indicates the minimum unit size of the hidden layer (D = H/N).
Restrictions on sparseMode
- When the shapes of all attenMaskOptional are less than 2048 and are the same, the default mode is recommended to reduce memory usage.
- When the value is set to 1, 2, 3, or 5, the user-configured preTokens and nextTokens do not take effect.
- When the value is set to 0 or 4, ensure that the ranges of attenMaskOptional, preTokens, and nextTokens are consistent.
- If no specific value is required, 0 is recommended.
- For details about the sparse modes, see Sparse Mode Description.
- When the value is set to 3, computation on invalid rows is not supported, and Sq <= Skv must be satisfied for each batch.
- When the value is set to 7, realShiftOptional is not supported.
- When the value is set to 8, realShiftOptional is supported when the q and kv of each sequence have the same length. PSE generation is performed globally. Outer splitting in the q direction is supported. q and kv of each sequence must have the same length before outer splitting, and actualSeqQLenOptional is passed after outer splitting.
In some scenarios, if the computation load is too large, the operator execution may time out (an AI Core error is reported, and errorStr is timeout or trap error). In this case, you are advised to perform axis splitting. Note: The computation load is affected by parameters such as B, S, N, and D. Larger values indicate larger computation loads.
In the band scenario, the values of preTokens and nextTokens must overlap.
The prefixOptional sparse computing scenario is sparseMode=6. When Sq > Skv, the value range of N of prefix is [0, Skv]. When Sq ≤ Skv, the value range of N of prefix is [Skv – Sq, Skv]. [0] - actualSeqKvLenOptional[0] + qStartIdxOptional - kvStartIdxOptional == 0 (experimental feature)
The actualSeqQLenOptional input supports the S length of 0 in a batch. In this case, the realShiftOptional input is not supported. The length of actualSeqQLenOptional ranges from 1 to 2K. When prefixOptional is present, its maximum length is 1K.
The attenMaskOptional input does not support padding. That is, attenMaskOptional cannot contain a row of all 1s.
The S length of a batch in actualSeqQLenOptional can be 0. If the S length is 0, the pse input is not supported. If the actual S length is [2,2,0,2,2], the value of actualSeqQLenOptional is [2,4,4,6,8].
pseType can only be 0 or 1.
realShiftOptional must be null.
dropMaskOptional must be null.
attenMaskOptional cannot be null.

Example

The following example is for reference only. For details, see Compilation and Running Sample.

#include <iostream>
#include <vector>
#include "acl/acl.h"
#include "aclnnop/aclnn_flash_attention_score.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;
}

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

int Init(int32_t deviceId, aclrtStream* stream) {
  // (Boilerplate) Initialize resources.
  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 memory on the device.
  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 data from the host to the device.
  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 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 aclCreateTensor 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 and stream. For details, see the acl API manual.
  // Set the device ID in use.
  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 input and output based on the API.
  std::vector<int64_t> qShape = {256, 1, 128};
  std::vector<int64_t> qRopeShape = {256, 1, 64};
  std::vector<int64_t> kShape = {256, 1, 128};
  std::vector<int64_t> kRopeShape = {256, 1, 64};
  std::vector<int64_t> vShape = {256, 1, 128};
  std::vector<int64_t> attenmaskShape = {256, 256};

  std::vector<int64_t> attentionOutShape = {256, 1, 128};
  std::vector<int64_t> softmaxMaxShape = {256, 1, 8};
  std::vector<int64_t> softmaxSumShape = {256, 1, 8};

  void* qDeviceAddr = nullptr;
  void* qRopeDeviceAddr = nullptr;
  void* kDeviceAddr = nullptr;
  void* kRopeDeviceAddr = nullptr;
  void* vDeviceAddr = nullptr;
  void* attenmaskDeviceAddr = nullptr;
  void* attentionOutDeviceAddr = nullptr;
  void* softmaxMaxDeviceAddr = nullptr;
  void* softmaxSumDeviceAddr = nullptr;

  aclTensor* q = nullptr;
  aclTensor* qRope = nullptr;
  aclTensor* k = nullptr;
  aclTensor* kRope = nullptr;
  aclTensor* v = nullptr;
  aclTensor* pse = nullptr;
  aclTensor* dropMask = nullptr;
  aclTensor* padding = nullptr;
  aclTensor* attenmask = nullptr;
  aclTensor* attentionOut = nullptr;
  aclTensor* softmaxMax = nullptr;
  aclTensor* softmaxSum = nullptr;
  aclTensor* softmaxOut = nullptr;

  std::vector<float> qHostData(32768, 1);
  std::vector<float> qRopeHostData(16384, 1);
  std::vector<float> kHostData(32768, 1);
  std::vector<float> kRopeHostData(16384, 1);
  std::vector<float> vHostData(32768, 1);
  std::vector<uint8_t> attenmaskHostData(65536, 0);
  std::vector<float> attentionOutHostData(32768, 0);
  std::vector<float> softmaxMaxHostData(2048, 3.0);
  std::vector<float> softmaxSumHostData(2048, 3.0);

  ret = CreateAclTensor(qHostData, qShape, &qDeviceAddr, aclDataType::ACL_BF16, &q);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(qRopeHostData, qRopeShape, &qRopeDeviceAddr, aclDataType::ACL_BF16, &qRope);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(kHostData, kShape, &kDeviceAddr, aclDataType::ACL_BF16, &k);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(kRopeHostData, kRopeShape, &kRopeDeviceAddr, aclDataType::ACL_BF16, &kRope);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(vHostData, vShape, &vDeviceAddr, aclDataType::ACL_BF16, &v);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(attenmaskHostData, attenmaskShape, &attenmaskDeviceAddr, aclDataType::ACL_UINT8, &attenmask);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(attentionOutHostData, attentionOutShape, &attentionOutDeviceAddr, aclDataType::ACL_BF16, &attentionOut);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(softmaxMaxHostData, softmaxMaxShape, &softmaxMaxDeviceAddr, aclDataType::ACL_FLOAT, &softmaxMax);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(softmaxSumHostData, softmaxSumShape, &softmaxSumDeviceAddr, aclDataType::ACL_FLOAT, &softmaxSum);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  std::vector<int64_t> prefixOp = {0};
  aclIntArray *prefix = aclCreateIntArray(prefixOp.data(), 1);
  std::vector<int64_t> qStartIdxOp = {0};
  std::vector<int64_t> kvStartIdxOp = {0};
  aclIntArray *qStartIdx = aclCreateIntArray(qStartIdxOp.data(), 1);
  aclIntArray *kvStartIdx = aclCreateIntArray(kvStartIdxOp.data(), 1);
  std::vector<int64_t>  acSeqQLenOp = {256};
  std::vector<int64_t>  acSeqKvLenOp = {256};
  aclIntArray* acSeqQLen = aclCreateIntArray(acSeqQLenOp.data(), acSeqQLenOp.size());
  aclIntArray* acSeqKvLen = aclCreateIntArray(acSeqKvLenOp.data(), acSeqKvLenOp.size());
  double scaleValue = 0.088388;
  double keepProb = 1;
  int64_t preTokens = 65536;
  int64_t nextTokens = 65536;
  int64_t headNum = 1;
  int64_t innerPrecise = 0;
  int64_t sparseMode = 0;
  int64_t pseType = 1;

  char layOut[5] = {'T', 'N', 'D', 0};

  // 3. Call the CANN operator library API, which needs to be replaced with the actual API.
  uint64_t workspaceSize = 0;
  aclOpExecutor* executor;

  // Call the first-phase API aclnnFlashAttentionVarLenScoreV3.
  ret = aclnnFlashAttentionVarLenScoreV3GetWorkspaceSize(
            q, qRope, k, kRope, v, pse, dropMask, padding, attenmask, prefix, acSeqQLen, acSeqKvLen, qStartIdx, kvStartIdx,
            scaleValue, keepProb, preTokens, nextTokens, headNum, layOut, innerPrecise,
            sparseMode, pseType, softmaxMax, softmaxSum, softmaxOut, attentionOut, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFlashAttentionVarLenScoreV3GetWorkspaceSize failed. ERROR: %d\n", ret);
            return ret);

  // Allocate device memory based on the computed workspaceSize.
  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 aclnnFlashAttentionVarLenScoreV3.
  ret = aclnnFlashAttentionVarLenScoreV3(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFlashAttentionVarLenScoreV3 failed. ERROR: %d\n", ret); return ret);

  // 4. (Boilerplate) Wait until the task execution is 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 to the host. Modify the code based on the API definition.
  PrintOutResult(attentionOutShape, &attentionOutDeviceAddr);
  PrintOutResult(softmaxMaxShape, &softmaxMaxDeviceAddr);
  PrintOutResult(softmaxSumShape, &softmaxSumDeviceAddr);

  // 6. Release aclTensor and aclScalar. Modify the configuration based on the API definition.
  aclDestroyTensor(q);
  aclDestroyTensor(qRope);
  aclDestroyTensor(k);
  aclDestroyTensor(kRope);
  aclDestroyTensor(v);
  aclDestroyTensor(attenmask);
  aclDestroyTensor(attentionOut);
  aclDestroyTensor(softmaxMax);
  aclDestroyTensor(softmaxSum);

  // 7. Release device resources.
  aclrtFree(qDeviceAddr);
  aclrtFree(qRopeDeviceAddr);
  aclrtFree(kDeviceAddr);
  aclrtFree(kRopeDeviceAddr);
  aclrtFree(vDeviceAddr);
  aclrtFree(attenmaskDeviceAddr);
  aclrtFree(attentionOutDeviceAddr);
  aclrtFree(softmaxMaxDeviceAddr);
  aclrtFree(softmaxSumDeviceAddr);
  if (workspaceSize > 0) {
    aclrtFree(workspaceAddr);
  }
  aclrtDestroyStream(stream);
  aclrtResetDevice(deviceId);
  aclFinalize();

  return 0;
}