QServe: W4A8KV4 Quantization and System Co-design for Efficient LLM Serving

[Paper] [LMQuant Quantization Algorithm Library] [Website]

QServe: Efficient and accurate LLM serving system on GPUs with W4A8KV4 quantization (4-bit weights, 8-bit activations, and 4-bit KV cache). Compared with leading industry solution TensorRT-LLM, QServe achieves 1.2x-1.4x higher throughput when serving Llama-3-8B, and 2.4x-3.5x higher throughput when serving Qwen1.5-72B, on L40S and A100 GPUs. QServe also allows users to achieve A100-level throughput on 3x cheaper L40S GPUs.

Introduction

Quantization can accelerate large language model (LLM) inference. Going beyond INT8 quantization, the research community is actively exploring even lower precision, such as INT4. Nonetheless, state-of-the-art INT4 quantization techniques only accelerate low-batch, edge LLM inference, failing to deliver performance gains in large-batch, cloud-based LLM serving. We uncover a critical issue: existing INT4 quantization methods suffer from significant runtime overhead (20-90%) when dequantizing either weights or partial sums on GPUs. To address this challenge, we introduce QoQ, a W4A8KV4 quantization algorithm with 4-bit weight, 8-bit activation, and 4-bit KV cache. QoQ stands for quattuor-octo-quattuor, which represents 4-8-4 in Latin. QoQ is implemented by the QServe inference library that achieves measured speedup. The key insight driving QServe is that the efficiency of LLM serving on GPUs is critically influenced by operations on low-throughput CUDA cores. Building upon this insight, in QoQ algorithm, we introduce progressive quantization that can allow low dequantization overhead in W4A8 GEMM. Additionally, we develop SmoothAttention to effectively mitigate the accuracy degradation incurred by 4-bit KV quantization. In the QServe system, we perform compute-aware weight reordering and take advantage of register-level parallelism to reduce dequantization latency. We also make fused attention memory-bound, harnessing the performance gain brought by KV4 quantization. As a result, QServe improves the maximum achievable serving throughput of Llama-3-8B by 1.2× on A100, 1.4× on L40S; and Qwen1.5-72B by 2.4× on A100, 3.5× on L40S, compared to TensorRT-LLM. Remarkably, QServe on L40S GPU can achieve even higher throughput than TensorRT-LLM on A100. Thus, QServe effectively reduces the dollar cost of LLM serving by 3×.

The current release supports:

• Blazingly fast system support for QoQ W4A8KV4 quantization (Algorithim release: LMQuant);
• Pre-quantized QServe model zoo with W4A8KV4 QoQ for mainstream LLMs;
• Fully PyTorch-based runtime and user interface for LLM serving, with TensorRT-LLM-level efficiency and PyTorch-level flexibility;
• Full support for in-flight batching and paged attention;
• Efficient fused CUDA kernels for W4A8/W8A8 GEMM and KV4/KV8 attention;
• Easy-to-use examples on speed benchmarking and large-scale end-to-end content generation (with W4A8KV4, in-flight batching and paged attention).

News

• [2024/05] 🔥 QServe is publicly released! Check our paper here.

Contents

• QServe: W4A8KV4 Quantization and System Co-design for Efficient LLM Serving
  • Introduction
  • News
  • Contents
  • Installation
  • QServe Model Zoo
  • Usage and Examples
  • Results
  • Accuracy Evaluation
  • Efficiency Benchmarks
  • Reference
  • Related Projects
  • Acknowledgement

Installation

1. Clone this repository and navigate to the corresponding folder:

   git clone https://github.com/mit-han-lab/qserve
   cd qserve

2. Install QServe

   conda create -n QServe python=3.10 -y
   conda activate QServe
   pip install --upgrade pip  # enable PEP 660 support
   pip install -e .
   pip install flash-attn --no-build-isolation

   We recommend starting an interactive python CLI interface and run

   import flash_attn

   to check whether FlashAttention-2 is installed successfully. If not, we recommend downloading pre-built wheels from here. Please notice:

• PyTorch version needs to exactly match with the version specified in the .whl name;
• Check out both cxx11abiTRUE and cxx11abiFALSE wheels if one of them does not work;
• It's recommended to match CUDA version specified in the .whl filename, but minor mismatches (e.g. 12.1 vs 12.2, or even 11.8 vs 12.2) usually do not matter.

3. Compile the CUDA kernels.

   cd kernels
   python setup.py install

