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$是上下文的长度。

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
query	输入	公式中的输入Q。 不支持空tensor和非连续。 layout_query为BSND时，shape为(B,S1,N1,D)；layout_query为TND时，shape为(T1,N1,D)。	FLOAT16、BFLOAT16	ND
key	输入	公式中的输入K。 不支持空tensor和非连续。 layout_key为PA_BSND时，shape为(block_num, block_size, N2, D)，其中block_num为PageAttention时block总数、block_size为一个block的token数；layout_key为BSND时，shape为(B, S2, N2, D)；layout_key为TND时，shape为(T2, N2, D)。	FLOAT16、BFLOAT16	ND
weights	输入	公式中的输入W。 不支持空tensor和非连续。 layout_query为BSND时，shape为(B,S1,N1)；layout_query为TND时，shape为(T1,N1)。	FLOAT16、BFLOAT16、FLOAT	ND
actual_seq_lengths_query	输入	每个Batch中Query的有效token数。 不支持空tensor和非连续。 可传入None表示与query的S长度相同。 支持长度为B的一维tensor，且每个Batch的有效token数不超过query中的维度S大小且不小于0。layout_query为TND时该入参必须传入，并以元素数量作为B值。 每个元素表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须大于等于前一个元素的值。	INT32	ND
actual_seq_lengths_key	输入	每个Batch中Key的有效token数。 不支持空tensor和非连续。 可传入None表示与key的S长度相同。 支持长度为B的一维tensor，且每个Batch的有效token数不超过key/value中的维度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,S2/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, S1, N2, sparseCount]；layout_query为TND时输出shape为[T1, N2, sparseCount]。	INT32	-
sparse_values	输出	公式中的Indices输出对应的value值。 不支持空tensor和非连续。 shape与sparseIndicesOut保持一致。	FLOAT16、BFLOAT16	ND

约束说明

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

调用示例

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