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

本文主要是介绍From self-attention 2 flash-attention 数学原理与 cuda 实现优化，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

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

1，self attention 的数学提炼

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

input 矩阵： $\mathbf{Q}, \mathbf{K}, \mathbf{V} \in \mathbf{R}^{N \times d}$

output 矩阵： $\mathbf{O} \in \mathbf{R}^{N \times d}$

$\mathbf{self\ attention\ algorithm:}$

step1: $\mathbf{S} = \mathbf{Q}*\mathbf{K}^t$

step2: $\mathbf{P} = \mathbf{softmax_{column}(S)}$

step3: $\mathbf{O} = \mathbf{P}*\mathbf{V}$

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) rowMjfloat *B, int ldb,		//B(K x N) rowMjfloat *C, int ldc,		//C(M x N) rowMjint 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) rowMjfloat *B, int ldb,		//B(N x K) rowMjfloat *C, int ldc,		//C(M x N) rowMjint 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;}}
}