GPUmachines

Does MXFP8 FlashAttention Make Blackwell Training 1.6x Faster?

PyTorch reports up to 1.6x over BF16 for an MXFP8 FlashAttention-4 path on Blackwell. The system-level gain depends on shapes, fusion, accuracy and the rest of the training step.

Does MXFP8 FlashAttention Make Blackwell Training 1.6x Faster?

The claim that catches attention is a 1.6 times speed-up over BF16. The purchasing question is narrower: does low-precision FlashAttention-4 make a Blackwell training system faster for your model, or only for the shapes and fused pipeline that the authors measured?

On 16 September 2026, PyTorch published a technical account of an MXFP8 implementation for FlashAttention-4 forward and backward passes. The team reports 2.85 PF/s forward and 2 PF/s backward on LLM shapes. On its internal shapes, it reports 2.54 PF/s forward and 1.58 PF/s backward, with gains of up to 1.6 times and 1.52 times over BF16. Meta uses the resulting zero-gather jagged module for its GEM training workload, and the code is open source.

Those are useful results. They are not a general promise that every Blackwell training job will run 60% faster. The gain depends on attention shapes, precision policy, surrounding operators, kernel integration, sequence layout, memory movement and whether the rest of the step can keep up.

What MXFP8 changes inside attention

Attention contains two large matrix multiplications. Queries multiply transposed keys to produce scores; after softmax, probabilities multiply values to produce the output. FlashAttention reduces memory traffic by tiling this work so intermediate matrices do not have to be written to and read from high-bandwidth memory in full.

BF16 stores each value in 16 bits. MXFP8 stores values in eight-bit formats and attaches scale information to small blocks, allowing Blackwell tensor cores to use block-scaled matrix instructions. In principle, less data moves through the chip and more low-precision matrix operations fit into each unit of time.

The paper's implementation is interesting because it does not stop at swapping one datatype. Scale factors need storage and scheduling; online values such as attention probabilities and backward intermediates need quantisation; conversions can erase the arithmetic gain if they run as separate passes.

The authors address those costs by:

  • fitting scale factors into an already full tensor-memory allocation;
  • calculating some quantisation information while nearby work is already happening;
  • fusing RMSNorm, quantisation and GEMM output handling into producer or epilogue stages;
  • keeping variable-length data compact while scattering only the smaller scale-factor tensors into aligned layouts;
  • using MXFP8 in both forward and backward paths rather than accelerating inference alone.

This is the difference between a format capability and a useful training kernel. A data sheet may say the tensor core supports MXFP8, but the training stack must feed it without spending the saving on conversion and layout work.

Why Blackwell's peak ratio is not the training ratio

The PyTorch article says Blackwell's block-scaled matrix instructions can provide two to four times BF16 MMA throughput for MXFP8, MXFP6, MXFP4 and NVFP4 formats. That is an arithmetic-unit comparison. End-to-end training includes operations that do not accelerate by the same amount.

Softmax remains a clear example. Lower-precision matrix multiplication finishes sooner, so time previously hidden behind those operations can become visible. The authors describe adjusting the unroll and barrier sequence because softmax work began to expose bubbles at tile boundaries. Faster tensor cores had shifted the bottleneck rather than removed it.

At a larger scale, a training step also contains embedding work, feed-forward layers, normalisation, optimiser updates, collective communication, input loading and checkpoint activity. If attention consumes 30% of step time and its usable speed doubles, Amdahl's law caps the ideal whole-step gain at about 18%:

``text old step time = 70 non-attention units + 30 attention units = 100 new step time = 70 + (30 / 2) = 85 ideal speed-up = 100 / 85 = 1.176x ``

That calculation is only an example, not a forecast. The actual attention share and speed-up must come from the target model's profile. It shows why a 1.6 times kernel result cannot be pasted into a cluster business case.

Accuracy is part of the performance result

Block-scaled FP8 represents a narrower numerical range and less detail than BF16. Scaling small groups of values helps preserve useful information, but a team still needs evidence that the chosen precision path maintains acceptable loss, convergence and downstream model quality.

The right test is not whether one training step completes. Run enough of the intended workload to compare loss curves and detect instability. Keep seeds, data order, optimiser settings and checkpoint behaviour fixed where possible. Then evaluate the resulting model on the metrics that matter to the application.

Some workloads tolerate low precision readily. Others contain outliers or sensitive operations that need BF16 or FP32. Mixed precision is a design choice, not a badge of modernity.

Ask where the implementation accumulates in higher precision, where it quantises, how scales are selected and whether fallback paths exist for unsupported shapes. If the response is only “Blackwell supports FP8”, the software qualification has not started.

What the published figures actually cover

PyTorch reports two sets of figures. The LLM-shape result reaches 2.85 PF/s forward and 2 PF/s backward. The internal-shape result reaches 2.54 PF/s forward and 1.58 PF/s backward, with the quoted gains over BF16. The source describes a cross-attention module for advertising-model training with jagged sequence lengths and fused surrounding operators.

That context limits the conclusion in a useful way. The result demonstrates that MXFP8 can improve production training on Blackwell when the kernel and its neighbours are designed together. It does not establish the same gain for dense decoder-only self-attention, every head dimension, every sequence length or every PyTorch training stack.

The code lives in Meta's Ads Model Kernel Library under the lp_fa4 path. Buyers and engineering teams can inspect the implementation, supported shapes, build requirements and tests. Open source reduces guesswork, but it does not make integration free. A custom kernel still needs version control, security review, compilation, correctness tests and an upgrade plan.

