aclnnFlashAttentionScoreV2

Supported Products

Description

• API function: Uses the FlashAttention algorithm to perform self-attention computation in training scenarios. Compared with the aclnnFlashAttentionScore API, this API has a psetype parameter.

  • When psetype is set to 1, the implementation is the same as that of aclnnFlashAttentionScore.
  • When psetype is set to other values, multiplication is required before addition.

• Formula:

  The forward propagation formula for attention is as follows:

  • When psetype is set to 1, the calculation formula is the same as that of aclnnFlashAttentionScore.

  • When psetype is set to other values, the formula is as follows:

    $$attention\_out=Dropout(Softmax(Mask(scale*(query*key^T) + pse),atten\_mask),keep\_prob)*value$$

Prototype

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

aclnnStatus aclnnFlashAttentionScoreV2GetWorkspaceSize(
  const aclTensor   *query,
  const aclTensor   *key,
  const aclTensor   *value,
  const aclTensor   *realShiftOptional,
  const aclTensor   *dropMaskOptional,
  const aclTensor   *paddingMaskOptional,
  const aclTensor   *attenMaskOptional,
  const aclIntArray *prefixOptional,
  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 aclnnFlashAttentionScoreV2(
  void             *workspace, 
  uint64_t          workspaceSize, 
  aclOpExecutor    *executor, 
  const aclrtStream stream)

aclnnFlashAttentionScoreV2GetWorkspaceSize

• Parameters

  Name	Input/Output	Description	Usage Notes	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.	FLOAT16, BFLOAT16, FLOAT32	ND	[BNSD], [BSND], [BSH], [SBH]	√
  key	Input	key in the formula.	The data type must be the same as that of query and value.	FLOAT16, BFLOAT16, FLOAT32	ND	[BNSD], [BSND], [BSH], [SBH]	√
  value	Input	value in the formula.	The data type must be the same as that of query and key.	FLOAT16, BFLOAT16, FLOAT32	ND	[BNSD], [BSND], [BSH], [SBH]	√
  realShiftOptional	Optional input	pse in the formula.	The data type must match query. Use this parameter with pseType.	FLOAT16, BFLOAT16, FLOAT32	ND	[B,N,Sq,Skv], [B,N,1,Skv], [1,N,Sq,Skv], [B,N,1024,Skv], [1,N,1024,Skv], [B,N], [N]	√
  dropMaskOptional	Input	Dropout in the formula.	-	UINT8	ND	0, 1	√
  attenMaskOptional	Input	atten_mask in the formula.	A value of 1 indicates that the position does not participate in the calculation, while a value of 0 indicates that it does.	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	-
  scaleValue	Input	scale in the formula, indicating the scale factor.	-	DOUBLE	-	-	-
  keepProb	Input	Proportion of 1s in dropMaskOptional.	The value range is (0, 1].	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 rank, that is, the length of the N axis of the input query.	-	INT64	-	-	-
  inputLayout	Input	Layout of the input query, key, and value.	BSH, SBH, BSND, and BNSD are supported.	String	-	-	-
  innerPrecise	Input	Used to improve precision.	-	INT64	-	-	-
  sparseMode	Input	Sparse mode.	The value can be 0, 1, 2, 3, 4, 5, or 6.	INT64	-	-	-
  pseType	Input	Calculation sequence of multiplication and addition. The value can be 0, 1, 2, or 3.	-	INT64	-	-	-
  softmaxMaxOut	Output	Intermediate result of the Max operation in Softmax, used for backward calculation.	-	FLOAT	ND	[B,N,Sq,8]	√
  softmaxSumOut	Output	Intermediate result of the Sum operation in Softmax, used for backward calculation.	-	FLOAT	ND	[B,N,Sq,8]	√
  attentionOutOut	Output	Final output of the formula.	The data type and shape must be the same as those of query.	FLOAT16, BFLOAT16, FLOAT32	ND	[BNSD], [BSND], [BSH], [SBH]	√
  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, key, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, softmaxMaxOut, softmaxSumOut, softmaxOutOut, or attentionOutOut is not supported.
  		The data format of query, key, value, realShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, softmaxMaxOut, softmaxSumOut, softmaxOutOut, or attentionOutOut is not supported.

aclnnFlashAttentionScoreV2

• Parameters

  Name	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 aclnnFlashAttentionScoreV2GetWorkspaceSize.
  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 computing:

  • aclnnFlashAttentionScoreV2 defaults to a deterministic implementation.

• When this API is used together with PyTorch, ensure that the CANN package versions match the PyTorch package versions.

• The constraints for input query, key, and value are as follows:

  • B: The batch sizes must be equal.
  • D: Head-Dim must satisfy (qD == kD && kD >= vD).
  • inputLayout must be consistent.

• The shapes of input key and value must be the same except D.

• You need to pay attention to constraints of the data shape. The following takes the inputLayout values BSND and BNSD as examples to describe the constraints (H = N*D in BSH and SBH):

  • B: The value ranges from 1 to 2M. When prefixOptional is passed, B supports a maximum of 2K.
  • N: The value ranges from 1 to 256.
  • S: The value ranges from 1 to 1M.
  • D: The value ranges from 1 to 768.

• The data format of query, key, and value can be interpreted from multiple dimensions. To be specific, B (Batch) indicates the size of an input sample batch, 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, and D (Head-Dim) indicates the minimum unit size of the hidden layer (D = H/N).

• realShiftOptional: If Sq is greater than 1024, Sq and Skv of each batch are of equal length, and it is a lower triangular mask scenario with sparseMode being 0, 2, or 3, ALiBi positional encoding compression can be enabled. In this case, only the last 1024 rows of the original PSE need to be input for memory optimization, that is, alibi_compress = ori_pse[:, :, -1024:, :]. Specifically:

  • If the parameters of each batch are different, the shape is BNHSkv (H=1024).
  • When each batch is the same, the shape is 1NHSkv (H=1024).
  • If pseType is 2 or 3, the data type must be FLOAT32, and the supported shapes are [B,N] and [N].
  • If this parameter is not enabled, pass a null pointer to realShiftOptional and 1 to pseType.

• innerPrecise: 0 and 1 are reserved, and 2 indicates that invalid row calculation is enabled. This function is used to prevent precision loss caused by the mask of the entire row during calculation. However, this configuration deteriorates the performance. If the operator can determine that invalid rows exist, the invalid row computation is automatically enabled, such as in scenarios where sparseMode is set to 3 and Sq is greater than Skv.

• Meanings of pseType values:

  pseType	Meaning	Remarks
  0	A value is externally passed to pse, and multiplication is required before addition.	-
  1	A value is externally passed to pse, and addition is required before multiplication.	The implementation is the same as that of FlashAttentionScore.
  2	A value is internally passed to pse, and multiplication is required before addition.	-
  3	A value is internally passed to pse, and multiplication and addition are required before square root operation.	-

• When pseType is set to 2 or 3, Sq and Skv must be of the same length.

• The constraints for sparseMode are as follows:

  • If the shape values of all attenMaskOptional are the same and less than 2048, you are advised to use the default mode 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, you are advised to set it to 0.
  • For details about the sparse modes, see Sparse Mode Description.

• In some scenarios, if the computation load is too large, the operator execution may time out (AI Core error, errorStr: 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=5 or 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].

• If Sq of realShiftOptional is greater than 1024, if BNHS and 1NHS are configured, Sq and Skv must have the same length.

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) {
  // (Fixed writing) 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 the data on the host to the memory on 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. (Fixed writing) 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 inputs and outputs based on the API definition.
  std::vector<int64_t> qShape = {256, 1, 128};
  std::vector<int64_t> kShape = {256, 1, 128};
  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 = {1, 1, 256, 8};
  std::vector<int64_t> softmaxSumShape = {1, 1, 256, 8};

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

  aclTensor* q = nullptr;
  aclTensor* k = 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> kHostData(32768, 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_FLOAT16, &q);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(kHostData, kShape, &kDeviceAddr, aclDataType::ACL_FLOAT16, &k);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(vHostData, vShape, &vDeviceAddr, aclDataType::ACL_FLOAT16, &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_FLOAT16, &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};
  std::vector<int64_t> qStartIdxOp = {0};
  std::vector<int64_t> kvStartIdxOp = {0};
  aclIntArray *prefix = aclCreateIntArray(prefixOp.data(), 1);
  aclIntArray *qStartIdx = aclCreateIntArray(qStartIdxOp.data(), 1);
  aclIntArray *kvStartIdx = aclCreateIntArray(kvStartIdxOp.data(), 1);
  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] = {'S', 'B', 'H', 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 of aclnnFlashAttentionScoreV2.
  ret = aclnnFlashAttentionScoreV2GetWorkspaceSize(
            q, k, v, pse, dropMask, padding, attenmask, prefix, qStartIdx, kvStartIdx, scaleValue,
            keepProb, preTokens, nextTokens, headNum, layOut, innerPrecise,
            sparseMode, pseType, softmaxMax, softmaxSum, softmaxOut, attentionOut, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFlashAttentionScoreV2GetWorkspaceSize 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 of aclnnFlashAttentionScoreV2.
  ret = aclnnFlashAttentionScoreV2(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFlashAttentionScoreV2 failed. ERROR: %d\n", ret); return ret);
  
  // 4. (Fixed writing) 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(k);
  aclDestroyTensor(v);
  aclDestroyTensor(attenmask);
  aclDestroyTensor(attentionOut);
  aclDestroyTensor(softmaxMax);
  aclDestroyTensor(softmaxSum);
  
  // 7. Release device resources.
  aclrtFree(qDeviceAddr);
  aclrtFree(kDeviceAddr);
  aclrtFree(vDeviceAddr);
  aclrtFree(attenmaskDeviceAddr);
  aclrtFree(attentionOutDeviceAddr);
  aclrtFree(softmaxMaxDeviceAddr);
  aclrtFree(softmaxSumDeviceAddr);
  if (workspaceSize > 0) {
    aclrtFree(workspaceAddr);
  }
  aclrtDestroyStream(stream);
  aclrtResetDevice(deviceId);
  aclFinalize();
  
  return 0;
}