MoQAE: Mixed-Precision Quantization for Long-Context LLM Inference via Mixture of Quantization-Aware Experts

TL;DR

MoQAE introduces a mixed-precision quantization method using a mixture of experts approach to optimize memory efficiency and accuracy in long-context LLM inference, outperforming existing quantization techniques.

Contribution

The paper proposes a novel mixture of quantization-aware experts method with efficient routing and fine-tuning for improved long-context LLM inference.

Findings

Outperforms state-of-the-art KV cache quantization methods in efficiency.

Effectively balances model accuracy and memory reduction.

Reduces inference overhead with routing freezing and sharing mechanisms.

Abstract

One of the primary challenges in optimizing large language models (LLMs) for long-context inference lies in the high memory consumption of the Key-Value (KV) cache. Existing approaches, such as quantization, have demonstrated promising results in reducing memory usage. However, current quantization methods cannot take both effectiveness and efficiency into account. In this paper, we propose MoQAE, a novel mixed-precision quantization method via mixture of quantization-aware experts. First, we view different quantization bit-width configurations as experts and use the traditional mixture of experts (MoE) method to select the optimal configuration. To avoid the inefficiency caused by inputting tokens one by one into the router in the traditional MoE method, we input the tokens into the router chunk by chunk. Second, we design a lightweight router-only fine-tuning process to train MoQAE with a comprehensive loss to learn the trade-off between model accuracy and memory usage. Finally, we introduce a routing freezing (RF) and a routing sharing (RS) mechanism to further reduce the inference overhead. Extensive experiments on multiple benchmark datasets demonstrate that our method outperforms state-of-the-art KV cache quantization approaches in both efficiency and effectiveness.