阅读Linear Attention官方实现的理解 本人转载于大佬哦~~
然后,在 https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/linear-attention/causal_product_cuda.cu#L661-L689 对应的lmha模板函数的实现中有两种不同的kernel dispatch逻辑:
如果blocks(batch 和 注意力头的乘积 即params.B * params.H)小于LOW_OCCUPANCY_THRESHOLD(=40)的时候,走的是lmha_low_occupancy_
这个kernel的实现,否则就会走到lmha_
的实现。另外,还会根据 query 的特征维度的大小 E 来设置kernel不同的模板参数。【BBuf的CUDA笔记】十,Linear Attention的cuda kernel实现解析详细解析了lmha_
这个kernel的实现,这篇文章就来详解一下lmha_low_occupancy_
的实现。
lmha_low_occupancy_ kernel实现解析
我们先从理论上来解释一下这个kernel的取名,cuda中occupancy指的是一个SM中实际活跃的warp与理论上可以最高可以活跃的warp的比值,然后如果occupancy太低直接带来的影响就是GPU没有足够多的warp来切换,就无法隐藏数据加载/计算的延时,直接导致了kernel B的算力下降。而触发这个lmha_low_occupancy_
kernel的条件就是blocks < LOW_OCCUPANCY_THRESHOLD
,且LOW_OCCUPANCY_THRESHOLD=40
,我们想一下如果blocks比较小的情况下我们也dispatch到lmha_
这个kernel会发生什么?
对于lmha_
这个kernel,它的Block数量就是上面blocks,所以对这个kernel来说,它只能启动少于40个Block,对于V100来说有80个sm,对于A100则有120个sm,如果在这两种显卡上只启动不到40个Block则显然SM只能用上一半不到,会导致GPU存在大量资源浪费的现象。因此,当blocks的数量小于40的时候,Linear Attention的官方实现选择实现了另外一个lmha_low_occupancy_
来尽量增加Block的数量,减少sm资源浪费。
感觉这里的40并没有设得很好,没有充分考虑到各种GPU的sm数量,感觉可以优化成根据SM的数量来自动选择。
接着来看lmha_low_occupancy_
kernel的具体逻辑:
这里还多了一层,会根据query的特征维度E,key的特征维度M以及LOW_OCCUPANCY_THRESHOLD=40
来决定lmha_low_occupancy_
kernel的第三个模板参数,也就是warp的数量。继续往下看lmha_low_occupancy_
真正的kernel启动部分:
再明确一下params的几个维度,批量大小是 B、头数是 H、序列长度是 L 和query的特征维度 E,key的特征维度M 。一般来说E和M是相等的。然后这里有个条件判断,当H和B有一个大于65536的时候就返回1,不执行这个kernel。都则就开一个三维线程网格,大小为(M, H, B),也就是有 B * H * M
这么多个 Block,而每个Block的线程数量为 WARPS*THREADS_PER_WARP
。其中WARPS是模板参数表示启动Kernel时用到多少个warp,由query的特征维度E,key的特征维度M以及LOW_OCCUPANCY_THRESHOLD=40
来共同决定,而THREADS_PER_WARP=32
表示一个warp有32个线程。
接下来就来到真正的cuda kernel实现了,代码和解释如下,这里假设E和M都是1024,WARPS=4:
请详细阅读上面的代码和注释,然后这里对这个kernel的核心计算逻辑做一个总结。
- 使用 warp 作为基本的计算单元,每个 warp 处理特定的数据列。通过 THREADS_PER_WARP 和 WARPS 参数来控制每个 block 的线程数量。然后,每个线程加载加载它对应的Q,K,V矩阵部分。并使用
offset_q
,offset_k
, 和offset_v
来计算每个线程的数据偏移量。对应:https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/linear-attention/causal_product_cuda.cu#L152-L310 - 在每个 warp 内部进行前缀和计算以获得部分 K*V 乘积。这是通过逐行累加来完成的。https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/linear-attention/causal_product_cuda.cu#L152-L261 ,这里对V的读取使用了共享内存。
- 使用共享内存进行跨 warp 的规约操作,以计算 K*V^T 的总和。https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/linear-attention/causal_product_cuda.cu#L263-L310
- 利用
__shfl_xor_sync
指令进行 warp 内部的并行规约,以合并每个 warp 的计算结果。https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/linear-attention/causal_product_cuda.cu#L312-L319 - 最终的输出结果被写回到共享内存,然后存储到全局内存中。https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/linear-attention/causal_product_cuda.cu#L321-L349
这里涉及到的技能主要是使用warp(32个线程)为基本单位来处理这个任务,Linear Attention的cuda kernel实现解析 中的lmha_kernel
以单个线程为单位。也就是说对于之前的lmha_kernel
来说,每个线程都需要完整的算一遍K(T*M
个元素)和V(T)的外积,再和Q(T*E
,E一般等于M)进行内积计算,由于这个kernel里面开了比较多的线程round_up(max(E, M*THREADS_PER_HEAD), 32);
,并且Block数量为B*H一般也是多于SM数量的,这样算是容易打满GPU的。而对于这里的lmha_low_occupancy_kernel
kernel来说,T维度已经被切开了,所以应当尽量减少线程的数量让每个线程做尽量多的工作来避免频繁的线程切换开销,所以这里使用warp为单位来处理。这种技巧也是应用得比较多了,例如oneflow的softmax和layernorm算子优化中,对于列数比较小的矩阵就是采用一个warp处理一行或者两行这种技巧。
可以把这里介绍的cuda kernel模型抽象成向量外积加内积,我们可以轻易的把它修改为单纯外积,内积用于我们自己的场景。此外,这个lmha_low_occupancy_kernel
并没有做数据向量化读取,double buffering等等,可能还有进一步优化空间。
总结
Linear Attention的cuda kernel实现解析就是阅读Linear Attention官方实现的理解。欢迎关注 https://github.com/BBuf/how-to-optim-algorithm-in-cuda 获取更多后续cuda优化相关的知识。写一个这种cuda kernel难度是挺大的,可以帮助更多的没有很好基础的读者入门cuda kernel开发