4. If you want to clone our model zoo, please make sure that git-lfs is installed.

QServe Model Zoo

We provide pre-quantized checkpoints for multiple model families. For example, for Llama-3-8B model, please run the following commands to download:

# git lfs install  # install git lfs if not already
mkdir -p qserve_checkpoints && cd qserve_checkpoints
git clone https://huggingface.co/mit-han-lab/Llama-3-8B-QServe 

For other models, please refer to the detailed support list for the links to download:

For flagship datacenter GPUs such as the A100, it is recommended to use QServe-per-channel, while for inference datacenter GPUs like the L40S, QServe-per-group is the recommended approach.

If you are interested in generating the quantized checkpoints on your own, please follow the instructions in LMQuant to run QoQ quantization and dump the fake-quantized models. We then provide checkpoint converter to real-quantize and pack the model into QServe format:

python checkpoint_converter.py --model-path <hf-model-path> --quant-path <fake-quant-model-path> --group-size -1 --device cpu
# <fake-quant-model-path> is a directory generated by LMQuant, including model.pt and scale.pt

We also provide a script to run the checkpoint converter. The final model will be automatically stored under qserve_checkpoints.

Usage and Examples

We support both offline benchmarking and online generation (in-flight-batching) in QServe.

1. Offline speed benchmarking (Batched input sequences, fixed context length = 1024 and generation length = 512). We take Llama-3-8B (per-channel quant) as an example here. Please make sure that you have already downloaded the QoQ-quantized QServe model.

export MODEL_PATH=./qserve_checkpoints/Llama-3-8B-QServe # Please set the path accordingly

GLOBAL_BATCH_SIZE=128 \
python qserve_benchmark.py \
  --model $MODEL_PATH \
  --benchmarking \
  --precision w4a8kv4 \
  --group-size -1

If you hope to use larger batch sizes such as 256, you may need to change NUM_GPU_PAGE_BLOCKS to a larger value than the automatically-determined value on A100. For example:

export MODEL_PATH=./qserve_checkpoints/Llama-3-8B-QServe # Please set the path accordingly

GLOBAL_BATCH_SIZE=256 \
NUM_GPU_PAGE_BLOCKS=6400 \
python qserve_benchmark.py \
  --model $MODEL_PATH \
  --benchmarking \
  --precision w4a8kv4 \
  --group-size -1

2. This is an online demonstration of batched generation, showcasing in-flight batching, paged attention of W4A8KV4 QoQ LLMs. We will randomly sample a set of safety-moderated conversations from the WildChat dataset and process them efficiently through in-flight batching.

export MODEL_PATH=./qserve_checkpoints/Llama-3-8B-Instruct-QServe # Please set the path accordingly

python qserve_e2e_generation.py \
  --model $MODEL_PATH \
  --ifb-mode \
  --precision w4a8kv4 \
  --quant-path $MODEL_PATH \
  --group-size -1

3. Argument list in QServe

   Below are some frequently used arguments in QServe interface:

• --model: Path to the folder containing hf model configs. Can be the same as --quant-path if you directly download the models from QServe model zoo.
• --quant-path: Path to the folder containing quantized LLM checkpoints.
• --precision: The precision for GEMM in QServe, please choose from the following values: w4a8kv4, w4a8kv8, w4a8 (means w4a8kv8), w8a8kv4, w8a8kv8, w8a8 (means w8a8kv8). Default: w4a8kv4.
• --group-size: Group size for weight quantization, -1 means per-channel quantization. QServe only supports -1 or 128. Please make sure your group size matches the checkpoint.
• --max-num-batched-tokens: Maximum number of batched tokens per iteration. Default: 262144.
• --max-num-seqs: Maximum number of sequences per iteration. Default: 256. Remember to increase it if you want larger batch sizes.
• --ifb-mode: Enable in-flight batching mode. Suggest to activate in e2e generation.

• --benchmarking: Enable speed profiling mode. Benchmark settings aligned with TensorRT-LLM.

  Environment variables in QServe:

• GLOBAL_BATCH_SIZE: Batch size used in offline speed benchmarking.
• NUM_GPU_PAGE_BLOCKS: Number of pages to be allocated on GPU. If not specified, it will be automatically determined based on available GPU memory. Note that the current automatic GPU page allocation algorithm is very conservative. It is recommended to manually set this value to a larger number if you observe that GPU memory utilization is relatively low.

4. One-line scripts:

We also provide sample scripts in QServe.

