LightningIndexer

产品支持情况

产品	是否支持
Ascend 950PR/Ascend 950DT	√
Atlas A3 训练系列产品/Atlas A3 推理系列产品	√
Atlas A2 训练系列产品/Atlas A2 推理系列产品	√
Atlas 200I/500 A2 推理产品	×
Atlas 推理系列产品	×
Atlas 训练系列产品	×

功能说明

• 算子功能：lightning_indexer基于一系列操作得到每一个token对应的Top-$k$个位置。

• 计算公式：

  $$Indices=\text{Top-}k\left\{[1{]}_{1\times g}@\left[(W@[1{]}_{1\times S_k})\odot \text{ReLU}\left(Q_{index}@K_{index}^T\right)\right]\right\}$$

  对于某个token对应的Index Query $Q_{index}\in {\mathbb{R}}^{g\times d}$，给定上下文Index Key $K_{index}\in {\mathbb{R}}^{S_k\times d},W\in {\mathbb{R}}^{g\times 1}$，其中$g$为GQA对应的group size，$d$为每一个头的维度，$S_k$是上下文的长度。

参数说明

说明：
参数维度含义：B表示Batch Size、Q_S和K_S分别表示query和key的Sequence Length、Q_N和K_N分别表示query和key的Head Num、D表示Head Dim（Q_D和K_D取值相等为128）、Q_T和K_T分别表示query和key的Total Tokens、sparse_count表示最后选取的索引个数（topK）、block_num和block_size分别表示PageAttention场景下的block总数和每个block的token数。K_N仅支持1。

参数名	输入/输出/属性	描述	数据类型	数据格式
query	输入	公式中的输入Q。 不支持空tensor和非连续。 layout_query为BSND时，shape为[B, Q_S, Q_N, D]；layout_query为TND时，shape为[Q_T, Q_N, D]。	FLOAT16、BFLOAT16	ND
key	输入	公式中的输入K。 不支持空tensor和非连续。 layout_key为PA_BSND时，shape为[block_num, block_size, K_N, D]，其中block_num为PageAttention时block总数、block_size为一个block的token数；layout_key为BSND时，shape为[B, K_S, K_N, D]；layout_key为TND时，shape为[K_T, K_N, D]。	FLOAT16、BFLOAT16	ND
weights	输入	公式中的输入W。 不支持空tensor和非连续。 layout_query为BSND时，shape为[B, Q_S, Q_N]；layout_query为TND时，shape为[Q_T, Q_N]。	FLOAT16、BFLOAT16、FLOAT	ND
actual_seq_lengths_query	输入	每个Batch中Query的有效token数。 不支持空tensor和非连续。 可传入None表示与query的Q_S长度相同。 支持长度为B的一维tensor，且每个Batch的有效token数不超过query中的Q_S大小且不小于0。layout_query为TND时该入参必须传入，并以元素数量作为B值。 每个元素表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须大于等于前一个元素的值。	INT32	ND
actual_seq_lengths_key	输入	每个Batch中Key的有效token数。 不支持空tensor和非连续。 可传入None表示与key的K_S长度相同。 支持长度为B的一维tensor，且每个Batch的有效token数不超过key中的K_S大小且不小于0。 layout_key为TND或PA_BSND时该入参必须传入；其中layout_key为TND时，每个元素表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须大于等于前一个元素的值。	INT32	ND
block_table	输入	表示PageAttention中KV存储使用的block映射表。 不支持空tensor和非连续。 PageAttention场景下，block_table必须为二维，第一维长度需要等于B，第二维长度不能小于maxBlockNumPerSeq（每个batch中最大actual_seq_lengths_key对应的block数量）。 shape支持[B, K_S/block_size]。	INT32	ND
layout_query	属性	用于标识输入Query的数据排布格式。 默认值为"BSND"，当前支持BSND、TND。	STRING	-
layout_key	属性	用于标识输入Key的数据排布格式。 默认值为"BSND"，当前支持PA_BSND、BSND、TND。	STRING	-
sparse_count	属性	topK阶段需要保留的block数量。 支持[1, 2048]，以及3072、4096、5120、6144、7168、8192。 默认值为2048。	INT32	-
sparse_mode	属性	表示sparse的模式。 sparse_mode为0时代表defaultMask模式。 sparse_mode为3时代表rightDownCausal模式的mask，对应以右顶点为划分的下三角场景。 默认值为3。	INT32	-
pre_tokens	属性	用于稀疏计算，表示attention需要和前几个Token计算关联，仅支持默认值2^63-1。	INT64	-
next_tokens	属性	用于稀疏计算，表示attention需要和后几个Token计算关联，仅支持默认值2^63-1。	INT64	-
return_values	属性	表示是否输出sparseValuesOut。 True表示输出，False表示不输出，默认值为False。仅在训练且layout_key不为PA_BSND场景支持。	BOOL	-
sparse_indices	输出	公式中的Indices输出。 不支持空tensor和非连续。 layout_query为BSND时输出shape为[B, Q_S, K_N, sparse_count]；layout_query为TND时输出shape为[Q_T, K_N, sparse_count]。	INT32	-
sparse_values	输出	公式中的Indices输出对应的value值。 不支持空tensor和非连续。 shape与sparseIndicesOut保持一致。	FLOAT16、BFLOAT16	ND

约束说明

• 该接口支持图模式。
• 参数key中的K_N支持1。
• headdim支持128。
• block_size取值为16的倍数，最大支持1024。
• 参数query、key的数据类型应保持一致。
• 参数weights不为float32时，参数query、key、weights的数据类型应保持一致。
• Ascend 950PR/Ascend 950DT：
  • query Q_N仅支持8、16、24、32、64。
  • 参数weights不支持float32类型。
• Atlas A3训练系列产品/Atlas A3推理系列产品、Atlas A2训练系列产品/Atlas A2推理系列产品：
  • query Q_N支持小于等于64。

调用示例

调用方式	样例代码	说明
aclnn接口	test_aclnn_lightning_indexer	通过 aclnnLightningIndexer 接口方式调用算子