From self-attention 2 flash-attention 数学原理与 cuda 实现优化

self attension 是transformer 编码器和解码器中共同的一个计算环节，在整个transformer 网络体系中耗费的算力比例占主导。所以节省self attention 的正向和反向的计算时间，就可以加速 transormer 的训练和推理过程。

1，self attention 的数学提炼

两个矩阵乘法，加入一个列向的softmax

input   矩阵：

output 矩阵：

        step1:       

        step2:       

        step3:       

2，cpu 实现self attention

这里的数据类型使用了 float，实际网络中一般采用 fp16，数学过程是相同的；

cpu_self_attention.cpp

#include <stdio.h>
#include <string.h>

#include "cpu_gemm.h"
#include "utils.h"
#include "soft_max.h"
//all matrices are row major.

void cpu_self_attention(float* Q, int ldq,
						float* K, int ldk,
						float* V, int ldv,
						float* S, int lds,
						float* P, int ldp,
						float* O, int ldo,
						int N, int d)
{
	gemm_nt(Q, ldq, K, ldk, S, lds, N, N, d);// S = Q*K^t     (NxN) = (Nxd) * (dxN)
					printf("\nS =\n");	print_matrix(S, N, N, lds);
	soft_max_column(P, ldp, S, lds, N, N);// P(NxN) = softmax(S(NxN))
					printf("\nP =\n");	print_matrix(S, N, N, lds);
	gemm_nn(P, ldp, V, ldv, O, ldo, N, d, N);// O = P*V     (Nxd) = (NxN) * (Nxd)
}

cpu_gemm.cpp

#include "cpu_gemm.h"

void gemm_nn(float *A, int lda,		//A(M x K) rowMj
	     	 float *B, int ldb,		//B(K x N) rowMj
	     	 float *C, int ldc,		//C(M x N) rowMj
	      	 int M,
			 int N,
			 int K)
{
	for(int i=0; i<M; i++)
	{
		for(int j=0; j<N; j++)
		{
			float sigma = 0.0;

			for(int k=0; k<K; k++)
			{
				sigma += A[i*lda + k] * B[k*ldb + j];
			}

			C[i*ldc + j] = sigma;
		}
	}
}

void gemm_nt(float *A, int lda,		//A(M x K) rowMj
	     	 float *B, int ldb,		//B(N x K) rowMj
	     	 float *C, int ldc,		//C(M x N) rowMj
	      	 int M,
			 int N,
			 int K)
{
	for(int i=0; i<M; i++)
	{
		for(int j=0; j<N; j++)
		{
			float sigma = 0.0;

			for(int k=0; k<K; k++)
			{
				sigma += A[i*lda + k] * B[k + j*ldb];
			}

			C[i*ldc + j] = sigma;
		}
	}
}

cpu_softmax_column.cpp

这里使用的是未数值优化的方式，直接按照原始公式计算：

#include "soft_max.h"
void soft_max_column(float *P, int ldp, float* S, int lds, int M, int N)//P = softmax(S)  P(i,j) = exp(S(i,j))/sigma(exp(S(r,j)));  r=0,1,..,n-1 ;
{
    for(int j=0; j<N; j++){
        float sigma = 0.0f;

        for(int i=0; i<M; i++){
            sigma += exp(S[i*lds + j])
        }

        for(int i=0; i<M; i++){
            P[i*ldp + j] = S[i*lds + j]/sigma;
        }
    }
}

3， gpu 实现 self attention 正向

cuda 实现上述过程：

gpu_self_attention.cu

gpu_gemm.cu

gpu_softmax_column.cu

4，为什么不需要gpu 实现self attention 反向

融合上述过程

5， gpu 实现 flash attention 反向

融合算子

数学原理

cuda 实现

挖坑，未完待续 。。。