aclnnScatterPaCache

📄 View Source Code

Applicable Products

Function

• Operator function: Updates the key at the specified position in KCache.

• Formula:

  • Scenario 1:

    key:[batch * seq_len, num_head, k_head_size]
    keyCache:[num_blocks, block_size, num_head, k_head_size]
    slotMapping:[batch * seq_len]
    cacheMode:"Norm"

    $$keyCache = slotMapping(key)$$

  • Scenario 2:

    key:[batch, seq_len, num_head, k_head_size]
    keyCache:[num_blocks, block_size, 1, k_head_size]
    slotMapping:[batch, num_head]
    compressLensOptional:[batch, num_head]
    compressSeqOffsetOptional:[batch * num_head]
    seqLensOptional:[batch]
    cacheMode:"Norm"

    $$\begin{aligned} keyCache =\ & slotMapping(key[: compressSeqOffset], \\ & ReduceMean(key[compressSeqOffset : compressSeqOffset + compressLens]), \\ & key[compressSeqOffset + compressLens : seqLens]) \end{aligned}$$

  • Scenario 3:

    key:[batch, seq_len, num_head, k_head_size]
    keyCache:[num_blocks, block_size, 1, k_head_size]
    slotMapping:[batch, num_head]
    compressLensOptional:[batch * num_head]
    seqLensOptional:[batch]
    cacheMode:"Norm"

    $$keyCache = slotMapping(key[seqLens - compressLens : seqLens])$$

The above scenarios are distinguished based on the constructed parameters. The first input parameter construction corresponds to Scenario 1, the second construction corresponds to Scenario 2, and the third construction corresponds to Scenario 3. Scenario 1 does not include the three optional inputs: compressLensOptional, seqLensOptional, and compressSeqOffsetOptional. Scenario 3 does not include the compressSeqOffsetOptional optional input.

Function Prototype

Each operator is divided into a two-phase API. The "aclnnScatterPaCacheGetWorkspaceSize" API must be called first to obtain the workspace size required for computation and the executor that encapsulates the operator computation flow. Then, the "aclnnScatterPaCache" API is called to perform the computation.

aclnnStatus aclnnScatterPaCacheGetWorkspaceSize(
  const aclTensor *key,
  aclTensor       *keyCacheRef,
  const aclTensor *slotMapping,
  const aclTensor *compressLensOptional,
  const aclTensor *compressSeqOffsetOptional,
  const aclTensor *seqLensOptional,
  char            *cacheMode,
  uint64_t        *workspaceSize,
  aclOpExecutor  **executor)

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

aclnnScatterPaCacheGetWorkspaceSize

