BNSD维度输入

功能

针对X1场景,使PA Op支持BNSD数据排布(原有的数据排布格式为BSND)。

开启方式

将参数“inputLayout”置为infer::inputLayout::TYPE_BNSD时开启BNSD功能。该参数默认值为infer::inputLayout::TYPE_BSND,表示默认数据排布格式为BSND。

若干tensor的规格发生了改变,改变的tensor如下表所示:

参数

维度

数据类型

格式

cpu or npu

q

[num_tokens, num_head, head_dim]

float16/bf16

ND

npu

kCache

[num_blocks, kv_head_num, block_size, head_size]

float16/bf16

ND

npu

vCache

[num_blocks, kv_head_num, block_size, head_size]

float16/bf16

ND

npu

out

[num_tokens, num_head, head_dim]

float16/bf16

ND

npu

num_tokens = batch * qseqlen,当前qseqlen = 1。

其余tensor未发生改变。

特殊约束