FlashAttention 论文解析：用 IO-Aware Exact Attention 革新 Transformer 训练与推理速度

import torch
import torch.nn.functional as F
import time

def standard_attention(Q, K, V, mask=None):
    """Standard Self-Attention: 整体 materialization 完整的 N x N 分数矩阵。"""
    d_k = Q.size(-1)
    # S = Q @ K^T -> (batch, heads, N, N) 整体存入内存
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)

    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))

    # softmax 也对整个 N x N 矩阵执行
    attn_weights = F.softmax(scores, dim=-1)  # (batch, heads, N, N) 存储

    # 计算最终输出
    output = torch.matmul(attn_weights, V)
    return output

# 基准测试
batch, heads, seq_len, d_head = 4, 32, 4096, 64
Q = torch.randn(batch, heads, seq_len, d_head, device='cuda', dtype=torch.float16)
K = torch.randn_like(Q)
V = torch.randn_like(Q)

# 预热
for _ in range(3):
    _ = standard_attention(Q, K, V)
torch.cuda.synchronize()

# 测量
start = time.time()
for _ in range(100):
    _ = standard_attention(Q, K, V)
torch.cuda.synchronize()
elapsed = (time.time() - start) / 100

print(f"Standard Attention: {elapsed*1000:.2f} ms/iter")
print(f"Peak memory: {torch.cuda.max_memory_allocated() / 1e9:.2f} GB")

import torch
import math

def flash_attention_forward(Q, K, V, block_size_q, block_size_kv):
    """FlashAttention v1 Forward Pass 伪代码。

    实际 CUDA 内核中，所有运算都在一个 GPU 内核内完成，
    中间结果只保留在 SRAM 中。
    """
    batch, heads, N, d = Q.shape
    O = torch.zeros_like(Q)           # 输出累加器
    m = torch.full((batch, heads, N, 1), float('-inf'), device=Q.device)  # 按行最大值
    l = torch.zeros((batch, heads, N, 1), device=Q.device)                # 按行总和

    # 外层循环：按块遍历 Q
    for i in range(0, N, block_size_q):
        qi = Q[:, :, i:i+block_size_q, :]  # 载入 SRAM

        # 内层循环：按块遍历 K、V
        for j in range(0, N, block_size_kv):
            kj = K[:, :, j:j+block_size_kv, :]  # 载入 SRAM
            vj = V[:, :, j:j+block_size_kv, :]  # 载入 SRAM

            # 1. 计算局部注意力分数（在 SRAM 内）
            sij = torch.matmul(qi, kj.transpose(-2, -1)) / math.sqrt(d)

            # 2. 局部 softmax 统计量
            mij_local = sij.max(dim=-1, keepdim=True).values
            pij_local = torch.exp(sij - mij_local)
            lij_local = pij_local.sum(dim=-1, keepdim=True)

            # 3. 用 Online Softmax 更新全局统计量
            mi_old = m[:, :, i:i+block_size_q, :]
            li_old = l[:, :, i:i+block_size_q, :]
            oi_old = O[:, :, i:i+block_size_q, :]

            mi_new = torch.maximum(mi_old, mij_local)
            alpha = torch.exp(mi_old - mi_new)
            beta = torch.exp(mij_local - mi_new)

            li_new = alpha * li_old + beta * lij_local

            # 4. 更新累加输出
            O[:, :, i:i+block_size_q, :] = (
                alpha * li_old * oi_old + beta * torch.matmul(pij_local, vj)
            ) / li_new

            m[:, :, i:i+block_size_q, :] = mi_new
            l[:, :, i:i+block_size_q, :] = li_new

    return O, m, l  # m, l 会在反向传播的 recomputation 中用到

import torch
import torch.nn.functional as F
from torch.nn.attention import sdpa_kernel, SDPBackend

# 基本用法：PyTorch 自动选择最优后端
batch, heads, seq_len, d_head = 2, 16, 4096, 64
Q = torch.randn(batch, heads, seq_len, d_head, device='cuda', dtype=torch.float16)
K = torch.randn_like(Q)
V = torch.randn_like(Q)

# 自动后端选择（若支持 FlashAttention 则自动使用）
output = F.scaled_dot_product_attention(Q, K, V)

# 强制只使用 FlashAttention 后端
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    output_flash = F.scaled_dot_product_attention(
        Q, K, V,
        dropout_p=0.1,      # 训练时的 dropout
        is_causal=True,      # causal masking（用于解码器）
        scale=None,          # 使用默认 1/sqrt(d_k)
    )

# 确认当前使用的是哪种后端
print(f"Flash available: {torch.backends.cuda.flash_sdp_enabled()}")
print(f"Memory efficient available: {torch.backends.cuda.mem_efficient_sdp_enabled()}")

# pip install flash-attn --no-build-isolation
from flash_attn import flash_attn_func, flash_attn_qkvpacked_func
import torch

# 方法一：Q、K、V 分别输入
# 形状：(batch, seqlen, nheads, headdim) - 注意：与 PyTorch 的 head/seq 顺序不同
batch, seqlen, nheads, headdim = 2, 4096, 16, 64
Q = torch.randn(batch, seqlen, nheads, headdim, device='cuda', dtype=torch.float16)
K = torch.randn_like(Q)
V = torch.randn_like(Q)

