BlockCopyOperation
Applicable Products
Hardware Model |
Supported |
Remarks |
|---|---|---|
Atlas 350 accelerator cards |
x |
- |
√ |
- |
|
√ |
- |
|
x |
- |
|
√ |
Only float16 allowed. |
|
x |
- |
Description
Copies data in keyCache and valueCache from a series of source blocks to a series of destination blocks by specifying the source block index and destination block index (the same copy is performed on keyCache and valueCache).
Operator Context
Definition
1 2 3 | struct BlockCopyParam { uint8_t rsv[16] = {0}; }; |
Parameters
Member |
Type |
Default Value |
Description |
|---|---|---|---|
rsv[16] |
uint8_t |
{0} |
Reserved |
Input
Parameter |
Dimension |
Data Type |
Format |
Description |
|---|---|---|---|---|
keyCache |
|
float16/bf16/int8 |
ND/NZ |
Key matrix |
valueCache |
|
float16/bf16/int8 |
ND/NZ |
Value matrix. The data type and format must be the same as those of keyCache. |
srcBlockIndices |
[src_count]. The maximum value is [num_blocks]. |
int32 |
ND |
keyCache and valueCache source location block index. The value range is [0,num_blocks). |
dstBlockIndices |
[dst_count]. The maximum value is [num_blocks]. |
int32 |
ND |
keyCache and valueCache target location block index. dstBlockIndices[cumSum[i-1]:cumSum[i]-1] is the target location list corresponding to srcBlockIndices[i] (specially, dstBlockIndices[0:cumSum[0]-1] is the target location list corresponding to srcBlockIndices[0]). The value range is [0,num_blocks). For details, see Figure 2. |
cumSum |
[src_count] |
int32 |
ND |
cumSum[i] is an end location, in dstBlockIndices, of a target location list corresponding to the source location srcBlockIndices[i]. |
Output
Parameter |
Dimension |
Data Type |
Format |
Description |
|---|---|---|---|---|
keyCacheNew |
[num_blocks, block_size, num_head, head_size] |
float16/bf16/int8 |
ND/NZ |
Key matrix after block copy, which is used as the output. The output keyCacheNew and input keyCache point to the same address, that is, in-place modification is performed. The data type and format must be the same as those of keyCache. |
valueCacheNew |
[num_blocks, block_size, num_head, head_size] |
float16/bf16/int8 |
ND/NZ |
Value matrix after block copy, which is used as the output. The output valueCacheNew and input valueCache point to the same address, that is, the in-place modification is performed. The data type and format must be the same as those of valueCache. |
Constraints
- A block index cannot appear in srcBlockIndices and dstBlockIndices at the same time.
- The block indexes in srcBlockIndices must be unique. If there are duplicate block indexes, combine them into one and obtain the union set of the corresponding dstBlockIndices.
- The block indexes in dstBlockIndices must be unique. That is, two different blocks cannot be copied to the same block.
- Each srcBlockIndex corresponds to at least one dstBlockIndex (i > 0, cumSum[i] - cumSum[i - 1] ≥ 1; cumSum[0] ≥ 1).
- According to the preceding four rules, the size of srcBlockIndices and dstBlockIndices must meet the following condition: dst_count ≥ src_count, src_count + dst_count ≤ num_blocks.
- Only the ND format is supported for
Atlas A2 training products /Atlas A2 inference products andAtlas A3 inference products /Atlas A3 training products . - There are the following restrictions on
Atlas inference products :- When the format is ND, the product of the last three dimensions of keyCache or valueCache must be 16-byte aligned. (Each block must be 32-byte aligned. For float16 data, each piece of data occupies two bytes.)
- When the format is NZ, block_size must be 16-byte aligned.