Abstract
Quantizing Large Language Models (LLMs) often involves a trade-off between hardware efficiency and model accuracy. While uniform integer quantization (I-WAQ) is fast on standard hardware, it suffers from significant accuracy loss at low bit-widths. Conversely, non-uniform K-Means quantization (K-WAQ) preserves accuracy but is computationally expensive on traditional GPUs due to dequantization overhead. This paper presents KLLM, a hardware-software co-design that enables efficient K-Means quantization through an index-based computation scheme and a dedicated outlier detection engine, achieving significant speedups and energy efficiency gains over state-of-the-art GPUs.
Introduction
The industry standard for deploying LLMs on resource-constrained hardware is uniform quantization. However, because LLM weights and activations typically follow Gaussian-like distributions, uniform spacing wastes precision on empty ranges and fails to capture the density near the mean. K-Means quantization offers a solution by finding optimal centroids for the data, but standard GPUs spend roughly 63% of their time performing codebook lookups to dequantize data before calculation. KLLM addresses this "quantization tax" by rethinking matrix multiplication (MatMul) at the hardware level.
Methodology
KLLM introduces a specialized architecture designed to handle non-uniform quantization natively through three primary innovations:
1. Index-Based Computation
Instead of dequantizing indices back to floating-point values, KLLM operates directly on the indices. For a 4-bit K-Means codebook, there are only 256 possible product combinations between two indices. KLLM precomputes these products and shifts the computational burden from expensive multipliers to simple counters that track the frequency of each index pair. A final weighted sum is then performed, transforming MatMul into a histogram-based problem.
2. Orizuru Outlier Detection
Activation outliers—rare but extreme values—can significantly degrade quantization performance. KLLM utilizes Orizuru, a lightweight hardware engine that identifies the top-k largest and smallest values in the activation stream during online inference. These outliers are maintained in high-precision FP16, while the remaining "inliers" are safely quantized. This approach reduces the error (RMSE) between offline and online centroids from 0.14 to 0.01, allowing offline-trained centroids to remain effective.
3. Hardware Accelerator Design
The custom ASIC (simulated at 28nm) features 16 Processing Element (PE) lanes, dedicated index distribution counters, and functional units for nonlinear operations. Because inputs to functions like SiLU or GELU are quantized, KLLM implements these via simple codebook swaps, making nonlinear functions computationally inexpensive.
Evaluation and Results
The authors evaluated KLLM using the WikiText-2 dataset for perplexity (PPL) and several zero-shot reasoning tasks (PIQA, ARC, HellaSwag, WinoGrande). Calibration was performed using samples from the C4 dataset.
Key Findings
- Accuracy: On LLaMA-2-7B (W4A4), KLLM achieved a perplexity of 5.90, closely approaching the FP16 baseline of 5.47 and significantly outperforming SmoothQuant (702.3).
- Performance: KLLM demonstrates an 8.66× speedup over an NVIDIA A100 GPU and a 2.12× improvement over specialized accelerators like SpAtten.
- Energy Efficiency: By avoiding high-precision floating-point units and reducing memory shuffling, KLLM achieves a 214.76× improvement in energy efficiency compared to the A100.
Discussion and Future Work
While KLLM offers a breakthrough in quantization efficiency, it currently requires custom silicon to realize these gains; implementing the index-counting logic on general-purpose GPUs via CUDA remains a challenge. Additionally, while Orizuru mitigates the gap between offline and online data, extreme distribution shifts in real-world applications may eventually require periodic centroid updates. Future research will explore scaling this architecture to even larger models, such as Mixture-of-Experts (MoE), where activation sparsity is more prevalent.