BlockCopyOperation
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). It is supported only by the
Operator Context
Definition
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 |
[num_blocks, block_size, num_head, head_size] |
float16/bf16/int8 |
ND |
Key matrix |
valueCache |
[num_blocks, block_size, num_head, head_size] |
float16/bf16/int8 |
ND |
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). |
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 |
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 |
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. |
Restrictions
- 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.