Optimizing Jagged Flash Attention with TLX: The Road Toward SOTA FA4 on Blackwell

使用 TLX 优化锯齿状 Flash Attention：通向 Blackwell 上 SOTA FA4 的道路 – PyTorch

特色项目

TL;DR

在本文中，我们介绍了在 NVIDIA Blackwell（B200）上基于 TLX（Triton Low-level Extensions）构建的 Jagged Flash Attention（JFA）——这是 Meta 生成广告模型（GEM）背后的核心注意力内核。TLX 在 Triton 的高级分块编程模型之上，提供了显式的硬件感知控制。

注意力是 GEM 中最慢的内核，而要达到 Blackwell 的峰值性能，传统上需要手写 CuteDSL 或 CUDA 代码——开发速度慢且难以扩展到新变体。我们证明 TLX 在这两个方面都弥补了这一差距。在开发效率方面，TLX 注意力内核约有 3.2K 行简洁的 Triton 级代码，相比最先进的 FlashAttention-4（FA4）约 10K 行的 CuteDSL 内核，代码量减少了约 3 倍。在性能方面，它在 GEM 关键的锯齿状形状上超越了 FA4（2026 年 5 月版本）——前向传播快约 13%，反向传播快约 50%。通过保持 Triton 的高级模型，该内核对建模工程师（而不仅仅是内核专家）来说也足够易读、可扩展和可融合。

我们将工作分为两部分——TLX 重写带来的结构变化，以及我们在其基础上叠加的优化——所有内容均在 B200 上使用 bfloat16 对比 FA4（Blackwell 当前最先进的方案）进行了基准测试。

代码可在以下地址获取：https://github.com/facebookresearch/ads_model_kernel_library/tree/main/tlx_jfa

引言

Meta 的广告模型，包括生成广告模型（GEM）[2] 和昆仑架构 [5]，都需要在锯齿状（可变长度或不规则）用户序列上运行注意力机制。如我们在 GEM 训练深度解析中所述，将这些序列填充到固定长度可能会浪费高达 50% 的计算资源。GEM 采用连续打包序列的方式，并通过偏移量张量记录每个序列的边界。

锯齿状 Flash Attention 是实现这种无填充表示效率的关键机制：它直接对打包的 Q/K/V 张量及其偏移量应用 FlashAttention 算法，而无需显式生成填充标记。早期的 GEM 博客文章从训练系统层面描述了这一挑战；本文则聚焦于为 Blackwell 优化底层注意力内核。

在 Blackwell 上达到峰值吞吐量意味着要持续保持张量核心的负载。普通 Triton 将大部分决策交给编译器；TLX [1] 则将这些决策作为一等原语暴露出来（显式的 SMEM/TMEM 分配、异步任务线程束专业化、屏障、异步 TMA 和 MMA 以及 CLC）。本文其余部分将利用这些控制能力，将一个由编译器调度、内存受限的内核转化为紧密流水线化的、线程束专业化的内核。

1. 挑战：在 Blackwell 上实现高性能且可扩展的注意力机制

在 Blackwell 上优化注意力机制意味着同时解决两个问题。第一个是性能：注意力是 GEM 中最慢的内核，只有当全局内存加载、softmax 和矩阵乘法紧密重叠时才能达到峰值吞吐量。第二个是迭代速度：注意力机制是新建模思想落地的地方（如滑动窗口、块稀疏等变体），因此内核经常被修改——但目前实现该性能的最先进路径是难以阅读、扩展或融合的手写 CuteDSL/CUDA 内核。我们通过 TLX 的目标是实现平衡：在保持内核易于迭代的同时，实现具有竞争力的性能。

为何编译器调度的基线存在局限。我们的起点是一个成熟的基于Triton的JFA内核——算法上正确，但其将所有片上数据移动和调度都交给编译器处理：无法控制共享内存分配或流水线深度，没有显式屏障或线程束专业化，且单一的同构线程束同时执行加载、softmax和矩阵乘法，导致张量核心在未发出MMA指令时持续空转。突破这一限制正是TLX显式控制所实现的。

