FlashAttention-4 针对 Blackwell GPU 的非对称硬件 scaling（MMA 翻倍但 smem/exp 不变）重新设计 attention kernel pipeline：利用 TMEM 和异步 MMA 重叠 softmax 与矩阵乘，用多项式逼近软件模拟 exp 绕过 MUFU 瓶颈，2-CTA MMA 减半 smem 流量和 dQ 原子操作。B200 BF16 达 1613 TFLOPs/s (71%)，超 cuDNN 9.13 达 1.3×。
Blackwell B200 tensor core 吞吐翻倍至 2.25 PFLOPS (BF16)，但 shared memory 带宽（128 bytes/clock/SM）和 exponential 单元（MUFU, 16 ops/clock/SM）不变。这造成 "asymmetric hardware scaling"：在 Hopper 上 MMA、smem、exp 近乎平衡，但在 Blackwell 上 smem + exp 联合超过 MMA 时间 25-60%，成为新瓶颈。FlashAttention-3 针对 Hopper 设计，在 Blackwell 上无法达到高利用率。
1. 新 pipeline 设计（§3.1.2）
Ping-pong schedule：两个 Q tile 交替执行。MMA 结果写入 TMEM（非寄存器），使 softmax warpgroup 与 MMA warpgroup 真正解耦。关键布局：
TMEM 分区：两个 $128 \times 128$ output tile + 两个 S tile（可覆盖 P 的存储位置，因为 P 存入后立即被 MMA 消费）。
2. 软件模拟 exp（§3.1.3）
MUFU.EX2 吞吐 16 ops/clock vs MMA 8192 ops/clock → 512× 差距。用 FMA 单元（与 MUFU 并行）运行 Cody-Waite 分解 + 多项式逼近：
$$2^x = 2^{\lfloor x \rfloor} \cdot 2^{x - \lfloor x \rfloor}$$
Degree-3 的 FP32 最大相对误差 $8.77 \times 10^{-5}$，但 BF16 量化误差（$3.89 \times 10^{-3}$）是主导项——软件模拟在 BF16 精度下与硬件 MUFU 功能等价。
部分模拟：仅 10-25% 的 exp 使用软件模拟（避免寄存器压力导致 spill），其余仍用 MUFU。比例经验调优。
3. 条件 softmax rescaling（§3.1.4）
Online softmax 每个 block 更新 $m_j = \max(m_{j-1}, \text{rowmax}(S_j))$ 后需 rescale output $O_{j-1}$。当 $m_j - m_{j-1} \leq \tau = \log_2(256) = 8.0$ 时跳过 rescaling——减少非 matmul 操作。最终 normalization 修正累计偏差。
4. 2-CTA MMA 减半 smem 流量（§3.2.3）
Backward 5 个 MMA 中 smem 流量（3328 cycles @ 1-CTA）超过 MMA 计算（2560 cycles）30%。2-CTA 模式：CTA pair 各自 stage 半个 operand B → smem 流量降至 2688 cycles（仅超 MMA 5%）。
dQ 步骤：通过 DSMEM 交换半个 dS tile，使每个 CTA 持有 $(M/2, 2N)$ operand 做 CTA-pair UMMA，减半 atomic reduction 写次数。
5. Deterministic backward（§3.2.4）
Semaphore-based serialization + CTA swizzling + causal-specific SPT (Shortest-Processing-Time-first) schedule → 确定性反向达到非确定性 75% 速度。
LPT (Longest-Processing-Time-first) for causal masking + varlen：按 head/batch swizzle，每 section 装满 L2 cache。Causal 4-8% FLOPS gain for MHA，7-14% for MQA。
核心技术壁垒：TMEM 的利用——将 MMA 结果写入 TMEM 而非寄存器，使 softmax warpgroup 能独立于 MMA pipeline 执行。这是 Blackwell 特有的硬件特性，是实现 MMA-softmax overlap 的前提。在 Hopper 上 MMA 结果在寄存器中，softmax 必须等待 MMA 完成才能访问。
| Resource | B200 Throughput | Hopper H100 | Scaling |
|---|---|---|---|
| Tensor core (BF16 MMA) | 8192 ops/clock/SM | 4096 | 2× |
| MUFU exp | 16 ops/clock/SM | 16 | 1× |
| SMEM read | 128 bytes/clock/SM | 128 | 1× |
| TMEM | 256 KB/SM | N/A (new) | — |
| MMA tile | 128×128 | 64×128 | 2× in M |
| SMs | 148 | 132 | 1.12× |
| Peak BF16 | 2.25 PFLOPS | 1.0 PFLOPS | 2.25× |
| Resource | Cycles |
|---|---|
| MMA compute | 1024 |
| Shared memory | 768 |
| Exponential unit | 1024 |
MMA 和 exp 是共同瓶颈；smem 略低。这驱动了 (1) overlap MMA+softmax，(2) 软件模拟 exp，(3) 条件 rescaling。
| Resource | 1-CTA (cycles) | 2-CTA (cycles) |
|---|---|---|
| MMA compute | 2560 | 2560 |
| Total shared memory | 3328 | 2688 |
| Exponential unit | 1024 | 1024 |
Backward 中 smem 是主瓶颈（超 MMA 30% / 5%）。2-CTA 将差距从 30% 压缩到 5%。
| 符号 | 含义 |
|---|---|
| $M, N$ | Tile dimensions along sequence length |
| $d$ | Head dimension (typically 128) |
| $\alpha$ | Scaling factor $1/\sqrt{d}$ |
| $T_{\text{MMA}}$ | MMA compute time (cycles) |
| $T_{\text{smem}}$ | Shared memory traffic time (cycles) |
| $T_{\text{exp}}$ | Exponential unit time (cycles) |
| $\tau$ | Conditional rescaling threshold ($\log_2(256) = 8.0$) |
Forward：$T_{\text{MMA}} = 4MNd/8192$ — 两个 MMA（$QK^\top$ 和 $PV$），每个 $2MNd$ FLOPs，tensor core 8192 FLOPs/cycle。$T_{\text{exp}} = MN/16$ — softmax 需要对 $MN$ 个元素做 exp，MUFU 16 ops/cycle。当 $d = 128$ 时 $T_{\text{MMA}} = T_{\text{exp}}$ = 1024 cycles——恰好 MMA 和 exp 等强度。smem 768 cycles 略低。
关键 insight：$T_{\text{exp}}$ 与 $d$ 无关而 $T_{\text{MMA}} \propto d$。对于 $d < 128$（如 64），exp 将超过 MMA 成为单独瓶颈。对于 $d > 128$（如 256），MMA 主导而 exp 可忽略。$d = 128$ 恰好是平衡点。
Backward：5 个 MMA → $T_{\text{MMA}} = 10MNd/8192 = 2.5\times$ forward。smem 额外包含 dS write + dQ write/read → 1-CTA 总 3328 cycles 超过 2560 cycles MMA。2-CTA 通过每 CTA 仅 stage 半个 operand B 降至 2688。
Forward ($d=128$): arithmetic intensity = $4MNd / (2 \times (M+N) \times d \times 2)$ bytes from HBM ≈ $N$ (for large $N$). 这是 compute-bound regime（大 $N$ 时）。kernel 目标是达 MMA peak——但被 smem 和 exp 分走时间。
| Axis | Value Chosen | Rejected Alternatives | Constraint |
|---|---|---|---|
| MMA tile size | 128×128 (M×N) | 64×128 (Hopper-style) | Blackwell native 128×128 tile; smaller wastes TMEM |
| Exp emulation fraction | 10-25% | 0% (all hardware) / 100% (all software) | >25% causes register spills; 0% leaves MUFU bottleneck |
| Rescaling threshold $\tau$ | $\log_2(256) = 8.0$ | Lower / higher | Lower → more rescaling ops; higher → risk numerical drift |
| Backward CTA mode | 2-CTA | 1-CTA | 1-CTA smem exceeds MMA by 30%; 2-CTA reduces to 5% |
| Softmax warpgroup size | 128 threads (1 per row) | Fewer threads | Each thread handles full row → no inter-warp shuffle needed |
| TMEM allocation (fwd) | 2×O + 2×S (overlap P) | 2×O + 1×S + 2×P | Overlap allows immediate dual S computation; P consumed fast |
| Seqlen | FA4 TFLOPs/s (non-causal) | cuDNN 9.13 | Triton | Speedup vs cuDNN |
|---|---|---|---|---|
| 1k | ~1100 | ~1000 | ~550 | 1.1× |
| 4k | ~1400 | ~1250 | ~600 | 1.12× |
| 8k | ~1500 | ~1300 | ~600 | 1.15× |
| 16k | ~1580 | ~1350 | ~620 | 1.17× |
| 32k | ~1613 | ~1250 | ~600 | 1.3× |
长序列优势更大（amortize pipeline startup + LPT scheduling 效果显著）。
FA4 2-CTA backward 在 8k+ 序列上一致优于 cuDNN。
SPT schedule 的确定性 backward 达到非确定性 1-CTA backward 速度的 75%——实用级别的确定性/性能 tradeoff。
| Forward | Backward | |
|---|---|---|
| FA-3 (C++ templates) | 55s | 45s |
| FA-4 (CuTe-DSL) | 2.5s | 1.4s |
| Speedup | 22× | 32× |
| Technique | Target Bottleneck | Hardware Primitive | Contribution |
|---|---|---|---|
| Ping-pong pipeline w/ TMEM | MMA-softmax overlap | TMEM (256 KB/SM), async MMA | Core pipeline structure |
| Software exp emulation | MUFU exp throughput (16 ops/clock) | FMA units (parallel to MUFU) | Degree-3 poly, 10-25% of exp ops |
| Conditional rescaling | Non-matmul ops in online softmax | Branch + FMA | Skip rescaling when $m_j - m_{j-1} \leq 8.0$ |
| 2-CTA MMA (backward) | SMEM bandwidth (128 B/clock) | CTA-pair cooperative MMA | Each CTA stages half operand B |
| DSMEM exchange (dQ) | dQ atomic reduction overhead | Distributed shared memory | Halves atomic add count |
| LPT scheduling | Causal / varlen load imbalance | CTA grid swizzling | 4-14% FLOPS gain |
| CuTe-DSL JIT | C++ compile time | Python → PTX → SASS | 22-32× faster builds |
| Step | Claim | Evidence | Strength |
|---|---|---|---|
| 1 | Blackwell MMA 翻倍但 smem/exp 不变导致新瓶颈 | §2.2 硬件 spec (8192 vs 16 vs 128) | Strong — 硬件事实 |
| 2 | Forward roofline: MMA 和 exp 是共同瓶颈 | Table 1 cycle analysis | Strong — 可手算验证 |
| 3 | 新 pipeline 利用 TMEM 实现 MMA-softmax 真正异步重叠 | §3.1.2 pipeline 设计 + Fig 1 | Strong — 架构设计清晰 |
| 4 | 软件 exp 模拟在 BF16 下与硬件功能等价 | Table 2 (rel err dominated by BF16 quantization) | Strong — 数值验证 |
| 5 | Backward smem 瓶颈 (30% > MMA) 被 2-CTA 压缩到 5% | Table 3 (3328→2688 vs 2560) | Strong — 定量分析 |
| 6 | 端到端: 1613 TFLOPs/s (71%) on B200，1.3× cuDNN | Fig 4, 5, 6 | Strong — 多配置验证 |
| 7 | CuTe-DSL 保持性能同时 22-32× 编译加速 | Table 4 + benchmark parity | Strong |
关键实现细节：
| GPU Family | Support | Key Dependency | Porting Effort |
|---|---|---|---|
| Blackwell (sm_100, B200/B100) | Native | TMEM, 128×128 MMA, 2-CTA MMA, async MMA | N/A |
| Hopper (sm_90, H100/H200) | Partial | Has wgmma but no TMEM; 64×128 MMA | FA-3 remains better fit |
| Ada (sm_89) | No | No wgmma/TMEM | Not applicable |
| CDNA 3/4 (gfx942/gfx950) | No | No TMEM equivalent; MFMA has different semantics | Separate design needed |
核心优化（TMEM-based pipeline、2-CTA cooperative MMA）是 Blackwell-specific——不可直接移植。但算法 insight（exp emulation、conditional rescaling、LPT scheduling）可移植到任何 architecture。