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

The work introduces Jagged Flash Attention (JFA), an attention kernel for Meta’s Generative Ads Model (GEM) optimized for NVIDIA Blackwell GPUs using Triton Low-level Extensions (TLX). The research addresses the challenge of efficiently processing variable-length sequences inherent in GEM by packing sequences contiguously, which is achieved via JFA. The TLX approach provides explicit, hardware-aware control over memory allocation and scheduling, moving beyond standard compiler scheduling to achieve peak Blackwell performance on the attention kernel.
The optimizations are categorized into structural changes that reorganize the kernel execution—such as warp specialization for tasks like TMA loads and matmuls, and persistent execution structures—and specific performance optimizations that eliminate stalls in the optimized structure. Key optimizations include software load balancing for jagged tile distribution across SMs, multi-stage staging for handling the heavily contended backward pass reduction, early release of tensor-memory buffers, loop peeling to reduce softmax overhead, and a two-CTA collaborative MMA strategy for the backward pass.
Overall performance benchmarks show that JFA outperforms FlashAttention-4 (FA4) on jagged shapes for GEM by approximately 13% in the forward pass and 50% in the backward pass on B200. Furthermore, TLX enables flexibility to fork the kernel into variants like low-precision attention (MXFP8) and block-sparse attention without rewriting the core structure, demonstrating a significant gain in development efficiency alongside performance improvement.

* Jagged Flash Attention (JFA) is the attention kernel for Meta’s Generative Ads Model (GEM), built on NVIDIA Blackwell (B200) using TLX.
* Attention is the slowest kernel in GEM, historically requiring hand-written CuteDSL or CUDA.
* TLX provides explicit control over hardware primitives like SMEM/TMEM allocation, warp specialization, and barriers.
* The baseline Triton JFA leaves data movement and scheduling to the compiler.
* JFA efficiently handles jagged sequences by applying FlashAttention directly to packed Q/K/V tensors and offsets without materializing padded tokens.
* Structural changes involve splitting the CTA into role-specialized async tasks (e.g., TMA loads, matmuls) and managing explicit memory for pipeline depth.
* Optimizations include scheduling jagged tiles across SMs using host-side sorting and Cluster Launch Control (CLC).
* Multi-stage dQ staging uses double-buffered SMEM to overlap HBM reductions with TMEM data movement.
* Early tensor-memory release involves releasing buffers before stores to allow concurrent MMA execution.
* Loop peeling splits KV loops into a branch-free bulk pass and a masked tail to reduce register spills during softmax computation.
* A two-CTA collaborative MMA scheme is adopted for the backward pass, involving two CTAs working together on wider matmuls over adjacent K/V blocks.
* On jagged shapes, JFA outperforms FA4 by ~13% in the forward pass and ~50% in the backward pass on B200.
* The TLX structure allows for easy porting to variants like low-precision (MXFP8) attention and block-sparse attention.

The core innovation lies in decoupling high-level algorithmic expression from low-level hardware scheduling by introducing explicit control via TLX. This shift moves the optimization challenge from merely finding faster math (the domain of existing FA4 research) to explicitly managing the complex, interdependent memory access and execution pipelines inherent in modern accelerator architectures like Blackwell. The observed performance gains, especially the 50% improvement in the backward pass on jagged data, suggest that the bottlenecks in attention kernels are less about raw FLOPs during the core matrix multiplication and more about synchronization overhead, HBM latency hiding, and scheduling inefficiency when managing non-uniform workloads across streaming multiprocessors.
The pattern of optimization—identifying stalls caused by compiler abstraction (softmax overhead, register spills) and imposing explicit control to resolve them—reflects a persistent tension in high-performance kernel development: the trade-off between generality/ease of iteration (the realm of high-level frameworks) and fine-grained control necessary for peak hardware utilization. The subsequent ability to re-target this control into new mathematical variants like MXFP8 or sparse attention shows that the established structural foundation provided by TLX is not merely an optimization tactic, but a platform for architectural evolution.
The implication for future AI kernel development is that abstracting complexity through explicit, composable scheduling primitives (like CLC and shared memory management) allows researchers to focus on the novel algorithmic structure rather than fighting the underlying hardware plumbing. The cost paid here is increased initial implementation complexity, which is offset by enabling superior performance on specific data layouts critical for real-world applications like broadcast-Q scenarios. Future inquiry should focus on how these explicit scheduling models can be generalized beyond attention to other complex memory-bound primitives, testing whether this structured control becomes a fundamental paradigm shift rather than an attention-specific optimization strategy.