output = flash_attn_func(
    Q, K, V,
    dropout_p=0.0,
    softmax_scale=None,   # 默认 1/sqrt(headdim)
    causal=True,
    window_size=(-1, -1), # (-1, -1) 表示完整注意力，(w, 0) 表示 sliding window
    return_attn_probs=False,
)

# 方法二：QKV packed 格式
# 形状：(batch, seqlen, 3, nheads, headdim)
qkv = torch.randn(batch, seqlen, 3, nheads, headdim, device='cuda', dtype=torch.float16)
output_packed = flash_attn_qkvpacked_func(
    qkv,
    dropout_p=0.0,
    causal=True,
)

# 方法三：Sliding Window Attention（FlashAttention-2+）
output_sliding = flash_attn_func(
    Q, K, V,
    causal=True,
    window_size=(256, 0),  # 左侧 256 token 的窗口
)

print(f"Output shape: {output.shape}")  # (batch, seqlen, nheads, headdim)

import triton
import triton.language as tl
import torch

@triton.jit
def flash_attention_kernel(
    Q_ptr, K_ptr, V_ptr, O_ptr,
    stride_qb, stride_qh, stride_qm, stride_qk,
    stride_kb, stride_kh, stride_kn, stride_kk,
    stride_vb, stride_vh, stride_vn, stride_vk,
    stride_ob, stride_oh, stride_om, stride_ok,
    N_CTX: tl.constexpr,
    D_HEAD: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    """FlashAttention Forward Kernel 的 Triton 实现示意。

    实际生产环境中的内核会包含更多优化，
    这里展示的是核心 tiling 逻辑的结构。
    """
    # 当前 program 负责的 Q 块索引
    pid_m = tl.program_id(0)
    pid_bh = tl.program_id(1)

    # 偏移量计算
    off_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    off_n = tl.arange(0, BLOCK_N)
    off_k = tl.arange(0, D_HEAD)

    # 载入 Q 块（HBM -> SRAM）
    q = tl.load(Q_ptr + pid_bh * stride_qh + off_m[:, None] * stride_qm + off_k[None, :] * stride_qk,
                mask=off_m[:, None] < N_CTX)

    # 初始化累加变量
    m_i = tl.full([BLOCK_M], value=float('-inf'), dtype=tl.float32)
    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
    o_i = tl.zeros([BLOCK_M, D_HEAD], dtype=tl.float32)

    # 遍历 K/V 块（内层循环）
    for start_n in range(0, N_CTX, BLOCK_N):
        curr_n = start_n + off_n

        # 载入 K、V 块（HBM -> SRAM）
        k = tl.load(K_ptr + pid_bh * stride_kh + curr_n[:, None] * stride_kn + off_k[None, :] * stride_kk,
                    mask=curr_n[:, None] < N_CTX)
        v = tl.load(V_ptr + pid_bh * stride_vh + curr_n[:, None] * stride_vn + off_k[None, :] * stride_vk,
                    mask=curr_n[:, None] < N_CTX)

        # 计算局部注意力分数（在 SRAM 内）
        s = tl.dot(q, tl.trans(k)) * (D_HEAD ** -0.5)

        # Online Softmax 更新
        m_ij = tl.max(s, axis=1)
        m_new = tl.maximum(m_i, m_ij)
        alpha = tl.exp(m_i - m_new)
        beta = tl.exp(m_ij - m_new)
        l_new = alpha * l_i + beta * tl.sum(tl.exp(s - m_ij[:, None]), axis=1)

        # 更新累加输出
        p = tl.exp(s - m_new[:, None])
        o_i = alpha[:, None] * l_i[:, None] * o_i + tl.dot(p.to(tl.float16), v)
        o_i = o_i / l_new[:, None]

        m_i = m_new
        l_i = l_new

    # 把结果存回 HBM
    tl.store(O_ptr + pid_bh * stride_oh + off_m[:, None] * stride_om + off_k[None, :] * stride_ok,
             o_i.to(tl.float16), mask=off_m[:, None] < N_CTX)

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 加载使用 FlashAttention-2 的模型
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B",
    torch_dtype=torch.float16,
    attn_implementation="flash_attention_2",  # 关键设置
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B")

# 基准对比
import time

text = "FlashAttention 是 " * 2000  # 长序列
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=4096).to("cuda")

# FlashAttention-2 推理
torch.cuda.reset_peak_memory_stats()
start = time.time()
with torch.no_grad():
    outputs = model(**inputs)
torch.cuda.synchronize()
flash_time = time.time() - start
flash_mem = torch.cuda.max_memory_allocated() / 1e9

print(f"FlashAttention-2: {flash_time:.3f}s, Peak Memory: {flash_mem:.2f} GB")

# 与 SDPA eager 模式对比（需要重新加载）
model_eager = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B",
    torch_dtype=torch.float16,
    attn_implementation="eager",
    device_map="auto",
)

torch.cuda.reset_peak_memory_stats()
start = time.time()
with torch.no_grad():
    outputs_eager = model_eager(**inputs)
torch.cuda.synchronize()
eager_time = time.time() - start
eager_mem = torch.cuda.max_memory_allocated() / 1e9

print(f"Eager Attention: {eager_time:.3f}s, Peak Memory: {eager_mem:.2f} GB")
print(f"Speedup: {eager_time/flash_time:.2f}x, Memory Savings: {eager_mem/flash_mem:.2f}x")