实现 FlashAttention-2 前向

题目描述

实现 FlashAttention-2 的前向计算（单头，简化版）。给定 Q（seq_len × d）、K、V，用分块 + 在线 softmax计算输出，不实例化完整的 seq_len × seq_len 注意力矩阵。

要求

• 实现右侧代码模板中给出的函数，签名以模板为准
• 使用 online softmax，维护 running max 和 running sum
• 缩放因子为

提示

FlashAttention-2 相比 v1 减少了非矩阵乘的 rescale 操作，并把并行维度放在 seq_len 上。