概述

本文档详细walkthrough fa4中在SM100/SM90下实现的forward code (serving暂不考虑backward)。

Walkthrough

flash_fwd_sm100.py

flashattention-4采用了warp specialization的技术，具体如

FlashAttention4 精读
概述

本文档详解flashattention4的具体实现和设计，其中，将核心分为paper read和code walkthrough两个部分；

Paper Read

相比fa3主要为Hopper GPU进行优化，通过异步的执行和warp特殊调度，fa4主要针对Blackwell架构进行优化，其中主要面对不对称的硬件性能scaling，分别是，tensor core的吞吐翻倍了，然而其他的功能性组件如shared memory带宽、指数单元增长却很慢，因此，fa4主要提出了如下三个方面的核心创新，分别是

重新设计pipeline来完全利用异步的MMA操作以及更大的tile size；

软件模拟的指数运算与条件性 softmax 缩放重算，以减少非矩阵乘法操作；

利用Tensor内存和2-CTA MMA模式来降低共享内存冲突和backward pass中的atomic add操作；

整个fa4是利用CuTe-DSL (Python)实现的，比传统基于C++的方案快20-30倍；

1/2 Introduction

Transformer是近期LLM的核心架构，其中attention非常重要，然后实现高效的attention面对一个不对称的硬件升级，即tensor core的计算能力成倍增长，然而其他组件如shared memory大小/特殊计算单元增长却很慢，导致了pipeline的设计需要很careful；
最开始TriDao提出Flashattention，通过tile和kernel fusion来消除中间到global memory的read/write；在基础上提出的fa2，核心在seq_len维度并行，fa3通过实现细粒度的warp specialization执行、并支持FP8，然而fa3主要优化的还是Hopper架构，在Blackwell上并不是最优，下面提供一些具体的Blackwell和Hopper的差别

Tensor core：Blackwell相比Hopper两倍了tensor core吞吐，BF16 MMA 吞吐8192 ops/clock/SM，Hopper 4096，这个值可以从理论最大FLOPS中计算得出，2.25PFLOPS/1850Mhz clock speed / 148SMs = 8192 ops/clock/SM；

Exponential unit，B200/GB200和Hopper一致，都是16 opus/clock/SM，B300/GB300将这个值提升到32 ops/clock/SM；

SMEM：Blackwell和Hopper保持一致，shared memory读写吞吐都是128 bytes/clock/SM；
同时，Blackwell提供额外的架构创新，保持扩每个SM新增256KB的TMEM (Tensor memory) 来直接储存中间的tensor core结果，MMA的tiles从Hopper的64 x 128提升到了 128 x 128，并直接给TMEM新增了完全异步的Tensor core操作；这些值为后续理论分析attention计算的bottleneck提供重要支撑；
具体的Attention forward和backward计算示意图如下：

3 Algorithm

3.1 Algorithm Forward Pass

本节系统分析目前在Blackwell架构上的系统pipeline bottleneck，并详细介绍为什么会提出一系列的flashattention的优化改进；

3.1.1 Feeds and speeds

在attention计算时，设计算Q和K的tile dimension为M x N，设置head dimension为d，随后分别分析MMA计算 /读写SMEM流量 /指数单元计算 (Softmax)的总时间开销

MMA计算：attention计算时核心存在两个MMA操作，分别是 $Q K^{T}$ (MxN output, Mxd and dxN)和 $P V$ (Mxd output, MxN and Nxd)，每个MMA需要2MND浮点数操作 (乘一次，加一次)，考虑到前面介绍的8192 FLOPS/cycle，因此执行时间为 == $T_{M M A} = (4 M N D) /8192 cycles$ ；

SMEM流量：对于两个MMA操作，第一个操作Q和K都需要从Shared memory读，因此具体要read M/128 x N/128 x 256d，对于PV操作，P在TMEM内，因此没有访寸开销，此时总开销位 M/128 x d/128 x 128N，因此考虑到每元素为BF16，读写速度为128 bytes/clock/SM，因此总时间为 3MND/8192 cycles；

指数单元：指数单元需要在M x N总操作上做softmax，考虑到B200上该数值为 MN/16 cycles；
对于常见的head dimension=128，MxN为256x128和128x128，具体的forward pass的数据如下

