@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 183 lines 6.1 KB
Rendered Raw
1 # Quantized KV Cache
2
3 ## FP8 KV Cache Overview
4
5 Efficient memory usage is crucial for working with large language models. Quantizing the KV (Key-Value) cache to FP8 format can significantly reduce its memory footprint. This optimization enables you to store more tokens in memory, leading to improved throughput and support for longer context windows.
6
7 > **Note:** When using the Flash Attention 3 backend with FP8 KV cache, attention operations are also performed in the quantized (FP8) domain. In this configuration, queries are quantized to FP8 in addition to keys and values.
8
9 ### Supported FP8 KV-Cache Quantization Schemes
10
11 vLLM supports two main quantization strategies for the FP8 KV-cache:
12
13 - **Per-tensor quantization:**
14 A single scale is applied for each Q, K, and V tensor individually. (`q/k/v_scale = [1]`)
15 - **Per-attention-head quantization:**
16 Each scale corresponds to an attention head: `q_scale = [num_heads]`, `k/v_scale = [num_kv_heads]`.
17
18 > **Note:**
19 > Per-attention-head quantization is currently available **only with the Flash Attention backend** and requires the calibration pathway provided by **llm-compressor**.
20
21 ### Scale Calibration Approaches
22
23 You can configure how the quantization scales are computed in vLLM using three different approaches:
24
25 1. **No calibration (default scales):**
26 All quantization scales are set to `1.0`.
27 _Configure with:_
28 ```python
29 kv_cache_dtype="fp8"
30 ```
31
32 2. **[Recommended] Calibration with a dataset (via llm-compressor):**
33 Scales are estimated using a curated calibration dataset for maximum accuracy.
34 This requires the [llm-compressor](https://github.com/vllm-project/llm-compressor) library.
35 _See example below!_
36
37 #### Additional `kv_cache_dtype` Options
38
39 - `kv_cache_dtype="auto"`: Use the model's default data type
40 - `kv_cache_dtype="fp8_e4m3"`: Supported on CUDA 11.8+ and ROCm (AMD GPUs)
41 - `kv_cache_dtype="fp8_e5m2"`: Supported on CUDA 11.8+
42
43 ### Skipping Specific Layers from KV-Cache Quantization
44
45 Some attention layer types (e.g. sliding-window) are more sensitive to KV-cache quantization. The `--kv-cache-dtype-skip-layers` flag leaves the specified layers at the model's native dtype while keeping the rest of the layers under the chosen quantized dtype. The flag accepts either layer indices or layer-type names:
46
47 ```bash
48 # Skip every sliding-window attention layer.
49 vllm serve <model> \
50 --kv-cache-dtype fp8 \
51 --kv-cache-dtype-skip-layers sliding_window
52
53 # Skip specific layer indices.
54 vllm serve <model> \
55 --kv-cache-dtype fp8 \
56 --kv-cache-dtype-skip-layers 0 1 23
57 ```
58
59 Programmatic usage:
60
61 ```python
62 llm = LLM(
63 model="meta-llama/Llama-3.1-8B-Instruct",
64 kv_cache_dtype="fp8",
65 kv_cache_dtype_skip_layers=["sliding_window"],
66 )
67 ```
68
69 ---
70
71 ## Examples
72
73 ### 1. No Calibration (`kv_cache_dtype="fp8"`)
74
75 All quantization scales are set to 1.0.
76
77 ```python
78 from vllm import LLM, SamplingParams
79
80 sampling_params = SamplingParams(temperature=0.7, top_p=0.8)
81 llm = LLM(
82 model="meta-llama/Llama-2-7b-chat-hf",
83 kv_cache_dtype="fp8",
84 )
85 prompt = "London is the capital of"
86 out = llm.generate(prompt, sampling_params)[0].outputs[0].text
87 print(out)
88 ```
89
90 ---
91
92 ### 2. **[Recommended] Calibration Using a Dataset (with `llm-compressor`)**
93
94 For the highest-quality quantization, we recommend calibrating against a dataset using `llm-compressor`. This enables advanced strategies such as per-attention-head quantization.
95
96 #### Install the required package
97
98 ```bash
99 pip install llmcompressor
100 ```
101
102 #### Example: Quantize Llama Attention & KV Cache to FP8
103
104 ```python
105 """
106 Quantize Llama attention + KV cache to FP8 (choose either 'tensor' or 'attn_head' strategy)
107 using llm-compressor one-shot calibration.
108 """
109
110 from datasets import load_dataset
111 from transformers import AutoModelForCausalLM, AutoTokenizer
112
113 from llmcompressor import oneshot
114 from llmcompressor.modifiers.quantization import QuantizationModifier
115 from compressed_tensors.quantization import QuantizationScheme, QuantizationArgs
116
117 # -----------------------------
118 # Config
119 # -----------------------------
120 MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"
121 DATASET_ID = "HuggingFaceH4/ultrachat_200k"
122 DATASET_SPLIT = "train_sft"
123 STRATEGY = "tensor" # or "attn_head"
124 NUM_CALIB_SAMPLES = 512 # Good starting value
125 MAX_SEQ_LEN = 2048
126
127 # -----------------------------
128 # Helpers
129 # -----------------------------
130 def process_and_tokenize(example, tokenizer: AutoTokenizer):
131 """Convert chat messages to tokens."""
132 text = tokenizer.apply_chat_template(example["messages"], tokenize=False)
133 return tokenizer(
134 text,
135 padding=False,
136 max_length=MAX_SEQ_LEN,
137 truncation=True,
138 add_special_tokens=False,
139 )
140
141 def build_recipe(strategy: str) -> QuantizationModifier:
142 fp8_args = QuantizationArgs(num_bits=8, type="float", strategy=strategy)
143 return QuantizationModifier(
144 config_groups={
145 "attention": QuantizationScheme(
146 targets=["LlamaAttention"], # Quantize queries: q_scale
147 input_activations=fp8_args,
148 )
149 },
150 kv_cache_scheme=fp8_args, # Quantize KV cache: k/v_scale
151 )
152
153 # -----------------------------
154 # Main
155 # -----------------------------
156 def main():
157 model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto")
158 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
159 ds = load_dataset(DATASET_ID, split=f"{DATASET_SPLIT}[:{NUM_CALIB_SAMPLES}]")
160 ds = ds.shuffle(seed=42)
161 ds = ds.map(
162 lambda ex: process_and_tokenize(ex, tokenizer),
163 remove_columns=ds.column_names,
164 )
165
166 recipe = build_recipe(STRATEGY)
167 oneshot(
168 model=model,
169 dataset=ds,
170 recipe=recipe,
171 max_seq_length=MAX_SEQ_LEN,
172 num_calibration_samples=NUM_CALIB_SAMPLES,
173 )
174
175 save_dir = f"{MODEL_ID.rstrip('/').split('/')[-1]}-kvattn-fp8-{STRATEGY}"
176 model.save_pretrained(save_dir, save_compressed=True)
177 tokenizer.save_pretrained(save_dir)
178
179 if __name__ == "__main__":
180 main()
181 ```
182
183 For more detailed and up-to-date examples, see the [`llm-compressor` official examples](https://github.com/vllm-project/llm-compressor/tree/main/examples/quantization_kv_cache).