SoftmaxFlashV2

产品支持情况

产品

是否支持

Ascend 950PR/Ascend 950DT

√

Atlas A3 训练系列产品 / Atlas A3 推理系列产品

√

Atlas A2 训练系列产品 / Atlas A2 推理系列产品

√

Atlas 200I/500 A2 推理产品

√

Atlas 推理系列产品 AI Core

√

Atlas 推理系列产品 Vector Core

x

Atlas 训练系列产品

x

功能说明

SoftmaxFlash增强版本，对应FlashAttention-2算法。将输入tensor[m0, m1, ...mt, n]（t大于等于0）的非尾轴长度相乘的结果看作m，则输入tensor的shape看作[m, n]。对输入tensor[m,n]按行做如下计算，不同的update值对应不同的计算公式，其中x、inmax和insum为输入，M、S、E均为输出。

当输入shape为ND格式时，内部的reduce过程按last轴进行；当输入shape为NZ格式时，内部的reduce过程按照last轴和first轴进行，reduce过程可以参考SoftMax中的图示说明。

为方便理解，通过Python脚本实现的方式，表达其计算公式如下，其中src、inmax、 insum、update为输入，dst、x_sum、x_max、exp_max为输出。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
 
def softmax_flash_2(src, inmax=None, insum=None, update=None):
    if update is None:
        x_max = np.max(src, axis=-1, keepdims=True)
        x_sub = src - x_max   
        dst = np.exp(x_sub) 
        x_sum = np.sum(dst, axis=-1, keepdims=True)
        exp_max = None
        return dst, x_max, x_sum, exp_max
    else:
        x_max = np.max(np.concatenate((inmax, src), axis=-1), axis=-1, keepdims=True)
        dst = np.exp(src - x_max)
        exp_max = np.exp(inmax - x_max)
        x_sum = np.sum(dst, axis=-1, keepdims=True)
        x_sum = exp_max * insum +  x_sum
        return dst, x_max, x_sum, exp_max

实现原理

以float类型，ND格式，shape为[m, k]的输入Tensor为例，描述SoftmaxFlashV2高阶API内部算法框图，如下图所示。

图1 SoftmaxFlashV2算法框图

计算过程根据isUpdate是否开启分为两个分支处理，均在Vector上进行。

函数原型

由于该接口的内部实现中涉及复杂的计算，需要额外的临时空间来存储计算过程中的中间变量。临时空间支持接口框架申请和开发者通过sharedTmpBuffer入参传入两种方式。

接口框架申请的方式，开发者需要预留临时空间；通过sharedTmpBuffer传入的情况，开发者需要为tensor申请空间。临时空间大小BufferSize的获取方式如下：通过SoftmaxFlashV2 Tiling接口中提供的GetSoftMaxFlashV2MinTmpSize/GetSoftMaxFlashV2MaxTmpSize接口获取所需最小和最大临时空间大小，最小空间可以保证功能正确，最大空间用于提升性能。

另外提供了一个kernel侧tiling计算的接口，当kernel侧的输入shape与通过host侧TilingData传入的shape不一致时，可使用该接口在kernel侧重新计算tiling。该接口的参数含义请参考SoftmaxFlashV2 Tiling接口。

参数说明

表1 模板参数说明

参数名

描述

T

操作数的数据类型。支持的数据类型为：half、float。

isUpdate

是否开启update部分中的计算。

isReuseSource

该参数预留，传入默认值false即可。

isBasicBlock

srcTensor和dstTensor的shape信息和Tiling切分策略满足基本块要求的情况下，可以开启该参数用于提升性能，默认为false表示不开启。是否满足基本块的要求，可以采用如下两种方式之一判断：

  • srcTensor和dstTensor的shape信息[m,n]需要满足如下条件：
    • 尾轴长度n小于2048并且大于等于256/sizeof(T)（即half场景下n最小为128，float场景下n最小为64），同时n是64的倍数；
    • 非尾轴长度的乘积m为8的倍数。
  • 在Tiling实现中，通过调用IsBasicBlockInSoftMax判断Tiling切分策略是否满足基本块的切分要求。

针对 Atlas 200I/500 A2 推理产品 ，该参数为预留参数，暂未启用，为后续的功能扩展做保留，保持默认值即可。

isDataFormatNZ

当前输入输出的数据格式是否为NZ格式，默认数据格式为ND，即默认取值为false。

针对 Atlas 200I/500 A2 推理产品 ，不支持配置为NZ格式。

config

结构体模板参数，此参数可选配，SoftmaxConfig类型，具体定义如下方代码所示，其中参数的含义为：

  • isCheckTiling：是否需要检查shape和tiling的一致性；若不一致，API内会根据shape重新计算所需tiling。默认取值true：API内部会检查一致性。
  • oriSrcM：原始非尾轴长度的乘积。设置该参数后，将shape常量化，编译过程中使用常量化的shape。
  • oriSrcK：原始尾轴长度。设置该参数后，将shape常量化，编译过程中使用常量化的shape。
  • mode：输出shape的处理模式。当输入输出的数据格式为NZ格式时，不支持配置mode参数。SoftmaxMode类型，取值如下：
    • SOFTMAX_NORMAL ：默认值，常规模式，对输出数据做Broadcast，使得输出shape由(m, 1)拓展成(m, 8)（输出为float数据类型）或者(m, 16)（输出为half数据类型）。
    • SOFTMAX_OUTPUT_WITHOUT_BRC ：非拓展模式，不对输出数据做Broadcast，输出shape均为(m, 1)，相应的输入参数（例如inExpSumTensor、inMaxTensor），shape也均为(m, 1) 。

此参数一般用于配合kernel侧tiling计算的接口使用。