因此可以看到，这给我们kernel设计的时候有三个如下的motivation

需要更大的tile size和最大化overlap MMA计算操作和Softmax 操作；

需要增大指数单元的吞吐，利用其他的硬件单元；

减少不必要的non-matmual操作

新pipeline来overlap matmul和softmax

由于tensor core能力翻倍，因此在Blackwell上，如何更好的把tensor core操作和softmax操作做overlap比Hopper更重要，整体上仍然采用类似FA3的ping-pong的方式，two tile的output互相overlap。一个tile执行MMA操作，另一个tile则执行softmax操作，Hopper会把accumulator结果放在寄存器中，而Blackwell会把他放在TMEM中，另外，Blackwell的tile是128x128大小，而Hopper的tile是64x128大小；
因此，一个自然的方法也就是利用2个总共有128 threads的warpgroup来处理softmax的每一行，我们这里叫他softmax warpgroup，和FA3，这里在两个softmax warpgroup的关键段落处同步来保证他们不要同时去执行softmax操作，因为指数单元能力固定，两个并行执行不如串行执行，能够让某个tile尽快进入下一个stage。
和FA3最大的区别在于，P是放在TMEM内，而不是放在寄存器里，因而采用单独的一个correction warp来对O做rescale，因为P的存储不再处于critical path上了；
Softmax WG:
计算新 row max
      ↓
得到 scale α
      ↓
继续计算并写出 P_j
          ──────────────────────┐
                                │ 并行
Correction WG:                  │
TMEM 读取旧 O                   │
      ↓                         │
寄存器中分块执行 O ← αO         │
      ↓                         │
写回 TMEM                       │
          ──────────────────────┘
                    ↓
MMA WG 执行 P_j V_j，并累加到 TMEM 中的 O
为了实现这样的pipeline overlap，如何管理TMEM的内存分配非常重要，因为每个SM只有256KB的TMEM大小，首先Output必须得放在TMEM里，格式为FP32，元素大小为128x128x4x2=128KB，因而，O就要占据TMEM中一半大小的空间，剩下的S和P在TMEM中的存储只能放剩下的一半，S格式是FP32，P是BF16，因此每个S要占据64KB，每个P要占据32KB，因此128KB可以放得下2个S或者4个P；我们有两种做法来存放S和P，一种是放一个S和2个P，一种是放2个S和P做overlap，显然采用第二种，一个可能的TMEM示例图如下，由于一部分S算完以后就能腾出空间放P，而S的大小要大于P，因而显然，TMEM中还能剩余一部分空间来给correction warpgroup通信rescale信息。
TMEM:
┌──────────┬──────────┬────────────────┬────────────────┐
│ Oᴴ FP32  │ Oᴸ FP32  │ Sᴴ → Pᴴ 复用区 │ Sᴸ → Pᴸ 复用区 │
└──────────┴──────────┴────────────────┴────────────────┘
 
启动：
MMA WG        : QᴴKᵀ → Sᴴ     QᴸKᵀ → Sᴸ
 
