分享
Deploy028 Flash attn
输入“/”快速插入内容
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为例:
•
HBM
40-80GB,带宽为1.5-2.0TB/s,
•
SRAM内存
分布在108个流式多处理器上,每个处理器的大小为192K,合计为20M左右,带宽约为19TB/s
1.1
CUDA的软件和硬件架构
50%
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%
50%