Search papers, labs, and topics across Lattice.
This paper addresses the attention-bound bottleneck in long-context LLM inference on NVIDIA Blackwell GPUs, where FP4 Tensor Cores offer high matmul throughput but are hindered by non-matmul costs. The authors introduce MpFA, a training-free FlashAttention kernel employing mixed precision (NVFP4 for QK, FP8 for PV) to leverage FP4 throughput while mitigating conversion and scaling overheads. MpFA achieves significant speedups and accuracy recovery through rank-one smoothing compensation, asynchronous pipelining, and adaptive partitioning.
Blackwell GPUs' FP4 Tensor Cores unlock 2.81x output throughput for long-context LLMs by strategically using mixed precision and rank-one smoothing, bypassing attention bottlenecks.
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.
