Deploy028 Flash attn

用户5190

用户3543

2024年11月19日修改

attention存在的问题：
的结果非常巨大，当sequence length增加时，显存占用以O(N^2)比例增加。受显存访问速度的限制，导致transformer大部分是显存访问瓶颈。因而flash attention系列通过一系列等价替换的计算方式减少数据从SRAM中搬运至HBM的次数，实现较优的运行效率。​

1.
GPU硬件背景介绍​

以下引自：FlashAttention2详解（性能比FlashAttention提升200%）

GPU主要包含计算单元（如浮点运算单元）和内存层次结构。大多数现代GPU包含专用的低精度矩阵乘法单元（如Nvidia GPU的Tensor Core用于FP16/BF16矩阵乘法）。GPU的内存由多个不同大小和不同读写速度的内存组成，内存越小，读写速度越快。内存层次结构分为高带宽内存（High Bandwidth Memory, HBM）和片上SRAM（也称为shared memory）。​

以A100 GPU为例：

•
HBM40-80GB，带宽为1.5-2.0TB/s，​

•
SRAM内存分布在108个流式多处理器上，每个处理器的大小为192K，合计为20M左右，带宽约为19TB/s​

common.docs_name - LarkCCM_Docs_Menu_Image

1.1
CUDA的软件和硬件架构​

50%

1.1.1
硬件架构​

从Hardware角度来看：

•
Streaming Processor（SP）：是最基本的处理单元，从fermi架构开始被叫做CUDA core。​

•
Streaming MultiProcessor（SM）：一个SM由多个CUDA core（SP）组成，每个SM在不同GPU架构上有不同数量的CUDA core，例如Pascal架构中一个SM有128个CUDA core。​

SM还包括特殊运算单元(SFU)，共享内存(shared memory)，寄存器文件(Register File)和调度器(Warp Scheduler)等。register和shared memory是稀缺资源，这些有限的资源就使每个SM中active warps有非常严格的限制，也就限制了并行能力。​

1.1.2
软件架构​

CUDA软件示例

•
thread：一个CUDA并行程序由多个thread来执行​
◦
thread是最基本的执行单元（the basic unit of execution）。​

•
warp：一个warp通常包含32个thread。每个warp中的thread可以同时执行相同的指令，从而实现SIMT（单指令多线程）并行。​
◦
warp是SM中最小的调度单位（the smallest scheduling unit on an SM），一个SM可以同时处理多个warp​

•
thread block：一个thread block可以包含多个warp，同一个block中的thread可以同步，也可以通过shared memory进行通信。​
◦
thread block是GPU执行的最小单位（the smallest unit of execution on the GPU）。​
◦
一个warp中的threads必然在同一个block中，如果block所含thread数量不是warp大小的整数倍，那么多出的那个warp中会剩余一些inactive的thread。也就是说，即使warp的thread数量不足，硬件也会为warp凑足thread，只不过这些thread是inactive状态，但也会消耗SM资源。​

•
grid： 在GPU编程中，grid是一个由多个thread block组成的二维或三维数组。grid的大小取决于计算任务的规模和thread block的大小，通常根据计算任务的特点和GPU性能来进行调整。​

1.2
硬件和软件的联系​

SM采用的是Single-Instruction Multiple-Thread（SIMT，单指令多线程）架构，warp是最基本的执行单元，一个warp包含32个并行thread，这些thread以不同数据资源执行相同的指令。​

当一个kernel被执行时，grid中的thread block被分配到SM上，大量的thread可能被分到不同的SM上，但是一个线程块的thread只能在一个SM上调度，SM一般可以调度多个block。每个thread拥有自己的程序计数器和状态寄存器，并且可以使用不同的数据来执行指令，从而实现并行计算，这就是所谓的Single Instruction Multiple Thread。​

一个CUDA core可以执行一个thread，一个SM中的CUDA core会被分成几个warp，由warp scheduler负责调度。GPU规定warp中所有thread在同一周期执行相同的指令，尽管这些thread执行同一程序地址，但可能产生不同的行为，比如分支结构。一个SM同时并发的warp是有限的，由于资源限制，SM要为每个block分配共享内存，也要为每个warp中的thread分配独立的寄存器，所以SM的配置会影响其所支持的block和warp并发数量。​

GPU执行模型小结：

GPU有大量的threads用于执行操作（an operation，也称为a kernel）。这些thread组成了thread block，接着这些blocks被调度在SMs上运行。在每个thread block中，threads被组成了warps（32个threads为一组）。一个warp内的threads可以通过快速shuffle指令进行通信或者合作执行矩阵乘法。在每个thread block内部，warps可以通过读取/写入共享内存进行通信。每个kernel从HBM加载数据到寄存器和SRAM中，进行计算，最后将结果写回HBM中。​

2.
标准attention读写次数​

50%

Deploy028 Flash attn​

Deploy028 Flash attn