FlashAttention-4 Hits 1,605 TFLOPS on NVIDIA Blackwell GPUs

NVIDIA has launched FlashAttention-4, the newest optimization for transformer neural networks that squeezes 1,605 TFLOPS out of its Blackwell structure—capturing 71% of the {hardware}’s theoretical most efficiency.

The announcement issues for anybody watching AI infrastructure investments. As giant language fashions push towards longer context home windows, the eye mechanism’s quadratic reminiscence complexity turns into a brutal bottleneck. FlashAttention-4 assaults this downside straight, and the benchmark numbers recommend significant beneficial properties for manufacturing AI workloads.

What the Numbers Present

On the B200 GPU, FA4 delivers a 3.6x speedup over FlashAttention-2 throughout ahead passes at 32,768 sequence size. Backward cross efficiency hits 3.15x sooner than FA2 below the identical situations. Towards present frameworks, FA4 posts 1.3x enchancment over cuDNN and a pair of.4x over Triton Inference Server implementations.

The reminiscence effectivity beneficial properties are equally important. Normal consideration scales at O(N²) with sequence size—that means doubling your context window quadruples reminiscence necessities. FA4 brings this all the way down to O(N) by means of tiling and incremental softmax normalization. NVIDIA claims 20x decrease reminiscence utilization in comparison with PyTorch baselines.

{Hardware}-Software program Co-Design

FA4 was constructed particularly for Blackwell’s quirks. The structure presents an uneven scaling downside: compute energy roughly doubles whereas reminiscence bandwidth would not preserve tempo. Conventional approaches depart tensor cores sitting idle whereas ready for information.

The answer leverages Blackwell’s devoted Tensor Reminiscence (TMEM)—256 KB of on-chip reminiscence per streaming multiprocessor. By storing intermediate calculations straight in TMEM as a substitute of shared reminiscence, FA4 sidesteps the bandwidth bottleneck that will in any other case throttle the sooner compute items.

Bigger tile sizes (as much as 128×128) and deeper pipelines preserve the {hardware} busy. The backward cross—usually the slower half of coaching—advantages from bypassing register accumulation completely.

Manufacturing Integration

Main inference frameworks together with SGLang and vLLM already assist FA4 prefill operations. NVIDIA has integrated these strategies into cuDNN 9.14, making the optimizations accessible to builders with out customized kernel work.

For AI corporations burning by means of compute budgets, the effectivity beneficial properties translate on to price financial savings. A 3x+ speedup on coaching passes means both sooner iteration cycles or the flexibility to coach bigger fashions inside present infrastructure constraints.

The broader pattern right here: as transformer fashions develop, algorithmic effectivity on the kernel degree turns into as essential as uncooked {hardware} functionality. FlashAttention-4 represents the present frontier of that optimization work.

Picture supply: Shutterstock