---
title: "Meta's TLX Block Attention Beats Flash Attention v2 by 2.3x on Blackwell GPUs"
subtitle: "Meta's open-source TLX Block Attention kernel delivers 2.3x faster attention and 3.5x fused backward speedup on NVIDIA Blackwell GPUs by exploiting compile-time knowledge of block-diagonal attention patterns."
url: "https://alphasignal.ai/news/meta-s-tlx-block-attention-beats-flash-attention-v2-by-2-3x-on-blackwell-gpus"
author: "AlphaSignal Newsroom"
author_url: "https://alphasignal.ai/editorial-team"
company: "PyTorch"
company_url: "https://alphasignal.ai/company/pytorch"
published: "2026-05-26T15:13:06.000Z"
updated: "2026-06-16T20:58:58.068Z"
type: "News"
topics:
  - "Gpus"
  - "Open Source"
subtopics:
  - "Kernels"
upvotes: 265
read_time: "7 min"
source_url: "https://pytorch.org/blog/tlx-block-attention-a-warp-specialized-blackwell-kernel-for-fixed-block-sparse-self-attention/"
comments: 0
---

# Meta's TLX Block Attention Beats Flash Attention v2 by 2.3x on Blackwell GPUs

> Meta's open-source TLX Block Attention kernel delivers 2.3x faster attention and 3.5x fused backward speedup on NVIDIA Blackwell GPUs by exploiting compile-time knowledge of block-diagonal attention patterns.

