EXAQ: Exponent Aware Quantization For LLMs Acceleration

TL;DR

This paper introduces EXAQ, a method for optimizing the softmax layer in quantized LLMs, achieving sub-4-bit quantization and approximately 4x acceleration with minimal accuracy loss.

Contribution

It proposes an analytical approach to determine optimal clipping for softmax inputs, enabling ultra-low bit quantization and significant acceleration in LLM inference.

Findings

Achieves 2-bit quantization with baseline performance on LLaMA1-30B.

Realizes approximately 4x acceleration in the softmax accumulation phase.

Attains a 36.9% overall acceleration in the softmax operation.

Abstract

Quantization has established itself as the primary approach for decreasing the computational and storage expenses associated with Large Language Models (LLMs) inference. The majority of current research emphasizes quantizing weights and activations to enable low-bit general-matrix-multiply (GEMM) operations, with the remaining non-linear operations executed at higher precision. In our study, we discovered that following the application of these techniques, the primary bottleneck in LLMs inference lies in the softmax layer. The softmax operation comprises three phases: exponent calculation, accumulation, and normalization, Our work focuses on optimizing the first two phases. We propose an analytical approach to determine the optimal clipping value for the input to the softmax function, enabling sub-4-bit quantization for LLMs inference. This method accelerates the calculations of both $e^x$ and $\sum(e^x)$ with minimal to no accuracy degradation. For example, in LLaMA1-30B, we achieve baseline performance with 2-bit quantization on the well-known "Physical Interaction: Question Answering" (PIQA) dataset evaluation. This ultra-low bit quantization allows, for the first time, an acceleration of approximately 4x in the accumulation phase. The combination of accelerating both $e^x$ and $\sum(e^x)$ results in a 36.9% acceleration in the softmax operation.