270 lines
8.9 KiB
Markdown
270 lines
8.9 KiB
Markdown
# LightningIndexer
|
||
|
||
## 产品支持情况
|
||
|
||
|产品 | 是否支持 |
|
||
|:----------------------------|:-----------:|
|
||
|<term>Ascend 950PR/Ascend 950DT</term>| √ |
|
||
|<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>| √ |
|
||
|<term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>| √ |
|
||
|<term>Atlas 200I/500 A2 推理产品</term>| × |
|
||
|<term>Atlas 推理系列产品</term>| × |
|
||
|<term>Atlas 训练系列产品</term>| × |
|
||
|
||
## 功能说明
|
||
|
||
- 算子功能：`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\R^{g\times d}$，给定上下文Index Key $K_{index}\in\R^{S_{k}\times d},W\in\R^{g\times 1}$，其中$g$为GQA对应的group size，$d$为每一个头的维度，$S_{k}$是上下文的长度。
|
||
|
||
## 参数说明
|
||
|
||
<table style="undefined;table-layout: fixed; width: 1080px"><colgroup>
|
||
<col style="width: 200px">
|
||
<col style="width: 150px">
|
||
<col style="width: 480px">
|
||
<col style="width: 200px">
|
||
<col style="width: 100px">
|
||
</colgroup>
|
||
<thead>
|
||
<tr>
|
||
<th>参数名</th>
|
||
<th>输入/输出/属性</th>
|
||
<th>描述</th>
|
||
<th>数据类型</th>
|
||
<th>数据格式</th>
|
||
</tr></thead>
|
||
<tbody>
|
||
<tr>
|
||
<td>query</td>
|
||
<td>输入</td>
|
||
<td>
|
||
<ul>
|
||
<li>公式中的输入Q。</li>
|
||
<li>不支持空tensor和非连续。</li>
|
||
<li>layout_query为BSND时，shape为(B,S1,N1,D)；layout_query为TND时，shape为(T1,N1,D)。</li>
|
||
</ul>
|
||
</td>
|
||
<td>FLOAT16、BFLOAT16</td>
|
||
<td>ND</td>
|
||
</tr>
|
||
<tr>
|
||
<td>key</td>
|
||
<td>输入</td>
|
||
<td>
|
||
<ul>
|
||
<li>公式中的输入K。</li>
|
||
<li>不支持空tensor和非连续。</li>
|
||
<li>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)。</li>
|
||
</ul>
|
||
</td>
|
||
<td>FLOAT16、BFLOAT16</td>
|
||
<td>ND</td>
|
||
</tr>
|
||
<tr>
|
||
<td>weights</td>
|
||
<td>输入</td>
|
||
<td>
|
||
<ul>
|
||
<li>公式中的输入W。</li>
|
||
<li>不支持空tensor和非连续。</li>
|
||
<li>layout_query为BSND时，shape为(B,S1,N1)；layout_query为TND时，shape为(T1,N1)。</li>
|
||
</ul>
|
||
</td>
|
||
<td>FLOAT16、BFLOAT16、FLOAT</td>
|
||
<td>ND</td>
|
||
</tr>
|
||
<tr>
|
||
<td>actual_seq_lengths_query</td>
|
||
<td>输入</td>
|
||
<td>
|
||
<ul>
|
||
<li>每个Batch中Query的有效token数。</li>
|
||
<li>不支持空tensor和非连续。</li>
|
||
<li>可传入None表示与query的S长度相同。</li>
|
||
<li>支持长度为B的一维tensor，且每个Batch的有效token数不超过query中的维度S大小且不小于0。layout_query为TND时该入参必须传入，并以元素数量作为B值。</li>
|
||
<li>每个元素表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须大于等于前一个元素的值。</li>
|
||
</ul>
|
||
</td>
|
||
<td>INT32</td>
|
||
<td>ND</td>
|
||
</tr>
|
||
<tr>
|
||
<td>actual_seq_lengths_key</td>
|
||
<td>输入</td>
|
||
<td>
|
||
<ul>
|
||
<li>每个Batch中Key的有效token数。</li>
|
||
<li>不支持空tensor和非连续。</li>
|
||
<li>可传入None表示与key的S长度相同。</li>
|
||
<li>支持长度为B的一维tensor，且每个Batch的有效token数不超过key/value中的维度S大小且不小于0。</li>
|
||
<li>layout_key为TND或PA_BSND时该入参必须传入；其中layout_key为TND时，每个元素表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须大于等于前一个元素的值。</li>
|
||
</ul>
|
||
</td>
|
||
<td>INT32</td>
|
||
<td>ND</td>
|
||
</tr>
|
||
<tr>
|
||
<td>block_table</td>
|
||
<td>输入</td>
|
||
<td>
|
||
<ul>
|
||
<li>表示PageAttention中KV存储使用的block映射表。</li>
|
||
<li>不支持空tensor和非连续。</li>
|
||
<li>PageAttention场景下，block_table必须为二维，第一维长度需要等于B，第二维长度不能小于maxBlockNumPerSeq（每个batch中最大actual_seq_lengths_key对应的block数量）。</li>
|
||
<li>shape支持(B,S2/block_size)。</li>
|
||
</ul>
|
||
</td>
|
||
<td>INT32</td>
|
||
<td>ND</td>
|
||
</tr>
|
||
<tr>
|
||
<td>layout_query</td>
|
||
<td>属性</td>
|
||
<td>
|
||
<ul>
|
||
<li>用于标识输入Query的数据排布格式。</li>
|
||
<li>默认值为"BSND"，当前支持BSND、TND。</li>
|
||
</ul>
|
||
</td>
|
||
<td>STRING</td>
|
||
<td>-</td>
|
||
</tr>
|
||
<tr>
|
||
<td>layout_key</td>
|
||
<td>属性</td>
|
||
<td>
|
||
<ul>
|
||
<li>用于标识输入Key的数据排布格式。</li>
|
||
<li>默认值为"BSND"，当前支持PA_BSND、BSND、TND。</li>
|
||
</ul>
|
||
</td>
|
||
<td>STRING</td>
|
||
<td>-</td>
|
||
</tr>
|
||
<tr>
|
||
<td>sparse_count</td>
|
||
<td>属性</td>
|
||
<td>
|
||
<ul>
|
||
<li>topK阶段需要保留的block数量。</li>
|
||
<li>支持[1, 2048]，以及3072、4096、5120、6144、7168、8192。</li>
|
||
<li>默认值为2048。</li>
|
||
</ul>
|
||
</td>
|
||
<td>INT32</td>
|
||
<td>-</td>
|
||
</tr>
|
||
<tr>
|
||
<td>sparse_mode</td>
|
||
<td>属性</td>
|
||
<td>
|
||
<ul>
|
||
<li>表示sparse的模式。</li>
|
||
<li>sparse_mode为0时代表defaultMask模式。</li>
|
||
<li>sparse_mode为3时代表rightDownCausal模式的mask，对应以右顶点为划分的下三角场景。</li>
|
||
<li>默认值为3。</li>
|
||
</ul>
|
||
</td>
|
||
<td>INT32</td>
|
||
<td>-</td>
|
||
</tr>
|
||
<tr>
|
||
<td>pre_tokens</td>
|
||
<td>属性</td>
|
||
<td>用于稀疏计算，表示attention需要和前几个Token计算关联，仅支持默认值2^63-1。</td>
|
||
<td>INT64</td>
|
||
<td>-</td>
|
||
</tr>
|
||
<tr>
|
||
<td>next_tokens</td>
|
||
<td>属性</td>
|
||
<td>用于稀疏计算，表示attention需要和后几个Token计算关联，仅支持默认值2^63-1。</td>
|
||
<td>INT64</td>
|
||
<td>-</td>
|
||
</tr>
|
||
<tr>
|
||
<td>return_values</td>
|
||
<td>属性</td>
|
||
<td>
|
||
<ul>
|
||
<li>表示是否输出sparseValuesOut。</li>
|
||
<li>True表示输出，False表示不输出，默认值为False。仅在训练且layout_key不为PA_BSND场景支持。</li>
|
||
</ul>
|
||
</td>
|
||
<td>BOOL</td>
|
||
<td>-</td>
|
||
</tr>
|
||
<tr>
|
||
<td>sparse_indices</td>
|
||
<td>输出</td>
|
||
<td>
|
||
<ul>
|
||
<li>公式中的Indices输出。</li>
|
||
<li>不支持空tensor和非连续。</li>
|
||
<li>layout_query为BSND时输出shape为[B, S1, N2, sparseCount]；layout_query为TND时输出shape为[T1, N2, sparseCount]。</li>
|
||
</ul>
|
||
</td>
|
||
<td>INT32</td>
|
||
<td>-</td>
|
||
</tr>
|
||
<tr>
|
||
<td>sparse_values</td>
|
||
<td>输出</td>
|
||
<td>
|
||
<ul>
|
||
<li>公式中的Indices输出对应的value值。</li>
|
||
<li>不支持空tensor和非连续。</li>
|
||
<li>shape与sparseIndicesOut保持一致。</li>
|
||
</ul>
|
||
</td>
|
||
<td>FLOAT16、BFLOAT16</td>
|
||
<td>ND</td>
|
||
</tr>
|
||
</tbody>
|
||
</table>
|
||
|
||
## 约束说明
|
||
|
||
- 该接口支持图模式。
|
||
- 参数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。
|
||
|
||
## 调用示例
|
||
|
||
<table class="tg"><thead>
|
||
<tr>
|
||
<th class="tg-0pky">调用方式</th>
|
||
<th class="tg-0pky">样例代码</th>
|
||
<th class="tg-0pky">说明</th>
|
||
</tr></thead>
|
||
<tbody>
|
||
<tr>
|
||
<td class="tg-9wq8" rowspan="6">aclnn接口</td>
|
||
<td class="tg-0pky">
|
||
<a href="./examples//test_aclnn_lightning_indexer.cpp">test_aclnn_lightning_indexer
|
||
</a>
|
||
</td>
|
||
<td class="tg-lboi" rowspan="6">
|
||
通过
|
||
<a href="./docs/aclnnLightningIndexer.md">aclnnLightningIndexer
|
||
</a>
|
||
接口方式调用算子
|
||
</td>
|
||
</tr>
|
||
</tbody></table>
|