PsumQuant: In-line Post-training Partial Sum Quantizer for Energy Efficient NPU Inference

The rapid growth of deep neural networks (DNNs) has intensified the demand for efficient hardware acceleration under quantization. While prior research has successfully reduced weight and activation precision, partial sums generated during accumulation often retain high precision, resulting in significant energy overhead. In this work, we analyze psum distributions in tiled architectures and reveal that within-tile outliers are input-dependent. We propose PsumQuant, a post-training, input-aware quantization that predicts psum scales on-the-fly. By leveraging the crest factor of input activations, our learnable scale predictor effectively bounds the psum bit-width while handling the extreme outliers in DNNs. Experimental results on a $128 \times 128$ systolic array demonstrate that PsumQuant compresses psum precision down to 8-bit within only a 1\% accuracy drop on ResNet-18 and a marginal 0.04 perplexity increase on Llama-3.1. Furthermore, bit-width reduction with PsumQuant results in a 45\% reduction in total energy with minimal accuracy loss, demonstrating that PsumQuant provides a highly efficient solution for actual NPU architectures.

Artificial intelligence (AI) models require substantial computing power and energy to run. To make them compute efficiently, engineers further shrink (quantize) the AI's data. However, even if the number of AI's data is reduced, the intermediate summation generated during processing remains large and bulky. Moving and storing these bulky numbers inside the computer chip creates a severe energy bottleneck. To solve this, we developed PsumQuant, a hardware-friendly method for on-the-fly compression of these intermediate summations. Our technique looks ahead at the incoming data and dynamically adjusts how it shrinks the numbers. We ensure that rare but important extreme values are preserved without slowing down the hardware. By efficiently compressing these intermediate summations, PsumQuant reduces the total energy consumption of AI chips by 45% with almost no loss in AI accuracy. It reduces the power drain of AI systems, paving the way for advanced, sustainable, and energy-efficient AI.