Basic Functions

When the default values of maskType, batchRunStatusEnable, quantType, calcType, compressType, and scaleType are used, the basic functions of PagedAttention are used. In this case, the input and output parameters are as follows.

Parameter	Dimension	Data Type	Format	cpu or npu	Description
query	[num_tokens, num_head, head_size]	float16/bf16	ND	npu	query of each batch is combined along the num_tokens axis.
keyCache	Atlas A2 training products/Atlas A2 inference products and Atlas A3 inference products/Atlas A3 training products: [num_blocks, block_size, kv_head_num, head_size_k]; Atlas inference products: [num_blocks, head_size *num_heads / 16 ,block_size, 16]	float16/bf16	Atlas A2 training products/Atlas A2 inference products and Atlas A3 inference products/Atlas A3 training products: ND Atlas inference products: NZ	npu	Cached key.
valueCache	Atlas A2 training products/Atlas A2 inference products and Atlas A3 inference products/Atlas A3 training products: [num_blocks, block_size, kv_head_num, head_size_v]; Atlas inference products: [num_blocks, head_size *num_heads / 16 ,block_size, 16]	float16/bf16	Atlas A2 training products/Atlas A2 inference products and Atlas A3 inference products/Atlas A3 training products: ND Atlas inference products: NZ	npu	Cached value.
blockTables	[num_tokens, max_num_blocks_per_query]	int32	ND	npu	Block table of kvcache of each query. The first dimension is the token index, and the second dimension indicates the block index.
contextLens	[batch]	int32	ND	cpu	Number of key/value tokens of each query.
attnOut	[num_tokens, num_head, head_size_v]	float16/bf16	ND	npu	Query output after computation.

Parent topic: Functions