BlockCopyOperation

功能

通过指定源block index和目的block index,将keyCache和valueCache中的数据从一系列源block拷贝到一系列目的block(keyCache和valueCache做相同的拷贝处理)。

算子上下文

图1 BlockCopyOperation

定义

1
2
3
struct BlockCopyParam {
    uint8_t rsv[16] = {0};
};

参数列表

成员名称

类型

默认值

描述

rsv[16]

uint8_t

{0}

预留参数。

输入

参数

维度

数据类型

格式

描述

keyCache

[num_blocks, block_size, num_head, head_size]

float16/bf16/int8

ND

key矩阵。

valueCache

[num_blocks, block_size, num_head, head_size]

float16/bf16/int8

ND

value矩阵。

数据类型和格式应该与keyCache保持一致。

srcBlockIndices

[src_count],最长为[num_blocks]

int32

ND

keyCache、valueCache源位置block索引,值域范围为[0,num_blocks)。

dstBlockIndices

[dst_count],最长为[num_blocks]

int32

ND

keyCache、valueCache目标位置block索引,dstBlockIndices[cumSum[i-1]:cumSum[i]-1]为srcBlockIndices[i]对应的目标位置List(特殊的,dstBlockIndices[0:cumSum[0]-1]为srcBlockIndices[0]对应的目标位置List),值域范围为[0,num_blocks)。

cumSum

[src_count]

int32

ND

cumSum[i]为源位置srcBlockIndices[i]对应的目标位置List在dstBlockIndices中的结束位置。

输出

参数

维度

数据类型

格式

描述

keyCacheNew

[num_blocks, block_size, num_head, head_size]

float16/bf16/int8

ND

block拷贝后的key矩阵,作为输出。输出keyCacheNew与输入keyCache指向同一地址,即进行原地修改。

数据类型和格式应该与keyCache保持一致。

valueCacheNew

[num_blocks, block_size, num_head, head_size]

float16/bf16/int8

ND

block拷贝后的value矩阵,作为输出。输出valueCacheNew与输入valueCache指向同一地址,即进行原地修改。

数据类型和格式应该与valueCache保持一致。

规格约束