• End to end generation: ./scripts/run_e2e.sh;
• Speed benchmarking: ./scripts/benchmark/benchmark_a100.sh or ./scripts/benchmark/benchmark_l40s.sh.

These scripts are expected to be executed in the QServe project folder (not in the scripts folder). Please note that git-lfs is needed for downloading QServe benchmark config files from huggingface before running the benchmark scripts.

Results

We evaluate QServe W4A8KV4 quantization on a wide range of mainstream LLMs. QServe consistently outperforms existing W4A4 or W4A8 solutions from the accuracy perspective, while providing State-of-the-Art LLM serving efficiency.

Efficiency Benchmarks

When serving the large language models Llama-3-8B and Qwen1.5-72B on L40S and A100 GPUs, QServe demonstrates superior performance, achieving 1.2x-1.4x higher throughput compared to the leading industry solution, TensorRT-LLM, for Llama-3-8B, and a 2.4x-3.5x higher throughput for Qwen1.5-72B. It is also able to deliver higher throughput and accomodate the same batch size on L40S compared with TensorRT-LLM on A100 for six of eight models benchmarked, effectively saving the dollar cost of LLM serving by around 3x.

Benchmarking setting: the criterion is maximum achieveable throughput on NVIDIA GPUs, and the input context length is 1024 tokens, output generation length is 512 tokens. For all systems that support paged attention, we enable this feature. In-flight batching is turned off in the efficiency benchmarks.

The absolute token generation throughputs of QServe and baseline systems (Unit: tokens/second. -- means unsupported). All experiments were conducted under the same device memory budget. Throughput increase of QServe is calculated with regard to the best baseline in each column. It is recommended to use QServe-per-channel on high-end datacenter GPUs like A100 and QServe-per-group is recommended on inference GPUs like L40S.

We recommend direcly setting the NUM_GPU_PAGE_BLOCKS environmental variable to 25 * batch size, since in our benchmarking setting we have a context length of 1024 and generation length of 512, which corresponds to 24 pages (each page contains 64 tokens). We leave some buffer by allocating one more page for each sequence.

Accuracy Evaluation

QServe also maintains high accuracy thanks to the QoQ algorithm provided in our LMQuant quantization library.

Below is the WikiText2 perplexity evaluated with 2048 sequence length. The lower is the better.

* SmoothQuant is evaluated with per-tensor static KV cache quantization.

Reference

If you find QServe useful or relevant to your research and work, please kindly cite our paper:

@article{lin2024qserve,
  title={QServe: W4A8KV4 Quantization and System Co-design for Efficient LLM Serving},
  author={Lin*, Yujun and Tang*, Haotian and Yang*, Shang and Zhang, Zhekai and Xiao, Guangxuan and Gan, Chuang and Han, Song},
  journal={arXiv preprint arXiv:2405.04532},
  year={2024}
}

Team

The QServe serving library is maintained by the following research team:

• Haotian Tang, system lead, MIT EECS;
• Shang Yang, system lead, MIT EECS;
• Yujun Lin, algorithm lead, MIT EECS;
• Zhekai Zhang, system evaluation, MIT EECS;
• Guangxuan Xiao, algorithm evaluation, MIT EECS;
• Chuang Gan, advisor, UMass Amherst and MIT-IBM Watson AI Lab;
• Song Han, advisor, MIT EECS and NVIDIA.

Related Projects

The following projects are highly related to QServe. Our group has developed full-stack application-algorithm-system-hardware support for efficient large models, receiving 9k+ GitHub stars and over 1M Huggingface community downloads.

You are also welcome to check out MIT HAN LAB for other exciting projects on Efficient Generative AI!

• [Algorithm] LMQuant: Large Foundation Models Quantization

• [Algorithm] AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration

• [System] TinyChat: Efficient and Lightweight Chatbot with AWQ

• [Application] VILA: On Pretraining of Visual-Language Models

• [Algorithm] SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models

• [Algorithm] StreamingLLM: Efficient Streaming Language Models with Attention Sinks

• [Hardware] SpAtten: Efficient Sparse Attention Architecture with Cascade Token and Head Pruning

Acknowledgement

We thank Julien Demouth, Jun Yang, and Dongxu Yang from NVIDIA for the helpful discussions. QServe is also inspired by many open-source libraries, including (but not limited to) TensorRT-LLM, vLLM, vLLM-SmoothQuant, FlashAttention-2, LMDeploy, TorchSparse++, GPTQ, QuaRot and Atom.