Progressive Mixed-Precision Decoding for Efficient LLM Inference

TL;DR

The paper tackles the challenge of running large language models on resource-constrained devices by showing that uniform quantization across prefill and decoding is suboptimal. It introduces Phase-Aware Precision Allocation and Progressive Mixed-Precision Decoding (PMPD) to adapt precision over phases and across the decoding sequence, supported by two schedulers (static and learned). Offline calibration selects per-phase precision pairs to meet a target quality, while PMPD progressively lowers precision during decoding to improve memory-bound throughput. Empirical results across GPU and NPU targets show substantial speedups with minimal or no loss in output quality, enabling practical edge deployment of diverse LLMs.

Abstract

In spite of the great potential of large language models (LLMs) across various tasks, their deployment on resource-constrained devices remains challenging due to their excessive computational and memory demands. Quantization has emerged as an effective solution by storing weights in reduced precision. However, utilizing low precisions (i.e.~2/3-bit) to substantially alleviate the memory-boundedness of LLM decoding, still suffers from prohibitive performance drop. In this work, we argue that existing approaches fail to explore the diversity in computational patterns, redundancy, and sensitivity to approximations of the different phases of LLM inference, resorting to a uniform quantization policy throughout. Instead, we propose a novel phase-aware method that selectively allocates precision during different phases of LLM inference, achieving both strong context extraction during prefill and efficient memory bandwidth utilization during decoding. To further address the memory-boundedness of the decoding phase, we introduce Progressive Mixed-Precision Decoding (PMPD), a technique that enables the gradual lowering of precision deeper in the generated sequence, together with a spectrum of precision-switching schedulers that dynamically drive the precision-lowering decisions in either task-adaptive or prompt-adaptive manner. Extensive evaluation across diverse language tasks shows that when targeting Nvidia GPUs, PMPD achieves 1.4$-$12.2$\times$ speedup in matrix-vector multiplications over fp16 models, while when targeting an LLM-optimized NPU, our approach delivers a throughput gain of 3.8$-$8.0$\times$ over fp16 models and up to 1.54$\times$ over uniform quantization approaches while preserving the output quality.

Figures (7)

• Figure 1: Illustration of the conventional paradigm of quantized LLM inference (left) and our method that comprises phase-aware precision allocation and progressive mixed-precision decoding (right). Our approach is motivated by i) the distinct error resilience observed during the prefill and decoding phases, and ii) the increasing fault tolerance as decoding progresses to longer tokens.
• Figure 2: Phase-aware Precision Allocation. Here prefilling was performed by 2- (top) or 3-bit (bottom) Vicuna-7B, while decoding employed the 2-bit model for both cases. We observe three key improvements. i) enhanced reasoning abilities. The model demonstrated higher accuracy for numerical tasks, such as calculating Fibonacci numbers (left). ii) improved instruction-following. The model understood and responded to user instructions (middle). iii) strengthened emerging abilities like creativity. The model successfully handled open-ended tasks such as poem writing (right). More quantitative results are shown in Appendix \ref{['app:error_res']} and Figure \ref{['fig:ablation_phase_aware']}.
• Figure 3: Different schedules mixing 3- and 4-bit models for decoding. We examined text summarization on CNN/DM using Phi-1.5. To ensure fair comparison, the same number of tokens are generated by the 4-bit baseline. Moreover, the 4-bit model was used for prefilling for all schedules.
• Figure 4: Precision-switching scheduler.
• Figure 5: End-to-end NPU throughput. Speedup ratios are obtained in comparison with fp16 model.