op融合
原始方法：
痛点：多次读取、写入显存。
解决：中间结果不保存，1个kernel顺序算完多个操作。
反向传播时用到这些中间结果要求导，怎么办？
答：类似activation checkpointing，重新计算一遍。
通过矩阵的分块计算：1. 达到了融合多个操作，减少了IO加快了速度；2. 消除了句子长度平方项矩阵，减少了显存占用。
计算步骤
K、V是外循环，Q、O是内循环。
1.1
1.2
1.3
2.1
2.2
2.3
S的softmax怎么办
为了不让e^x溢出，safe_softmax的做法是，分子分母同时除以e^m，m为最大的x：
每次重新做归一化即可：
l和m占用显存非常小。
原文算法：
M/4d里的4怎么来的？
SharedMemory里，只需要保存4个子矩阵：Q、K、V、O。（S*V可以用乘加来跳过中间结果矩阵的存储。但是S矩阵不需要占空间吗？）IO分析：
原始方法，读写几次Q、K、V、O矩阵，复杂度是O(Nd)；读写各1次S矩阵，复杂度是O(N^2)。
FlastAttention，K、V矩阵大小是Nd，SharedMemory大小是M，则外层循环次数是Nd/M；内存循环读写Q、O矩阵，矩阵大小是O(Nd)；总的读写复杂度是:前者*后者=O(Nd/M*Nd)=O(N^2*d^2/M)
因为，Multi-head attention里，一般d是64或128，M一般是100KB，因此d^2/M远小于1，因此，FlashAttention的复杂度远小于原始方法。
FlashAttention-2
第3条，说明训练时，也要加Mask的，句子里，靠前面的token的Q，不能和靠后面token的K和V进行计算。