aclnnScatterPaKvCache

📄 View Source Code

Applicable Products

Function

• Description: Updates the key and value at specified positions in the KV cache.

• The input and output support the following scenarios:

  • Scenario 1:

    key:[batch * seq_len, num_head, k_head_size]
    value:[batch * seq_len, num_head, v_head_size]
    keyCache:[num_blocks, num_head * k_head_size // last_dim_k, block_size, last_dim_k]/[num_blocks, num_head, k_head_size // last_dim_k, block_size, last_dim_k]
    valueCache:[num_blocks, num_head * v_head_size // last_dim_v, block_size, last_dim_v]/[num_blocks, num_head, v_head_size // last_dim_v, block_size, last_dim_v]
    slotMapping:[batch * seq_len]
    cacheMode:"PA_NZ"

  • Scenario 2:

    key:[batch * seq_len, num_head, k_head_size]
    value:[batch * seq_len, num_head, v_head_size]
    keyCache:[num_blocks, block_size, num_head, k_head_size]
    valueCache:[num_blocks, block_size, num_head, v_head_size]
    slotMapping:[batch * seq_len]
    cacheMode:"Norm"
    scatter_mode:"None"/"Nct"

    k_head_size and v_head_size can be different or the same.

  • Scenario 3:

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

  • Scenario 4:

    key:[num_tokens, num_head, k_head_size]
    value:[num_tokens, num_head, v_head_size]
    keyCache:[num_blocks, block_size, 1, k_head_size]
    valueCache:[num_blocks, block_size, 1, v_head_size]
    slotMapping:[batch * num_head]
    compressLensOptional:[batch * num_head]
    seqLensOptional:[batch]
    cacheMode:"Norm"
    scatter_mode:"Alibi"

  • Scenario 5:

    key:[num_tokens, num_head, k_head_size]
    value:[num_tokens, num_head, v_head_size]
    keyCache:[num_blocks, block_size, 1, k_head_size]
    valueCache:[num_blocks, block_size, 1, v_head_size]
    slotMapping:[batch * num_head]
    compressLensOptional:[batch * num_head]
    seqLensOptional:[batch]
    compressSeqOffsetOptional:[batch * num_head]
    cacheMode:"Norm"
    scatter_mode:"Rope"/"Omni"

    • Scenario 6:

    key:[batch * seq_len, num_head, k_head_size]
    value:[]
    keyCache:[num_blocks, block_size, num_head, k_head_size]
    valueCache:[]
    slotMapping:[batch * seq_len]
    cacheMode:"Norm"
    scatter_mode:"None"/"Nct"

• The above scenarios are distinguished based on the constructed parameters. The first parameter construction corresponds to Scenario 1, the second to Scenario 2, the third to Scenario 3, the fourth to Scenario 4, the fifth to Scenario 5, and the sixth to Scenario 6. Scenarios 1, 2, and 6 do not have the three optional parameters: compressLensOptional, seqLensOptional, and compressSeqOffsetOptional. Scenario 4 does not have the compressSeqOffsetOptional optional parameter.

• Atlas A3 training products/Atlas A3 inference products, Atlas A2 training products/Atlas A2 inference products: Only Scenarios 1, 2, 4, 5, and 6 are supported.

Function Prototype

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

aclnnStatus aclnnScatterPaKvCacheGetWorkspaceSize(
  const aclTensor   *key,
  aclTensor         *keyCacheRef,
  const aclTensor   *slotMapping,
  const aclTensor   *value,
  aclTensor         *valueCacheRef,
  const aclTensor   *compressLensOptional,
  const aclTensor   *compressSeqOffsetOptional,
  const aclTensor   *seqLensOptional,
  char              *cacheModeOptional,
  char              *scatterModeOptional,
  const aclIntArray *stridesOptional,
  const aclIntArray *offsetsOptional,
  uint64_t          *workspaceSize,
  aclOpExecutor    **executor)

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

aclnnScatterPaKvCacheGetWorkspaceSize

• Parameters

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
  key (aclTensor*)	Input	Key values to be updated, representing the keys of multiple tokens at the current step.		FLOAT16, FLOAT, BFLOAT16, INT8, UINT8, INT16, UINT16, INT32, UINT32, HIFLOAT8, FLOAT8_E5M2, FLOAT8_E4M3FN	ND	3-4	x
  keyCacheRef (aclTensor*)	Input/Output	Key cache to be updated, i.e., the key cache of the current layer.	Supports 4D or 5D. When a null pointer or "Norm" is passed, only the ND memory layout format is supported. When "PA_NZ" is passed, only the FRACTAL_NZ memory layout is supported.	Same as key.	ND	4-5	x
  slotMapping (aclTensor*)	Input	Storage offset of each token key or value in the cache.		INT32, INT64	ND	1	x
  value (aclTensor*)	Input	Value to be updated, i.e., the values of multiple tokens at the current step.	Supports 0, 3, or 4 dimensions. For non-zero dimensions, the shape is the same as that of key.	Must be consistent with key.	ND	0, 3, 4	x
  valueCacheRef (aclTensor*)	Input/Output	Specifies the value cache to be updated, which is the value cache of the current layer.	Supports 0D, 4D, or 5D. For non-0D, the shape is the same as keyCacheRef. When a null pointer or "Norm" is passed, only the ND memory layout format is supported. When "PA_NZ" is passed, only the FRACTAL_NZ memory layout is supported.	Consistent with key.	ND	0, 4, 5	x
  compressLensOptional (aclTensor*)	Optional input.	Compression lengths.	-	Consistent with slotMapping.	ND	1	x
  compressSeqOffsetOptional (aclTensor*)	Optional input	Specifies the compression start point for each head in each batch.	-	Must be consistent with slotMapping.	ND	1	x
  seqLensOptional (aclTensor*)	Optional input	Actual seqLens of each batch.	-	Consistent with slotMapping.	ND	1	x
  cacheModeOptional (char*)	Input	Indicates the memory layout format of keyCacheRef and valueCacheRef.	When a null pointer or "Norm" is passed, only the ND memory layout format is supported. When "PA_NZ" is passed, only the FRACTAL_NZ memory layout is supported.	-	-	-	-
  scatterModeOptional (char*)	Input	Indicates the state of the key and value to be updated.	When a null pointer or "None" is passed, it indicates that the key and value to be updated are in an uncompressed and contiguous state. When "Alibi" is passed, it indicates that the key and value to be updated are in a compressed state based on the Alibi structure. When "Rope" is passed, it indicates that the key and value to be updated are in a compressed state based on the Rope structure. When "Omni" is passed, it indicates that the key and value to be updated are in a compressed state based on the Omni structure. When "Nct" is passed, it indicates that the key and value to be updated are in an uncompressed but non-contiguous state.	-	-	-	-
  stridesOptional (aclIntArray*)	Input	Specifies the strides of key and value in a non-contiguous state.	The array length is 2. Its values must be greater than 0. It takes effect only when scatterModeOptional is set to "Nct", representing strideK and strideV respectively.	-	-	-	-
  offsetsOptional (aclIntArray*)	Input	Specifies the offsets of key and value in the non-contiguous state.	The array length is 2. Its values must be greater than 0. This parameter takes effect only when scatterModeOptional is "Nct", representing offsetK and offsetV respectively.	-	-	-	-
  workspaceSize (uint64_t*)	Output	Returns the workspace size to be requested on the Device side.	-	-	-	-	-
  executor (aclOpExecutor**)	Output	Returns the op executor, which contains the operator computation flow.	-	-	-	-	-

  • Atlas A2 training products/Atlas A2 inference products, Atlas A3 training products/Atlas A3 inference products:

    • The input key, keyCacheRef, value, and valueCacheRef do not support the FLOAT, UINT8, INT16, UINT16, INT32, UINT32, HIFLOAT8, FLOAT8_E5M2, and FLOAT8_E4M3FN data types.

• 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, slotMapping, value, and valueCacheRef are null pointers.
  ACLNN_ERR_PARAM_INVALID	161002	The data type of the key or value parameter is not within the supported range.
  		The data types of key, keyCacheRef, value, and valueCacheRef 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 3 or 4, or the number of dimensions of value is not 0, 3, or 4.

aclnnScatterPaKvCache

• Parameters

  Parameter	Input/Output	Description
  workspace	Input	Workspace memory address applied for on the Device side.
  workspaceSize	Input	Workspace size applied for on the Device side, obtained through the first-phase API aclnnScatterPaKvCacheGetWorkspaceSize.
  executor	Input	Operator executor that contains the operator computation flow.
  stream	Input	Specifies the stream for task execution.

• Return Value

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

Constraints

• Deterministic computation:

  • aclnnScatterPaKvCache defaults to a deterministic implementation.

  • The data types of key, value, keyCacheRef, and valueCacheRef must be consistent.

  • The data types of slotMapping, compressLensOptional, compressSeqOffsetOptional, and seqLensOptional must be consistent.

  • The value range of slotMapping is [0, num_blocks * block_size - 1], and the elements in slotMapping must be unique. Correctness is not guaranteed if duplicates exist.

  • When both key and value are 3D, the shapes of the first two dimensions of key and value must be the same.

  • When both key and value are 4D, the shapes of the first three dimensions of key and value must be the same, and the third dimension of keyCacheRef and valueCacheRef must be 1.

  • When key and value are 4D, compressLensOptional and seqLensOptional are mandatory parameters. When key and value are 3D, compressLensOptional, compressSeqOffsetOptional, and seqLensOptional are optional parameters.

  • When both key and value are 4D, slotMapping is 2D, where the first dimension of slotMapping equals the first dimension of key (batch), and the second dimension of slotMapping equals the third dimension of key (num_head) (corresponding to Scenario 3).

  • When both key and value are 4-dimensional, seqLensOptional is 1-dimensional, and the value of seqLensOptional equals the first dimension of key, which is batch (corresponding to Scenario 3);

  • When key and value are 3-dimensional and seqLensOptional exists, the sum of all values in seqLensOptional equals the first dimension of key, which is num_blocks (corresponding to Scenarios 4 and 5);

  • Each element value in seqLensOptional and compressLensOptional must satisfy the formula: reduceSum(seqLensOptional[i] - compressLensOptional[i]) <= num_blocks * block_size (corresponding to Scenarios 3, 4, and 5).

  • When cacheModeOptional is "PA_NZ", the second-to-last dimension of keyCacheRef and valueCacheRef must be less than UINT16_MAX (corresponding to Scenario 1).

Example

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

• x1 is INT8, x2 is INT8, x1Scale is FLOAT32, and x2Scale is FLOAT32.

  #include <iostream>
  #include <vector>
  #include "acl/acl.h"
  #include "aclnnop/aclnn_scatter_pa_kv_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 for 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, aclFormat format) {
    auto size = GetShapeSize(shape) * sizeof(T);
    // Call aclrtMalloc to apply for memory on the device side.
    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 side 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);
  
    // 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, format,
                              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 inputs and outputs. Custom construction is required based on the API.
    std::vector<int64_t> keyShape = {2, 2, 32};
    std::vector<int64_t> keyCacheShape = {1, 4, 32, 16};
    std::vector<int64_t> slotMappingShape = {2};
    std::vector<int64_t> valueShape = {2, 2, 32};
    std::vector<int64_t> valueCacheShape = {1, 4, 32, 16};
  
    void* keyDeviceAddr = nullptr;
    void* valueDeviceAddr = nullptr;
    void* slotMappingDeviceAddr = nullptr;
    void* keyCacheDeviceAddr = nullptr;
    void* valueCacheDeviceAddr = nullptr;
    void* compressLensDeviceAddr = nullptr;
    void* compressSeqOffsetDeviceAddr = nullptr;
    void* seqLensDeviceAddr = nullptr;
  
    aclTensor* key = nullptr;
    aclTensor* value = nullptr;
    aclTensor* slotMapping = nullptr;
    aclTensor* keyCache = nullptr;
    aclTensor* valueCache = nullptr;
    aclTensor* compressLens = nullptr;
    aclTensor* compressSeqOffset = nullptr;
    aclTensor* seqLens = nullptr;
    char * cacheMode = "PA_NZ";
    char * scatterMode = "None";
  
    std::vector<int16_t> hostKey(128, 1);
    std::vector<int16_t> hostValue(128, 1);
    std::vector<int32_t> hostSlotMapping(2, 1);
    std::vector<int16_t> hostKeyCacheRef(2048, 1);
    std::vector<int16_t> hostValueCacheRef(2048, 1);
    std::vector<int64_t> hostStrides(2, 1);
    std::vector<int64_t> hostOffsets(2, 0);
  
    // Create the key aclTensor.
    ret = CreateAclTensor(hostKey, keyShape, &keyDeviceAddr, aclDataType::ACL_FLOAT16, &key, aclFormat::ACL_FORMAT_ND);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the value aclTensor.
    ret = CreateAclTensor(hostValue, valueShape, &valueDeviceAddr, aclDataType::ACL_FLOAT16, &value, aclFormat::ACL_FORMAT_ND);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the slotMapping aclTensor.
    ret = CreateAclTensor(hostSlotMapping, slotMappingShape, &slotMappingDeviceAddr, aclDataType::ACL_INT32, &slotMapping, aclFormat::ACL_FORMAT_ND);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the keyCache aclTensor.
    ret = CreateAclTensor(hostKeyCacheRef, keyCacheShape, &keyCacheDeviceAddr, aclDataType::ACL_FLOAT16, &keyCache, aclFormat::ACL_FORMAT_ND);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the valueCache aclTensor.
    ret = CreateAclTensor(hostValueCacheRef, valueCacheShape, &valueCacheDeviceAddr, aclDataType::ACL_FLOAT16, &valueCache, aclFormat::ACL_FORMAT_ND);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    aclIntArray *strides = aclCreateIntArray(hostStrides.data(), 2);
    CHECK_RET(strides != nullptr, return ACL_ERROR_INTERNAL_ERROR);
    aclIntArray *offsets = aclCreateIntArray(hostOffsets.data(), 2);
    CHECK_RET(offsets != nullptr, return ACL_ERROR_INTERNAL_ERROR);
  
    // 3. Call the CANN operator library API. Replace with the specific API name.
    uint64_t workspaceSize = 0;
    aclOpExecutor* executor;
    // Call the first-phase API of aclnnScatterPaKvCache.
    ret = aclnnScatterPaKvCacheGetWorkspaceSize(key, keyCache, slotMapping, value, valueCache, compressLens, compressSeqOffset, seqLens, cacheMode, scatterMode, strides, offsets, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnScatterPaKvCacheGetWorkspaceSize 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 of aclnnScatterPaKvCache.
    ret = aclnnScatterPaKvCache(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnScatterPaKvCache 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 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. Release aclTensor and aclScalar. Modify based on the specific API definition.
    aclDestroyTensor(key);
    aclDestroyTensor(value);
    aclDestroyTensor(slotMapping);
    aclDestroyTensor(keyCache);
    aclDestroyTensor(valueCache);
    aclDestroyTensor(compressLens);
    aclDestroyTensor(compressSeqOffset);
    aclDestroyTensor(seqLens);
    aclDestroyIntArray(strides);
    aclDestroyIntArray(offsets);
    // 7. Release device resources. Modify parameters based on the specific API definition.
    aclrtFree(keyDeviceAddr);
    aclrtFree(valueDeviceAddr);
    aclrtFree(slotMappingDeviceAddr);
    aclrtFree(keyCacheDeviceAddr);
    aclrtFree(valueCacheDeviceAddr);
     if (workspaceSize > 0) {
      aclrtFree(workspaceAddr);
    }
    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    aclFinalize();
    return 0;
  }