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

TL;DR

This paper introduces BAPS, a low-precision softmax scheme using 8-bit floating-point format and block-aware rescaling, significantly improving Transformer inference speed and hardware efficiency without accuracy loss.

Contribution

The paper presents a novel low-precision softmax method with 8-bit floating-point format and block-aware rescaling, enabling faster inference and reduced hardware costs.

Findings

Halves data movement bandwidth in softmax operations.

Reduces exponentiation unit area by computing in 8-bit precision.

Validates effectiveness on language and multi-modal models.

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.