实现 FlashAttention-2 前向

题目描述

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

其中 是因果掩码：被屏蔽的位置为 ，即每个 token 只能看到自己和它之前的 token —— 这是自回归语言模型的基本约束，不加掩码会让训练时"看到未来"。

因果掩码的下标规则

设 context_len = len(K) - len(Q)，则第 m 行 query 只能看到 K/V 的第 0 .. context_len + m 个位置（含），后面一律屏蔽：

• prefill（len(K) == len(Q)，context_len = 0）：标准下三角，第 m 行看 0..m。第 0 行只能看到第 0 个 key，所以 O[0] 必然等于 V[0]。
• decode（len(Q) == 1，context_len = len(K) - 1）：唯一那行 query 能看到全部历史 KV。

要求

• 实现右侧代码模板中给出的函数，签名以模板为准
• 必须实现因果掩码，被掩掉的位置不允许参与 softmax
• 使用 online softmax，维护 running max 和 running sum，逐个位置增量更新，不要先算出完整的分数矩阵再做 softmax
• 缩放因子为 ，其中 d = len(Q[0])
• V 的宽度可能与 d 不同，输出宽度以 len(V[0]) 为准

提示

• online softmax 的核心是三个累加器：running max m、running sum l、加权累加 acc。遇到更大的分数时，用 exp(m_old - m_new) 把已有的 l 和 acc 重标定一遍。
• 因果掩码在这里不需要真的构造 矩阵：顺序遍历 key，越过上界直接 break 即可。
• 数值稳定性：分数很大时（如 ）若不减 running max 会 exp 溢出。
• FlashAttention-2 相比 v1 减少了非矩阵乘的 rescale 操作，并把并行维度放在 seq_len 上。