FlashAttention-2 相比 FlashAttention-1，到底改了什么

一句话区别

源码地图：先知道看哪里

第一处源码差异：launch grid 怎么铺

// v1: fmha_fwd_launch_template.h
if (params.num_splits <= 0) {
  params.num_splits = num_splits_heuristic_fwd(
    params.b * params.h, nSMs, ctas_per_sm, max_splits);
}

dim3 grid(params.b, params.h, params.num_splits);
kernel<<<grid, Kernel_traits::THREADS, smem_size, stream>>>(params);

// v2: flash_fwd_launch_template.h
const int num_m_block =
  (params.seqlen_q + Kernel_traits::kBlockM - 1) / Kernel_traits::kBlockM;

dim3 grid(num_m_block, params.b, params.h);
kernel<<<grid, Kernel_traits::kNThreads, smem_size, stream>>>(params);

怎么讲这段源码

v1 会根据 batch * heads 和 SM 数量估计 num_splits，目的是让 occupancy 不要太低，但 split 太多会带来额外 HBM 读写和中间结果合并。 v2 直接把 query sequence 切成 m_block，让 grid.x 覆盖 query block；这对小 batch、少 head、长序列的训练和推理更友好。

第二处源码差异：kernel 主循环怎么组织

// v1 源码讲解用压缩版
gmem_k.load(); gmem_q.load(); gmem_v.load();
gmem_q.commit(smem_q);
gmem_v.commit(smem_v);
gmem_k.commit(smem_k);
__syncthreads();

for (int l = blockIdx.z; l < steps; l += step_stride) {
  gemm_q_k(acc_p);                  // P score = Q * K^T
  softmax.unpack_noscale(acc_p);
  softmax.apply_mask(mask);
  softmax.reduce_max(...);
  softmax.scale_apply_exp(...);
  softmax.reduce_sum_before_sync_(...);
  softmax.pack(frag_p);

  gemm_cl(acc_o, frag_p, frag_v);    // O partial = P * V
  smem_o.store(acc_o, 0);
  __syncthreads();
  softmax.reduce_sum_after_sync_(...);
  smem_o.load(out);
  store O or temporary O;
}

// v2 源码讲解用压缩版
clear(acc_o);

for (masking_step...) {
  flash::gemm(acc_s, Q, K);          // S = QK^T
  apply_mask_if_needed(scores);
  softmax_rescale_o(scores, max, sum, acc_o, scale);
  rP = convert_type(scores);
  flash::gemm_A_in_regs(acc_o, rP, V);
}

for (; n_block >= 0; --n_block) {
  flash::gemm(acc_s, Q, K);
  softmax_rescale_o(scores, max, sum, acc_o, scale);
  rP = convert_type(scores);
  flash::gemm_A_in_regs(acc_o, rP, V);
}

// epilogue: normalize acc_o, write O and LSE

for each query block:
  load Q block
  m = -inf; l = 0; acc = 0

  for each K/V block:
    S = Q @ K.T
    S = apply_mask(S)
    m_new = max(m, rowmax(S))
    P = exp(S - m_new)
    acc = acc * exp(m - m_new) + P @ V
    l = l * exp(m - m_new) + rowsum(P)
    m = m_new

  O = acc / l

怎么讲这段源码

v1 和 v2 都遵循同一个 online softmax 数学流程。源码层面最容易讲清楚的区别是： v1 中 smem_o.store、__syncthreads、reduce_sum_after_sync_ 这类操作体现了跨 warp 归并路径； v2 中 acc_o 在主循环中持续累加，最后 epilogue 统一归一化和写回，warp 之间的 output 责任边界更清楚。

先用矩阵图理解 block 和 warp

这里有两个不同层级的 “block”：矩阵里的 Q block / K block / V block 是算法上的 tile； CUDA 里的 thread block / CTA 是执行这些 tile 的线程组。FA2 中可以粗略理解为：一个 CUDA block 负责一个 Q block，然后循环扫描多个 K/V block。

第三处源码差异：warp 内 work partitioning

// v1: mapping changes, so recompute row max
softmax.reduce_max_after_sync_(p_max_o, rows);
smem_o.store(acc_o, 0);
__syncthreads();
softmax.reduce_sum_after_sync_(p_sum_o, rows);
smem_o.load(out);

// v2: kernel_traits.h
static constexpr int kNWarps = kNWarps_;
static constexpr int kNThreads = kNWarps * 32;
static constexpr int kBlockM = kBlockM_;
static constexpr int kBlockN = kBlockN_;

using TiledMma = TiledMMA<
  MMA_Atom_Arch,
  Layout<Shape<Int<kNWarps>,_1,_1>>,
  ValLayoutMNK>;

第四处差异：v2 减少非 matmul 开销

FA2 论文和博客把一个重要点说得很直接：A100 上 Tensor Core 的 matmul 很快，非 matmul 的 FP32 操作、mask、rescale、shared memory 同步相对更贵。 所以 v2 不是只盯着 FLOPs，而是尽量把主循环拆成更规整的路径。

录视频时的讲解脚本

面试速答

• FA2 相比 FA1 的本质区别：不是数学公式变化，而是更好的并行度和 work partitioning。
• 源码第一眼看哪里：先看 launch grid，再看 forward kernel 主循环，再看 kernel traits 中的 block/warp/MMA 配置。
• v1 为什么还不够：IO 降下来了，但 threadblock/warp 分工导致 occupancy 和 shared memory 通信仍有瓶颈。
• v2 为什么更快：sequence 维度增加 CTA 并行，warp 按 Q 行切分，减少跨 warp 通信，减少非 matmul FLOPs。
• 什么时候收益明显：长序列、小 batch、少 head、大模型训练或推理中 attention 占比高时更明显。

练习题与追问