当参数maskType、batchRunStatusEnable、quantType、calcType、compressType、scaleType均为默认值时,使用PagedAttention的基础功能,此时输入输出参数如下:
参数 |
维度 |
数据类型 |
格式 |
cpu or npu |
描述 |
---|---|---|---|---|---|
query |
[num_tokens, num_head, head_size] |
float16/bf16 |
ND |
npu |
各batch的query在num_tokens轴合并。 |
keyCache |
|
float16/bf16 |
|
npu |
cache好的key。 |
valueCache |
|
float16/bf16 |
|
npu |
cache好的value。 |
blockTables |
[num_tokens, max_num_blocks_per_query] |
int32 |
ND |
npu |
每个query的kvcache的block table,第一维是token索引,第二维表示block索引。 |
contextLens |
[batch] |
int32 |
ND |
cpu |
每个query对应的key/value的token数量。 |
attnOut |
[num_tokens, num_head, head_size_v] |
float16/bf16 |
ND |
npu |
经过计算输出的query。 |