flash attention的CUDA实现探讨-V4

之前关于flash attention的实现参考添加链接描述，添加链接描述，添加链接描述
lash attention的数学变换：给定三个矩阵Q,K,V，形状都是[N,d]，计算S=QK.T，然后针对dim=1做softmax，然后和V继续做矩阵乘法得到形状为[N,d]的输出矩阵O，即O=softmax(QK.T,dim=1)V。
下面本人的实现以下图为基础，从下图也可以看出attention算法的关键在于对Q,K,V的分块，因此我们需要首先考虑清楚编程过程中一个线程块到底处理哪些数据。

blockDim=(Br,Bc,1),gridDim=(N/Br, d/Bc, 1)

我们考虑基本的baseline实现，根据上面这个算法流程，我们设置线程块和线程网格为blockDim=(Br,Bc,1),gridDim=(N/Br, N/Bc, 1)，即一个block需要处理矩阵Q的Br行，矩阵V的Br列以及全部的矩阵K。也就是说，对于固定的block，处理的Qi的形状是[Br,d]