• Parameters

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
  key (aclTensor*)	Input	Key value to be updated, which is the key in the formula.	Nullable. HIFLOAT8, FLOAT8_E5M2, FLOAT8_E4M3FN, FLOAT4_E2M1, and FLOAT4_E1M2 are supported only when key is 3-dimensional. The shape satisfies [batch * seq_len, num_head, k_head_size] or [batch, seq_len, num_head, k_head_size]. For FLOAT4_E2M1 and FLOAT4_E1M2, k_head_size must be an even number.	FLOAT16, FLOAT, BFLOAT16, INT8, UINT8, INT16, UINT16, INT32, UINT32, HIFLOAT8, FLOAT8_E5M2, FLOAT8_E4M3FN, FLOAT4_E2M1, FLOAT4_E1M2	ND	3-4	x
  keyCacheRef(aclTensor*)	Input/Output	The keyCache to be updated, i.e., keyCache in the formula.	Nullable. When key is 3D, the shape satisfies [num_blocks, block_size, num_head, k_head_size]; when key is 4D, the shape satisfies [num_blocks, block_size, 1, k_head_size].	Consistent with key.	ND	4	x
  slotMapping (aclTensor*)	Input	Storage offset of each token of key in the cache, i.e., slotMapping in the formula.	Nullable tensor. When key is 3D, the shape satisfies [batch * seq_len]; when key is 4D, the shape satisfies [batch, num_head]. The value range is [0, num_blocks * block_size - 1], and element values must not be duplicated; correctness is not guaranteed when duplicates exist.	INT32, INT64	ND	1-2	x
  compressLensOptional(aclTensor*)	Optional input.	Compression amount, compressLens in the formula.	Nullable. When key is 4-dimensional and compressSeqOffsetOptional is not a null pointer, the shape satisfies [batch, num_head]; when key is 4-dimensional and compressSeqOffsetOptional is a null pointer, the shape satisfies [batch * num_head]. Pass a null pointer in Scenario 1.	Consistent with slotMapping.	ND	1-2	x
  compressSeqOffsetOptional(aclTensor*)	Optional input	Compressed start offset of each head in each batch, which is compressSeqOffset in the formula.	Nullable tensor. The shape satisfies [batch * num_head]. Pass a null pointer in Scene 1 and Scene 3.	Consistent with slotMapping.	ND	1	x
  seqLensOptional(aclTensor*)	Optional input	Actual seqLens of each batch, i.e., seqLens in the formula.	Nullable tensor. The shape satisfies [batch]. Pass a null pointer in Scenario 1.	Consistent with slotMapping.	ND	1	x
  cacheMode(char*)	input	Specifies the memory layout format of keyCacheRef.	Reserved parameter. Currently only the ND data format is supported, so this parameter does not take effect.	-	-	-	-
  workspaceSize (uint64_t*)	Output	Returns the workspace size to be allocated on the Device side.	-	-	-	-	-
  executor(aclOpExecutor**)	Output	Returns the operator executor, which includes the operator computation flow.	-	-	-	-	-

• 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	The input key, keyCacheRef, and slotMapping are null pointers.
  ACLNN_ERR_PARAM_INVALID	161002	The data type of the key parameter is not within the supported range.
  		The data types of key and keyCacheRef are inconsistent.
  		The data types of slotMapping, compressLensOptional, compressSeqOffsetOptional, and seqLensOptional are inconsistent.
  ACLNN_ERR_PARAM_INVALID	561002	The number of dimensions of key is not equal to 3 or 4.

aclnnScatterPaCache

• Parameters

  Parameter	Input/Output	Description
  workspace	Input	Side-allocated workspace memory address on the Device.
  workspaceSize	Input	Side-allocated workspace size on the Device, obtained from the first-phase API aclnnScatterPaCacheGetWorkspaceSize.
  executor	Input	Operator executor, which includes the computation process of the operator.
  stream	Input	Specifies the stream for executing the task.

• Return Value

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

Constraints

• Deterministic computation: aclnnScatterPaCache defaults to a deterministic implementation.

• The variables used in shape descriptions in Parameters are as follows:

  • batch: the number of sequences in the current input (the number of samples processed at a time), which is a positive integer.
  • seq_len: the length of the sequence, which is a positive integer.
  • num_head: Number of "heads" in multi-head attention, which is a positive integer.
  • k_head_size: Feature dimension of key per attention head (length of a single-head key), which is a positive integer.
  • num_blocks: Total number of pre-allocated blocks in keyCache for storing key data of all sequences, which is a positive integer.
  • block_size: Number of tokens contained in each cache block, which is a positive integer.

• Input value range restriction: Each element value in seqLensOptional and compressLensOptional must satisfy the formula: reduceSum(seqLensOptional[i] - compressLensOptional[i] + 1) <= num_blocks * block_size (corresponding to scenarios 2 and 3).

Example

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

