TLX-Optimized Jagged Flash Attention: Approaching FA4-Level Performance on Blackwell

Why Jagged Flash Attention Deserves Attention

Since its introduction, Flash Attention has become the de facto standard for attention computation in modern large-scale models. By combining tiled computation with online softmax, it eliminates the need to explicitly store the full attention matrix, dramatically reducing memory overhead and improving throughput. In recommendation and advertising systems, however, sequence lengths vary wildly — and standard fixed-length Flash Attention requires heavy padding to align them, wasting significant compute.

Jagged Flash Attention (JFA) was built specifically to solve this problem. It operates directly on variable-length (jagged) sequences, avoiding the need to pad with invalid tokens for alignment. According to the original work, JFA is the core attention kernel powering Meta's Generative Ads Model (GEM), underscoring its real-world value in large-scale industrial recommendation and advertising systems.

The core innovation of Flash Attention lies in its tiling + online softmax paradigm. Traditional attention requires writing out the full N×N attention score matrix to HBM (High Bandwidth Memory) before applying softmax normalization — an O(N²) memory complexity. Flash Attention instead loads Q/K/V tiles into SRAM (on-chip cache), accumulates intermediate statistics (row-wise maximum and exponential sum) using numerically stable online softmax, and never materializes the full matrix — compressing memory complexity to O(N) while dramatically reducing HBM read/write operations, which are the primary bottleneck in GPU-based attention. The FA1 → FA2 → FA3 → FA4 progression has been driven by increasingly efficient use of hardware asynchrony (Tensor Core pipelines, async memory copies) to push real-world throughput closer to theoretical peak FLOPS.

Why Rewriting the Kernel for Blackwell (B200) Matters

NVIDIA's Blackwell architecture (B200) introduces next-generation tensor cores and higher memory bandwidth — but fully unlocking that compute power requires redesigning kernel scheduling and data flow for the new architecture. The central contribution of this work is porting and optimizing JFA for Blackwell, with the goal of approaching the current state-of-the-art: Flash Attention 4 (FA4).

Two points are worth unpacking here:

Architecture Migration Is Not Just a Recompile

Moving from the previous GPU generation to Blackwell isn't as simple as recompiling with a new target. The memory hierarchy, async copy mechanisms, and tensor core instruction sets have all changed — tiling strategies, pipeline scheduling, and register allocation all need to be re-tuned from scratch. For a kernel like JFA, which already handles variable-length sequences with more complex branching logic, the migration challenge is even steeper.

Targeting FA4 Means Targeting the Highest Bar

Each generation of Flash Attention pushes closer to the hardware's theoretical peak. Getting a specialized variable-length kernel to approach FA4-level efficiency means preserving all the compute savings that come from avoiding padding, while ensuring the additional control-flow overhead doesn't drag down peak throughput — a genuinely demanding engineering tradeoff.

Flash Attention 4 (FA4) is the latest iteration, designed specifically for NVIDIA Hopper/Blackwell architectures. Its key improvements include: full exploitation of the Tensor Memory Accelerator (TMA) introduced in Hopper for asynchronous data movement, complete pipelining of compute and data loading via producer-consumer warp separation, and dedicated optimizations for FP8 low-precision paths. On H100, FA4's forward pass can achieve MFU (Model FLOP Utilization) north of ~75% of theoretical peak. Getting a JFA kernel — one that introduces additional variable-length branching logic — to approach that efficiency means the control-flow overhead from variable-length handling must be fully hidden within the async pipeline. That is the core challenge of the Blackwell rewrite.

The Role TLX Plays in Kernel Optimization

The original work explicitly uses TLX to build and optimize this kernel. TLX, as a low-level GPU kernel development tool and language layer, allows developers to describe tensor computations at a higher level of abstraction while retaining fine-grained control over scheduling and memory access.

Compared to writing raw CUDA, tools like this offer several advantages:

• Lower engineering cost when rewriting kernels for new architectures like Blackwell
• Easier expression of irregular computation patterns like variable-length sequences
• Better maintainability without sacrificing performance close to hand-tuned levels

For teams that need to keep pace with hardware generations by frequently rewriting kernels, this combination of portability and expressiveness often determines how quickly they can capture the performance benefits of new silicon.

TLX (Tensor Language X, or a similarly positioned DSL) belongs to the kernel compiler / domain-specific language category — alongside tools like Triton (OpenAI) and Hidet (Microsoft). The shared philosophy is to provide a higher-level tensor operation abstraction above CUDA, letting the compiler handle the mapping to architecture-specific instructions (such as Blackwell's new-generation async warpgroup MMA instructions), so developers don't have to manually manage register allocation, bank conflict avoidance, or pipeline slot scheduling. The performance gap between these tools and hand-written CUDA continues to narrow as compiler maturity improves. But for control-flow-heavy scenarios like variable-length or sparse patterns, the engineering efficiency gains from high-level abstractions are especially pronounced — because hand-tuning correctness and performance for irregular memory access patterns is extremely costly.

Broader Industry Trends in Attention Optimization

This work reflects several signals worth tracking. First, the competition in attention kernels has expanded beyond general-purpose LLMs into structured sparse and variable-length domains like recommendation and advertising — variable-length optimization is becoming a major front. Second, the hardware generation transition to Blackwell is driving a full round of kernel rewrites across the industry; teams that migrate their specialized kernels to new architectures fastest will gain the earliest advantages in cost and latency. Third, kernel development tools like TLX are becoming critical infrastructure connecting algorithmic innovation to hardware capability.

For engineering teams working on large-scale inference and training systems, the takeaway is direct: don't treat attention kernels as a black box. Specializing for the sequence length distribution of your actual workload can yield substantial cost savings. At the same time, investing early in a portable kernel development toolchain pays dividends at every hardware upgrade cycle.

Summary

This work ports the Jagged Flash Attention kernel — the engine behind Meta's GEM — to NVIDIA's Blackwell platform, optimizing it with TLX to approach SOTA-level FA4 performance. It demonstrates both the real-world demand for variable-length attention in industrial settings and the continuous evolution of kernel optimization alongside hardware generations. Since the source material is a preview-style summary, specific performance numbers, implementation details, and benchmark comparisons await disclosure in a full technical blog post.