MpFA boosts long-context AI speed on NVIDIA Blackwell GPUs
MpFA: Hardware-Efficient Train-Free QK4V8 FlashAttention Kernels on Blackwell GPUs
Distributed, Parallel, and Cluster Computing
Summary
Long-context AI models require a lot of memory and computing power, slowing down their operations on GPUs. The authors introduced MpFA, a new method that mixes different precisions of math operations to speed up a key part of these models without needing retraining. They also found a way to fix accuracy lost from this change using a clever smoothing step. Overall, MpFA makes AI models run faster when handling very long inputs on certain NVIDIA GPUs.
What this means in practice
- •For ai model deployers: Accelerate long-context large language model inference on NVIDIA Blackwell GPUs by using MpFA for faster attention kernels.
- •For gpu performance engineers: Optimize GPU kernels for complex matrix operations in AI workloads leveraging mixed precision and adaptive pipelining for better throughput.
Authors
Chencheng Deng, Jianbin Fang, Dezun Dong
Abstract
Long-context LLM inference pushes modern GPU serving stacks into an attention-bound regime, where both compute and memory are dominated by the softmax-GEMM pipeline. On NVIDIA Blackwell GPUs, FP4 Tensor Cores offer high matmul throughput, but we find that fully FP4 attention often fails to translate this throughput into end-to-end speedups due to non-matmul costs: online quantization after softmax, tensor/shared-memory data movement, and contention on the softmax path. We present MpFA, a training-free FlashAttention kernel optimized for Blackwell. Guided by hardware characterization, MpFA uses mixed precision: NVFP4 for QK and FP8 for PV (QK4PV8). This preserves low-bit QK throughput while avoiding the conversion and scaling overheads of FP4 PV. To recover accuracy without further stressing the softmax pipeline, MpFA introduces rank-one smoothing compensation implemented as an additional Tensor Core MMA. MpFA further improves performance with a fine-grained asynchronous pipeline, tensor-memory reuse, and adaptive parallel partitioning across prefill and decode. On an NVIDIA B200 and across 16K-128K contexts, MpFA improves prefill throughput over state-of-the-art BF16/FP8 baselines and increases end-to-end output throughput by 2.81$\times$ over BF16 FA4 across Llama-3.1-8B and Qwen3-14B. Across five benchmark suites and two models, rank-one compensation recovers 62.5% of the accuracy loss with about 2.0% kernel overhead.