BlockCopyOperation

Applicable Products

Hardware Model	Supported	Remarks
Atlas 350 accelerator cards	x	-
Atlas A3 inference products/Atlas A3 training products	√	-
Atlas A2 training products/Atlas A2 inference products	√	-
Atlas training products	x	-
Atlas inference products	√	Only float16 allowed.
Atlas 200I/500 A2 inference products	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

Figure 1 BlockCopyOperation

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	ND: [num_blocks, block_size, num_head, head_size] NZ: [num_blocks,num_head*head_size/16, block_size, 16]	float16/bf16/int8	ND/NZ	Key matrix
valueCache	ND: [num_blocks, block_size, num_head, head_size] NZ: [num_blocks,num_head*head_size/16, block_size, 16]	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.

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

Figure 2 kvCache target position diagram

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 and Atlas 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.

Parent topic: atb/infer_op_params.h