Source-linked AI summary
Efficient FlashAttention on Blackwell via Fixed-Shift Softmax and Persistent Scheduling
Oleksandr Stashuk, Hongtao Yu, Jay Shah
TL;DR
On Blackwell, normalization and operand movement can limit attention kernels whose matrix multiplications are already deeply pipelined. This work combines fixed-shift dense softmax, a saved backward reciprocal, guarded causal handling, and scheduling changes in a TLX pipeline. Across 24 qualified dimension-128 cases on B200, the combined implementation improves geometric-mean throughput over FA4 by 7.3 percent.
Problem
On Blackwell, normalization and operand movement can constrain attention kernels even when their matrix multiplications are deeply pipelined.
Source: §1 · p. 1Method
The TLX pipeline combines fixed-shift dense softmax, a saved inverse denominator for backward, guarded causal recovery, and scheduling and data-movement changes.
Source: §1 · p. 1, §3.5 · p. 5, §4.4 · p. 8Results
7.3 percent higher geometric-mean throughput than FA4 across 24 numerically qualified dimension-128 cases on B200 at 1,000 W, using prepared launches and the original timer.
Source: §8 · p. 14Takeaways & Limitations
The combined implementation improves throughput over FA4 while targeting normalization and operand-movement costs in the specified Blackwell benchmark setting.
Source: §1 · p. 1, §8 · p. 14The packed exponential has a finite encoded range, and the benchmark checks use one seeded Gaussian fixture per length; broader output and gradient guarantees require further bounds.
Source: §5.5 · p. 12
Abstract
from arXiv · showhide
On Blackwell, normalization and operand movement can limit attention kernels whose matrix multiplications are already deeply pipelined. We implement an FA4-style pipeline in Triton TLX with fixed-shift dense softmax and a saved inverse denominator for backward. The fixed shift removes recurrent accumulator corrections; the saved reciprocal moves row normalization from the quadratic backward loop into linear preprocessing. The guarded causal path uses a first-block anchor with selective recovery, and fully masked diagonal contractions are skipped. A packed BF16 exponential approximation and pipelined dQ publication reduce per-element work and on-chip movement. On B200 at 1,000 W, using prepared launches and the original FA4 timer, TLX improves geometric-mean throughput by 7.3 percent across the 24 cases in FA4's BF16 variable-batch grid through 32,768 tokens at head dimension 128. Dense forward and backward improve by 13.0 and 10.4 percent, respectively. Causal backward improves by 8.0 percent, while causal forward is 1.5 percent slower. The TritonBench comparison has geometric- mean gains of 13.4 percent in forward and 24.4 percent in backward
1 Introduction
On Blackwell, operand preparation and normalization can constrain attention pipelines even as matrix multiplications become deeply pipelined. The introduction presents TLX changes to dense and causal softmax, scheduling, and backward data movement to address these constraints.
1 Introduction
- The work targets forward accumulator corrections, backward TMEM reuse hazards, and causal tile scheduling that must balance work against K/V locality.These constraints are addressed together through TLX asynchronous tasks, resource allocation, barriers, and on-chip storage views.Source: §1 · p. 1
- Dense dimension-128 forward holds a fixed row shift and uses packed BF16 exponentials to eliminate recurrent denominator and output-accumulator corrections.The packed approximation generates BF16 weights for the PV operation.Source: §1 · p. 1
- Guarded causal forward anchors each row’s shift to its first visible key block, selectively recomputes tiles that fail checks, and skips fully masked diagonal contractions.The skipped work is avoided before issuing QK or PV.Source: §1 · p. 1
- Dense forward distributes work nearly evenly across persistent two-CTA clusters, while causal query tiles run in reverse order within L2-sized head sections.The sectioned reverse order retains the causal scheduling strategy described for the implementation.Source: §1 · p. 1
- Dense forward saves an FP32 inverse denominator that backward uses to scale dO once and reconstruct unnormalized weights, avoiding LSE loads and per-score normalization.This moves row normalization out of the quadratic backward loop.Source: §1 · p. 1
- Backward overlaps dQ publication with later matrix multiplications, using paired BF16 reductions for dense cases and bounded packed fixed-point reductions for specified causal lengths.The stated causal range is 1,024 through 16,384 tokens.Source: §1 · p. 1
2 Attention and Blackwell Resource Limits
The section defines the attention and gradient notation, then describes Blackwell resource limits that shape the kernel schedule. Its discussion covers exponential evaluation, work distribution, shared-memory traffic, and backward operand readiness.
2.1 Forward and backward equations
- The forward and backward equations use Q, K, and V for sequence inputs, P for exact softmax probabilities, O for exact output, and dX for a loss gradient.Causal attention sets scores for j > i to −∞ before normalization; eP, m, ℓ, and L denote shifted exponentials, row shift, row sum, and log-sum-exp.Source: §2.1 · p. 2
- The fused backward pass retains five MMA contractions, targeting normalization, operand formats, and storage reuse to determine when contractions can begin.dK and dQ consume opposite orientations of dS.Source: §2.1 · p. 2
2.2 Resource limits after the FlashAttention-4 schedule
- For a forward tile, FA4 models exponential evaluation at the same 2,048 cycles as tensor work, motivating overlap; TLX constructs dense BF16 operands with packed instructions that remove MUFU.EX2 from the score path.The cited tile has 256 query rows, 128 key rows, and head dimension 128.Source: §2.2 · p. 2
- After dense transformations, one profiled configuration recorded 2,048 tensor-active cycles out of 2,125 elapsed per SM, leaving about 3.7 percent maximum throughput improvement from eliminating remaining idle cycles.The profile used B = 4, H = 48, N = 8,192, and d = 128 at 750 W.Source: §2.2 · p. 2
- Dense cluster counts are chosen to distribute output tiles nearly evenly, while causal tiles run longest-first within L2-sized head sections to reduce the completion tail and preserve K/V locality.The causal sections are sized using the K/V working-set estimate.Source: §2.2 · p. 2
- In FA4’s backward model, BF16 dQ staging lowers shared-memory cost from 2,688 to 2,432 cycles per SM, below the 2,560-cycle tensor-work estimate.At N = 8,192, amortized row-state and dK/dV publication raise the estimate to 2,456 cycles; these terms grow at shorter sequences.Source: §2.2 · p. 3
- An earlier backward kernel’s counters motivate scheduling contractions around operand readiness and shared-storage reuse: dV waits for QKT probabilities, while dK and dQ also need dOV T to form dS.The intermediates share 512 TMEM columns.Source: §2.2 · p. 3
3 Algorithmic Changes
The algorithm removes recurrent normalization corrections in dense attention and redesigns causal normalization and query-gradient reductions. It also reuses a saved reciprocal in backward to avoid per-score normalization.
3.1 Fixed-shift streaming softmax
- A fixed dense row shift makes every rescaling factor one, removing recurrent accumulator correction and its synchronization, while requiring finite-range weights and a finite, nonzero row sum.The shift trades dynamic-range adaptation for a simpler recurrence; guarded causal attention uses a row-dependent anchor and recovery instead.Source: §3.1 · p. 3
3.2 Guarded causal normalization
- Guarded causal rows take their shift from the first visible block; failed normalization checks trigger recomputation of the full query tile with conditionally rescaled online normalization.The route saves the final online state for backward; inputs outside admitted shapes use online normalization.Source: §3.2 · pp. 3–4
3.3 Packed BF16 exponential evaluation
- The dense path uses a packed BF16 exponential evaluator for every score, constructing approximate weights from paired FP32 scores and leaving final normalization to the epilogue.This replaces FA4’s mixed native-exponential and polynomial evaluation; the paper notes that the available benefit depends on hardware exponential throughput.Source: §3.3 · p. 4
3.4 Reducing causal diagonal work
- Skipping fully masked causal diagonal halves removes one QK and two PV calls per tile, and is enabled through N = 4,096 and at N = 8,192 on the guarded route.The skip also omits associated score and probability work while preserving the remaining producer–consumer sequence.Source: §3.4 · pp. 4–5
3.5 Saving the inverse denominator for backward
- Backward applies the saved reciprocal to dO in its linear preprocessing pass, then reconstructs weights without LSE loads or per-score normalization; rounded weights may not sum exactly to one.The dense route forms the denominator from BF16-rounded fragments accumulated in FP32, saves its FP32 reciprocal, and folds score scaling into the FP32 QK result without scaling K.Source: §3.5 · pp. 5–6
3.6 Bounded query-gradient reduction
- Dense dQ partials use paired BF16 reductions, while causal lengths 1,024–16,384 use bounded packed fixed-point reductions; at N = 32,768, the route uses FP32 reduction.Bounds depend on conversion and recovered sums staying in range; runtime checks enforce budgets, and failures trigger FP64 dQ recomputation.Source: §3.6 · p. 6
4 TLX Dataflow and Scheduling
The implementation organizes Blackwell attention around constrained on-chip storage, asynchronous forward and backward pipelines, and scheduling tailored to dense and causal workloads. It retains FA4’s contraction structure while changing data ownership, normalization paths, and publication to improve overlap and balance work.
4.1 On-chip storage
- At d = 128, each forward tile owns 256 query rows and advances through 128 keys, using TMA for tiles and TMEM for FP32 scores and accumulators.The dense non-causal two-CTA route covers 512 query rows per cluster; storage pressure leaves room for one CTA per SM.Source: §4.1 · p. 6
4.2 Measured configurations
- Measured binaries use different normalization and work-acquisition paths: only dense d = 128 forward saves the reciprocal, while other routes save LSE.Dense d = 128 forward uses a fixed-stride persistent grid; other main forward kernels use cluster launch control, and sparse recovery maps to flagged tiles.Source: §4.2 · p. 7, Table 1 · p. 7
4.3 Forward pipeline
- Interleaving H and L query groups lets a group’s PV work overlap the other group’s QK and softmax, while fixed shifts remove rescaling and accumulator correction.Short causal tiles can still expose the QK–softmax–PV dependency, and on-chip storage and registers limit resident CTAs.Source: §4.3 · p. 8
4.4 Causal scheduling
- Causal scheduling processes query blocks in reverse within L2-sized head sections, reducing the completion tail while preserving reuse of each head’s K/V operands.Section size depends on head dimension, sequence length, and the number of heads sharing the cache budget.Source: §4.4 · p. 8
4.5 Backward contractions and accumulator flow
- Backward retains FA4’s five-contraction order and reuses TMEM according to operand lifetimes, with DSMEM exchange required for dQ and publication overlapped with later contractions.Incremental dQ publication stages partials through a shared-memory ring; α is applied once during publication, and dK is scaled after accumulation.Source: §4.5 · p. 10
5 Numerical Behavior
The numerical evaluation examines approximation and rounding effects across exponential evaluation, benchmarked outputs and gradients, and causal recovery. The tested dimension-128 kernels pass the stated tolerance, while the paper identifies input-range, generalization, determinism, and evaluation-scope limits.
5.1 Finite-precision operations
- Forward’s packed BF16 exponentials and backward’s reciprocal scaling and residual rounding add gradient error beyond final BF16 conversion, while direct dQ publication and packed causal reduction add further numerical effects.For short causal prefixes, the authors recompute D for the first 128 query rows and add selected low-precision probability-residual contributions to dV.Source: §5.1 · p. 10
5.2 Intrinsic exponential accuracy
- On four million sampled exponents, packed outputs exceed the FP64 reference by factors of 1.0440–1.0687, and the varying ratio leaves input-dependent error after normalization.With the deployed nominal shift, maximum and mean relative errors are 6.93e−2 and 5.85e−2.Source: §5.2 · pp. 10–11, Table 3 · p. 11
5.3 Accuracy of the benchmarked kernels
- Both providers pass |a − b| ≤ 0.01 + 0.01|b| for outputs and gradients at all six tested d = 128 lengths through 32,768 tokens under both masks.At 32,768 tokens, TLX’s dense dQ relative ℓ2 error is 1.47e−2, compared with 2.37e−3 for FA4.Source: §5.3 · p. 11, Table 4 · p. 11
5.4 Recovery observations
- All guarded d = 128 cases recorded zero flagged groups, while separate d = 64 checks triggered recovery at 16,384 and 32,768 tokens.The recovery maps repeated across probes, and the timed callback includes the recovery launch even when no flags are set.Source: §5.4 · p. 11
5.5 Limitations
- Packed evaluation produces NaNs at effective encoded base-two exponents of −127 and −126.5; extending the seeded-fixture observations to general guarantees requires input-distribution error bounds.The paper also says full-pass determinism remains to be evaluated and that individual transformations, triggered-recovery cost, compilation times, and a CuTe implementation were not evaluated.Source: §5.5 · p. 12
6 Evaluation
The evaluation compares TLX with FA4 on a B200 variable-batch grid and with multiple implementations on a separate TritonBench workload. TLX improves overall throughput on the FA4 grid, though causal forward is slower; the TritonBench comparison also reports gains in both passes.
6.1 FA4 benchmark configuration
- The FA4 grid uses BF16 inputs with d = 128 and H = 16 on B200 at 1,000 W, covering dense and causal forward and backward across six lengths.There are 24 cases; each case uses a common sequence length across its batch, and the inputs are reused across masks, providers, and replays.Source: §6.1 · p. 12
6.2 FA4 variable-batch grid
- For dense forward, gains range from 12.8% to 13.4% across all six sequence lengths, including 1,024 through 32,768 tokens.At 32,768 tokens, TLX reaches 1,786.8 TFLOP/s; the current FA4 reference reaches 1,582.5 TFLOP/s.Source: §6.2 · p. 12, Table 5 · p. 13
- On the B200 at 1,000 W, TLX Attention shows geometric-mean gains of 11.7% for dense and 3.1% for causal attention, 7.3% overall across 24 cases; causal forward remains 1.5% slower.The benchmark uses the original timer.Source: §6.2 · p. 12, Table 6 · p. 13
6.3 TritonBench
- The TritonBench comparison is a separate dense B200 workload at 750 W, with B = 4, Hq = Hkv = 48, d = 128, and five sequence lengths.It includes benchmark-native CuTeDSL, public FA4, and cuDNN as comparators; its setup differs from the FA4 grid.Source: §6.3 · p. 12
- 13.4% forward and 24.4% backward are the minimum estimated geometric-mean gains over evaluated implementations on the dense TritonBench workload.TLX values are ratio-adjusted estimates for the August 2026 configuration, while comparator values come from profiler measurements.Source: §6.3 · p. 14, Table 8 · p. 14
7 Related Work
The related work situates the paper among approaches to attention memory traffic, normalization, arithmetic approximation, and scheduling. It distinguishes its training-focused combination of fixed shifts, packed BF16 exponentiation, and inverse-denominator reuse from prior methods.
Attention memory and normalization
Fixed-shift normalization
- FlashDecoding++ uses a precomputed unified max value for decoding [16]; this work applies fixed shifts to training forward and backward with packed BF16 evaluation and inverse-denominator reuse.VFA reaches a fixed shift through a different update schedule [17].Source: §7 · p. 14
Exponential approximation
- Exponent-field exponential approximations were introduced by Schraudolph; VEXP applies related BF16 constructions to softmax and FlashAttention-2 [5].Softermax combines low-precision base-two softmax with online normalization [8].Source: §7 · p. 14
Scheduling
- Prior attention-backward and persistent scheduling work includes DASH’s deterministic dQ reduction [20] and persistent attention planning in FlashInfer [19].The paper contrasts DASH’s ordered floating-point accumulation with its range-valid packed integer reduction.Source: §7 · p. 14
On-chip storage
- Colfax’s schedules use spare TMEM to release scores earlier or overlap the next QK with softmax, whereas this measured forward retains aliased score/probability storage.Source: §7 · p. 14
8 Conclusion
The combined implementation improves geometric-mean throughput over FA4 on the evaluated B200 grid, while causal forward remains slower. The conclusion emphasizes implementation-specific results and identifies future work integrating the ideas into CuTe.
8 Conclusion
- 7.3% geometric-mean throughput improvement over FA4 across all 24 numerically qualified dimension-128 cases on B200 at 1,000 W.The evaluation uses the original timer and prepared launches.Source: §8 · p. 14
- Causal forward remains 1.5% slower because PV still waits for probability publication after rescaling is removed.The authors identify overlapping more probability work with independent tensor operations as a remaining scheduling opportunity within available registers and on-chip storage.Source: §8 · p. 14
- A public artifact release has not been prepared for this draft.Source code, binary manifests, benchmark records, and numerical-check scripts are retained in an internal archive.Source: §8 · p. 15
- The results compare specific implementations and do not establish Triton or TLX as inherently better languages.The authors state that the algorithmic and scheduling ideas are transferable to other GPU languages and frameworks, with future work planned in CuTe.Source: §8 · p. 15