flash attention的CUDA编程和二维线程块实现softmax

本文参考了链接添加链接描述

flash attention介绍

flash attention的介绍可以参考论文:FlashAttention: Fast and Memory-Efficient Exact Attention
with IO-Awareness,具体的数学公式参考下面这个图片:其中注意关于矩阵S有两个维度,softmax的操作维度是dim=1,用pytorch表示就是torch.softmax(S, dim=1)
flash attention的CUDA编程和二维线程块实现softmax_第1张图片
对于flash attention来说,里面有两次矩阵乘法,对于这样的二维数组矩阵乘法,一般来说都会考虑使用二维线程块,但是我们之前实现的softmax都是以一维线程块来处理,其中专门用到了一个cub库的函数BlockReduce,经过本人测试,发现这个函数只能针对一维线程块做线程块内部的规约,不能用于二维线程块内部针对某个维度规约,因此在实现flash attention之前,我们需要编写一个二维线程块实现softmax的算法,其中注意BLOCK_DIM_x和BLOCK_DIM_y都必须要选取2的幂次方。

二维线程块实现softmax

之前我们实现一维线程块处理softmax的时候,参考链接添加链接描述

你可能感兴趣的:(高性能计算,算法)