---
title: FastSoftMaxOperation
description: "| 产品 | 是否支持 |"
url: https://www.hiascend.com/document/detail/zh/canncommercial/latest/API/ascendtb/ascendtb_01_0088.html
sourcePath: /source/zh/canncommercial/900/API/ascendtb/ascendtb_01_0088.html
indexId: 9927fb32a3314d2dc01249c6a8b910d8d3181e9daf4898984eb442391bd3109464
---
# FastSoftMaxOperation

#### 产品支持情况

| 产品 | 是否支持 |
| --- | --- |
| Atlas 350 加速卡 | x |
| Atlas A3 推理系列产品 / Atlas A3 训练系列产品 | √ |
| Atlas A2 训练系列产品 / Atlas A2 推理系列产品 | √ |
| Atlas 训练系列产品 | x |
| Atlas 推理系列产品 | x |
| Atlas 200I/500 A2 推理产品 | x |


#### 功能说明

训练算子，与unpadOperation一起使用，将unpad处理后的Q矩阵和K矩阵相乘的结果做高性能的SoftMax处理。

SoftMax数学公式：


#### 使用场景

图1 模型流程图

Multi-Head Attention计算过程如图1所示。经过unpad处理后，Q矩阵和K矩阵相乘的结果实际上由batchSize个大小为(seqLen[i], seqLen[i])的矩阵拼接而成（如图2所示），需要对这个结果的末轴做SoftMax操作。采用FastSoftMax算子，接收seqLen和headNum的信息可以对这种排列方式的数据实现SoftMax操作。

图2 输入数据的排列方式


#### 定义

```
struct FastSoftMaxParam {
    int32_t headNum = 0;
    std::vector<int32_t> qSeqLen;
    uint8_t rsv[8] = {0};
};
```


#### 参数列表

| 成员名称 | 类型 | 默认值 | 描述 |
| --- | --- | --- | --- |
| headNum | int32\_t | 0 | Attention的head数量，需大于0。 |
| qSeqLen | std::vector<int32\_t> | \- | 每个batch的实际输入长度。元素个数为batchSize，最大不超过32。 |
| rsv[8] | uint8\_t | {0} | 预留参数。 |


#### 输入

| 参数 | 维度 | 数据类型 | 格式 | 描述 |
| --- | --- | --- | --- | --- |
| inTensor | [nSquareTokens] | float16 | ND | 输入tensor，是batch个(headNum, qSeqLen[i], qSeqLen[i])大小的矩阵按照ND排布拼接成一维的结果。 |


#### 输出

| 参数 | 维度 | 数据类型 | 格式 | 描述 |
| --- | --- | --- | --- | --- |
| outTensor | [nSquareTokens] | float16 | ND | 输出tensor，数据范围在[0, 1]之间。 |


#### 约束说明

- qSeqLen数组长度不超过32，且要求各元素大于0。
- 输入tensor（inTensor）、输出tensor（outTensor）的维度大小nSquareTokens与参数中的headNum和qSeqLen有关。须满足
  ：表示qSeqLen中的第k个元素。


#### 功能列表

当参数列表中

时，

其对应的输入tensor的shape大小为。

例如参数中设置


此时的输入tensor的shape大小为


输入和输出tensor的shape为[5528144]。
