Flash Attention

# 标准实现（简化）
S = Q @ K.T / sqrt(d)      # N×N矩阵，必须存储
P = softmax(S)             # N×N矩阵，必须存储
O = P @ V                  # 最终输出

# 内存需求：O(N²)
# IO复杂度：O(N²)

# 避免存储完整的注意力矩阵
# 使用稳定的在线算法计算softmax
m_new = max(m_old, row_max)  # 更新最大值
exp_sum = exp_sum * exp(m_old - m_new) + row_sum

[Tensor Core计算] || [数据传输]
     ↓                    ↓
[softmax计算]    || [预取下一块]

# 自动选择最优实现
import torch.nn.functional as F

output = F.scaled_dot_product_attention(
    query, key, value,
    attn_mask=mask,
    dropout_p=0.0,
    is_causal=True
)

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "model_name",
    torch_dtype=torch.float16,
    attn_implementation="flash_attention_2"
)

from flash_attn import flash_attn_func

# q, k, v: (batch, seq_len, n_heads, head_dim)
output = flash_attn_func(q, k, v, causal=True)