2. 实现标准Attention的朴素版本

任务目标

在深入FlashAttention之前，先实现一个标准Attention的朴素版本作为性能基线。该版本将完整实现Softmax(QK^T)V的计算流程，使用PyTorch的einsum与显式Softmax操作，不引入任何分块或近似优化。后续所有FlashAttention案例都将与此基线进行正确性对比和性能对照。

实现代码

import torch
import torch.nn.functional as F

def standard_attention(Q, K, V):
    """
    标准Attention朴素实现。
    
    Args:
        Q: (batch, seq_len, num_heads, head_dim)
        K: (batch, seq_len, num_heads, head_dim)
        V: (batch, seq_len, num_heads, head_dim)
    
    Returns:
        output: (batch, seq_len, num_heads, head_dim)
        attn_weights: (batch, num_heads, seq_len, seq_len)  # 可选返回
    """
    # QK^T: (batch, num_heads, seq_len, seq_len)
    scores = torch.einsum('bqhd,bkhd->bhqk', Q, K)
    
    # 缩放
    head_dim = Q.shape[-1]
    scores = scores / (head_dim ** 0.5)
    
    # Softmax
    attn_weights = F.softmax(scores, dim=-1)
    
    # 加权求和
    output = torch.einsum('bhqk,bkhd->bqhd', attn_weights, V)
    
    return output, attn_weights

# ========== 最小验证闭环 ==========
if __name__ == "__main__":
    torch.manual_seed(42)
    batch, seq_len, num_heads, head_dim = 2, 4, 3, 8
    
    Q = torch.randn(batch, seq_len, num_heads, head_dim)
    K = torch.randn(batch, seq_len, num_heads, head_dim)
    V = torch.randn(batch, seq_len, num_heads, head_dim)
    
    output, attn_weights = standard_attention(Q, K, V)
    
    # 形状验证
    assert output.shape == (batch, seq_len, num_heads, head_dim), f"输出形状错误: {output.shape}"
    assert attn_weights.shape == (batch, num_heads, seq_len, seq_len), f"注意力权重形状错误: {attn_weights.shape}"
    
    # Softmax行和为1验证
    row_sums = attn_weights.sum(dim=-1)
    assert torch.allclose(row_sums, torch.ones_like(row_sums), atol=1e-6), "Softmax行和不为1"
    
    # 与PyTorch官方SDPA对比（需PyTorch ≥ 2.0）
    try:
        sdpa_output = F.scaled_dot_product_attention(Q, K, V)
        diff = (output - sdpa_output).abs().max().item()
        assert diff < 1e-5, f"与官方SDPA差异过大: {diff}"
        print(f"✅ 与PyTorch官方SDPA一致，最大差异: {diff:.2e}")
    except AttributeError:
        print("⚠️ 当前PyTorch版本不支持F.scaled_dot_product_attention，跳过官方对比")
    
    print(f"✅ 输出形状: {tuple(output.shape)}")
    print(f"✅ 注意力矩阵形状: {tuple(attn_weights.shape)}")
    print("✅ 朴素Attention实现验证通过")

运行与验证

环境依赖：Python 3.8+，PyTorch 1.10+（若需官方SDPA对比则需2.0+），无需其他第三方库。

运行方式：将代码保存为 naive_attention.py，直接执行 python naive_attention.py。

验证判定标准：

• 程序无异常退出，输出形状断言全部通过；
• 注意力权重矩阵每行求和为1（Softmax性质）；
• 若PyTorch版本≥2.0，与官方 F.scaled_dot_product_attention 的输出最大差异小于1e-5，证明实现正确；
• 控制台打印三行✅标记即表示验证通过。

关键说明

该实现将Q、K、V按 (batch, seq_len, num_heads, head_dim) 布局传入，通过einsum完成批量矩阵乘法。缩放因子 1/√d_k 在Softmax之前应用，这是Attention的标准做法。此版本显式物化了完整的 (batch, num_heads, seq_len, seq_len) 注意力矩阵，内存复杂度为O(n²)，这正是FlashAttention后续要解决的核心瓶颈。本实现仅用于正确性基线与性能对照，不包含任何IO优化或分块策略。