enginex-ascend-910-vllm/csrc/attention/vllm_quant_lightning_indexer_metadata

VllmQuantLightningIndexerMetadata

产品支持情况

产品	是否支持
Ascend 950PR/Ascend 950DT	√
Atlas A3 推理系列产品	√
Atlas A2 推理系列产品	√

功能说明

• API功能：VllmQuantLightningIndexerMetadata是VllmQuantLightningIndexer的前置算子，通过AICPU为VllmQuantLightningIndexer算子生成分核结果，包括每个核需要处理的数据的起始点、结束点等内容，随后，VllmQuantLightningIndexer根据该分核结果进行实际计算。

• 主要计算过程为：

  1. 获取每个batch的基本块大小，并计算负载。
  2. 计算所有batch的总负载和总的基本块个数。
  3. 为每个核分配负载，并记录分核结果，分核结果包括每个核需要处理的数据的起始点、结束点等内容。

参数说明

• 参数维度含义：B（Batch Size）表示输入样本批量大小、S（Sequence Length）表示输入样本序列长度、H（Head Size）表示hidden层的大小、N（Head Num）表示多头数、D（Head Dim）表示hidden层最小的单元尺寸，且满足D=H/N、T表示所有Batch输入样本序列长度的累加和。
• 使用S1和S2分别表示query和key的输入样本序列长度，N1和N2分别表示query和key对应的多头数，k表示最后选取的索引个数。

参数名	输入/输出/属性	描述	数据类型	数据格式
num_heads_q	属性	Q的多头数。	INT32	-
num_heads_k	属性	K的多头数。	INT32	-
head_dim	属性	注意力头的维度。	INT32	-
query_quant_mode	属性	用于标识query的量化模式，当前支持Per-Token-Head量化模式，当前仅支持传入0	INT32	-
key_quant_mode	属性	用于标识输入key的量化模式，当前支持Per-Token-Head量化模式，当前仅支持传入0。	INT32	-
actual_seq_lengths_query	可选输入	表示不同Batch中`query`的有效token数，数据类型支持`int32`。如果不指定seqlen可传入None，表示和`query`的shape的S长度相同。该入参中每个Batch的有效token数不超过`query`中的维度S大小且不小于0。支持长度为B的一维tensor。 当`layout_query`为TND时，该入参必须传入，且以该入参元素的数量作为B值，该入参中每个元素的值表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须大于等于前一个元素的值。不能出现负值。特殊说明：actual\_seq\_lengths\_query和actual\_seq\_lengths\_key至少传入一个。	INT32	-
actual_seq_lengths_key	可选输入	表示不同Batch中压缩前原始`key`的有效token数，数据类型支持`int32`。如果不指定seqlen可传入None，表示和key的shape的S长度相同。该参数中每个Batch的原始有效token数除以压缩率后不超过`key`中的维度S大小且不小于0，支持长度为B的一维tensor。 当`layout_kv`为TND或PA_BSND时，该入参必须传入，`layout_kv`为TND，该参数中每个元素的值表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须大于等于前一个元素的值。特殊说明：actual\_seq\_lengths\_query和actual\_seq\_lengths\_key至少传入一个。	INT32	-
batch_size	可选属性	输入样本批量大小。	INT32	-
max_seqlen_q	可选属性	当layout_query为BSND时，表示每个Batch中的q的有效token数。	INT32	-
max_seqlen_k	可选属性	当layout_kv为BSND时，表示每个Batch中的k的有效token数。	INT32	-
cmp_ratio	可选属性	用于稀疏计算，表示key的压缩倍数。Atlas A3 推理系列产品支持1/2/4/8/16/32/64/128，Ascend 950PR/Ascend 950DT支持1/4/128。数据类型支持int32，默认值1。	INT32	-
layout_query	可选属性	用于标识query的数据排布格式，当前支持BSND、TND，默认值"BSND"。	STRING
layout_key	可选属性	用于标识key的数据排布格式，当前仅支持PA_BSND。	STRING	-
sparse_count	可选属性	代表topK阶段需要保留的block数量，Atlas A3推理系列产品支持[1, 2048]，Ascend 950PR/Ascend 950DT支持512。	INT32	-
sparse_mode	可选属性	表示sparse的模式，支持0/3，数据类型支持int32。为0时，代表defaultMask模式。为3时，代表rightDownCausal模式的mask，对应以右顶点为划分的下三角场景.	INT32	-
pre_token	可选属性	用于稀疏计算，表示attention需要和前几个Token计算关联,仅支持默认值2^63-1。	INT64	-
next_token	可选属性	用于稀疏计算，表示attention需要和前几个Token计算关联，仅支持默认值2^63-1。	INT64	-
device	可选属性	npu的ID。	STRING	-
metadata	输出	VllmQuantLightningIndexerMetadata算子传入的分核信息，包括每个Cube核上FlashAttention计算任务的Batch、Head以及Q和K分块的索引，以及每个Vector核上FlashDecode的规约任务索引。数据类型支持`int32`，shape大小为[1024]	INT32	-

约束说明

• 该接口支持推理场景下使用。
• 该接口支持aclgraph模式。