手撕 Multi-Head Self-Attention

手撕多头自注意力算子 (Live Coding: Multi-Head Self-Attention in Pure Numpy / PyTorch) 是机器学习工程师 (MLE) 白板编码与现场代码考核中最高频、最核心的压轴题目；考察从零编写标准的 Scaled Dot-Product Multi-Head Attention：1) 线性投影：将输入张量 $X \in \mathbb{R}^{B \times L \times d_{\text{model}}}$ 分别通过线性权重映射为 $Q, K, V$；2) 维度变换与切分：将张量 Reshape 并 Transpose 为 $(B, H, L, d_k)$（其中 $d_k = d_{\text{model}} / H$）；3) 点积注意力得分：计算 $S = \frac{Q K^T}{\sqrt{d_k}}$ 并根据需要施加因果掩码 (Causal Mask / Upper Triangular Mask with $-\infty$)；4) 稳定 Softmax 与加权求和：$\text{Attention}(Q, K, V) = \text{softmax}(S) V$；5) 多头拼接与最终线性投影输出。

顶级大厂 MLE 面试现场 20 分钟手撕代码考核、Transformer 核心算子自研与调试。

考察候选人是否真正理解多头机制的张量维度流转 (Tensor Shape Transitions)、广播机制、缩放因子 $\sqrt{d_k}$ 的数学根因以及掩码加法的数值稳定性，而非单纯调用黑盒库。

现场写出包含 Batch 维度的 Scaled Dot-Product Attention 核心 Pure Numpy 代码，并标注每一步的 Tensor Shape？

数学推导：为什么必须除以缩放因子 $\sqrt{d_k}$（证明当 $q_i, k_i \sim \mathcal{N}(0, 1)$ 独立同分布时，内积点积的方差为 $d_k$）？

手写因果自回归掩码 (Causal Mask: `np.triu(np.ones(...), k=1)`) 的生成与为什么掩码值要填充为 `-1e9` 或 `-np.inf`？

对比 `torch.einsum('bhqd,bhkd->bhqk', q, k)` 与标准 `torch.matmul(q, k.transpose(-2, -1))` 的可读性与底层算子执行？

FlashAttention 核心优化原理手撕阐述：如何通过 Tiling 分块计算与 Online Softmax 消除中间注意力矩阵 $O(N^2)$ 的 HBM 读写？