In Meta's ads model training, attention is the slowest kernel, and reaching peak on Blackwell B200 has traditionally required hand-written CuteDSL or CUDA kernels: fast, but slow to extend when a modeling team introduces sliding windows, block sparsity, or new numeric formats. The team working on Jagged Flash Attention for GEM has published an alternative path: rewriting the kernel with TLX, Triton's low-level extensions, and exceeding FlashAttention-4 on jagged broadcast-Q workloads, with about 3,200 lines of code versus the 10,000 lines of FA4's CuteDSL kernel.

The result is not marginal. In bfloat16 on B200, the forward is on average 13% faster on the jagged workloads relevant for ads, while the backward gains about 50%. On dense LLM-style workloads the forward remains competitive, around 87% of FA4, and the backward wins again at +17%. The advantage comes not from any isolated trick but from a combination of warp specialization, explicit pipelines, manual shared memory and TMEM allocation, persistent execution, software load balancing, and Cluster Launch Control. All of these were previously expressible only in low-level languages, whereas TLX brings them into readable Python-level code.

The turning point: iteration moves from kernel engineers to modeling engineers

The most interesting figure is not the benchmark, but who can touch the kernel. With CuteDSL you need specialists who understand barriers, TMA, MMA, and register allocation; with TLX the structure is written once and variants become forks of the math. The two cited experiments prove it: the MXFP8 version (microscaling-precision matmuls, with P quantized to FP8 on the fly and scale factors kept in tensor memory) reuses the same warp-specialized skeleton and in the forward surpasses FA4's FP8 kernel, with backward at parity on dense workloads. The block-sparse variant only adds iteration over selected blocks, and at a 0.5 selection ratio runs 1.3 to 1.5 times faster than dense. In both cases, most of the kernel remains unchanged.

This changes incentives. Companies managing Blackwell clusters for self-hosted models no longer have to choose between generic libraries and hiring a CUDA expert team to get acceptable performance. They can start from a TLX kernel, profile with Nsight Compute, and modify only the part they need, retaining control over precision, sparsity, and variable-length sequence handling. That is a concrete advantage for those with data sovereignty constraints or non-standard workloads that do not justify a dedicated kernel team.

The downside

There is, however, a second-order cost. TLX productivity shifts part of the complexity onto the toolchain: the compiler, the Triton runtime, and low-level primitives become critical infrastructure. Adopters depend on PyTorch and the community keeping TLX stable across GPU generations, and they still need to read hardware counters to understand where cycles are lost. This is not the elimination of specialization but its redistribution: from a single hand-written kernel to the shared maintenance of a framework.

The code is available in the facebookresearch/ads_model_kernel_library repository, under tlx_jfa. The lesson for those watching the attention kernel market is that the barrier to peak performance on Blackwell is shifting: no longer how many lines of CuteDSL you can write, but how well you can express a scheduling intent in a high-level language. For Meta's ads models, this has already produced a faster, more maintainable kernel.