Causal Self-Attention

题面

实现因果（下三角掩码）自注意力：softmax(mask(QK^T/√d)) V，对位置 i 仅允许关注 j≤i，掩码对 j>i 置 −∞，按行做 softmax。

Implementation Requirements

• Use only native features (external libraries are not permitted)
• The solve function signature must remain unchanged
• 输出写入 output

Examples

见页面两组示例（2×4 与 2×2）。

Constraints

• Q,K,V 形状均为 M×d；1 ≤ M ≤ 10000；1 ≤ d ≤ 128
• 值域 −100..100，float32；Performance: M=5,000