CANN/ops-transformer密集闪电索引Softmax算子
DenseLightningIndexerSoftmaxLse
【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer
产品支持情况
| 产品 | 是否支持 |
|---|---|
| Ascend 950PR/Ascend 950DT | × |
| Atlas A3 训练系列产品 | √ |
| Atlas A2 训练系列产品 | √ |
| Atlas 200I/500 A2 推理产品 | × |
| Atlas 推理系列产品 | × |
| Atlas 训练系列产品 | × |
功能说明
算子功能：DenseLightningIndexerSoftmaxLse算子是DenseLightningIndexerGradKlLoss算子计算Softmax输入的一个分支算子。
计算公式：
$$ \text{res}=\text{AttentionMask}\left(\text{ReduceSum}\left(W\odot\text{ReLU}\left(Q_{index}@K_{index}^T\right)\right)\right) $$
$$ \text{maxIndex}=\text{max}\left(res\right) $$
$$ \text{sumIndex}=\text{ReduceSum}\left(\text{exp}\left(res-maxIndex\right)\right) $$
maxIndex，sumIndex作为输出传递给算子DenseLightningIndexerGradKlLoss作为输入计算Softmax使用。
参数说明
| 参数名 | 输入/输出/属性 | 描述 | 数据类型 | 数据格式 |
|---|---|---|---|---|
| queryIndex | 输入 | lightningIndexer结构的输入queryIndex。 | FLOAT16、BFLOAT16 | ND |
| keyIndex | 输入 | lightningIndexer结构的输入keyIndex。 | FLOAT16、BFLOAT16 | ND |
| weights | 输入 | 权重。 | FLOAT16、BFLOAT16、FLOAT32 | ND |
| actualSeqLengthsQuery | 输入 | 每个Batch中，Query的有效token数。 | INT64 | ND |
| actualSeqLengthsKey | 输入 | 每个Batch中，Key的有效token数。 | INT64 | ND |
| layout | 输入 | layout格式。 | - | - |
| sparseMode | 输入 | sparse的模式。 | INT64 | - |
| preTokens | 输入 | 用于稀疏计算，表示Attention需要和前几个token计算关联。 | INT64 | - |
| nextTokens | 输入 | 用于稀疏计算，表示Attention需要和后几个token计算关联。 | INT64 | - |
| softmaxMaxOut | 输出 | softmax计算使用的max值。 | FLOAT32 | ND |
| softmaxSumOut | 输出 | softmax计算使用的sum值。 | FLOAT32 | ND |
约束说明
无
调用说明
| 调用方式 | 调用样例 | 说明 |
|---|---|---|
| aclnn调用 | test_aclnn_dense_lightning_indexer_softmax_lse | 通过aclnnDenseLightningIndexerSoftmaxLse接口方式调用dense_lightning_indexer_softmax_lse算子。 |
【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考