生产场景中的broadcast-Q进一步塑造了设计：单个密集Q被广播到批次中每个序列，因此其梯度dQ必须在整个批次中求和——这使得dQ尾部处理变成一个竞争激烈的跨程序规约，而我们的多项优化正是为加速这一过程而设计。除非另有说明，本文其余部分均假设broadcast_q=True。

2. 使用TLX优化Jagged Flash Attention

我们的工作可分为两类。结构化变更（§2.1）是TLX使能的内核底层重组——它们改变了CTA线程束间工作和内存的布局方式。优化（§2.2）则保持该结构不变，主要通过消除停顿和冗余工作来提升性能。我们通过每个优化针对的瓶颈而非实现机制来阐述其动机。

2.1 结构化变更

线程束专业化、显式共享/张量内存管理以及屏障流水线化是任何高性能Blackwell注意力内核的标准要素，因此此处简要说明——关键在于TLX让我们能用高级Triton代码而非手写CuteDSL或CUDA来表达这些特性，这正是使§2.2优化能快速原型化的关键。具体而言，我们将CTA拆分为角色专用的异步任务——专用线程束分别负责TMA加载、张量核心矩阵乘法、softmax/校正计算以及尾部存储，反向传播中还设有专用dQ规约线程束，从而使张量核心在softmax运行时能持续接收连续的MMA指令，而非像基线那样由单一线程组阻塞矩阵乘法。

这种显式控制同样延伸至内存和调度：我们手动分配每个片上缓冲区并选择其流水线深度（例如三缓冲K/V使加载线程束提前运行），将TMEM缓冲区与非重叠生命周期的内存区域别名（QK得分、P和softmax统计信息共享一个分配，而PV累加器保留独立分配），通过显式生产者/消费者屏障在不同线程束间传递数据以实现确定性流水线，并使两次遍历都保持持久化——每个SM上的CTA循环处理瓦片。这些特性都无法用纯Triton表达，而持久化结构尤其为我们提供了自主调度决策的空间——这是§2.2优化的基础。

2.2 优化

在确立线程束专业化、持久化结构后，剩余工作就是消除导致张量核心无法持续繁忙的停顿和冗余工作。以下优化基本相互独立，且每个都针对特定瓶颈。

我们通过一致的方法找到了这些优化点：使用 NVIDIA Nsight Compute (NCU) 读取硬件计数器（如张量核心/TMEM 流水线利用率、SM-busy 和寄存器溢出（表现为本地内存流量））；通过 ptxas 倾倒日志确认溢出；以及使用 TritonBench 的 profiler-measure 模式进行干净的前后 A/B 延迟对比。这些优化分为两组。一些优化来自 FA4 实现——这本身表明 TLX 可以用可读的 Python 表达与手写 CuteDSL 内核相同的技术。其他优化则是我们通过对自身内核的分析和性能分析得出的，专门针对锯齿状广播-Q 工作负载。

简要来看，以下优化包括：

• 在 SM 之间调度锯齿状瓦片——软件负载均衡加上 Cluster Launch Control 使每个 SM 保持忙碌，即使锯齿状输入（前向和后向）的长度差异达到数量级。
• 多阶段 dQ 分阶段处理——通过双缓冲 SMEM 分阶段流水线隐藏了高度竞争的广播-Q dQ reduce-add，这是后向阶段的首要瓶颈。
• 早期张量内存释放——在最终存储之前释放 dQ 张量内存缓冲区，使 MMA warp 能够更早开始下一个瓦片的 dQ 矩阵乘法。
• 循环剥离——将 KV 循环拆分为无分支的批量处理和微小的掩码尾部，消除了每次迭代的掩码开销及其引起的寄存器溢出。
• 2-CTA 协作 MMA——两个 CTA 在单个矩阵乘法上协作，提高矩阵乘法密集型后向阶段的张量核心利用率（从 FA4 采用）。