并行：
Softmax WG-H  : load Sᴴ → max → exp → 分段写 Pᴴ
Softmax WG-L  : load Sᴸ → max → exp → 分段写 Pᴸ
Correction WG : 根据统计量 rescale O
MMA WG        : P 前 3/4 ready 后提前开始 PV
另一个Blackwell tile size的issue是，显然，计算max时，寄存器需要把完整的128个元素放在寄存器里 (每个thread)，由于同时存在四个warpgroup，两个softmax warpgroups、1个correction warpgroups和一个触发tensor core和TMA单元的warpgroup，因此分配充足的寄存器给softmax很重要；对于BF16的input，需要128个input寄存器和64个output寄存器 (S是FP32，P是BF16），每个thread最多只有256个寄存器，此外还要保存别的如旧 max 和 rescale factor、exponentiation 临时变量、多项式 exp 的中间值等一系列结果到寄存器中，因而寄存器很容易逼近上限。如果寄存器不足发生 spill，数据会被放到 local memory，最终通常落入 L1/L2，性能会严重下降。
因此，fa4采用stage out storing P，具体做法是，P未按照dimension维度，切成3/4和1/4两个部分，前面3/4先算P，算完以后可以直接触发MMA操作，把P从寄存器搬到TMEM中，因此，不需要全量在寄存器中放所有128个输出P，也不需要完整的放128个S，而是流式的把S从TMEM搬到寄存器，然后算P，算完96个，就先搬走到TMEM中算MMA；

3.1.3 Emulation of the exponential function

指数操作一般是MUFU单元来做的，指数单元能力 (16)比MMA能力 (8192)差特别多，因此，指数操作是attention kernel中的一个bottleneck。本节的核心是一部分元素继续使用硬件 MUFU.EX2，另一部分元素使用普通浮点 FMA 指令，通过多项式近似计算指数，从而同时利用两类执行单元。GPU 的硬件指令通常是 MUFU.EX2，直接计算： $2^{x}$ ，因此，实际上是先计算 $x = (S_{ij} - m_{i}) lo g_{2} e$ ，再计算 $2^{x}$ ，注意此处softmax已经减去了行最大值，因此不需要担心正方向溢出。
本论文采用经典的cody-waite range reduction，核心分解非常简单，即 $x = n + f$ ，其中 $n = ⌊ x ⌋, f = x - ⌊ x ⌋ \in [0, 1)$ ，于是有 $2^{x} = 2^{n + f} = 2^{n} 2^{f}$ ，因此原问题转化了两个更好计算的子问题 1）计算整数次幂；2）计算 $2^{f}$ ：只需要在固定区间 $f \in [0, 1)$ 上近似。IEEE 754 FP32 的正规浮点数可以写为： $(- 1)^{s} \times 2^{E - 127} \times (1. M)$ ，其中E是8位指数域，127是FP32的exponent bias，M是尾数域，对于纯粹的 $2^{n}$ ，有效数字正好是1.0，尾数全0，因此其指数域只需要写为E = n + 127，对应的FP32 bit pattern位 (n+127) << 23 (23即最后的尾数全0)，例如 $2^{3} = 8.0$ ，FP32指数域是 $3 + 127 = 130$ ，因此不需要真正执行指数函数，只需要构造指数位。而对于浮点数位，由于 $f \in [0, 1)$ ，因而计算目标 $2^{f}$ 一定 $\in [1, 2)$ ，所以q的指数位一定是127，因而bit pattern可以表示为 $bits (q) = (127 ≪ 23) + mantissa (q)$ ，于是有 $bits (2^{n} q) = ((127 + n) ≪ 23) + mantissa (q)$ ，于是就有 $bits (2^{n} q) = bits (q) + (n ≪ 23)$ ，因而，最终的浮点数计算其实不需要额外的浮点乘法，可以采用如下三步骤来解决：

把q reinterperet 成整数

给exponent field加上n

在reinterpret回FP32
为了快速计算整数部分n，引入了magic bias： $C = 2^{23} + 2^{22} = 1.5 \times 2^{23}$ ，即 $C = 12582912$ ，然后通过指定向下舍入模式完成 $n = (x + C) - C$ ，这里利用了FP32尾数有23个显式fraction bits的特点，当给一个较小的x加上很大的C后，结果所在数量级的ULP约为1，因此自然就保留下来的就是n了 (小数点被23位給mask掉了)；
关于小数部分计算，核心采用多项式： $2^{f} \approx p_{0} + p_{1} f + p_{2} f^{2} + \dots + p_{n} f^{n}$ ，其中p0为1，且为了减小误差，这些系数都不是简单的泰勒展开，而是使用Sollya工具求出的优化系数，目标是在整个 $[0, 1)$ 区间最小化相对误差，即，它更接近minimax polynomial。
在此基础上，进一步采用Horner方法优化计算，直接计算 $p_{0} + p_{1} f + p_{2} f^{2} + p_{3} f^{3}$ 较慢，需要乘法构造 $f^{2}, f^{3}$ ，而Horner形式可以写成 $p (f) = ((p_{3} f + p_{2}) f + p_{1}) f + p_{0}$ ，这样相当于只用做3次FMA操作就行了。
论文通过实际数据证明degree-3的误差已经够用：具体的FP32原始误差如下
| 方法 | 最大相对误差 | 平均相对误差 |
| ---------- | ------------------: | ------------------: |
| MUFU.EX2 | $(1.41 \times 1 0^{- 7})$ | $(3.04 \times 1 0^{- 8})$ |
| degree 3 | $(8.77 \times 1 0^{- 5})$ | $(5.43 \times 1 0^{- 5})$ |
| degree 4 | $(3.05 \times 1 0^{- 6})$ | $(1.84 \times 1 0^{- 6})$ |
| degree 5 | $(1.44 \times 1 0^{- 7})$ | $(5.48 \times 1 0^{- 8})$ |
可以看到，degree3确实远不如MUFU.EX2，然而，Flashattention的需要把P转换成BF16，此时光光BF16的误差就达到了3.9x10^{-3}，因此degree3的误差比BF16自身的量化误差要小44倍，因此在degree3下，degree的结果再约99%输入上与硬件结果相差不超过1个BF16 ULP；
Partial emulation：虽然走多项式估计能够节约MUFU指令，然而，他会带来额外的寄存器开销、更大的寄存器带宽花费、更高的延迟，因而，只对10%-25%的entry做多项式估计，剩余的entry仍然通过MUFU.EX2进行计算，具体的fraction根据经验进行微调，在给定的tile配置下根据MMA和指数单元的吞吐。

3.1.4 Skipping online softmax rescaling

直接切换rescale为如下的公式，仅当rescale差距超过256的时候，才进行更新一次，不然就延迟更新，能够大幅度减少rescaling操作。在实际运行中，为了避免warp内的threads多样性，不让他们走不同分支，只要32个threads有任何一个线程需要resacle，整个warp的的所有线程都执行rescale。

3.2 Attention backward pass

不用特别关注，毕竟推理没有backward pass

3.2.1 Feeds and Speeds

和forward pass类似，我们首先提供指导关于我们的kernel设计和优化的motivation，基于具体计算出的tensor core, smem访问和指数单元的所需时间；

MMA计算：backward pass需要5个MMA操作，每个MMA包括一个MxN的matrix，一个Mxd的matrix和一个dxN的matrix，需要2MND的浮点数操作，总共10MND的浮点数操作，因此需要时间为 10MND / 8192 cycles。

SMEM流量：其中三个操作= KQ⊤, dP⊤ = VdO⊤, and dQ = dSK是shared-shared操作，而剩下两个操作dV = P⊤dO and dK = dS⊤Q是Tensor-shared操作，因而Shared memory总带宽为 (4md+3ND+MN) / 64 cycles，同时，考虑到算法还有额外把立即数梯度dDS (MxN)以BF16写到Shared memory，dQ (Mxd)要以FP32写到shared memory，随后利用TMA读回来做reduction，总共8MD的SMEM read/write，因此总时间为 (4MD+3ND+MN)/64 + MN/64 + Md/16 cycles。

指数单元：仍然为MN/16 cycles。

可以看到，在M=N=d=128的场景下，SMEM流量占据3328个cycle，超过了MMA计算时间 (2560 cycles)和指数计算单元 (1024 cycles)，显示SMEM读写成为主要的bottleneck。

3.2.2 New pipeline to overlap matmul and softmax

skip

3.2.3 2-CTA backward pass

需要特意提一下2-CTA backward pass，他能够降低一半的SMEM流量，

3.2.4 Deterministic backward pass

skip

3.3 Scheduling

这一节非常重要，本质是回答一个问题：不同的attention tile按照顺序分配给SM，fa4直接借鉴经典的longest-processing-time-first (LPT)调度的思路；
一个worktile通常由坐标表示： $(mb l oc k, h e a d, ba t c h)$ ，其中mblock代表分配给CTA/block的那个Query tile，head代表attention head，batch代表batch item，在该CTA中，固定处理一块Q_j，和所有的K算。
而这天然带来了一个问题，即不同的worktile的运行时间是不一样的：

普通non-causal、固定长度attention，每个Q tile遍历所有的K/V tile，因此，所有worktile的mainloop长度基本一致

causal attention：假设Q和K block size相同，第m个query block只能看到他之前的key blocks，因此 $T (m) \propto m + 1$ ，假设共有8个query blocks，那么就有如下的表格

mblock 需要处理的 KV blocks 相对工作量
0 $(K_{0})$ 1
1 $(K_{0}, K_{1})$ 2
2 $(K_{0}, K_{1}, K_{2})$ 3
… … …
7 $(K_{0}, \dots, K_{7})$ 8
所以causal attention的tile计算量天然呈三角形，对于varlen attention，情况更复杂：

不同batch的Q长度不同

不同batch的KV长度不同

有的batch是短prefill

有的batch是长context decode

causal和non-causal可能混合

不同tile的inner-loop次数差异很大
因此，这本质上是一个load-imbalanced scheduling problem。假设有4个SM，每个SM不断领取新tile，假设有4个SM，任务耗时分别为 $1, 2, 3, 4, 5, 6, 7, 8$ ，如果按照从短到长分发，那就是
第一轮：
SM0: 1
SM1: 2
SM2: 3
SM3: 4
 
第二轮：
SM0: +5 → 总计 6
SM1: +6 → 总计 8
SM2: +7 → 总计 10
SM3: +8 → 总计 12
总执行时间由最慢的SM决定，也即12s，在最后的尾部，大量的SM已经没事可做，这就是tail effect/wave quantization tail；
使用LPT能够先执行长任务，然后按需给不同的SM分配新任务，这样能保证先开始执行的SM拿到的是更慢的tile计算目标，然而，简单的LPT排序并不优，核心是他可能会影响L2 cache locality。不同的tile可能复用相同的K/V数据，调度顺序如果过于随机，虽然SM load balance了，但却可能破坏KV的L2 cache reuse，因此，真正的优化目标应该是 $减少 SM 尾部不均衡 + 保持 K/V 的 L2 locality$ .
不同batch的KV实际上是完全不同的数据，这取决于sequence具体的token，这也会导致随意的交错可能会导致L2 cache miss，一个示意图如下：
处理 batch 0 的 tile
→ K0/V0 进入 L2
 
立即处理 batch 1
→ K1/V1 进入 L2，挤掉部分 K0/V0
 
又回到 batch 0
→ 之前的 K0/V0 可能已经被驱逐
如果全局LPT把不同batch按任务长度任意混合，那么几乎无法有效复用L2中的KV，因此论文采用了一个naive但有效的策略，即 $batch 始终作为最外层维度$ ，也就是大致如下的伪代码：
for batch in batches:
    process_tiles_of_this_batch()
这样能够在处理一个batch时持续使用它的KV，之后再切换到下一个batch。
对MHA model来说，每个head有自己的KV，如果一个batch有很多heads，所有heads的KV总容量可能大于L2，因此如果按照如下的格式来进行计算的话，最早加载的KV head仍然会被后面的heads挤出L2：
mblock 15:
    head 0, head 1, ..., head 31
 
mblock 14:
    head 0, head 1, ..., head 31
当运行到mblock14时，head0的KV可能早就已经被别的heads给驱逐了，因此，fa4没有一次性处理全部heads，而是把heads划分为若干section：
section 0: heads 0–7
section 1: heads 8–15
section 2: heads 16–23
section 3: heads 24–31
目标是让 $W_{heads in section} ≲ C_{usable L2}$ ，这样同一个section内的KV数据，在处理多个mblocks时仍然有机会留在L2；
最终使用的循环结构大致如下伪代码：
for batch in batches:                         # 最外层
    for head_section in head_sections:
        for mblock in reversed(mblocks):      # LPT：长 tile 先
            for head in head_section:         # section 内 heads
                run_tile(batch, head, mblock)
类似的遍历顺序如下：
batch 0, section 0:
    mblock 3: head 0, head 1
    mblock 2: head 0, head 1
    mblock 1: head 0, head 1
    mblock 0: head 0, head 1
 
batch 0, section 1:
    mblock 3: head 2, head 3
    mblock 2: head 2, head 3
    mblock 1: head 2, head 3
    mblock 0: head 2, head 3
 
然后才进入 batch 1
这个顺序兼顾了load balance，同一个head，mblock从大到小执行，且同时考虑了L2 locality，不随意跨batch，不执行过多的heads导致不同的mblock无法复用KV，因此他不是纯LPT，而是 $cache-aware hierarchical LPT$ 。
对于MQA和GQA，需要做特殊的处理，因为，不同的q head会共享完全相同的k head，按照上面的逻辑，每个head各算各的，在MQA/GQA场景，显然应该先遍历同一个KV head对应的所有相关的query head，概念上的顺序类似，这样能保证，同一个KV head能够尽可能的被共享，这也是为什么MQA8 (4%-8%)的scheduling收益比MQA8 (7%-14%)的收益更大。
for batch in batches:
    for kv_head_section in kv_head_sections:
        for mblock in reversed(mblocks):
            for kv_head in kv_head_section:
                for q_head in query_heads_sharing(kv_head):
                    run_tile(batch, q_head, kv_head, mblock)
对于variable-length attention，不同的batch item长度不同，例如常见的一些case如下：

默认情况下，attention metadata中的batch顺序可能就是用户传入的顺序，这个顺序与计算量没有任何的关系，例如可能出现
很多短 prefill 先执行
最后才执行几个长 context decode
这与casual mblock从短到长的问题相同，长问题被留到最后，产生严重长尾；
论文针对varlen实施LPT的做法是先启动一个preprocessing kernel，这个预处理kernel根据每个batch的

query length；

KV length；

causal状态；

每个worktile的最大mainloop长度；
估算该batch中最重的worktile的执行时间，可以具体抽象为 $w_{b} = max_{m \in tiles of batch b} T (b, m)$ ，然后按照 $w_{b_{0}} \geq w_{b_{1}} \geq \dots$ 对batch进行排序，这里论文特意写的是maximum per-worktile execution time，而不是简单按照总token数或总FLOPS排序。原先是调度单元是tile，kernel的拖尾往往由最慢的tile决定。
当然，嗯的不可能真的把Q/K/V tensor在内存中重新排序，因为那会产生巨大的搬运开销，所以preprocessing kernel只会生成一个permutation：
virtual batch index → actual batch index
例如原始batch顺序是：0，1，2，3，根据预计任务长度排序后是1，2，3，0，那么就生成如下的映射：
virtual 0 → actual 1
virtual 1 → actual 2
virtual 2 → actual 3
virtual 3 → actual 0
attention kernel调度时遍历的是virtual batch 0, 1, 2, 3，真正读取数据时，通过映射找到actual batch 1, 2, 3, 0；
Link to original

mblock	需要处理的 KV blocks	相对工作量
0	$(K_{0})$	1
1	$(K_{0}, K_{1})$	2
2	$(K_{0}, K_{1}, K_{2})$	3
…	…	…
7	$(K_{0}, \dots, K_{7})$	8
所以causal attention的tile计算量天然呈三角形，对于varlen attention，情况更复杂：

中所讲，采用ping-pong的方式来进行MMA和softmax的overlap，其中，共存在两个softmax warpgroup，每个softmax warpgroup包含128个threads也即4个warp，每个thread负责一行softmax计算，因此前0-7号共8个warp作为两个softmax_warp_group，同时，由于TMEM的存在，rescale不再存在关键路径上，P算完以后不用一直放在寄存器内等待rescale操作完成后再累加到O上，而是可以直接把P丢到TMEM以后，立刻释放register压力，因而，维护correction_warpgroup来专门负责rescale操作，同样的，每一行由一个thread来负责，由于rescale操作并不那么频繁，因而不用额外维护两个warpgroup；

self.softmax0_warp_ids = (0, 1, 2, 3)
self.softmax1_warp_ids = (4, 5, 6, 7)
self.correction_warp_ids = (8, 9, 10, 11)
self.mma_warp_id = 12
self.epilogue_warp_ids = (13,)
self.load_warp_ids = (14,)
self.empty_warp_ids = (15,)
self.tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_100") # 512 coloum

init详解

首先详细拆解FlashAttentionForwardSm100的__init__(…)中的输入参数以及具体的设置，具体的输入变量如下：

head_dim: int, # q/k dimension
head_dim_v: Optional[int] = None, # v dimension, needed for MLA
qhead_per_kvhead: cutlass.Constexpr[int] = 1, # GQA
is_causal: bool = False, # causal mask, needed for prefill
is_local: bool = False, # sliding window, local attention
is_split_kv: bool = False, # flash decoding，把超长的kv切成多端，每个CTA算一部分，随后靠着scatter聚合到一起
pack_gqa: bool = False, # 对于decode和短prefill，没法切成tile，此时支持把同一个KV head对应的q head聚合到一起算
q_subtile_factor: int = 1, # 块稀疏 mask 的粒度细分因子
kv_subtile_factor: int = 1, # 块稀疏 mask 的粒度细分因子
m_block_size: int = 128, # tile M
n_block_size: int = 128, # tile N
q_stage: cutlass.Constexpr[int] = 2, # 2代表两个tile ping-pong做overlap
is_persistent: bool = True, # persistent kernel，SM直接按照偏移去拿tile进行计算
score_mod: cutlass.Constexpr | None = None, # 用户自定义分数修改回调
mask_mod: cutlass.Constexpr | None = None, # flex-attention 风格的用户自定义 mask 回调
has_aux_tensors: cutlass.Constexpr = False, # score_mod/mask_mod 是否携带辅助张量(如 ALiBi 的 slopes)。影响 score_vec_size
paged_kv_non_tma: bool = False, # 是否支持用TMA来拷贝KV数据
is_varlen_q: bool = False, # 是否是varlen Q，needed for prefill
use_2cta_instrs: bool = False, # 使用2-CTA MMA单元
use_clc_scheduler: bool = False,

在__init__函数中，会把head_dim和head_dim_v都padding到16的倍数，其中，核心通过调用self.arch = BaseDSL._get_dsl().get_arch_enum()来获取具体的硬件结构，随后获取具体三个tiler的mma的形状，分别如下，其中对于cta的单个tile而言，由于分为2个stage，因此输出是2MxN，最后一维是d，对于mma_tiler_qk而言，结果是MxN，中间维度是d，对于pv的mma，输出是Mxd，中间维度是N；

self.cta_tiler = (self.q_stage * m_block_size, n_block_size, self.head_dim_padded)
# With 2CTA, the MMA tiler M covers both CTAs, so it's cta_group_size * m_block_size.
# Each CTA owns m_block_size rows; the 2CTA MMA instruction spans both.
self.mma_tiler_qk = (self.cta_group_size * m_block_size, n_block_size, self.head_dim_padded)
self.mma_tiler_pv = (self.cta_group_size * m_block_size, self.head_dim_v_padded, n_block_size)

其中，核心注意use_tma_O这个变量，代表是否允许是有TMA来搬运O矩阵，他的开启条件很严苛，如下所示，如果开启pack_gqa，那么必须M可以被qhead_per_kvhead整除，这也是为了方便TMA的索引，同时，如果开启MQA，那么必须不能使用split_kv，同时，不能使用is_varlen_q，因为Q每条序列边界任意,TMA box 罩不住不许越界写的行边界)。放弃后走 gmem_tiled_copy_O:先 smem→寄存器→再逐元素带谓词写 gmem(_store_O_to_gmem, :2863,每行判断 < seqlen_q 才写)。其中，如果能用TMA把O从shared memory搬回global memory，那么只用epilogue_warp_ids也即第13号warp也就够了，不然，则直接使用correction的四个warp来负责搬运，并把原先的13号warp放入empty_warp_ids中。

