第 12 章 · FlashAttention

学习目标

• 理解 FlashAttention 为什么"既快又省显存"
• 把第 10 章的 online softmax 推广到 tile-wise（block 内分块）
• 看懂 v1 → v2 的差异（循环顺序换、warp 切分更友好）

前置知识

已完成 Ch06、Ch10 与 Ch11，理解 tile、online softmax、attention 数学和 HBM 流量分析。

核心概念

12.1 朴素 Attention 的根本问题

第 11 章看到：S, P 都是 T×T 张量，必须写回 HBM 才能下一步用。HBM 读写带宽是 attention 瓶颈，FLOPs 反而不是问题。

FlashAttention 的核心洞察：

12.2 算法推导

把 Q 沿行切成 Br × D 的块，K/V 沿行切成 Bc × D 的块。 对每个 Q 块，扫一遍所有 K/V 块，维护 running output O_i 和 (m_i, l_i)：

graph TD
    Init["初始化:
O_i = 0,  l_i = 0,  m_i = -∞"]
    Init --> Loop["for j in K/V tiles:"]
    Loop --> S["S_ij = Q_i @ K_j^T * scale"]
    S --> Mnew["m_new = max(m_i, rowmax(S_ij))"]
    Mnew --> Pij["P_ij = exp(S_ij - m_new)"]
    Pij --> Lnew["l_new = exp(m_i - m_new) * l_i + rowsum(P_ij)"]
    Lnew --> Onew["O_i = exp(m_i - m_new) * O_i + P_ij @ V_j
(暂不除 l_new)"]
    Onew --> Update["m_i = m_new,  l_i = l_new"]
    Update --> Loop
    Loop --> Out["最终: O_i /= l_i,  写回 HBM"]
    style S fill:#f3f1e8,stroke:#8b1538
    style Onew fill:#f3f1e8,stroke:#2f5d3a

关键公式：online rescale

每次新 K/V 块来了，行最大值 m 可能变大。已经累积的 l_old 与 O_old 都是相对旧 m_old 算的，必须乘 exp(m_old - m_new) 修正：

l_new = exp(m_old - m_new) * l_old + rowsum(exp(S_ij - m_new))
O_new = exp(m_old - m_new) * O_old + exp(S_ij - m_new) @ V_j

最后才一次除 l 得真正的归一化输出。这就是不需要看到完整行就能算 softmax的秘密。

关键代码

12.3 实现要点

源码：flash_attn_v1.cu（教学简化版）。 为方便读 D=64 固定、Br=Bc=64、单 head、fp32、无 Tensor Core。

// 简化: blockDim.x = Br = 64, 每 thread 处理 Q 的 1 行
// 每 thread 在 shared mem 上维护 O_s[Br][Dm], m_s[Br], l_s[Br]
__shared__ float Qs[Br][Dm], Ks[Bc][Dm], Vs[Bc][Dm], Ss[Br][Bc];

for (int kt = 0; kt < n_kv_tiles; ++kt) {
    /* 1) load K_j, V_j to shared (256 thread 协作) */
    __syncthreads();

    if (tid < Br) {
        /* 2) S_ij = Q_i @ K_j^T * scale */
        for (int j = 0; j < Bc; ++j) {
            float s = 0;
            for (int d = 0; d < Dm; ++d) s += Qs[tid][d] * Ks[j][d];
            Ss[tid][j] = s * scale + causal_mask(tid, j);
        }
        /* 3) online softmax merge */
        float m_old = m_s[tid], l_old = l_s[tid];
        float m_new = fmaxf(m_old, rowmax(Ss[tid]));
        float l_new = exp(m_old - m_new) * l_old + rowsum(exp(Ss[tid] - m_new));
        float rescale = exp(m_old - m_new);
        /* 4) O update */
        for (int d = 0; d < Dm; ++d) {
            float acc = O_s[tid][d] * rescale;
            for (int j = 0; j < Bc; ++j) acc += exp(Ss[tid][j] - m_new) * Vs[j][d];
            O_s[tid][d] = acc;
        }
        m_s[tid] = m_new; l_s[tid] = l_new;
    }
    __syncthreads();
}
/* 5) O /= l, 写回 */

12.4 v1 → v2 的进化

本章教学版接近 v1。生产实现（FA v2/v3、xformers、CUTLASS-FA）都用：

• Tensor Core (mma.sync) 跑 S 和 PV 的两个 GEMM
• cp.async 重叠 K/V 加载和计算
• warp specialization（FA v3, Hopper）让一组 warp 做 load 另一组做 compute
• swizzled shared layout 消除 bank conflict

运行结果

先验证 teaching FlashAttention 与朴素 attention 的逐元素误差，再采集内存与时间指标。

性能数据

12.5 性能采集表（需在目标 GPU 上实测）

本章教学版只追上"FlashAttention 的思想"，没追上"FlashAttention 的工程"。真要做工业用版，去看 官方仓库。

12.6 与 vLLM PagedAttention 的关系

FlashAttention 优化的是单个序列的 attention 计算。 vLLM 的 PagedAttention 优化的是多请求并发时 KV cache 的显存管理——它把 KV cache 按"页"管理，避免显存碎片化。 两者正交，可叠加：vLLM 的 attention kernel 实际是 PagedAttention 调用 FlashAttention 风格 fused kernel。

自检清单

练习题

1. 把本章 kernel 改成支持非整数倍 T（最后一个 Q tile 不满）。
2. 实现 v2 风格的"外循环改 Q tile, 内循环 K/V"，并对比 GFLOPS。
3. fp16 输入 fp32 累加版本（提示：把 Q/K/V 改 __half，shared 也改 __half，acc 仍 float）。
4. 用 Triton 实现同一个 kernel（参考 Triton 官方 tutorial 06），对比代码量。

12.9 工业实战：FA v2/v3、KV cache、PagedAttention、长上下文

12.9.1 FA v1 → v2 的三个核心改动

1. 外循环换成 Q tile：v1 外循环 K/V → 每个 Q tile 反复加载；v2 外循环 Q → Q tile 加载一次保留在寄存器，K/V tile 流过。每个 Q tile 独立 → block 间并行更高。
2. 减少 rescale 次数：v1 每个 K/V tile 都把已累加的 O 乘 exp(m_old - m_new)；v2 把这步推迟到 K/V 循环结束统一做一次。
3. warp 切分更细：v1 每 warp 算 S 的一行 → warp 间需要 row-reduce；v2 每 warp 算 S 的一列 → warp 间几乎独立。

评估 v2/v3 时不要只看 TFLOPS，还要记录有效 HBM bytes、L2 hit、Tensor pipe active、barrier stall 和端到端 tokens/s。公开论文数字只能作为复现目标，不能替代本机 profile。

12.9.2 FA v3 (Hopper)：TMA + warp specialization + fp8

• TMA 一条指令加载 64×64 tile，硬件自动 swizzle
• warpgroup mma (wgmma)：4 个 warp 协作算一个大 mma
• warp specialization：4 warp 分两组——一组 producer（只 load），一组 consumer（只算），mbarrier 同步。软件版乱序
• fp8：降低输入/中间字节并提高 Tensor Core 理论吞吐，但要额外处理量化 scale、amax 和精度回归

12.9.3 与 KV cache 集成（decode 路径）

推理 decode 时 attention 输入：

Q_new     : (1, D)             新一步 query
K_cache   : (T_so_far, D)      历史 K
V_cache   : (T_so_far, D)      历史 V

// FA decode kernel 骨架: Q 只有 1 行, K/V tile 流过
__global__ void flash_attn_decode(
    const __half* Q,                // (1, n_head, D_head)
    const __half* K_cache,          // (T_max, n_head_kv, D_head)
    const __half* V_cache, __half* O,
    int T_so_far, int n_head, int n_head_kv, int D_head)
{
    int h    = blockIdx.x;
    int h_kv = h / (n_head / n_head_kv);   // GQA mapping

    __shared__ __half Qs[D_head];           // Q 整行直接进 shared
    if (threadIdx.x < D_head) Qs[threadIdx.x] = Q[h * D_head + threadIdx.x];
    __syncthreads();

    float m = -INFINITY, l = 0;
    float O_acc[D_head] = {0};
    for (int t0 = 0; t0 < T_so_far; t0 += BLOCK_T) {
        // load K/V tile of BLOCK_T tokens, compute QK, online softmax, update O
    }
    // 写回
}

12.9.4 PagedAttention — 降低 KV cache 内部碎片

批量推理 N 个请求长度各异，contiguous 分配 max_seq_len 会产生大量内部碎片。PagedAttention 借鉴 OS 分页：

• KV cache 切成固定大小 block（典型 16 token / block）
• 每请求维护 block_table（逻辑 → 物理 block 映射）
• 新 token 来了 alloc 一个 block
• attention kernel 通过 block table indirection 读 K/V

// 朴素:
const __half* k = K_cache + t * stride;

// Paged:
int block_id    = block_table[req_id * max_blocks + t / BLOCK_SIZE];
int offset_in_b = t % BLOCK_SIZE;
const __half* k = K_blocks + block_id * (BLOCK_SIZE * D_head) + offset_in_b * D_head;

PagedAttention 的收益来自更高的 KV pool 利用率和更稳的 continuous batching；代价是 kernel 多一次 block table indirection。实际吞吐要同时记录 batch size、有效 token 数、页大小、L2 hit 和 indirection 带来的 stall。

12.9.5 Sliding Window Attention (Mistral)

每个 Q 只 attend 最近 W=4096 个 K/V。FA kernel 在外循环跳过 out-of-window tile：

if (kt_end <= max(0, qi - W)) continue;   // tile 完全在 window 外
if (kt_start > qi)             continue;   // causal 跳过

T=32K 时复杂度从 O(T²) 降到 O(T·W)。Mistral 7B 用 8K window 实现 32K 上下文。

12.9.6 100K+ 长上下文的额外手段

• Ring Attention：多卡分担 K/V，环形传递，1M context 可行
• YaRN / NTK-aware RoPE：把 4K 训练模型外推到 100K（不重训）
• State-space (Mamba)：放弃 attention，O(T) 复杂度
• KV cache 量化：fp16 → int8 / int4，显存减半到 1/4
• Cache 卸载：冷 KV 放 CPU memory，需要时拉回

12.9.7 用现成库 vs 自己写

12.10 研究前沿（2025-2026）：FA v3/v4、FlashMLA、FlashInfer、FlexAttention

12.10.1 FlashAttention v3（2024.07）— Hopper 完整解法

论文标题就是"Fast and Accurate Attention with Asynchrony and Low-precision"。三大武器：

1. TMA bulk async load：替代手写 cp.async 循环。一条 PTX 指令拷 64×64 tile
2. warp specialization：producer warp 跑 TMA，consumer warps 跑 wgmma（见 Ch6.10.1）
3. FP8 + 块量化：Q、K、V、S 全 fp8，accumulator fp32

12.10.2 FlashMLA（DeepSeek 2025）— MLA 专用

MLA（见 11.9.1）的 attention 不一样：cache 的是低维 c_kv（d_c=512），attention 内部要做 c_kv 的 up-projection。FlashMLA 是为 MLA 设计的 FA 变种：

• 主 loop 直接读 c_kv 而非完整 K/V
• Q 端预先吸收 W_UK，attention 内不再 up-project K
• V 的 up-project 留到 PV 阶段，融合进 mma 链
• tile shape 针对 d_c=512 优化（比标准 d_head=128 高 4×）

开源仓库：github.com/deepseek-ai/FlashMLA。复现时重点看 decode 阶段的 Tensor pipe、HBM throughput、TMA load 效率和 MLA up-projection 是否被融合。

12.10.3 FlashInfer — attention kernel dispatch 层

vLLM、SGLang、Llama-stack 越来越倾向不自己写 attention，统一通过 FlashInfer 调用。它的设计：

import flashinfer
# 用户传 KV layout + 模型配置, FlashInfer 选最优 kernel
o = flashinfer.batch_decode_with_paged_kv_cache(
    q, paged_k_cache, paged_v_cache, block_tables, kv_lens,
    pos_encoding_mode="ROPE_LLAMA",
    kv_layout="NHD",       # paged 布局风格
    use_tensor_cores=True,
    sm_scale=None,
)

支持矩阵：

• MHA / MQA / GQA / MLA
• Causal / sliding window / sparse mask
• RoPE / ALiBi / 无 PE
• fp16 / bf16 / fp8 / fp4
• Paged KV / contiguous KV / radix-cached

2025-2026 工业事实：自己手写 attention 的越来越少，除非有特别的研究 motivation。生产团队都接入 FlashInfer 或类似 dispatch 层。

12.10.4 PyTorch FlexAttention（2024.10）

Meta 发布的 PyTorch 2.5+ 新 API，把"自定义 attention mask + score_mod"编译成高效 FA-style kernel：

from torch.nn.attention.flex_attention import flex_attention, create_block_mask

def causal_mask(b, h, q_idx, kv_idx):
    return q_idx >= kv_idx

def alibi_score(score, b, h, q_idx, kv_idx):
    bias = (q_idx - kv_idx) * alibi_slope[h]
    return score + bias

mask = create_block_mask(causal_mask, B, H, M, N)
out = flex_attention(q, k, v, score_mod=alibi_score, block_mask=mask)
# 编译器自动产出 FA 风格 fused kernel

对研究者最大价值：试一个新 mask / 新 score 变体不用写 CUDA。性能是否接近手写 FA 取决于 mask 规则、动态 shape、编译器生成的 tile 和目标 GPU，需要对照生产 kernel 测。

12.10.5 Ring Attention & Stripe Attention — 跨 GPU 长 context

单 GPU 装不下 1M token 的 K/V。Ring Attention（Liu et al., 2023, 2024 工业化）：

1. 把 K, V 沿 T 维切到 N 张 GPU, 每张持一段
2. Q 也切到 N 张
3. K/V 在环上传递: GPU_i 收完 GPU_(i-1) 的 K/V → 算一部分 attention → 传给 GPU_(i+1)
4. 一圈后所有 GPU 都看到全部 K/V, 但峰值显存只是 1/N

延迟: 依赖通信/计算 overlap；用 Nsight Systems 记录 NVLink/NCCL row 与 attention kernel 的重叠

Stripe Attention（Brandon et al., 2023）：跟 Ring 同思路但 chunk 设计对 causal mask 更友好，长 context 训练用得多。

实战：Llama 4 / Gemini 1.5 系列长 context（1-10M token）用的就是这类技术。NVSHMEM（见 8.9.5）替代 NCCL 让通信粒度更细。

12.10.6 Striped + Chunked + Prefill 混合调度

当前 vLLM V1 用 unified scheduler 在固定 token budget 下统一调度 prompt 与 output token，并由此组合 chunked prefill、 prefix caching 和 speculative decoding；SGLang 与 TensorRT-LLM 有各自的调度实现，不能只凭版本号假设行为相同：

graph TB
    A["用户请求队列"] --> B["调度器"]
    B --> C["chunked prefill
(64-512 token 一片)"]
    B --> D["decode batch
(已 prefill 的请求)"]
    C --> E["在 prefill chunk 末用 online softmax
merge 之前 prefix"]
    D --> F["FlashInfer paged decode"]
    E --> G["统一 attention kernel"]
    F --> G
    G --> H["输出 token"]
    style C fill:#f3f1e8,stroke:#a86420
    style D fill:#f3f1e8,stroke:#2f5d3a

核心组件：

• Chunked Prefill：长 prompt 分多 chunk，每 chunk 跟 decode 同 step 跑，消除 prefill 阻塞 decode
• Prefix Caching：相同 system prompt 命中 KV cache, 跳过 prefill
• Disaggregated Serving：prefill 节点 + decode 节点分离, 各自优化（见 Ch14.6）

12.10.7 2026 attention kernel 推荐栈

常见坑

• online softmax 的 rescale 漏掉 → output 看起来正常但量级错
• causal mask 在 tile 边界算错（应该 j_global > i_global 而非 j_local）
• shared 数组超 48 KB → 需 cudaFuncSetAttribute 解锁
• fp16 累加 → 长序列下数值崩
• PagedAttention block_table 越界 → illegal memory access；block_size 必须是 16 倍数对齐
• GQA 时 K/V head 索引算错 (n_head / n_head_kv 整除) → 输出乱
• sliding window + causal mask 双重判断写颠倒 → 边界 attend 错误
• FA v3 warp specialization 时 producer/consumer warp 数比例错 → producer 太少跑不饱 consumer
• FlashInfer 配 KV layout 时 NHD vs HND 选错 → 访问模式异常或正确性失败；用文档、shape assert 与 profile 定位
• FlexAttention score_mod 用 Python 闭包变量 → 每次 launch 都 recompile, 慢死

12.12 CUDA 官方手册精讲（CUDA Programming Guide 13.2（核验：2026-07-20））

FA v1/v2/v3 数学与 Hopper 实现、Producer-Consumer、PagedAttn

FA v3 producer / consumer 完整最小骨架

12.9.2 把 FA v3 描述成"4 warp 分两组 producer/consumer"——下面给出基于 CUDA Programming Guide 13.2 的 Pipelines 与 Asynchronous Data Copies 章节（cuda::barrier + cp.async.bulk.tensor）的最小可读骨架（约 80 行，省略 wgmma 内联汇编细节）。

#include <cuda/barrier>
    #include <cuda/ptx>
    namespace cg  = cooperative_groups;
    namespace ptx = cuda::ptx;

    constexpr int Br = 64, Bc = 64, D = 128;
    constexpr int N_STAGE = 2;                  // double buffer
    constexpr int N_WARPS = 12;                 // 4 producer + 8 consumer

    __global__ __launch_bounds__(N_WARPS * 32)
    void fa3_skeleton(
        const __grid_constant__ CUtensorMap q_tmap,    // (T, D) fp16
        const __grid_constant__ CUtensorMap k_tmap,
        const __grid_constant__ CUtensorMap v_tmap,
        float* O_lse,                                   // (T, 2) logsumexp 输出
        __half* O, int T, float scale)
    {
        // ---------- shared memory: double-buffered K/V ----------
        __shared__ alignas(1024) __half Qs[Br][D];
        __shared__ alignas(1024) __half Ks[N_STAGE][Bc][D];
        __shared__ alignas(1024) __half Vs[N_STAGE][Bc][D];

        #pragma nv_diag_suppress static_var_with_dynamic_init
        __shared__ cuda::barrier<cuda::thread_scope_block> kv_full[N_STAGE];
        __shared__ cuda::barrier<cuda::thread_scope_block> kv_empty[N_STAGE];

        int warp_id = threadIdx.x / 32;
        int lane    = threadIdx.x % 32;
        bool is_producer = (warp_id < 4);                // warps 0..3 = producer

        if (threadIdx.x == 0) {
            for (int s = 0; s < N_STAGE; ++s) {
                init(&kv_full[s],  N_WARPS * 32);        // 全 block 等
                init(&kv_empty[s], 1);                   // 1 个 producer warp 触发
                kv_empty[s].arrive();                    // 初始 empty=1
            }
        }
        __syncthreads();

        int qi = blockIdx.x * Br;
        int n_kv_tiles = (T + Bc - 1) / Bc;

        // ---------- Q 加载一次, 整个 K/V loop 复用 ----------
        if (warp_id == 0 && lane == 0) {
            int crd[2] = {0, qi};
            ptx::cp_async_bulk_tensor(
                ptx::space_shared, ptx::space_global,
                &Qs, &q_tmap, crd,
                cuda::device::barrier_native_handle(kv_full[0]));
        }

        // ---------- producer 路径 ----------
        if (is_producer) {
            for (int t = 0; t < n_kv_tiles; ++t) {
                int s = t % N_STAGE;
                // 等 consumer 把上一份 K/V 用完
                kv_empty[s].arrive_and_wait();

                if (warp_id == 0 && lane == 0) {
                    int crd[2] = {0, t * Bc};
                    ptx::cp_async_bulk_tensor(
                        ptx::space_shared, ptx::space_global,
                        &Ks[s], &k_tmap, crd,
                        cuda::device::barrier_native_handle(kv_full[s]));
                    ptx::cp_async_bulk_tensor(
                        ptx::space_shared, ptx::space_global,
                        &Vs[s], &v_tmap, crd,
                        cuda::device::barrier_native_handle(kv_full[s]));
                    cuda::device::barrier_arrive_tx(
                        kv_full[s], 1, sizeof(Ks[0]) + sizeof(Vs[0]));
                }
            }
        }
        // ---------- consumer 路径 ----------
        else {
            float m_i = -INFINITY, l_i = 0.f;
            float O_acc[D / 8] = {0};                    // 跨 8 个 consumer warp 切分

            for (int t = 0; t < n_kv_tiles; ++t) {
                int s = t % N_STAGE;
                kv_full[s].arrive_and_wait();            // 等 producer 装好

                // 1) S = Q @ K^T  via wgmma (省略 PTX 内联)
                //    每 consumer warp 处理 S 的 1/8 列
                float S_tile[Br / 8][Bc];
                wgmma_qkt(Qs, Ks[s], S_tile, scale);

                // 2) online softmax merge (10.2 节算法)
                float m_new = max(m_i, rowmax(S_tile));
                float P_tile[Br / 8][Bc];
                float scale_factor = __expf(m_i - m_new);
                l_i = scale_factor * l_i + rowsum_exp(S_tile, m_new, P_tile);
                // 3) O = scale_factor * O + P @ V  via wgmma
                wgmma_pv(P_tile, Vs[s], O_acc, scale_factor);
                m_i = m_new;

                kv_empty[s].arrive();                    // 通知 producer: 这一 slot 可重写
            }

            // ---------- epilogue ----------
            float inv = 1.f / l_i;
            for (int d = 0; d < D / 8; ++d) O_acc[d] *= inv;
            // O 写回 (省略, 用 TMA store)
            if (lane == 0) O_lse[qi + warp_id - 4] = m_i + __logf(l_i);
        }
    }
    

实际 FA v3 源码（github.com/Dao-AILab/flash-attention/csrc/flash_attn_v3）还做了：epilogue overlap、async wgmma issue 提前 2 个 K/V tile、TMA store 也异步——总共 1500+ 行 PTX 内联，但骨架就是上面这 80 行。

FlashAttention-4：Blackwell 异步流水与软件调度

FlashAttention-4 论文 已于 2026 年 3 月公开，不再是预告或 RFC。论文针对 Blackwell 中 Tensor Core 吞吐增长快于 shared-memory bandwidth 与 exponential units 的非对称扩展，联合调整算法和 kernel pipeline。

Paged + TMA：FlashInfer 在 Hopper 上的 KV indirection 新方案

12.9.4 给的 PagedAttention indirection 写法（每 token 算 block_id = block_table[t / B]）在 Ampere/Hopper 都能跑，但 没法配 TMA——TMA 要求一整个 64×64 tile 在 global 是连续的，paged KV 偏偏是分散的。

FlashInfer 2024 末的方案：使用 CUDA Programming Guide 13.2 Asynchronous Data Copies / Tensor Maps 中的 device-side tensor map modification，让 kernel 内部按需修改 TMA descriptor 的 global_address 字段：

// nvcc -arch sm_90a   (仅 Hopper+)
    #include <cuda/ptx>
    namespace ptx = cuda::ptx;

    __launch_bounds__(32)
    __global__ void paged_fa_load(
        const __grid_constant__ CUtensorMap template_tmap,    // 通用模板
        const void* const* page_ptrs,                          // 每 page 的物理地址
        const int* block_table,                                // 逻辑 → page id
        __half* dst_smem, int t_logical)
    {
        __shared__ alignas(128) CUtensorMap tmap;
        if (threadIdx.x == 0) {
            tmap = template_tmap;                              // copy template
            int page_id = block_table[t_logical / 16];
            // 把 global_address 字段就地改成这个 page 的物理地址
            ptx::tensormap_replace_global_address(
                ptx::space_shared, &tmap, page_ptrs[page_id]);
        }
        __syncwarp();

        // release-acquire fence (PG §4.11.2.2.3 强制)
        ptx::n32_t<128> bytes_128;
        ptx::tensormap_cp_fenceproxy(
            ptx::sem_release, ptx::scope_cta, &tmap, &tmap, bytes_128);

        // 现在可以用 tmap 触发 TMA, 像普通 dense 加载一样
        if (threadIdx.x == 0) {
            int crd[2] = {0, (t_logical % 16) * D_head};
            ptx::cp_async_bulk_tensor(
                ptx::space_shared, ptx::space_global,
                dst_smem, &tmap, crd, /*barrier*/ ...);
        }
    }
    

比较 vLLM 原版 paged attention 与 FlashInfer paged-TMA 时，固定模型、batch、context、KV layout 和 page size，记录每 step 时间、L2 hit、TMA sector、register spill 与 page table load 数。

    graph LR
        BT["block_table[req][t/16]
(int32 per token)"] --> Mod["device-side
tensormap_replace_global_address"]
        Mod --> Fence["release-acquire
fence_proxy_tensormap"]
        Fence --> TMA["cp.async.bulk.tensor
(连续 16 token 一次搬)"]
        TMA --> Smem["Ks[i] in shared (swizzle 128B)"]
        Smem --> WGMMA["consumer warpgroup
wgmma + online softmax"]
        style Mod fill:#f3f1e8,stroke:#a86420
        style WGMMA fill:#f3f1e8,stroke:#2f5d3a
    

这套机制也解释了为什么 MLA 的 c_kv tile 更适合 TMA：tile 更宽时，单次搬运的事务字节更大，固定开销更容易被摊薄。真实带宽利用率仍然需要 profile。

Cluster Launch Control：独立的 Blackwell 调度原语

CUDA Programming Guide 13.2 的 Work Stealing with Cluster Launch Control 介绍了 cancellable thread block 和 work stealing。它适合单独学习与测量，但不是 FlashAttention-4 论文所列机制的替代描述。本仓库不把未经 B200/B300 验证的 sliding-window cancel 示例标成 FA4 实现；相关收益统一保留为 TODO(on GPU)。

下一章导览

第 13 章补齐 RoPE、SwiGLU、KV cache 与 sampling，为完整推理循环准备最后一组零件。