标签：PyTorch

• 
  PyTorch 实现 multi-head attention 逐行讲解

  2026-06-02 · ai

  注意力机制动手实现PyTorch

  用 nn.Module 手写一个 multi-head attention，每行标注 shape 变化，讲清 fused QKV 投影、head 维折叠、scaled_dot_product_attention 的三种后端选择，以及和 nn.MultiheadAttention 的对应关系。