🤖 machine learning

Search Your Block Floating Point Scales!

This paper introduces ScaleSearch, a novel strategy that optimizes Block Floating Point scale factors through a fine-grained search of mantissa bits to significantly reduce quantization error and improve the performance of low-precision generative models, including a specialized ScaleSearchAttention algorithm that achieves near-zero performance loss.

Original authors: Tanmaey Gupta, Hayden Prairie, Xiaoxia Wu, Reyna Abhyankar, Qingyang Wu, Austin Silveria, Pragaash Ponnusamy, Jue Wang, Ben Athiwaratkun, Leon Song, Tri Dao, Daniel Y. Fu, Chris De Sa

Published 2026-05-13

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Tanmaey Gupta, Hayden Prairie, Xiaoxia Wu, Reyna Abhyankar, Qingyang Wu, Austin Silveria, Pragaash Ponnusamy, Jue Wang, Ben Athiwaratkun, Leon Song, Tri Dao, Daniel Y. Fu, Chris De Sa

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Packing a Suitcase Efficiently

Imagine you are trying to pack a suitcase (a computer's memory) for a trip. You have a lot of items (data) to fit in, but the suitcase is small. To make everything fit, you have to squish the items down (this is called quantization).

In the past, when packing a group of items, you would look at the biggest item in the group and say, "Okay, I'll shrink everything else so the biggest item fits perfectly." This is the standard method used by modern AI chips.

The Problem:
The authors noticed that just because the biggest item fits, it doesn't mean the rest of the items are packed efficiently. Sometimes, shrinking everything based on the biggest item leaves a lot of empty space or squishes the smaller items too much, making them hard to recognize later. It's like trying to fit a giant teddy bear and a tiny button in a box; if you size the box for the bear, the button gets lost in the noise.

The Solution: "ScaleSearch" (The Smart Packager)

The paper introduces a new strategy called ScaleSearch. Instead of just looking at the biggest item and setting the rule once, the algorithm takes a tiny "searchlight" and checks a few different ways to pack the suitcase.

The Old Way: "The biggest item is 10 inches tall. I'll make my scale 10."
The New Way (ScaleSearch): "The biggest item is 10 inches. But wait... if I set my scale to 9.5, the big item still fits, but the medium items become much clearer. If I set it to 10.5, the small items look better. Let me quickly test these few options and pick the one that makes the whole group look the best."

The paper shows that by doing this tiny, quick search, the computer can pack the data much more accurately, reducing "quantization error" (the messiness caused by squishing data) by about 27%.

The Special Hardware: NVFP4

This trick works specifically because of new, super-fast computer chips (NVIDIA's Blackwell architecture) that use a format called NVFP4.

Think of the old packing method as using a ruler with only big markings (1, 2, 3). The new chips have a ruler with tiny, fine markings (1.0, 1.1, 1.2, etc.). ScaleSearch is the technique that uses those tiny, fine markings to find the perfect fit, rather than just guessing with the big markings.

The "Attention" Upgrade: ScaleSearchAttention

AI models (like chatbots) have to pay attention to words they've already said to understand the next word. This "memory" is called the KV Cache. Storing this memory takes up a lot of space and slows things down.

The authors created a special version called ScaleSearchAttention.

What it does: It applies the "Smart Packager" logic to the AI's memory.
The Result: It allows the AI to store its memory in a very compact format (4-bit) without losing its ability to understand context.
The Analogy: Imagine a librarian who usually has to write down every book title in full detail (slow and bulky). With this new method, the librarian uses a clever shorthand code that is tiny and fast to write, but they double-check the code to make sure they didn't miss any important details.

What Did They Prove? (The Results)

The paper tested this on real AI models (like Llama and Qwen) and video generation tools (like Mochi).

Better Math & Reasoning: When they used ScaleSearch on math problems, the AI got significantly smarter. For one model, the score on a math test jumped by 15 points compared to the standard method.
Better Language: When the AI wrote stories or answered questions, it made fewer mistakes. The "perplexity" (a measure of how confused the AI is) dropped, meaning the AI sounded more natural and human.
Speed: The best part? This "search" is so fast that it barely slows down the computer. It's like checking three different ways to tie your shoes; it takes a split second but ensures you don't trip later. The system runs at 98% of the speed of the standard method.

Summary

The paper says: "Don't just guess how to shrink your data based on the biggest piece. Take a split second to check a few nearby options, and you will get a much clearer, more accurate result with almost no speed penalty."

They call this ScaleSearch, and when applied to the AI's memory, they call it ScaleSearchAttention. It makes AI models smarter and more efficient without needing new hardware.

Technical Summary: Search Your Block Floating Point Scales!

Problem Statement

Quantization is a standard technique for accelerating generative model inference, yet maintaining accuracy at ultra-low bit-widths remains challenging. While Block Floating Point (BFP) formats offer a balanced trade-off between dynamic range, memory efficiency, and throughput, standard BFP algorithms rely on a fixed scale determined by the maximum absolute magnitude within a block. The authors observe that this "max-scaling" heuristic is often suboptimal, leading to unnecessary quantization errors. This issue is particularly relevant with the advent of NVIDIA's Blackwell architecture, which introduces hardware support for microscaling BFP formats like NVFP4. Unlike traditional BFP, NVFP4 includes mantissa bits within its 8-bit block scale (E4M3 format), providing a finer granularity for scale selection that standard max-scaling fails to exploit.

Methodology

ScaleSearch

The core contribution is ScaleSearch, an algorithm that leverages the mantissa bits in microscaling formats to perform a fine-grained search for the optimal block scale.

Mechanism: Instead of rigidly selecting the scale based on the block's maximum value ( $s = \text{round}(\|x\|_\infty / 6.0)$ ), ScaleSearch explores a neighborhood of scales around this default value. It treats the scale as an integer offset $f$ applied to the default scale's bit representation: $s(f) = \text{reinterpret}(\text{reinterpret}(s, \text{int8}) + f, \text{fp8}_{UE4M3})$ .
Search Strategy: The algorithm exhaustively searches a small range of offsets (e.g., $f \in [-2, 6]$ ) to find the scale that minimizes the Mean Squared Error (MSE) for the given block distribution.
Rationale: Empirical analysis on both synthetic Gaussian data and real neural network activations (e.g., Llama 3.1 Key states) reveals a bimodal distribution of optimal offsets. While the default scale (offset 0) maps the largest magnitude to the FP4 value 6, a significant portion of blocks are better represented by mapping the largest magnitude to the FP4 value 4, requiring a scale factor approximately 1.5x larger. ScaleSearch identifies these optimal mappings.

ScaleSearchAttention

To extend these benefits to the attention mechanism and KV cache, the authors propose ScaleSearchAttention. This is an end-to-end FP4 attention algorithm designed for causal language modeling that avoids de-quantization overheads.

Native NVFP4 Execution: It performs QKT and PV matrix multiplications directly on Tensor Cores using NVFP4 precision with FP32 accumulation.
Incoherence Processing (IP) & Magnitude Reduction: To mitigate outliers and reduce quantization error in Q and K tensors, the method employs Hadamard transforms (following QuIP#) and an additional linear transformation pair ( $Q' = QR^{-T}, K' = KR$ ) that preserves attention scores while reducing the average squared magnitude of the projected tensors.
Attention-Sink-Aware Mixed-Precision Caching: Recognizing that attention scores concentrate on initial and recent tokens (Attention Sink), the method stores the first and last blocks of the KV cache in full precision, while the intermediate blocks are stored in compressed NVFP4. This ensures high precision for critical tokens without scaling memory usage with sequence length.

Key Contributions

ScaleSearch Algorithm: A scale-searching technique for Block Floating Point quantization that minimizes block-wise quantization error by exploring neighboring representable scales. It specifically leverages the mantissa resolution of the NVFP4 E4M3 scale format.
ScaleSearchAttention: A co-optimized FP4 attention mechanism that enables end-to-end inference with near-zero accuracy degradation. It integrates ScaleSearch, incoherence processing, and mixed-precision KV caching to perform QKT and PV operations directly in NVFP4 on Tensor Cores.
Empirical Validation: Extensive experiments demonstrating consistent gains in perplexity and benchmark accuracy over prior FP4 baselines (such as SageAttention3 and Naive-FP4) across large-scale language models and diffusion models.

Experimental Results

Post-Training Quantization (PTQ)

Language Models: On the Qwen3-8B model, ScaleSearch improved performance by up to 15 points on the MATH-500 benchmark compared to standard NVFP4 quantization. It also significantly closed the gap on GPQA and MMLU benchmarks.
Diffusion Models: When applied to the Mochi and CogVideoX models using SageAttention3, ScaleSearch improved VQA-a, VQA-t, and FScore metrics while matching full-precision alignment scores (CLIPSIM, CLIP-T).

Attention and Perplexity

Perplexity (Wikitext-2): ScaleSearchAttention reduced perplexity by up to 22% for Llama 3.1 70B compared to vanilla FP4 quantization. For Llama 3.1 70B, it achieved a PPL of 2.6348, a significant improvement over SageAttention3 (3.3899) and Naive-FP4 (3.4000).
GPQA Diamond: On the Llama 3.1 8B Instruct model, ScaleSearchAttention achieved 32.32% accuracy, outperforming SageAttention3 (26.26%) and matching the full-precision baseline (31.81%).

Efficiency and Overhead

Quantization Overhead: ScaleSearch introduces a minimal overhead of 1.74x during the FP32 to NVFP4 quantization phase (0.0449 ms vs. 0.0258 ms baseline) for a restricted search range.
Throughput: In attention kernels, ScaleSearch achieves 98.3% of the baseline throughput (SageAttention3) for non-causal attention at 32K sequence lengths, demonstrating negligible runtime impact.
Latency: End-to-end latency for text-to-video generation increased only marginally (e.g., from 353.40s to 364.68s for Mochi), confirming the method's practicality for real-world workloads.

Significance

The paper claims that ScaleSearch addresses a fundamental suboptimality in standard BFP scaling strategies by utilizing the previously underutilized mantissa bits in microscale formats. By systematically searching for the optimal scale rather than relying on a heuristic maximum, the method significantly reduces quantization error. The introduction of ScaleSearchAttention demonstrates that these gains can be extended to the memory- and compute-intensive attention mechanism, enabling high-throughput, low-precision inference with near-zero accuracy loss. The authors position this work as a step toward fully utilizing emerging 4-bit hardware capabilities (like NVIDIA Blackwell) for efficient generative AI.