Warm-Up: Accessing Tensors (3 Points)

想象成多个立方体，x索引第几个立方体，y索引立方体的第几层平面，z索引平面的第几行，b索引行内第几列。

这样排布方便顺序访问。

Part 1: A Simple (But Not So Efficient) Implementation of Attention (10 Points)

Part 2: Blocked Matrix Multiply and Unfused Softmax (20 Points)

在这里踩了一个坑就是第一个QxK的时候这个K的base应该也是独立的所以就是Q有一个BaseX，K有一个BaseY，然后是块乘的中间BaseInner，然后里面就是对Q的i，对K的j，然后是乘法的中间迭代k。所以也挺好记了。

A1:因为tileSize=16的时候，刚好是一个tile的一行对应一行cache line，相比其他的，cache利用率会更高

该代码不能跑inference

主要是每个b,h都复用QK_t所以要清零，以及这个j应该跟d对应。

A2:

part1:

every 16 float has 1 memory access.

for an element of result, it has load of Q $\frac{d}{16}$$\frac{d}{16}$times of reading. the same for K, so total

of reading.

and the reading of result is $n^2$$n^2$ times, also writing for $n^2$$n^2$ times, because everytime it reads a line into cache, it becomes evicted by the loading of the coming K,V arrays.

part2:

assuming tileSize = 16

for every 16*16 of element, it has load of Q

and the same for K

so total

of reading.

the reading times of the result is

the evict or write back times is the same(it would not be evicted by the loading of Q and K)

so total part1:

total part2:

if d is very large, ratio approximately is 16x.

lets expand to normal conditions

$tileSize$$tileSize$, $cacheLineSize$$cacheLineSize$

我们先忽略写的部分，因为乘积里面没有d所以是小项

但是tile不是越大越好，如果一个tile不能容纳进一个cache，那就存在evict。此时就存在退化趋近于无分块的性能。

and above, the writeback of evcition has'nt been considered yet.

but we can think it as the same.

Part 3: Fused Attention (25 Points)

对Q的i行，跟K的每一行进行点乘，得出一个行向量，然后进行softmax，然后这一行跟V的所有列点乘就可以得到O的i行了。

A1

Because after fusion, we only needs to store Cxn of temp that stores $Q\ast K^T$$Q*K^T$, the C depends on the parallelism of openmp

A2

when not openmp accelarated it slows down

fused attention has row-granularity, each thread can compute one row of the output,

but part1 is a head-granularity, each thread can only compute one head.

Note that when computing one row $Q{K}^T$$QK^T$ with $V$$V$, the loop order can much influence cache behaviour

the upper one total loads Nxd of V, and the down is $N\frac{d}{cachelinesize}$$N\frac{d}{cachelinesize}$

so we can exceed the reference performance, optimizing performance.

Part 4 : Putting it all Together - Flash Attention (35 Points)

在实现的时候各种神秘情况没有考虑到导致调了很久最后让GPT看bug才调好。。

bug包括但不限于：

1. 将N分块时不考虑最后一块，即Tc=N/Bc but corrected is Tc=(N+Bc-1)/Bc

2. 行索引乘回乘到数量而不是大小。即rowIdx = i * Tr but corrected is rowIdx = i * Br

3. 四维数组用二维写。

4. 不考虑最后一块时，块内读写的索引上界

optimized way?:

fusion:

1. 大量的即时计算，比如计算出块内元素就将其乘到PV中。属于一种fusion了

2. 计算出一个l元素就将其与PV累加到O，这样O就直接出来了。成立的原因是当行计算完毕，用的PV也是行级别的。那么行PV已经计算完毕

然后居然就可以达到将近3x的speedup

A1

less than part1,2 similar but less than part3.

part3 mainly has threadNum * N float space

part4 mainly has Const+Const*d + N float space, so maybe less than part3.

A2

NO.

similar to part3, we can parallyze batches and heads loop.

and we can also parallel the row blocks.

因为算法是先取一个K块，然后对所有行计算O

然后随着K块步进取，最后才计算出最后正确的O

那么每个块内的行O就可以并行。

Other Optimizations

1 Matrix Multiply loop orderings to be cache-friendly

优化前（即上面那个not cache-friendly的实现）：

优化后：

Extra Credit: Optimize Further (12 Total Points - 3 Points Per Part)

Part1

加上Other Optimizations对$Q{K}^{T}V$$QK^{T}V$的顺序的优化，一共优化了3处：

1. $Q{K}^T$$QK^{T}$的点积SIMD

2 $SoftMax$$SoftMax$的SIMD

3. $Q{K}^T$$QK^{T}$与$V$$V$的SIMD

这样做由于粒度太小导致overhead增加。

这样粗粒度会更好。每个lane执行的工作多一点。同时用uniform限制j保证每个lane在同一个时间只加载j行对应lane列的V

这样，假设O/V被分为tile个列SIMD块

对每列SIMD块每个lane会计算行-列对应的完整输出

然后才进到下一个SIMD块

最上层是行级iteration。

将part1的uniform float改成每个lane的，最后reduce，然后由于一次启动有overhead所以让他一次计算一行。

Part1到这里就结束了。

Part2

1. $Q{K}^T$$QK^{T}$的SIMD

块内SIMD即可。不过粒度还是有点细。如果是粗粒度的话可能要维护一个 $tileSize\ast tileSize$$tileSize * tileSize$的中间矩阵。即每次调用SIMD kernel的时候计算出这整个tile的结果

2. $SoftMax$$SoftMax$的SIMD

与Part1完全一致

3. $Q{K}_{T}V$$QK_{T}V$的SIMD

与Part1的cache优化一致。然后再进行SIMD

类似的，复习了一下ISPC的semantics。

每到一节loop的时候，每个lane作为一个instance，负责对应的work

比如loop_i=0的时候，如果有8个SIMD-lane的话，那么0 lane的programIndex就是0，1就是1，类推。

然后类似于分块，进行里面的工作。

所以这样再看上面的这个解法，就是将输出的列按SIMD宽度进行分块，每个宽度内计算这个tile的tile输出。相比于没有分tile的乘法仅仅是变成了部分和，而没有tile的做法是这个SIMD列块一次性计算出整个列块元素结果的。

而且访存其实是按块访存，而且是不连续的。GPT骗我说什么连续的。实际上就是不懂的人描述出来就不懂了结果大家都不懂了。

不过也可能GPT理解成为，每次外部foreach的时候，也就是类似于一整个SIMD lane在同时加载数据的时候是否连续吧。

而循环间连续的要求属于访存优化的属性了。

让我们看看列级别accumulate和考虑行cache连续性的区别

列accumulate（即当前版本）

行连续访问

所以看来确实是ispc启动的overhead更大一些，毕竟外部循环都是N*N了（）

Part3

似乎跟前两个part很相似

第一个QK几乎一样的

然后后面那个的策略也是跟part1一样

因为本身fused就相当简单吧

Part4