TLX Rewrite of Jagged Flash Attention Beats FA4 on Blackwell B200

Meta's Jagged Flash Attention (JFA) kernel, built with TLX (Triton Low-level Extensions), outperforms the May 2026 version of FlashAttention-4 (FA4) on NVIDIA Blackwell B200 for the jagged shapes used in the Generative Ads Model (GEM). The forward pass is about 13% faster and the backward pass about 50% faster, all benchmarked in bfloat16. The TLX kernel is roughly 3.2K lines of Triton-level code, about three times less than the ~10K-line CuteDSL kernels of FA4. The code is available in the ads_model_kernel_library repository on GitHub.

Attention is the single slowest kernel in GEM, and reaching peak Blackwell performance has traditionally required hand-written CuteDSL or CUDA, which is slow to develop and hard to extend to new variants. TLX adds explicit, hardware-aware control on top of Triton's high-level, tile-based programming model, closing the gap on both development efficiency and performance. The kernel stays approachable enough for modeling engineers, not just kernel specialists, to read, extend, and fuse.

Structural changes and optimizations

The rewrite splits the CTA into role-specialized async tasks: dedicated warps for TMA loads, tensor-core matmuls, softmax and correction math, and the epilogue store, plus a dedicated dQ-reduction warp in the backward pass. This keeps the tensor cores fed with back-to-back MMAs while softmax runs concurrently. The kernel hand-allocates every on-chip buffer, chooses pipeline depth (triple-buffered K/V), aliases TMEM buffers with non-overlapping lifetimes, and uses explicit producer and consumer barriers. Both passes are persistent, with one CTA per SM looping over tiles.

Several optimizations target specific bottlenecks. For jagged inputs, sequence lengths vary by orders of magnitude, causing severe load imbalance across SMs. In the forward pass, the host sorts all tiles by KV workload descending and hands them to SMs in a zigzag pattern, recovering roughly 20% on the forward kernel. The backward pass uses a different strategy: precomputing valid tiles on the host and round-robinning them across SMs with a block-innermost ordering for L2 reuse. Cluster Launch Control (CLC), a Blackwell feature, handles residual runtime variance by handing out the next tile index on demand.

Backward pass improvements

The dQ epilogue was the number one backward bottleneck, accounting for about 9 to 11% of lost tensor-core utilization. Because dQ is summed across the whole batch, every SM issues reduce-adds into the same dQ locations, causing heavy contention. The high-level reduce-add API serializes them with no pipelining. The fix makes the staging explicit and double-buffered, so one slice's reduce-add to HBM overlaps the next slice's copy out of TMEM. Early tensor-memory release frees the dQ buffer before its final stores, letting the MMA warp start the next tile's dQ matmul sooner. Loop peeling splits the KV loop into a branch-free bulk pass and a tiny masked tail, removing per-iteration mask overhead and register spills. A 2-CTA collaborative MMA, adopted from FA4, raises tensor-core utilization in the matmul-heavy backward pass.

The optimizations were found using NVIDIA Nsight Compute to read hardware counters such as tensor-core and TMEM pipeline utilization, SM busy, and register spills, plus ptxas dump logs and TritonBench in profiler-measure mode for clean before and after A/B latency. Some techniques are adopted from FA4, showing that TLX can express in readable Python the same methods a hand-written CuteDSL kernel uses. Others are derived from profiling, targeting the jagged broadcast-Q workload specifically.