Amazon LLM：FlashAttention 为什么省显存，但不是线性 Attention？

题目重点：分块与在线softmax合并，避免完整score/probability写回，解释精确性、重计算与测量范围。

题意与假设

这道练习题保留了面经中值得训练的核心约束，并将需要澄清的条件单独列出。现场拿到变体时，先和面试官确认这些条件。

朴素QKᵀ把T×T score写显存，再读softmax、写probability，再读乘V，多次大矩阵IO。分块将K/V和query tiles尽量在快存计算，避免完整T²中间矩阵在显存物化。
对单query在线保m=已处理score最大值、l=Σexp(score−m)、o=Σexp(score−m)V。新块max为b，更新m'=max(m,b)、l'=exp(m−m')l+Σ新exp(score−m')，o'同样重缩旧o再加新加权V，最终o/l。
旧和按新max重缩是关键：不能每块各softmax后简单平均输出，因为块内归一化丢了全局概率质量。mask也在对应tile保留，future/padding不能重新变成候选。
反传可重计算部分score/probability换取少存大中间量；实际性能受tile大小、head维、长度、dtype、硬件与调度影响。数学exact不是bitwise一致，也不等于每配置都有同样加速。
密集所有允许token对仍被处理，算术主项O(T²d)；减少IO与峰值中间内存不等于把算法近似稀疏化。滑窗/低秩/线性attention改变交互或目标，应另讨论精度/可达性。
测量分prefill/训练/decode、同步计时、warmup、batch/长度与峰值allocated/reserved，比较同精度/同mask输出与梯度。短decode只有一条query时与长prefill瓶颈不同，不用一个benchmark推全场景。

自拟scores1000/1001、V2/4，稳定权重按exp(0)/exp(1)，输出(2+4e)/(1+e)约3.4621；分别每块只有一项时各softmax均1，简单平均得到3，错误。在线更新旧块质量后能恢复正确全局值。

dense算术仍O(T²d)，避免显存中的完整score/probability矩阵，使相关中间存储随T线性加tile快存；总峰值还含模型/激活/KV，不是所有显存都只O(T)。确切IO式依快存容量/tiling与实现条件。