Flash Attention is one of the most celebrated optimizations in modern deep learning , but it was designed for the general case. [TLX Block Attention](https://pytorch.org/blog/tlx-block-attention-a-warp-specialized-blackwell-kernel-for-fixed-block-sparse-self-attention/), a new open-source kernel from Meta's Ads AI team, asks a sharper question: what if you already know your attention pattern at compile time? The answer, it turns out, is a 2.3x speedup over Flash Attention v2 and a 3.5x speedup when rotary embeddings are fused into the backward pass.

## The Pattern That Changes Everything

Block-diagonal self-attention is a specific structural pattern where a sequence is split into fixed-size groups, and each token only attends to others within its own group. It sounds like a constraint, but it's actually a very common design in recommendation and feature-interaction models , the kind that power ads ranking systems at scale. Meta's production workloads run batch sizes of 1152, sequences up to ~4k tokens, and roughly 70% sparsity in the attention structure.

The problem is that general-purpose kernels like Flash Attention v2 don't know about this structure ahead of time. They're built to handle arbitrary-length causal attention, which means they carry a lot of algorithmic machinery that becomes pure overhead when the pattern is fixed. FlexAttention (FA4) supports block-sparse patterns but has a minimum tile size of 256 , incompatible with the 64-token blocks these models require.

![Comparison of Sliding Window vs Block Attention patterns showing attended and masked regions](https://pytorch.org/wp-content/uploads/2026/05/unnamed-4-1.png)

## One Constraint, Five Eliminations

The core insight is deceptively simple: in block-diagonal attention, every query (Q) tile attends to exactly one key/value (K/V) tile , its own block. That single constraint cascades through the entire Flash Attention algorithm and eliminates five categories of overhead:

- **No multi-tile loop.** The score matrix S = Q·Kᵀ is complete after a single matrix multiply. There is no iteration to maintain state across.
- **No online softmax correction.** Standard Flash Attention applies a correction factor at each tile step to stay numerically stable. With one tile, the row-wise max and sum are globally correct immediately , the correction factor is identically 1 and can be dropped.
- **No logsumexp (L) tensor.** Flash Attention saves a per-row log-sum-exp tensor to GPU memory (HBM) so the backward pass can reconstruct attention probabilities. With a single tile, the backward pass can recompute everything inline , eliminating an entire HBM write and read per forward/backward pair.
- **No Di preprocessing kernel.** Standard Flash Attention backward launches a separate kernel just to compute a row-wise sum before the main backward pass. TLX Block Attention computes this inline, killing the extra kernel launch entirely.
- **No output accumulation with rescaling.** The output O = P·V is a fresh result from a single GEMM, not an accumulation of rescaled partial results. This lets the kernel signal to the hardware that it doesn't need to preserve the tensor core accumulator across calls , freeing it for reuse.

These are not micro-optimizations. They are the elimination of entire algorithmic stages that Flash Attention needs for correctness in the general case but that become dead weight when the pattern is fixed.

## Built on Blackwell's Hardware Primitives

The kernel is written in [TLX (Triton Language Extensions)](https://github.com/triton-lang/triton-ext) , a set of low-level extensions to the Triton compiler that expose Blackwell-specific hardware control: warp specialization, asynchronous tensor core operations, and explicit memory hierarchy management. Think of it as Triton with the guardrails lowered to match what you'd get writing raw CUDA, but with Python-native ergonomics.

Warp specialization is the key architectural idea. In a standard GPU kernel, all threads run the same code. Here, different groups of threads (warps) are permanently assigned to different hardware units and run completely different code paths for the kernel's entire lifetime. The forward pass uses 15 warps across 5 specialized stages:

- **Load warp (1 warp):** Feeds Q, K, V tiles into shared memory via the TMA (Tensor Memory Accelerator) engine asynchronously.
- **QK MMA warp (1 warp):** Runs the Q·Kᵀ matrix multiply on Blackwell's tcgen05 tensor cores.
- **Softmax warps (4 warps):** Applies masking, scaling, exp2, and normalization to produce the attention probability matrix P.
- **PV MMA warp (1 warp):** Runs the P·V matrix multiply on tensor cores.
- **Epilogue warps (8 warps):** Moves results from tensor memory to registers, casts to BF16, and stores to global memory.

While the Epilogue warps write tile *i* to global memory, the MMA warps are computing tile *i+1*, and the Load warp is already fetching tile *i+2* , three tiles in flight simultaneously. The backward pass scales up to 20 warps across 7 stages and fully saturates the 256 KB of Tensor Memory on each streaming multiprocessor.

![Gantt chart showing the warp-specialized pipeline stages across time for the forward pass](https://pytorch.org/wp-content/uploads/2026/05/unnamed-3-1.png)

## The Rotary Fusion Bonus

Many transformer models use rotary position embeddings (RoPE) , a technique that encodes positional information by rotating the query and key vectors. Conventionally, the rotary backward pass runs as a separate kernel after the attention backward kernel, which means gradients get written to GPU memory in BF16, then immediately read back for the rotary transform, then written again. That's two quantization steps and a full memory round-trip on tensors that can span millions of sequence elements.

TLX Block Attention fuses the rotary conjugate directly into the attention backward epilogue, while gradients are still in FP32 registers. The result is striking:

| Configuration | Time (ms) |
| --- | --- |
| Attention backward (standalone) | 1.556 |
| Rotary backward (standalone) | 4.880 |
| Unfused total | 6.436 |
| Fused attention+rotary backward | 1.819 |
| **Speedup** | **3.54x** |

The standalone rotary kernel was more than 3x more expensive than the attention backward itself , it's purely memory-bandwidth bound with no meaningful compute. Fusing it eliminates the round-trip entirely. As a bonus, doing the rotary transform in FP32 instead of BF16 cuts the maximum gradient error in query gradients (dQ) by over 2x compared to Flash Attention v2.

## Handling Real-World Messiness

Production recommendation models don't process neatly uniform sequences. Traffic arrives as jagged, variable-length sequences packed into a single flattened buffer. TLX Block Attention handles this with a persistent kernel scheduling strategy: it launches exactly one thread block per streaming multiprocessor, then uses a precomputed prefix-sum array and a branchless binary search (fixed at 32 iterations) to let each SM figure out which sequence tiles it owns , with zero GPU synchronization points.

## Numbers on B200

All benchmarks were run on NVIDIA B200 GPUs in BF16 precision, with batch size 1152, head dimension 128, 4 heads, and 70% sparsity , representative of Meta's production traffic. The kernel-level results against Flash Attention v2 with block masking:

- **Forward pass:** 1.81 ms → 0.98 ms (**1.85x speedup**)
- **Backward pass:** 5.89 ms → 2.36 ms (**2.50x speedup**)
- **Combined:** 7.70 ms → 3.33 ms (**2.31x speedup**)
- **Fused attention+rotary backward:** 6.44 ms → 1.82 ms (**3.54x speedup**)
- **End-to-end MFU gain on production self-attention layers: +30.6%**

The backward speedup (2.50x) is larger than the forward (1.85x) because the backward pass benefits from two independent simplifications: eliminating the logsumexp HBM round-trip and removing the Di preprocessing kernel launch.

## Who Should Use This

The kernel is available now, open-source under Apache 2.0 at [facebookresearch/ads\_model\_kernel\_library](https://github.com/facebookresearch/ads_model_kernel_library), and it's free. But there are hard constraints on applicability:

- **Blackwell GPUs only.** The kernel uses tcgen05 MMA instructions, TMEM allocation, and Blackwell-era TMA descriptors that don't exist on Ampere or Hopper. You need sm\_100+ (B100, B200).
- **Block-diagonal attention patterns only.** If your model uses causal attention, sliding window, or any cross-block attention, this kernel is not a drop-in replacement.
- **HEAD\_DIM must be 64 or 128.** Other values require recompilation and new memory budget calculations.
- **Recommendation and feature-interaction models** are the primary target , ads ranking, retrieval systems, and similar architectures where block-local attention is a natural fit.

The broader signal here matters beyond the specific numbers. The field has spent years optimizing Flash Attention for the general case. TLX Block Attention is a proof of concept that domain-specific kernels , ones that bake structural assumptions directly into the algorithm , can deliver compounding gains that general-purpose implementations simply cannot reach. As Blackwell GPUs become more widely available, expect this pattern of compile-time-constrained kernel design to spread well beyond recommendation systems.

## Sources

- [X post by @pytorch](https://x.com/pytorch/status/2059291743632101697)
- [TLX Block Attention: A Warp-Specialized Blackwell Kernel for Fixed-Block Sparse Self-Attention – PyTorch Blog](https://pytorch.org/blog/tlx-block-attention-a-warp-specialized-blackwell-kernel-for-fixed-block-sparse-self-attention/)
- [facebookresearch/ads_model_kernel_library – GitHub](https://github.com/facebookresearch/ads_model_kernel_library)
- [triton-ext – TLX Blackwell Extensions](https://github.com/triton-lang/triton-ext)