BNSD Dimension Input
Description
For the X1 scenario, the PA Op supports the BNSD data layout (the original data layout format is BSND).
How to Enable
Set inputLayout to infer::inputLayout::TYPE_BNSD. The default value is infer::inputLayout::TYPE_BSND, indicating that the default data layout format is BSND.
The specifications of some tensors are changed. The following table lists the changed tensors.
Parameter |
Dimension |
Data Type |
Format |
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
Other tensors remain unchanged.
Constraints
None
Parent topic: Functions