#### 在 SM 之间调度锯齿状瓦片

锯齿状输入会在 SM 之间造成严重的负载不平衡：由于序列长度差异达到数量级，简单的瓦片到 SM 映射会导致一些 SM 空闲，而其他 SM 则处理长序列。SM 占用热图（Triton-MPP）使这种现象可视化——在其余 SM 空闲后，少数几行仍然保持高温。真正贪婪的重新平衡方法接近最优，但无法向量化，因此会增加主机端的启动开销。

在前向传递中，外层循环是密集的 Q，内层循环是锯齿状的 K/V，因此瓦片的成本与其批次的键长度成正比——瓦片高度不均匀。平衡这种不均匀性是我们针对锯齿状情况的原创想法：在主机端按 KV 工作负载（降序）对所有瓦片进行排序，并以之字形（boustrophedon）模式分发给 SM——偶数次传递从左到右填充 SM，奇数次传递从右到左填充 SM，这样每个 SM 都能以几乎可以忽略的 CPU 成本获得长瓦片和短瓦片的均衡混合。这种方法在前向内核中恢复了约 20% 的性能。

后向传递需要不同的策略。其 dK/dV 传递在外层循环 K/V 块，在内层循环密集的 Q，因此现在 K/V 块是瓦片：每个瓦片的成本大致相同（一个 K/V 块与共享的密集 Q 对比），但每个批次的瓦片数量会随着锯齿状键长度变化。按工作负载排序不再有效——我们改在主机端预计算所有有效（批次、头、K/V 块）瓦片列表，并以轮询方式分发到 SM，使用块最内层排序使连续瓦片在 L2 中重用相同的 K/V，并将每个头的块数填充为 2-CTA 路径的集群大小的倍数。以这种方式分发数量可变但成本相同的瓦片，使每个 SM 都保持忙碌，无需任何排序。

静态平衡处理我们可以提前预测的偏斜；而仅在运行时出现的残余方差则由集群启动控制（Cluster Launch Control，CLC）[3] 动态处理。作为 Blackwell 的一项特性，CLC 按需分配下一个瓦片索引，因此无论哪个 SM 首先完成任务，都会立即获取下一块工作内容，避免单个长序列导致整个网格停滞。

CTAs 工作负载热图（无 CLC）

CTAs 工作负载热图（有 CLC）

我们在前向和反向核函数中都应用了 CLC；在 TLX 中，它是一个基于共享调度上下文的轻量级生产者/消费者协议。

尽管软件负载均衡和 CLC 看起来可能重复——两者都跨 SM 分发工作——但它们是互补的：CLC 虽然能动态调度，但无法判断哪些瓦片为空，因此在锯齿状（尤其是稀疏）输入中，它仍会消耗周期来分发填充引入的哨兵/空瓦片，而我们主机端的合法瓦片预计算则提前剔除了这些内容。这就是为什么即使在支持 CLC 的核函数中，软件负载均衡仍然具有价值。

#### 多阶段 dQ 预分配

Triton-MPP 障碍分析将反向传播的 dQ 尾部处理标记为首要瓶颈：其将结果写入 HBM 的 reduce-add 操作导致约 9–11% 的张量核心利用率损失，这是最大的可操作性能缺口。导致其缓慢的有两个原因。首先，由于 dQ 需要对整个批次求和，每个 SM 都会将 reduce-add 操作写入相同的 dQ 位置——导致严重竞争。其次，高层的 reduce-add API 会将其序列化——每个列切片的加法必须在下一个开始前完成，无法进行流水线处理。

我们通过显式且双缓冲的预分配方式，使一个切片的 HBM 写入与下一个切片从 TMEM 的复制重叠，始终保持一次写入在飞行中：

NCOL   = BLOCK_D // (EPILOGUE_SUBTILE * 2)   # 窄列切片，尺寸适配 SMEM
STAGES = 2                                   # 双缓冲预分配
dq_smem = tlx.local_alloc((BLOCK_M, NCOL), dq.dtype, STAGES)

