当数据和并行策略都没有明显问题时,训练速度常常被单个高频算子卡住: attention、LayerNorm/RMSNorm、Adam 更新、dropout + residual、activation + bias。 算子层加速的核心是减少三件事:多余的 HBM 读写、多余的 kernel launch、多余的中间张量。

标准 attention 的数学形式没有变: softmax(QKᵀ / sqrt(d))V。 FlashAttention 变的是计算组织方式:不把完整 attention matrix 写回 HBM, 而是用 tile、online softmax 和重计算来降低显存占用并提高吞吐。 v2 主要改善工作划分与并行效率;v3 进一步针对 Hopper GPU 的异步能力和 Tensor Core 管线做优化。

FlashAttention 主要解决 attention 这条热路径。 xFormers 提供 memory-efficient attention、sparse/block-sparse attention 和若干 fused building blocks; Apex 则常用于 FusedAdam、FusedLayerNorm、FusedRMSNorm 等 CUDA 扩展。 这两类库的共同点是:替换看起来像 drop-in,但收益必须用 profiler 和 benchmark 验证。

import xformers.ops as xops

# q, k, v: [B, S, H, D] 或 xFormers 支持的 layout
out = xops.memory_efficient_attention(
    q, k, v,
    attn_bias=None,
    p=0.0,
    op=None,  # 让 xFormers 自动 dispatch
)

from apex.optimizers import FusedAdam
from apex.normalization import FusedLayerNorm

norm = FusedLayerNorm(hidden_size)
optimizer = FusedAdam(model.parameters(), lr=3e-4)

实操建议从“最少侵入”开始:先替换成 PyTorch SDPA,确认 dtype、mask、dropout 行为一致; 再尝试 FlashAttention 或 xFormers。每一步都保留 correctness check,否则速度再快也没有意义。

最小 attention wrapper

import torch
import torch.nn.functional as F

def attention_sdpa(q, k, v, *, is_causal=True, dropout_p=0.0):
    # q/k/v: [B, heads, S, head_dim]
    return F.scaled_dot_product_attention(
        q, k, v,
        attn_mask=None,
        dropout_p=dropout_p,
        is_causal=is_causal,
    )

def attention_naive(q, k, v, *, is_causal=True):
    scale = q.shape[-1] ** -0.5
    score = (q @ k.transpose(-2, -1)) * scale
    if is_causal:
        s = q.shape[-2]
        mask = torch.ones(s, s, device=q.device, dtype=torch.bool).tril()
        score = score.masked_fill(~mask, float("-inf"))
    prob = score.softmax(dim=-1)
    return prob @ v

FlashAttention 可选路径

def attention_flash_attn(q, k, v, *, causal=True, dropout_p=0.0):
    # flash_attn_func 常见 layout: [B, S, heads, head_dim]
    from flash_attn import flash_attn_func
    q = q.transpose(1, 2).contiguous()
    k = k.transpose(1, 2).contiguous()
    v = v.transpose(1, 2).contiguous()
    out = flash_attn_func(q, k, v, dropout_p=dropout_p, causal=causal)
    return out.transpose(1, 2).contiguous()

数值对齐检查

torch.manual_seed(0)
device = "cuda"
q = torch.randn(2, 16, 512, 64, device=device, dtype=torch.float16)
k = torch.randn_like(q)
v = torch.randn_like(q)

ref = attention_naive(q.float(), k.float(), v.float()).half()
out = attention_sdpa(q, k, v)

max_err = (ref - out).abs().max().item()
mean_err = (ref - out).abs().mean().item()
print({"max_err": max_err, "mean_err": mean_err})

单算子 benchmark 解决“这个 kernel 有没有更快”; MFU 解决“整个模型有没有更接近硬件峰值”。 真实训练里,attention 加速可能被数据等待、通信、pipeline bubble、checkpoint 重算抵消, 所以 Day26 的产出必须同时包含 micro benchmark 和 end-to-end step benchmark。

CUDA event 计时模板

import statistics
import torch

def bench(fn, *args, warmup=20, iters=100):
    for _ in range(warmup):
        fn(*args)
    torch.cuda.synchronize()

    times = []
    for _ in range(iters):
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()
        fn(*args)
        end.record()
        torch.cuda.synchronize()
        times.append(start.elapsed_time(end) / 1000.0)

    times.sort()
    return {
        "median_s": statistics.median(times),
        "p90_s": times[int(len(times) * 0.9)],
        "min_s": times[0],
    }

def attention_tflops(batch, heads, seq, head_dim, seconds):
    flops = 4 * batch * heads * seq * seq * head_dim
    return flops / seconds / 1e12

报告格式

shape: B=2, H=16, S=4096, D=64, dtype=bf16, causal=True
gpu: H100-SXM-80GB, torch=..., cuda=...

backend                 median_ms    p90_ms    peak_mem_gb    attention_tflops
naive eager             132.4        134.8     18.2           1.04
torch sdpa flash         12.8         13.1      3.6          10.76
flash-attn v2            11.9         12.2      3.4          11.57
xformers mem-eff         13.5         13.9      3.8          10.20

end-to-end:
baseline step_time=812 ms, tokens/s=20176, max_mem=71.4 GB
optimized step_time=674 ms, tokens/s=24305, max_mem=63.1 GB

1. FlashAttention 为什么能降低显存?它具体少保存了什么中间量?
2. 为什么 attention benchmark 必须固定 dtype、head_dim、causal、dropout 和 mask?
3. xFormers 的 memory-efficient attention 和 FlashAttention 的关系是什么?什么时候你会选 xFormers?
4. FusedAdam 为什么可能比普通 Adam 快?它减少的是算术量、内存访问,还是 kernel launch?
5. 如果 micro benchmark 快 2 倍,但 end-to-end step 只快 5%,你会如何定位原因?