self.use_tma_O = (
    not (self.pack_gqa and self.m_block_size % self.qhead_per_kvhead != 0)
    and not (self.pack_gqa and self.is_split_kv)
    and not is_varlen_q
)
self.use_correction_warps_for_epi = not self.use_tma_O
# ... ...
if self.use_correction_warps_for_epi:
    self.empty_warp_ids = self.empty_warp_ids + self.epilogue_warp_ids
    self.epilogue_warp_ids = self.correction_warp_ids

随后，下面值得注意的变量就是enable_ex2_emu，即利用多项式估计来算softmax，释放MUFU压力；同时，其中存在overlap_sO_sQ，仅在head_dim为192时打开，可以算笔账，hdim 192 时 sQ 要 2×128×192×2B = 96KB,K/V 每 stage 48KB,再单独留一块 sO(2×128×128×2B=64KB)的话流水级数就没了。而 Q 和 O 的生命周期其实错开:Q 在最后一次 QK GEMM 后就没用了,O 在最后才写 smem——所以让 sO 直接骑在 sQ 的地址上(:1039-1042 的 recast_ptr),smem_size_q_o 取 max 而不是 sum，同时，由于persistent kernel 会提前用TMA把下一个Q給load进来，可能和O冲突，因此强制关闭persistent kernel。注意，即使dim是128，如果开split KV，同样也需要overlap，因为
split KV的话会导致O需要存两份，因为后续需要进行累加，不复用 sO 的话:

