TLX JFA: Meta’s Blackwell Attention Kernel — ~13% Fwd / ~50% Bwd vs FA4 on Jagged

Last night’s Dist-MoE note was Meta’s Blackwell SM100 MoE kernels — fused dispatch, grouped expert GEMM, and combine with a named BF16 / MXFP8 / NVFP4 matrix. Yesterday AM’s Kolibri-1 was open serving MoE weights. Olmo-core 3 answered the open MoE trainer. Tonight is the complementary attention piece: on October 1, 2026, the PyTorch blog published Optimizing Jagged Flash Attention with TLX — Meta’s production jagged / broadcast-Q attention kernel for the Generative Ads Model (GEM), rewritten in TLX (Triton Low-level Extensions) for NVIDIA B200, with a named comparison against open-source FlashAttention-4.

Primary: PyTorch blog · GitHub tlx_jfa. Every number below is Meta’s, from the blog and linked repo, measured on their shapes. Meta did not publish tok/s or MFU, so I don’t either.

Spec sheet (named numbers)

The question this post answers is narrow and useful: can you beat FA4 on the jagged broadcast-Q shapes that dominate a production ads ranking model, in readable TLX rather than 10K lines of CuteDSL? Meta says yes, and shows the work.

What it is selling

Dist-MoE sits under routed experts once you already have IDs and scores. TLX JFA sits under attention on packed jagged sequences — the kernel Meta says is the single slowest piece inside GEM. Padding those sequences to a fixed length can waste up to 50% of compute (Meta’s earlier GEM writeup, cited from the Oct 1 post). JFA runs FlashAttention directly on contiguous packed Q/K/V plus an offsets tensor, without materializing padded tokens.

Plain Triton left on-chip movement and scheduling to the compiler. TLX exposes the knobs Meta needs on Blackwell: explicit SMEM/TMEM allocation, async_task warp specialization, barriers, async TMA / MMA, and Cluster Launch Control. The CTA splits into role-specialized warps — TMA loads, tensor-core matmuls, softmax/correction, epilogue store, plus a dedicated dQ-reduction warp in the backward — so tensor cores stay on back-to-back MMAs while softmax runs concurrently.

The production case shapes the design: a single dense Q is broadcast across every sequence in the batch, so the backward dQ must be summed across the entire batch. That turns the dQ epilogue into a contended cross-program reduction — the #1 backward bottleneck Meta’s Triton-MPP barrier analysis pinned at roughly 9–11% lost tensor-core utilization before the multi-stage staging fix.

Structural changes and layered opts (Meta’s list)

Structural: warp specialization + explicit SMEM/TMEM + barrier pipelining; persistent CTAs (one per SM); triple-buffered K/V; TMEM aliasing across non-overlapping lifetimes.

On top of that structure:

1. Jagged tile scheduling. Forward: sort tiles by KV workload, zigzag across SMs (~20% forward recovery claimed). Backward: precompute valid (batch, head, K/V-block) tiles, round-robin for L2 reuse, pad to cluster size for the 2-CTA path. Residual runtime variance handed to CLC.
2. Multi-stage dQ staging. Double-buffered SMEM pipeline so one slice’s HBM reduce-add overlaps the next slice’s TMEM→SMEM copy.
3. Early TMEM release. Pre-load the last 1–2 dQ slices into registers, free TMEM, let the MMA warp start the next tile’s dQ matmul while stores finish — waterfall tests recovered +8–11% TC utilization.
4. Loop peeling. Branch-free bulk KV pass + tiny masked tail; Meta’s point is register pressure, not the cheap branch itself.
5. 2-CTA collaborative MMA (backward). Two CTAs cooperate on a wider matmul over adjacent K/V blocks; Meta claims ~+12% throughput / −11% latency on broadcast-Q HEAD_DIM=128.

Where it sits next to earlier posts

This isn’t a Dist-MoE follow-up. Different layer: attention vs MoE. Dist-MoE assumes you already routed experts. TLX JFA assumes you already packed jagged sequences. Neither ships a checkpoint.

Related but separate: Meta’s earlier TLX Block Attention post (2026-07-21) targets fixed block-diagonal attention with different benches (FA2 baselines, fused rotary). The October 1 post is the jagged / broadcast-Q JFA vs FA4 comparison, and the two shouldn’t be confused.

Caveats

1. Jagged production shapes, not a universal FA4 killer. Headline wins are on jagged broadcast-Q. On dense equal-length shapes, forward is ~87% of FA4 while backward still wins (~+17%). The jagged +50% backward number is not a dense LLM result.
2. Vendor / shape-owned benches. All comparison percentages are Meta’s bf16 B200 measurements against May 2026 open-source FA4 on their shapes and sparsities (0.1 → 0.9 jagged sweep). Re-run on your shapes before you plan capacity.
3. Blackwell / TLX stack. Targets B200 with TLX primitives (async TMA/MMA, TMEM, CLC). Ampere / Hopper and plain-Triton-only stacks are outside the published comparison.
4. Not Dist-MoE. Dist-MoE is CuTe DSL MoE kernels (ARM64 + SM100, named precision matrix, no published throughput benches). TLX JFA is Triton/TLX attention with published FA4 comparison percentages. Different API, different layer.
5. Ads GEM context. Motivating workload is Meta’s Generative Ads Model (and Kunlun, cited). Uniform dense LLM decode may not share the jagged / broadcast-Q bottleneck.
6. MXFP8 / sparse forks are not the main comparison table. Main FA4 comparison is bf16. MXFP8 and block-sparse results are separate forks Meta describes qualitatively — treat as Meta claims on their forks.

How I would evaluate

1. Read the Oct 1 PyTorch post — especially the jagged vs dense charts and the broadcast-Q dQ discussion — before quoting any percentage.
2. Clone tlx_jfa and pin the TLX / Triton / CUDA stack the repo documents; do not assume Dist-MoE’s ARM64+cu130 pins apply.
3. Reproduce the jagged broadcast-Q sweep (sparsity 0.1 / 0.5 / 0.9) and the dense B=768 H=4 head_dim=128 case against the same May 2026 FA4 reference; log NCU tensor-core / TMEM counters the way the post describes.
4. Exercise MXFP8 and block-sparse forks only after the bf16 jagged path is clean.
5. Keep this separate from Dist-MoE’s MoE precision matrix, Kolibri’s serve residency, Olmo-core’s trainer tok/s, and Naive-N0.5-Flash’s zero-full-attn decode path — this is Blackwell jagged attention.

Bottom line

If you run Blackwell training stacks and your attention traffic is jagged or broadcast-Q, Meta published the comparison numbers: TLX JFA claims ~+13% forward and ~+50% backward vs May 2026 FA4 on those shapes, in ~3.2K lines of readable TLX instead of ~10K CuteDSL, with open code under tlx_jfa. Last night answered MoE kernels. Tonight answers attention. Same generation, different layer — and the layer most GEM-like ranking models actually stall on.

Create-only. Specs from Meta’s October 1, 2026 PyTorch blog on Optimizing Jagged Flash Attention with TLX and the facebookresearch/ads_model_kernel_library tlx_jfa tree. Update when third-party rematches, tok/s, or MFU numbers ship.

Sources: PyTorch blog · GitHub tlx_jfa.

Anpoo Sivanadi, Staff Software Engineer