Input and Output List

The shapes in the list are used only in common scenarios. For details about the shapes in special scenarios, see Functions.

Parameter	Dimension	Data Type	Format	cpu or npu	Description	Application Scenarios
query	[num_tokens, num_head, head_size]	float16/bf16/int8	ND	npu	query of each batch is combined on the num_tokens axis. When full quantization is enabled, only int8 is supported.	Basic scenario
keyCache	Atlas A2 inference 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/int8	Atlas A2 inference products: ND Atlas inference products: NZ	npu	Cached key. When the dequantization fusion and full quantization functions are enabled, only int8 is supported.	Basic scenario
valueCache	Atlas A2 inference 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/int8	Atlas A2 inference products: ND Atlas inference products: NZ	npu	Cached value. When the dequantization fusion and full quantization functions are enabled, only int8 is supported. This parameter is not passed when the function of outputting KVCache after mla combination is enabled.	Basic scenario
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.	Basic scenario
contextLens	[batch]	int32	ND	cpu	Number of key/value tokens of each query	Basic scenario
mask	For details, see mask types.	float16/bf16	Atlas A2 inference products: ND Atlas inference products: NZ	npu	Atlas A2 inference products: This parameter is required when maskType is not set to UNDEFINED.	Mask
batchRunStatus	[batch]	int32	ND	cpu	Flag bit of the batches involved in computable batch control computation	Computable batch control
kDescale	[k_head_num*head_size]	int64/float	ND	npu	Stride tensor for dequantization. Must have the same data type as vDescale.	Dequantization fusion or full quantization
kOffset	[k_head_num*head_size]	int32	ND	npu	Offset of k when the quantization type is dequantization and hasQuantOffset is true.	Asymmetric dequantization
vDescale	[v_head_num*head_size]	int64/float	ND	npu	Step tensor for dequantization	Dequantization fusion or full quantization
vOffset	[v_head_num*head_size]	int32	ND	npu	Offset of v when the quantization type is dequantization and hasQuantOffset is true.	Asymmetric dequantization
qSeqLens	[batch]	int32	ND	cpu	seqLen of each batch when parallel decoding is enabled.	Parallel decoding
razorOffset	[num_blocks, block_size]	float	ND	npu	This tensor is required when the Razor Rope function is enabled.	Multi-head adaptive compression (Rope)
pScale	[q_head_num]	float	ND	npu	Quantization scale of the intermediate p_score, which needs to be passed when offline full quantization is enabled.	Offline full quantization
logN	[batch]	float	ND	npu	Scaling coefficient of each batch, which is required when the scale type is logN.	LogN
attnOut	[num_tokens, num_head, head_size_v]	float16/bf16	ND	npu	Query output after computation	Basic scenario

In the preceding table, contextLens and qSeqLens are special. Although they are input tensors, they are used as parameters. Therefore, tensor binding is required to extract them from the device and bind them to the host. contextLens is on the CPU in the runner of the Atlas A2 inference products and is on the NPU in the runner of other products.

In the preceding table, basic scenario means that the corresponding tensor is required in PagedAttention, including five input tensors and one output tensor. The input tensor in other scenarios can be passed only in the corresponding scenario. The following are the scenario definitions:

Mask: maskType is not UNDEFINED.
Computable batch control: batchRunStatusEnable is set to True.
Dequantization fusion: quantType is set to TYPE_DEQUANT_FUSION.
Full quantization: quantType is set to TYPE_QUANT_QKV_OFFLINE (offline full quantization) or TYPE_QUANT_QKV_ONLINE (online full quantization).
Asymmetric dequantization: quantType is set to TYPE_DEQUANT_FUSION and hasQuantOffset is set to True.
Parallel decoding: calcType is set to CALC_TYPE_SPEC.
Multi-head adaptive compression (rope): compressType is set to COMPRESS_TYPE_KVHEAD_ROPE.
Offline full quantization: quantType is set to TYPE_QUANT_QKV_OFFLINE.
LogN: scaleType is set to SCALE_TYPE_LOGN.

Parent topic: PagedAttentionOperation