for s in range(BLOCK_D // NCOL):
    # TMEM -> 寄存器
    dq = tlx.local_load(dq_tmem[:, s * NCOL : (s + 1) * NCOL]) * LN2

    # 寄存器 -> SMEM（乒乓操作）
    tlx.local_store(dq_smem[s % STAGES], dq.to(dq.dtype))
    tlx.fence_async_shared()

    tlx.async_descriptor_store(desc_dq, dq_smem[s % STAGES],
                               offsets, store_reduce="add")  # SMEM -> HBM，累加

    # 保持最多一个写入在飞行中
    tlx.async_descriptor_store_wait(STAGES - 1)

该技术可推广到任何受共享内存限制的 reduce-store 尾部处理，而不仅限于注意力机制。

#### 早期释放张量内存

同样的分析也解释了为什么该尾部处理会从 TMEM 侧阻塞计算：归约线程块会持有 dQ 张量内存缓冲区，直到所有切片都已写入 HBM，因此 MMA 线程块必须在开始下一个瓦片的 dQ 矩阵乘法之前等待 dq-empty 障碍。一个简单的移除 dQ 写入的瀑布测试验证了这一点，恢复了 +8–11% 的张量核心利用率（跨序列长度）。

改进方法是提前将最后 1-2 个切片加载到寄存器中，并在发出它们的写入之前释放 TMEM 缓冲区，这样 MMA 线程块可以在归约线程块完成向 HBM 写入时立即重用该内存：

# 以正常方式排空大部分切片
# (TMEM -> SMEM -> 异步归约加到 HBM)
for s in range(N_SLICES - EARLY_RELEASE_SUBTILES):
    reduce_add_slice(s)
# 首先将最后 1-2 个切片预加载到寄存器中，然后立即释放
# dQ TMEM 缓冲区
# 以便 MMA 线程束可以开始下一个瓦片的 dQ 矩阵乘法。
dq_tail = tlx.local_load(dq_tmem[:, last_slice])
# TMEM -> 寄存器
tlx.barrier_arrive(dq_empties[buf])
# <-- 早期释放 TMEM
reduce_add_from_registers(dq_tail)

需要提前释放的切片数量是自动调整的（1 或 2 个）；最佳选择较小，因为释放更多会增加寄存器压力到足以影响性能的程度。结合多阶段分阶段处理，这使得 MMA 线程束可以在前一个瓦片的 dQ 写出过程中就开始下一个 dQ 矩阵乘法。

#### 循环剥离

NCU 分析显示前向传播是 MMA-issue 饥饿——而不是内存或计算受限——在张量核心（TMEM）流水线中利用率约为 57%，而 FA4 在相同形状下约为 82%，丢失的周期来自 softmax 线程束开销渗透到 MMA-issue 间隙中。任何能缩小或简化 softmax 主体的内容都会让 MMA 线程束更早开始工作。

两个信号指向了罪魁祸首。首先，ptxas 转储和 NCU 显示在寄存器密集型 softmax 线程束中存在寄存器溢出（本地内存流量）；在早期一轮中，仅通过为线程束分配更多寄存器就将该流量减少了约 40%，从而实现了约 6% 的前向传播性能提升，确认了压力是真实的。其次，最严重的罪魁祸首是 KV 循环中每个迭代的掩码分支，该分支仅在部分最后一个瓦片时才有意义。

这个分支代价高昂的原因微妙且特定于编译器如何分配寄存器：Triton/TLX 静态地将寄存器分配给变量，因此在热循环内部的掩码分支迫使编译器在整个循环中为仅在很少执行的掩码路径（列偏移、掩码张量、选择）中使用的变量保留寄存器。这种永久保留会增加寄存器压力并可能触发溢出——对寄存器密集型 softmax 线程束尤其有害。

解决方案是将循环剥离为无分支的主体部分加上一个微小的掩码尾部。由于掩码标志是编译时常量，主体部分完全不包含任何掩码变量，因此编译器可以释放这些寄存器并安排更紧凑的 MMA-issue 节奏：

# APPLY_MASK 是 constexpr：当 APPLY_MASK=False 时，编译器静态删除
# 比较、offs_n / 掩码张量和选择 -> 一个无分支的主体，该主体
# 不为（罕见的）掩码路径保留寄存器。
aligned = (klen // BLOCK_N) * BLOCK_N
for start_n in tl.range(lo, aligned, BLOCK_N):
    # 主体：直线执行，无掩码代码
    softmax_iter(start_n, ..., APPLY_MASK=False)
for start_n in tl.range(aligned, klen, BLOCK_N):
    # 尾部：0 或 1 次迭代
    softmax_iter(start_n, ..., APPLY_MASK=True)
    # 仅对部分最后一个瓦片应用掩码

我们在前向 softmax 循环和反向传播中寄存器密集型路径上应用相同的剥离技术，其中收益并非来自跳过一个廉价的运行时分支，而是来自它解锁的降低的寄存器压力；在反向传播中，它特别地回收了正确性掩码本会增加的约 9% 的延迟。

#### 2-CTA 协作 MMA（反向传播）

后向计算是矩阵乘法密集型的——计算dQ、dK和dV需要每个K/V块执行五次GEMM操作——单个CTA无法单独充分利用Blackwell张量核心。我们采用FA4的2-CTA（配对CTA tcgen05）方案：集群中的两个CTA协作处理同一（batch, head）下两个相邻K/V块的更宽协作矩阵乘法，将累加器行在两个SM之间分配，并通过分布式共享内存（DSMEM）在芯片内交换dS中间结果，而非通过全局内存。两个SM共同向单个MMA供能，提高了后向计算的张量核心利用率。

我们在此提出的是TLX实现：将该方案移植到我们锯齿状的广播Q布局中，并在我们的持久化和CLC多瓦片调度器下运行。

针对生产环境的广播Q场景（HEAD_DIM=128），与单CTA路径相比，后向计算的吞吐量提高了约+12%（延迟降低约11%）。

### 性能

我们在B200（bf16）上进行两种场景的基准测试：生产环境的锯齿状场景（分层种子池化（HSP）——广播密集Q与锯齿状K/V的对比）和LLM风格的密集场景（等长Q/K/V），均与最先进的开源CuteDSL FlashAttention-4 [6] 内核（2026年5月版本）FA4进行对比。在锯齿状形状——对广告场景至关重要的形状上，TLX在两种计算路径上均优于FA4：前向计算在大多数形状和稀疏度下更快（平均提升约+13%，仅在高密度最长序列上略逊），后向计算在所有情况下均更快（平均提升约+50%）。在密集形状上，前向计算与FA4竞争（约87%的FA4性能），后向计算再次胜出（约+17%）。锯齿状图表的稀疏度范围从0.1（高度可变长度）到0.5再到0.9（接近统一长度，1.0表示所有序列长度相同）——这与生产环境中的实际范围一致。

锯齿状（广播Q）——前向TFLOPS（bf16，B200）

锯齿状（广播Q）——后向TFLOPS（bf16，B200）

LLM密集——前向TFLOPS（bf16，B200，B=768，H=4，head_dim=128）

LLM密集——后向TFLOPS（bf16，B200，B=768，H=4，head_dim=128）

## 3. 注意力变体：TLX的灵活性

基于TLX构建的实用优势在于，内核的结构足够清晰，可以针对新需求进行分支扩展，而无需从零开始重写。由于warp专业化、内存分配、屏障和调度等逻辑以可读的Python级代码编写一次，将内核重新定位到新的数值格式或新的注意力模式时，主要只需更改数学计算部分，而非底层机制。两个变体案例说明了这一点。

#### 低精度（MXFP8）注意力

为了支持FP8容错训练，我们通过将BF16矩阵乘法替换为TLX的块缩放MMA（FP8 E4M3数据与E8M0块缩放因子），复用相同的warp专业化骨架、屏障和TMEM布局，构建了前向和后向的微缩FP8变体。softmax输出P在生成时直接量化为FP8，其缩放因子保留在张量内存中，以便块缩放MMA可直接使用。

经过调优后，MXFP8前向计算的性能超过了FA4自身的FP8内核，后向计算在密集场景下与FA4达到性能一致——两者均显著优于BF16基线，这一显著的吞吐量提升主要通过更改GEMM调用实现，而非内核结构。

#### 块稀疏注意力

对于那些需要长期处理复杂历史数据且完全依赖O(N²)注意力机制的场景，我们开发了一种两阶段稀疏变体[4]。该方法首先通过一个低成本的评分内核对每个Q和K块进行平均池化，然后根据可调节的选择比例（selection ratio）为每个Q块选择最相关的前k个KV块；接着，一个从密集注意力内核衍生而来的TLX注意力内核（新增了按瓦片迭代稀疏块的功能）仅对选定的块执行注意力计算。

该方案复用相同的瓦片专用结构、CLC调度机制和软件负载均衡（此处用于跳过未选中和空瓦片），保持PT2友好性，同时仍支持广播Q、GQA和窗口机制。在0.5的选择比例下，前向计算速度比密集实现快约1.3–1.5倍。

在两种实现中，内核主体代码基本保持不变，仅数学计算部分存在差异——包括GEMM精度或每个瓦片访问的KV块集合。这种复用性正是TLX的实际价值所在：相同的内核可以重新定位到新精度和稀疏模式，而无需降级到手动编写CUDA代码，整个过程仍可在Python层面完成。

## 4. 总结与未来工作

结构性改进——瓦片专用化、持久化执行、显式SMEM/TMEM分配与复用别名、以及屏障中介流水线——为性能提升奠定了基础；叠加的优化措施——软件负载均衡、集群启动控制、多阶段dQ分阶段、早期TMEM释放、循环剥离、以及2-CTA协作MMA——则回收了大部分剩余性能空间。这些改进使TLX内核在生产环境的复杂不规则形状上，无论是哪种执行路径都能匹配或超越FA4的性能，同时仍完整支持广播Q、滑动窗口和GQA机制。

除了原始性能，最大的优势在于开发效率。CuteDSL原本需要手动编写大量底层代码（如异步流水线、屏障管理、tcgen05 MMA初始化），这些现在均由TLX编译器自动处理。我们只需用少量代码表达意图而非具体实现，这种杠杆效应使团队能快速原型化、分析并部署所有优化方案，随后仅需少量额外代码即可将相同内核分叉为MXFP8和块稀疏变体。综合来看——在不规则广告形状上双路径都更快、在密集场景保持竞争力、且构建维护成本远低于FA4——TLX JFA比当前最先进的开源FA4更适合作为广告模型的解决方案。

## 参考文献

[1] TLX：通过Triton实现集群启动控制 —— https://pytorch.org/blog/enabling-cluster-launch-control-with-tlx/

[2] GEM：Meta的生成式广告模型 —— https://engineering.fb.com/2025/11/10/ml-applications/metas-generative-ads-model-gem-the-central-brain-accelerating-ads-recommendation-ai-innovation/

[3] NVIDIA Blackwell调优指南（线程块集群） —— https://docs.nvidia.com/cuda/blackwell-tuning-guide/index.html#thread-block-clusters

[4] TLX块注意力：面向固定块稀疏自注意力的瓦片专用Blackwell内核 —— https://pytorch.org/blog/tlx-block-attention-a-warp-specialized-blackwell-kernel-for-fixed-block-sparse-self-attention/

[5] Kunlun：通过统一架构设计建立大规模推荐系统的扩展定律 —— https://arxiv.org/abs/2602.10016

[6] FlashAttention-4：面向非对称硬件扩展的算法与内核流水线协同设计 —— https://arxiv.org/abs/2603.05451