项目	大小
sQ:2 stage × 128 × 192 × 2B	96 KB
sO:2 stage × 128 × 128 × 2B	64KB
剩给 K/V	224 − 160 = 64 KB
能买的槽位	64 ÷ 48 = 1 个

kv_stage = 1 意味着零重叠:MMA 每算一块都要干等 TMA——“流水级数就没了”说的就是这个。开 overlap_sO_sQ 后,sO 骑在 sQ 上(:353 取 max(96, 64) = 96KB 而非 sum):

项目	大小
sQ/sO 合体	96 KB
剩给 K/V	224 − 96 = 128 KB
槽位(按 48KB)	128 ÷ 48 = 2 个

2 已经能活,但这里还叠了第二层抠门(:361-376,就是之前提过的 uneven_kv_槽位装的东西固定——0、2 号槽装的总是 48KB 的 K,1 号槽装的总是 32KB 的V,那就不必每个槽都开 48KB。按"大-小-大"排列,平均 stride 取 (48+32)/2 = 40KB,3 个槽只要 120KB ≤ 128KB,塞下了。代价是寻址不再均匀:1 号槽的真实地址相对”基址 + 1×40KB”差了 ±8KB,要按相位补偿——这就是 offset_kv_smem(:3111)和 uneven,也是为什么 load_KV 里 stage 0 要额外等 stage 1 排空(:3085-3089,因为0 号 K 的 48KB 会侵入原本 1 号的地盘)。注意，对于128的head_dim和is_split KV而言，split kv后，O会从bf16变成fp32，因而会占据两倍的SMEM space，因而，为了保证掩盖KV的TMA搬运时间，仍然需要开启overlap_sO_sQ，其中虽然O在TMEM中是FP32，但正常从寄存器写会SMEM的时候会转为BF16；
随后，根据具体是否开启varlen_q, 是否persistent kernel，是否有causal mask等来初始化Scheduler，如果都不是，直接fallback成普通的SingleTileScheduler，这个Scheduler本质就是做映射，即每个Block获取哪个Tile进行计算；

随后，判断是否可以开启use_tma_Q，随后判断，是否能用TMA搬数据，如果不能，则设置搬数据的warp为14和15，不再保留empty warp

注意，其中存在use_correction_warps_for_epi，如果不能用TMA搬运O到SMEM中，那么单warp太慢，此时直接用correction_warp来搬运；
随后，来具体分配TMEM中各个组件的偏移，其中就和paper里所述的一样，低0位开始放S->P low，128位开始放S->P high，256位和384位放O low和O high；

随后根据具体的调参config获取最优的寄存器分配；

Kernel

随后具体查看核心的GPU device kernel的实现，其中输入如下：

其中，kernel核心实现5类型warp，load warp/MMA warp/Softmax Warp/Correction Warp/Epilogue warp的pipeline。

Quartz 5

Explorer

FlashAttention-4 Code walkthrough

概述

Walkthrough

flash_fwd_sm100.py

FlashAttention4 精读

概述

Paper Read

1/2 Introduction

3 Algorithm

3.1 Algorithm Forward Pass

3.1.1 Feeds and speeds

新pipeline来overlap matmul和softmax

3.1.3 Emulation of the exponential function

3.1.4 Skipping online softmax rescaling

3.2 Attention backward pass

3.2.1 Feeds and Speeds

3.2.2 New pipeline to overlap matmul and softmax

3.2.3 2-CTA backward pass

3.2.4 Deterministic backward pass

3.3 Scheduling

init详解

Kernel

Graph View

Table of Contents

Quartz 5

Explorer

FlashAttention-4 Code walkthrough

概述

Walkthrough

flash_fwd_sm100.py

FlashAttention4 精读

概述

Paper Read

1/2 Introduction

3 Algorithm

3.1 Algorithm Forward Pass

3.1.1 Feeds and speeds

新pipeline来overlap matmul和softmax

3.1.3 Emulation of the exponential function

3.1.4 Skipping online softmax rescaling

3.2 Attention backward pass

3.2.1 Feeds and Speeds

3.2.2 New pipeline to overlap matmul and softmax

3.2.3 2-CTA backward pass

3.2.4 Deterministic backward pass

3.3 Scheduling

__init__详解

Kernel

Graph View

Table of Contents

init详解