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

news2024/11/29 12:41:18

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) 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;
        }
    }
}