|
| 1 | +(quantized-kvcache)= |
| 2 | + |
| 3 | +# Quantized KV Cache |
| 4 | + |
| 5 | +## FP8 KV Cache |
| 6 | + |
| 7 | +Quantizing the KV cache to FP8 reduces its memory footprint. This increases the number of tokens that can be stored in the cache, improving throughput. |
| 8 | + |
| 9 | +### FP8 Formats |
| 10 | + |
| 11 | +[OCP (Open Compute Project)](https://www.opencompute.org) specifies two common 8-bit floating point data formats: |
| 12 | + |
| 13 | +- E5M2 (5 exponent bits and 2 mantissa bits) |
| 14 | +- E4M3FN (4 exponent bits and 3 mantissa bits, often shortened as E4M3) |
| 15 | + |
| 16 | +The E4M3 format offers higher precision compared to E5M2. However, due to its small dynamic range (±240.0), E4M3 typically requires a higher-precision (FP32) scaling factor alongside each quantized tensor. |
| 17 | + |
| 18 | +### Current Limitations |
| 19 | + |
| 20 | +For now, only per-tensor (scalar) scaling factors are supported. Development is ongoing to support scaling factors of a finer granularity (e.g. per-channel). |
| 21 | + |
| 22 | +### Performance Impact |
| 23 | + |
| 24 | +The current FP8 KV cache implementation primarily benefits throughput by allowing approximately double the amount of space for KV cache allocation. This enables either: |
| 25 | + |
| 26 | +- Processing longer context lengths for individual requests, or |
| 27 | +- Handling more concurrent request batches |
| 28 | + |
| 29 | +However, there are currently no latency improvements as the implementation does not yet include fused dequantization and attention operations. Future releases will support quantized attention with hardware acceleration, which should provide additional performance benefits. While the most recent silicon offerings (e.g. AMD MI300, NVIDIA Hopper or later) support native hardware conversion between FP8 and other formats (fp32, fp16, bf16), this benefit is not yet fully realized. |
| 30 | + |
| 31 | +Studies have shown that FP8 E4M3 quantization typically only minimally degrades inference accuracy, making it a practical choice for throughput optimization. |
| 32 | + |
| 33 | +## Usage Example |
| 34 | + |
| 35 | +Here is an example of how to enable FP8 quantization: |
| 36 | + |
| 37 | +```python |
| 38 | +from vllm import LLM, SamplingParams |
| 39 | + |
| 40 | +sampling_params = SamplingParams(temperature=0.7, top_p=0.8) |
| 41 | +llm = LLM(model="meta-llama/Llama-3.1-8B-Instruct", kv_cache_dtype="fp8") |
| 42 | +prompt = "London is the capital of" |
| 43 | +out = llm.generate(prompt, sampling_params)[0].outputs[0].text |
| 44 | +print(out) |
| 45 | + |
| 46 | +# output w/ scaling factors: England, the United Kingdom, and one of the world's leading financial, |
| 47 | +# output w/o scaling factors: England, located in the southeastern part of the country. It is known |
| 48 | +``` |
| 49 | + |
| 50 | +The `kv_cache_dtype` argument specifies the data type for KV cache storage: |
| 51 | +- `"auto"`: Uses the model's default "unquantized" data type |
| 52 | +- `"fp8"` or `"fp8_e4m3"`: Supported on CUDA 11.8+ and ROCm (AMD GPU) |
| 53 | +- `"fp8_e5m2"`: Supported on CUDA 11.8+ |
| 54 | + |
| 55 | +## Calibrated Scales for Better Accuracy |
| 56 | + |
| 57 | +For optimal model quality when using FP8 KV Cache, we recommend using calibrated scales tuned to representative inference data. [LLM Compressor](https://github.com/vllm-project/llm-compressor/) is the recommended tool for this process. |
| 58 | + |
| 59 | +### Installation |
| 60 | + |
| 61 | +First, install the required dependencies: |
| 62 | + |
| 63 | +```console |
| 64 | +pip install llmcompressor |
| 65 | +``` |
| 66 | + |
| 67 | +### Example Usage |
| 68 | + |
| 69 | +Here's a complete example using `meta-llama/Llama-3.1-8B-Instruct` (most models can use this same pattern): |
| 70 | + |
| 71 | +```python |
| 72 | +from datasets import load_dataset |
| 73 | +from transformers import AutoModelForCausalLM, AutoTokenizer |
| 74 | +from llmcompressor.transformers import oneshot |
| 75 | + |
| 76 | +# Select model and load it |
| 77 | +MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct" |
| 78 | +model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto", torch_dtype="auto") |
| 79 | +tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) |
| 80 | + |
| 81 | +# Select calibration dataset |
| 82 | +DATASET_ID = "HuggingFaceH4/ultrachat_200k" |
| 83 | +DATASET_SPLIT = "train_sft" |
| 84 | + |
| 85 | +# Configure calibration parameters |
| 86 | +NUM_CALIBRATION_SAMPLES = 512 # 512 samples is a good starting point |
| 87 | +MAX_SEQUENCE_LENGTH = 2048 |
| 88 | + |
| 89 | +# Load and preprocess dataset |
| 90 | +ds = load_dataset(DATASET_ID, split=DATASET_SPLIT) |
| 91 | +ds = ds.shuffle(seed=42).select(range(NUM_CALIBRATION_SAMPLES)) |
| 92 | + |
| 93 | +def process_and_tokenize(example): |
| 94 | + text = tokenizer.apply_chat_template(example["messages"], tokenize=False) |
| 95 | + return tokenizer( |
| 96 | + text, |
| 97 | + padding=False, |
| 98 | + max_length=MAX_SEQUENCE_LENGTH, |
| 99 | + truncation=True, |
| 100 | + add_special_tokens=False, |
| 101 | + ) |
| 102 | + |
| 103 | +ds = ds.map(process_and_tokenize, remove_columns=ds.column_names) |
| 104 | + |
| 105 | +# Configure quantization settings |
| 106 | +recipe = """ |
| 107 | +quant_stage: |
| 108 | + quant_modifiers: |
| 109 | + QuantizationModifier: |
| 110 | + kv_cache_scheme: |
| 111 | + num_bits: 8 |
| 112 | + type: float |
| 113 | + strategy: tensor |
| 114 | + dynamic: false |
| 115 | + symmetric: true |
| 116 | +""" |
| 117 | + |
| 118 | +# Apply quantization |
| 119 | +oneshot( |
| 120 | + model=model, |
| 121 | + dataset=ds, |
| 122 | + recipe=recipe, |
| 123 | + max_seq_length=MAX_SEQUENCE_LENGTH, |
| 124 | + num_calibration_samples=NUM_CALIBRATION_SAMPLES, |
| 125 | +) |
| 126 | + |
| 127 | +# Save quantized model |
| 128 | +SAVE_DIR = MODEL_ID.split("/")[1] + "-FP8-KV" |
| 129 | +model.save_pretrained(SAVE_DIR, save_compressed=True) |
| 130 | +tokenizer.save_pretrained(SAVE_DIR) |
| 131 | +``` |
| 132 | + |
| 133 | +The above script will create a folder in your current directory containing your quantized model (e.g., `Llama-3.1-8B-Instruct-FP8-KV`) with calibrated scales. |
| 134 | + |
| 135 | +When running the model you must specify `kv_cache_dtype="fp8"` in order to enable the kv cache quantization and use the scales. |
| 136 | + |
| 137 | +```python |
| 138 | +from vllm import LLM, SamplingParams |
| 139 | + |
| 140 | +sampling_params = SamplingParams(temperature=0.7, top_p=0.8) |
| 141 | +llm = LLM(model="Llama-3.1-8B-Instruct-FP8-KV", kv_cache_dtype="fp8") |
| 142 | +prompt = "London is the capital of" |
| 143 | +out = llm.generate(prompt, sampling_params)[0].outputs[0].text |
| 144 | +print(out) |
| 145 | +``` |
0 commit comments