Technical Summary: OASIS

Problem Statement

Large Language Models (LLMs) demand substantial memory and computational resources during inference. While quantization is a primary method for compression, existing approaches face a trade-off between efficiency and accuracy:

• Weight-Only Quantization (WOQ): Reduces memory but requires costly dequantization of weights to FP16 before execution, often consuming 20–90% of total runtime.
• Integer Weight-and-Activation Quantization (INT-WAQ): Enables fully low-precision computation but suffers from significant accuracy degradation due to limited representation capabilities.
• Non-Uniform Weight-and-Activation Quantization (NU-WAQ): Uses learned codebooks to better match data distributions, improving accuracy. However, its index-coded format is incompatible with standard low-precision hardware, necessitating dequantization to FP16 and negating computational benefits.

Furthermore, activation outliers in LLMs are more numerous and larger in magnitude than weight outliers, causing high quantization noise. Dynamically identifying these outliers during inference typically introduces latency overhead by placing detection on the critical path of the General Matrix Multiplication (GEMM) operation.

Methodology

The paper proposes OASIS (Outlier-Aware LUT-Based GEMM with Dual-Side Quantization), an algorithm-architecture co-design framework that enables efficient GEMM between non-uniformly quantized weights and activations without dequantization.

1. WAQ LUT-Based GEMM

OASIS leverages three specific opportunities in Non-Uniform WAQ (NU-WAQ) to optimize Lookup Table (LUT) based computation:

• Offline Pre-computation: Unlike WOQ methods where LUTs depend on streaming activations, NU-WAQ uses offline-learned centroids for both weights and activations. This allows the entire LUT to be pre-computed and stored on-chip before inference.
• Cartesian Product LUTs: Instead of storing full inner products (which require large LUTs dependent on reduction length), OASIS stores the Cartesian Product of weight and activation centroids. This drastically reduces LUT size.
• High Parallelism: The Cartesian Product LUT is independent of the reduction length (input channel count), allowing for large-granularity GEMMs and significantly higher computational parallelism compared to existing group-wise LUT methods.

2. Outlier-Aware Quantization with Look-Ahead

To handle activation outliers without runtime penalties, OASIS introduces a dual-branch computation scheme:

• Main Branch (Look-Ahead): Performs WAQ LUT-GEMM on the entire activation vector, treating outliers as if they were quantized normally. This branch ignores the quantization error of outliers initially.
• Outlier Branch (Error Compensation): Concurrently identifies outliers and calculates their quantization errors (residuals). These residuals are multiplied by dequantized weights and accumulated to compensate for the errors introduced in the main branch.
• Result: The final output is the sum of the look-ahead results and the error compensation, mathematically equivalent to conventional dynamic outlier detection but with the detection process removed from the critical path.

3. Hardware Architecture

The OASIS accelerator includes:

• Orizuru: A lightweight, tree-based top-k detection engine that identifies the largest and smallest activation values in real-time. It uses two-fold binary trees with shared leaf nodes to minimize comparisons and latency.
• Processing Elements (PE): Utilize "Concat Units" to merge weight and activation indices, feeding them into Index Counters to determine the frequency of each unique index pair.
• Reduction: A MAC Tree performs weighted sums using the index counts and pre-stored Cartesian Product values from the LUT.

Key Contributions

1. WAQ LUT-GEMM: A novel method enabling efficient GEMM between non-uniformly quantized weights and activations without dequantization, utilizing pre-computed Cartesian Product LUTs to reduce LUT size by 64× and increase parallelism by 1024× compared to existing WOQ LUT methods.
2. Outlier-Aware Scheme: A concurrent look-ahead computation and error compensation mechanism that handles activation outliers without incurring additional runtime overhead.
3. OASIS Accelerator: An architecture design that efficiently supports the proposed WAQ LUT-GEMM scheme.
4. Orizuru Engine: A lightweight top-k detection engine for real-time identification of activation outliers in data streams.

Experimental Results

The authors evaluated OASIS across various LLMs (OPT, LLaMA, Mistral) and tasks (perplexity, zero-shot accuracy) on an NVIDIA A100 GPU for algorithmic benchmarks and a cycle-accurate simulator for hardware metrics.

• Algorithmic Performance:

  • Compared to the FP16 baseline, OASIS incurs an average accuracy drop of only 1.94% (measured by perplexity), which is 6.34% lower than the Atom baseline.
  • In W4A4 precision, OASIS improves accuracy by 6.44% over Atom and 6.92% over QuaRot.
  • Dynamic outlier detection (OASIS) outperforms static thresholding (OASIS-S), reducing perplexity by 0.05 (W4A4) and 0.27 (W4A3).

• Hardware Performance:

  • Compared to the FIGLUT ASIC accelerator, OASIS achieves an average 3.00× speedup and a 1.44× improvement in energy efficiency.
  • Compared to the A100 GPU and QuaRot, OASIS achieves speedups of 5.41× and 3.12× respectively in throughput.
  • OASIS reduces LUT sizes by 62.1× to 994.2× and reduction FLOPs by 248.6× to 497.1× compared to existing WOQ LUT-GEMM designs (FIGLUT, LUT Tensor Core, LUT-GEMM).

Significance

The paper claims that OASIS bridges the gap between algorithmic efficiency and hardware implementation for Non-Uniform Weight-and-Activation Quantization. By eliminating the need for dequantization and enabling massive parallelism through Cartesian Product LUTs, OASIS achieves high model accuracy with substantial throughput and energy efficiency gains. The work demonstrates that NU-WAQ, previously hindered by hardware incompatibility, can be executed efficiently on custom accelerators while maintaining model quality comparable to FP16 baselines.