• Ascend 950PR/Ascend 950DT:

  #include <iostream>
  #include <vector>
  #include "acl/acl.h"
  #include "aclnnop/aclnn_scatter_pa_cache.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;
  }
  
  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 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 = 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.
    // 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 input and output. Custom construction is required based on the API.
    std::vector<int64_t> keyShape = {2, 2, 3, 7};
    std::vector<int64_t> keyCacheShape = {2, 2, 1, 7};
    std::vector<int64_t> slotMappingShape = {2, 3};
  
    std::vector<int64_t> compressLensShape = {2, 3};
    std::vector<int64_t> compressSeqOffsetShape = {6};
    std::vector<int64_t> seqLensShape = {2};
    void* keyDeviceAddr = nullptr;
    void* slotMappingDeviceAddr = nullptr;
    void* keyCacheDeviceAddr = nullptr;
    void* compressLensDeviceAddr = nullptr;
    void* compressSeqOffsetDeviceAddr = nullptr;
    void* seqLensDeviceAddr = nullptr;
  
    aclTensor* key = nullptr;
    aclTensor* slotMapping = nullptr;
    aclTensor* keyCache = nullptr;
    aclTensor* compressLens = nullptr;
    aclTensor* compressSeqOffset = nullptr;
    aclTensor* seqLens = nullptr;
    char* cacheMode = const_cast<char*>("Norm");
  
    std::vector<float> hostKey = {1};
    std::vector<int32_t> hostSlotMapping = {0, 3, 6, 9, 12, 15};
    std::vector<float> hostKeyCacheRef = {1};
    std::vector<int32_t> hostCompressLens = {1, 0, 0, 0, 1, 0};
    std::vector<int32_t> hostCompressSeqOffset = {0, 0, 1, 0, 1, 1};
    std::vector<int32_t> hostSeqLens = {2, 1};
  
    // Create the key aclTensor.
    ret = CreateAclTensor(hostKey, keyShape, &keyDeviceAddr, aclDataType::ACL_FLOAT, &key);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the slotMappitng aclTensor.
    ret = CreateAclTensor(hostSlotMapping, slotMappingShape, &slotMappingDeviceAddr, aclDataType::ACL_INT32, &slotMapping);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the keyCache aclTensor.
    ret = CreateAclTensor(hostKeyCacheRef, keyCacheShape, &keyCacheDeviceAddr, aclDataType::ACL_FLOAT, &keyCache);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the compressLens aclTensor.
    ret = CreateAclTensor(hostCompressLens, compressLensShape, &compressLensDeviceAddr, aclDataType::ACL_INT32, &compressLens);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the compressSeqOffset aclTensor.
    ret = CreateAclTensor(hostCompressSeqOffset, compressSeqOffsetShape, &compressSeqOffsetDeviceAddr, aclDataType::ACL_INT32, &compressSeqOffset);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the seqLens aclTensor.
    ret = CreateAclTensor(hostSeqLens, seqLensShape, &seqLensDeviceAddr, aclDataType::ACL_INT32, &seqLens);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
  
    // 3. Call the CANN operator library API. Replace with the actual API name.
    uint64_t workspaceSize = 0;
    aclOpExecutor* executor;
    // Call the first-phase API of aclnnScatterPaCache.
    ret = aclnnScatterPaCacheGetWorkspaceSize(key, keyCache, slotMapping, compressLens, compressSeqOffset, seqLens, cacheMode, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnScatterPaCacheGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
    // Allocate 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 of aclnnScatterPaCache.
    ret = aclnnScatterPaCache(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnScatterPaCache 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-side memory to the host side. Modify this based on the specific API definition.
    auto size = GetShapeSize(keyShape);
    std::vector<float> resultData(size, 0);
    ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), keyDeviceAddr,
                      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++) {
      LOG_PRINT("result[%ld] is: %f\n", i, resultData[i]);
    }
  
    // 6. Destroy aclTensor and aclScalar. Modify this based on the specific API definition.
    aclDestroyTensor(key);
    aclDestroyTensor(slotMapping);
    aclDestroyTensor(keyCache);
    aclDestroyTensor(compressLens);
    aclDestroyTensor(compressSeqOffset);
    aclDestroyTensor(seqLens);
    // 7. Release device resources. Configure parameters based on the specific API definition.
    aclrtFree(keyDeviceAddr);
    aclrtFree(slotMappingDeviceAddr);
    aclrtFree(keyCacheDeviceAddr);
    aclrtFree(compressLensDeviceAddr);
    aclrtFree(compressSeqOffsetDeviceAddr);
    aclrtFree(seqLensDeviceAddr);
    if (workspaceSize > 0) {
      aclrtFree(workspaceAddr);
    }
    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    aclFinalize();
    return 0;
  }