TikTok FlashAttention：分块在线 Softmax 与 IO 成本

题目重点：保持精确attention目标，解释跨块max/分母/分子重标定，不把显存优化当线性计算。

题意与假设

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

普通流程把QKᵀ写HBM，再softmax读写权重，最后乘V，n²中间张量增加显存与带宽。分块将部分Q/K/V载入快速存储，直接累计输出。
每query维护running max m、指数和l、加权分子o。新块最大值改变到m_new时，把旧l/o乘exp(m_old−m_new)，再加该块exp(score−m_new)贡献。
重标定不可漏：不同块各自softmax再简单平均不是全局softmax，因为块分母/最大值不同。最后输出o/l实现跨全部合法keys的归一化。
causal/padding mask在score贡献前应用，整行遮蔽仍需规定；分块边界不能让未来key参加分母。dropout/backward需保存或重建一致随机信息。
避免存全部attention概率，反向可重算局部scores而用额外算术换显存IO；并不是剪掉低权重key或近似核，所以与linear attention的目标不同。
实际加速取决于硬件、tile、dtype、序列长和head维；小shape或其他瓶颈不保证收益。应用应对照完整reference输出/梯度、峰值显存和端到端吞吐。

第一块score较低、第二块较高，新的max增大后必须把第一块旧指数权重按exp(old−new)缩小；各块独立归一化并拼接会让两块各自总权重1，最终错误。

经典dense算术仍O(n²d)，避免O(n²)权重驻留，工作张量常线性于n加块存储；IO规模依片上存储和tile，不能宣称所有维度无条件O(n)。