flash attention的CUDA编程（一个线程处理多个元素）-V5

本文接着之前的文章介绍添加链接描述

单线程处理多元素

之前文章提到的算法都有一个特点，即一个(threadIdx.x,threadIdx.y)只计算输出O矩阵的一个元素，在内部循环计算matmul的时候也是如此，循环体内部的两个矩阵乘法也严格遵循了一个(threadIdx.x,threadIdx.y)计算输出矩阵的一个元素，但是在里面的matmul里面，为了计算一个输出，需要访问A,B矩阵的一行一列，这样就导致计算访存比非常小，为了加速内部的矩阵乘法，我们希望能够做出修改，一个线程同时计算多个元素，具体操作如下：
在原本的算法里面，对于固定某个线程块，也就是说固定好blockIdx.x,blockIdx.y以后，该线程块内的threadIdx.x处理的是矩阵V的某一列元素
V[:,threadIdx.x+blockIdx.x×blockDim.x]，
threadIdx.y处理的是矩阵Q的某一行元素
Q[threadIdx.y+blockIdx.y×blockDim.y,:]，
循环体内threadIdx.x处理的是矩阵K的某一行元素
K[threadIdx.x+j×Bc,:]，
现在我们希望对于固定的线程块，一个threadIdx.y可以处理Q的Rq行，一个threadIdx.x可以处理V的Rv列，循环体内部也是如此。具体来说，此时线程块设置仍然是(Bc,Br,1)，但是线程网格是((d/(Rv×Bc),N/(Rq×Br),1)，那么对于固定的blockI