Sliding Window Self-Attention

原始题目：LeetGPU - Sliding Window Self-Attention

题目描述

实现滑动窗口自注意力。给定查询矩阵 $Q$、键矩阵 $K$ 和值矩阵 $V$（均为 $M \times d$）以及窗口大小 $W$，每个位置 $i$ 只关注窗口内的位置 $j \in [i - W, i]$（因果窗口）：

$$score_{i,j} = \frac{Q_i \cdot K_j}{\sqrt{d}}$$

$$output_i = \sum_{j=\max(0, i-W)}^{i} \text{softmax}(score_{i,*})_j \cdot V_j$$

窗口外的注意力分数设为 $-\infty$（softmax 前）。所有数据为 float32。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。

约束条件

• $1 \le M \le 10{,}000$，$1 \le d \le 128$，$1 \le W \le M$。
• 性能测试在 $M = 5{,}000$ 下进行。

解题思路

滑动窗口注意力将注意力限制在局部窗口内，将复杂度从 $O(M^2)$ 降为 $O(M \cdot W)$。这避免了长序列中 attention 矩阵的显存爆炸问题，被 Mistral 等模型采用。GPU 实现中可以利用每个 query 只读局部 window 内 KV 的特点，使用分块策略加载共享内存，显著减少全局内存访问。因果掩码（$j \le i$）+ 窗口掩码（$j \ge i - W$）的结合限制了每个 query 的有效 attention 范围。欢迎在 GitHub Discussions 分享你的解法。