Which buyers should care now?

Teams purchasing Blackwell for model training should care if attention occupies a material share of the step, sequences are long, the workload supports the required shapes and the software team can maintain specialised kernels. Recommender and advertising models with jagged cross-attention are an obvious fit because they resemble the disclosed production use.

Large-language-model teams should treat the release as a qualification target. Profile first. If attention dominates and the model can use the implementation, the saving may justify a test on B200 or B300 hardware. If feed-forward layers, communication or data loading dominate, buying extra GPUs to chase an attention-kernel result may solve the wrong problem.

Researchers running varied code with short project lives may prefer standard framework kernels even when they are slower. A highly tuned custom path can create maintenance work every time PyTorch, CUDA, the compiler or model code changes.

Inference-only buyers should not read a forward-and-backward training result as a serving benchmark. Their decision needs time to first token, inter-token latency, concurrency, KV-cache behaviour and request scheduling.

Choosing a server for this work

MXFP8 support starts with the accelerator, but server architecture determines whether the rest of the workload can feed it.

An HGX server is the natural route when training depends on fast communication across eight tightly connected GPUs. Check the exact B200 or B300 platform, NVLink and NVSwitch topology, host-memory bandwidth, local storage, network adapters, cooling and the software image the supplier will support.

A PCIe GPU server can make more sense for independent experiments, smaller models or teams that value flexible accelerator and NIC placement. The PCIe path may become the limit for communication-heavy training, but paying for a scale-up fabric provides little benefit when jobs remain on one GPU.

Neither route fixes a poorly fed training loop. Data staging, CPU preprocessing and checkpoint writes need enough capacity. For a cluster, the scale-out network must carry the chosen parallelism pattern without leaving the attention kernel waiting at collective barriers.

A qualification plan that answers the buying question

Use the target model and compare a current supported BF16 baseline with the MXFP8 path. Do not change hardware, data, batch size and precision at once.

1. Freeze the environment

Record GPU model, firmware, driver, CUDA, PyTorch, compiler, kernel commit, container digest and relevant environment variables. Save the model commit, dataset revision, sequence-length distribution and batch policy.

2. Profile the baseline

Measure complete training-step time and the share spent in attention forward, attention backward, other GEMMs, collectives, input processing and checkpointing. Capture GPU clocks, power, memory use and utilisation. A kernel test without the full-step profile cannot predict the purchase outcome.

3. Validate numerical behaviour

Compare outputs on fixed inputs, then run a representative training window. Record loss, gradient statistics, overflows, skipped updates and validation metrics. Define pass limits before seeing the faster result.

4. Measure more than the best case

Test the sequence lengths and batch sizes that occur in production, including awkward or sparse cases. Report median and tail step time. Note any shape that falls back to BF16 or another kernel.

5. Scale only after one GPU is understood

Repeat across the intended GPU count. Record compute time, exposed communication and scaling efficiency. If the low-precision kernel shortens compute, the same collective operation occupies a larger fraction of the step; network or NVLink limits can become visible sooner.

6. Run a long soak

Short tests miss memory growth, thermal behaviour, intermittent numerical problems and checkpoint delays. Run long enough to cross normal checkpoint and validation intervals, then restart from the saved state.

Use completed training work per wall-clock hour and per kilowatt-hour as decision metrics. Include quality. A faster step that requires more retries or converges to a worse result is not cheaper.

What to request from a supplier

Ask for a supported software matrix that names the exact accelerator, framework and kernel path. Request evidence for the model class and shape range you intend to run, not only a vendor slide about FP8 peak throughput.

The proposal should state whether the result comes from a public implementation, a private fork or a future software release. It should identify fallback behaviour and show how the configuration is reproduced after service or image rebuild.

For a cluster, ask for end-to-end training-step results at one node and the planned node count, with communication exposed separately. Include rack input power if tokens, samples or training steps per watt enter the financial case.

If the local facility is not ready for dense power and cooling, Buy & Host can keep the same owned-hardware qualification inside a data-centre environment. Hosting does not remove the need to pin and test the kernel stack.

Who should wait?

Wait when the workload has not been profiled, the kernel does not support the required shapes, or the team cannot own a specialised build. Keep BF16 when numerical qualification remains incomplete. Do not move a production training programme to a new precision path during a deadline simply because the microbenchmark is attractive.

Also wait if the intended hardware is not yet deliverable under a supported configuration. A reproducible BF16 job on available systems is more useful than a faster plan that depends on an unqualified software release.

Questions buyers ask

Is MXFP8 simply FP8 with a different name?

It is an eight-bit microscaling format that uses scale factors for small blocks of values. The block scaling helps adapt the limited numeric range to local data, and Blackwell provides native matrix instructions for it.

Does the 1.6 times result apply to all Blackwell training?

No. It is an “up to” result on the authors' internal shapes against their BF16 baseline. Your whole-model gain depends on shapes, attention share, fusion, communication and data movement.

Can an existing PyTorch model use this without changes?

Do not assume so. The open implementation has build, shape and integration requirements. Test the model, precision policy and environment against the repository documentation.

Does lower precision always reduce memory use?

The main data can be smaller, but scale factors, higher-precision accumulators, saved training state and fallback tensors still consume memory. Measure peak allocated and reserved memory in the real step.

Is BF16 now obsolete on Blackwell?

No. BF16 remains a useful baseline and may be required for sensitive operations, unsupported shapes or workloads whose engineering cost outweighs the gain.

Sources and further reading

← Back to blog