作业描述及代码参见:cs149gpt
Warm-Up:访问张量
张量/数组都是按行存储的,四维数组可以看作元素为三维数组的数组,元素大小即为三维数组内元素总数,以此类推。
第 1 部分:简单(但不太高效)的注意力机制实现
主要实现两个矩阵乘法和一个 softmax 运算。
第 2 部分:块矩阵乘法和 Unfused Softmax
通过对矩阵进行分块,有效提高缓存的利用率,减少 miss。
N=1024 时,块大小设为 \(8 \times 8\)、\(16 \times 16\)、\(32 \times 32\)、\(64 \times 64\) 对应的时间为 \(185.069ms\)、\(154.328ms\)、\(159.785ms\)、\(170.758 \text{ms}\),缓存行的大小为 \(64B\),float
类型的大小为 \(4B\),因此块大小为 \(16 \times 16\) 时,一个缓存行恰好可以装下块内的一行。
第 3 部分:Fused Attention
Fused Attention 使得 \(N \times N\) 的临时矩阵减小为长度为 \(N\) 向量,虽然因为多线程要使用的多个长度为 \(N\) 向量,但线程数往往远小于 \(N\),因此减少了内存占用。
注释掉 #pragma omp ...
语句后,时间为 \(217.351 \text{ms}\)。在单线程的情况下,Fused Attention 虽然减小了内存占用,但性能有所降低,同时对缓存的利用率也比第 2 部分的低很多。
与第 1 部分相比,由于使用 Fused Attention,可并行化的循环变成了 \(3\) 个:batch、head、row 易于并行计算。
第 4 部分:Flash Attention
第 4 部分的内存使用最少。Flash Attention 是对 Fused Attention 的改进,在减少内存使用的同时,通过分块,提高了对缓存的利用。第 4 部分的性能比之前各部分要慢,但比单线程的 Fused Attention 稍快。
目前的 Flash Attention 可以通过使用多线程、使用 CPU 向量化硬件单元等方式提高性能。
ISPC加速
使用 ISPC 对各部分加速效果如下:
- part1:运行时间在 \(60 \text{ms}\) 左右
- part2:运行时间在 \(85 \text{ms}\) 左右
- part3:运行时间大致在 $$15-55 \text{ms}$$
- part4:运行时间在 \(55 \text{ms}\) 左右
主要使用 ISPC 加速了矩阵乘法,实现主要参考 ISPC Examples。
标签:缓存,Fused,--,Assignment,Attention,times,NanoGPT149,ms,text From: https://www.cnblogs.com/bienboy/p/18466165