注意：设置了oriSrcM与oriSrcK后，模板参数isBasicBlock不生效，计算数据是否为基本块由API内部判断并处理。

针对Ascend 950PR/Ascend 950DT，支持该参数。

Atlas A3 训练系列产品 / Atlas A3 推理系列产品 ，支持该参数。

Atlas A2 训练系列产品 / Atlas A2 推理系列产品 ，支持该参数。

针对 Atlas 200I/500 A2 推理产品 ，该参数为预留参数，暂未启用，保持默认值即可。

Atlas 推理系列产品 AI Core，支持该参数，不支持配置mode。

 
1
2
3
4
5
6
 
struct SoftmaxConfig{
    bool isCheckTiling = true; 
    uint32_t oriSrcM = 0;
    uint32_t oriSrcK = 0;
    SoftmaxMode mode = SoftmaxMode::SOFTMAX_NORMAL;
};

配置示例如下。

1
 
constexpr SoftmaxConfig SOFTMAX_DEFAULT_CFG = {true, 0, 0, SoftmaxMode::SOFTMAX_NORMAL};
表2 接口参数说明

参数名

输入/输出

描述

dstTensor

输出

目的操作数。

类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。

dstTensor的shape和源操作数srcTensor一致。

outReduceMax

输出

目的操作数。用于保存softmax计算过程中reducemax第一次计算的结果。

类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。

outReduceMax的shape与目的操作数maxTensor一致。

对于输出该结果的接口：

  • 模板参数isUpdate为false时，不输出该结果。
  • 仅支持输入输出的数据格式为ND，模板参数isDataFormatNZ为预留参数，传入默认值false即可。
  • 模板参数config.isCheckTiling为预留参数，传入默认值false即可。
  • 模板参数config.mode仅支持配置为非拓展模式SoftmaxMode::SOFTMAX_OUTPUT_WITHOUT_BRC。

expSumTensor、outExpSum

输出

目的操作数。用于保存softmax计算过程中reducesum的结果。

类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。

  • 除模板参数config配置为非拓展模式（SoftmaxMode::SOFTMAX_OUTPUT_WITHOUT_BRC）的场景外，expSumTensor的last轴长度固定为32Byte，即一个datablock长度。该datablock中的所有数据为同一个值，比如half数据类型下，该datablock中的16个数均为相同的reducesum的值。
  • 非last轴的长度与dstTensor保持一致。

maxTensor、outMax

输出

目的操作数。用于保存softmax计算过程中reducemax的结果。

类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。

  • 除模板参数config配置为非拓展模式（SoftmaxMode::SOFTMAX_OUTPUT_WITHOUT_BRC）的场景外，maxTensor的last轴长度固定为32Byte，即一个datablock长度。该datablock中的所有数据为同一个值。比如half数据类型下，该datablock中的16个数均为相同的reducemax的值。
  • 非last轴的长度与dstTensor保持一致。

srcTensor

输入

源操作数。

类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。

last轴长度需要32Byte对齐。

expMaxTensor、outExpMax

输出

目的操作数。用于保存inmax与reducemax差值的e的指数幂的结果。

类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。

  • 除模板参数config配置为非拓展模式（SoftmaxMode::SOFTMAX_OUTPUT_WITHOUT_BRC）的场景外，expMaxTensor的last轴长度固定为32Byte，即一个datablock长度。该datablock中的所有数据为同一个值，比如half数据类型下，该datablock中的16个数均为相同的值。
  • 非last轴的长度需要与dstTensor保持一致。

inExpSumTensor、inExpSum

输入

源操作数。softmax计算所需要的sum值。

类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。

  • 除模板参数config配置为非拓展模式（SoftmaxMode::SOFTMAX_OUTPUT_WITHOUT_BRC）的场景外，inExpSumTensor的last轴长度固定为32Byte，即一个datablock长度。该datablock中的所有数据为同一个值，比如half数据类型下，该datablock中的16个数均为相同的值。
  • 非last轴的长度需要与dstTensor保持一致。

inMaxTensor、inMax

输入

源操作数。softmax计算所需要的max值。

类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。

  • 除模板参数config配置为非拓展模式（SoftmaxMode::SOFTMAX_OUTPUT_WITHOUT_BRC）的场景外，inMaxTensor的last轴长度固定为32Byte，即一个datablock长度。该datablock中的所有数据为同一个值，比如half数据类型下，该datablock里的16个数均为相同的值。
  • 非last轴的长度需要与dstTensor保持一致。

sharedTmpBuffer

输入

临时空间。

类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。

该操作数的数据类型固定uint8_t。

接口内部复杂计算时用于存储中间变量，由开发者提供。

临时空间大小BufferSize的获取方式请参考SoftmaxFlashV2 Tiling接口。

tiling

输入

softmaxflashv2接口计算所需tiling信息，Tiling信息的获取请参考SoftmaxFlashV2 Tiling接口。

softmaxShapeInfo

输入

srcTensor的shape信息。SoftMaxShapeInfo类型，具体定义如下方代码所示，其中参数的含义为：

  • srcM：非尾轴长度的乘积。
  • srcK：尾轴长度，必须32Byte对齐
  • oriSrcM：原始非尾轴长度的乘积。
  • oriSrcK：原始尾轴长度。

需要注意，当输入输出的数据格式为NZ格式时，尾轴长度为reduce轴长度即图2中的W0*W1，非尾轴为H0*H1。

 
1
2
3
4
5
6
 
struct SoftMaxShapeInfo {
  uint32_t srcM; 
  uint32_t srcK; 
  uint32_t oriSrcM;  
  uint32_t oriSrcK;  
};

返回值说明

无

约束说明

调用示例

父主题： SoftMax接口