Meta's TLX Jagged Flash Attention Beats FA4 on Blackwell With 3x Less Code

Why attention in GEM needed a rewrite

Meta's ads models, including the Generative Ads Model (GEM) and the Kunlun architecture, run attention over jagged user sequences, meaning variable-length or ragged inputs. The authors say padding such sequences to a fixed length can waste up to 50% of compute. GEM instead packs sequences contiguously and records their boundaries in an offsets tensor. Jagged Flash Attention applies the FlashAttention algorithm directly to these packed Q/K/V tensors, without materializing padded tokens.

According to the post, attention is the single slowest kernel in GEM. Reaching peak Blackwell performance has traditionally required hand-written CuteDSL or CUDA, which the authors describe as slow to develop and hard to extend to new variants such as sliding window or block-sparse attention. Their starting point was a mature Triton-based JFA kernel that left on-chip data movement and scheduling to the compiler, so the tensor cores stalled whenever the single warp group was not issuing a matrix multiply.

What TLX changes

TLX adds explicit, hardware-aware control on top of Triton's tile-based programming model. It exposes explicit shared-memory and tensor-memory allocation, warp specialization through async tasks, barriers, async TMA and MMA, and Cluster Launch Control (CLC).

With this control, the team split each CTA into role-specialized 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. They hand-allocate on-chip buffers and choose pipeline depth, for example triple-buffered K/V. They also alias tensor-memory buffers with non-overlapping lifetimes and use explicit producer/consumer barriers. Both passes are persistent, with one CTA per SM looping over tiles. The authors say none of this is expressible in plain Triton.

On development effort, the TLX kernel is about 3.2K lines of Triton-level code, roughly 3x less than the ~10K-line CuteDSL kernels of FA4, according to the post.

Optimizations layered on top

The team found its optimizations using NVIDIA Nsight Compute, ptxas dump logs and TritonBench. Some are adopted from FA4, while others were derived from profiling the jagged broadcast-Q workload, where one dense Q is shared across every sequence in the batch and its gradient dQ must be summed across the whole batch. The listed optimizations are:

• Scheduling jagged tiles across SMs: software load balancing plus Cluster Launch Control keep SMs busy despite large length skew.
• Multi-stage dQ staging: a double-buffered shared-memory pipeline hides the contended dQ reduce-add, which the authors call the number one backward bottleneck.
• Early tensor-memory release: freeing the dQ buffer before its final stores lets the MMA warp start the next tile's dQ matmul sooner.
• Loop peeling: splitting the KV loop into a branch-free bulk pass and a small masked tail removes per-iteration mask overhead and the register spills it causes.
• 2-CTA collaborative MMA: two CTAs cooperate on one matmul to raise tensor-core utilization in the backward pass, adopted from FA4.

For the forward pass, the team sorts tiles on the host by KV workload in descending order and hands them to SMs in a zigzag pattern. According to the post, this recovered roughly 20% on the forward kernel at negligible CPU cost.

What it means for kernel development

The authors argue that staying within Triton's high-level model makes the kernel approachable enough for modeling engineers, not only kernel specialists, to read, extend and fuse. All benchmarks were run in bfloat16 on B200 against FA4, which the post calls the current state of the art on Blackwell. The code is available in Meta's ads_model_kernel_library repository on GitHub under tlx_jfa.