实现因果自注意力（Flash Attention 数值版）

实现带因果掩码的自注意力，数值上等价于 Flash Attention 的输出。Flash Attention 的关键是通过分块在线 softmax 避免物化 S×S 矩阵，本题在小规模上验证你实现的**数值正确性**（与直接计算的结果完全一致）： scores = q @ k^T / sqrt(D) scores[i, j] = -inf 当 j > i（因果掩码，不允许看未来） out = softmax(scores, dim=-1) @ v 输入形状：q/k/v 均为 (B, H, S, D)，输出形状与 q 相同。 建议按 Flash Attention 的思路做块级在线 softmax（block_size 建议取 8~16）；判题只看数值（与标准实现的 allclose 偏差）。注意分块累加时用 max 值做平移校正，避免中间结果溢出。 要求： - 不要调用 torch.nn.functional.scaled_dot_product_attention 或 torch.nn.functional.softmax（请自己实现带在线 max 平移的 softmax，这也是 Flash Attention 的核心技巧）； - 计算必须可微。