BAPS: A Fine-Grained Low-Precision Scheme for Softmax in Attention via Block-Aware Precision reScaling

TL;DR

This work targets the softmax bottleneck in Transformer inference by proposing a fine-grained low-precision workflow that uses an 8-bit floating format HiF8 and block-aware precision reScaling. It advances two key ideas: (i) performing softmax operations in 8-bit precision and (ii) reducing data movement by constraining matrix outputs to 8-bit while computing exponentiations in low precision, aided by a dynamic restart mechanism to preserve accuracy. The approach combines 8-bit exponents for $P$, 8-bit attention scores for $S$, and integer powers of two to bound error, achieving substantial throughput gains with limited accuracy loss (typically ≤1% across NLP and multimodal benchmarks) and offering a concrete hardware-software co-design path for future low-precision accelerators. The results suggest near-doubling of end-to-end inference throughput without increasing chip area, and the framework supports integration with existing quantized GEMM backends and potential future extensions to PTQ and sparse attention methods.

Abstract

As the performance gains from accelerating quantized matrix multiplication plateau, the softmax operation becomes the critical bottleneck in Transformer inference. This bottleneck stems from two hardware limitations: (1) limited data bandwidth between matrix and vector compute cores, and (2) the significant area cost of high-precision (FP32/16) exponentiation units (EXP2). To address these issues, we introduce a novel low-precision workflow that employs a specific 8-bit floating-point format (HiF8) and block-aware precision rescaling for softmax. Crucially, our algorithmic innovations make low-precision softmax feasible without the significant model accuracy loss that hampers direct low-precision approaches. Specifically, our design (i) halves the required data movement bandwidth by enabling matrix multiplication outputs constrained to 8-bit, and (ii) substantially reduces the EXP2 unit area by computing exponentiations in low (8-bit) precision. Extensive evaluation on language models and multi-modal models confirms the validity of our method. By alleviating the vector computation bottleneck, our work paves the way for doubling end-to-end inference throughput without increasing chip area, and offers a concrete co-design path for future low-precision hardware and software.

Figures (12)

• Figure 1: The cumulative distribution function $\Phi$ of different data formats. The scores are dumped from Qwen3-30B-A3B-Instruct model.
• Figure 2: Illustration of numerical error in accumulative products: a comparison between exponents of 2 with integers and direct exponents of 2 against FP32 ground truth.
• Figure 3: The single frame generated from different setups using Wan-2.2. The illustration only shows the upper-left sub-frame for each setup. Only setup Naive E2E generated blurred frame and the other three frames are consistent with each other in a high quality.
• Figure 4: Probability density of restart rates for both NLP and T2V tasks. The solid line represents the density of NLP task and the dashed line represents the density of T2V task.
• Figure 5: Prompt 1: "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about."