@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 44f08af3a7
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 9, 18:15 UTC

Add llmcompressor fp8 kv-cache quant (per-tensor and per-attn_head) (#30141)

Signed-off-by: Eldar Kurtic <8884008+eldarkurtic@users.noreply.github.com> Signed-off-by: eldarkurtic <8884008+eldarkurtic@users.noreply.github.com>

Eldar Kurtić committed Jan 22, 2026 at 21:29 UTC 44f08af3a75edcd64e525ac619ebb3b52b99a9c5
18 files changed +538 -243
csrc/cache_kernels.cu
+33 -12
@@ -202,7 +202,8 @@ __global__ void reshape_and_cache_flash_kernel(
202 const int64_t block_stride, const int64_t page_stride,
203 const int64_t head_stride, const int64_t key_stride,
204 const int64_t value_stride, const int num_heads, const int head_size,
205 - const int block_size, const float* k_scale, const float* v_scale) {
205 + const int block_size, const float* k_scale, const float* v_scale,
206 + const int kv_scale_stride) {
207 const int64_t token_idx = blockIdx.x;
208 const int64_t slot_idx = slot_mapping[token_idx];
209 // NOTE: slot_idx can be -1 if the token is padded
@@ -226,21 +227,23 @@ __global__ void reshape_and_cache_flash_kernel(
227 // this is true for the NHD layout where `head_stride == head_size`
228 const bool is_contiguous_heads = (head_stride == head_size);
229
229 - float k_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *k_scale;
230 - float v_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *v_scale;
230 constexpr int VEC_SIZE = (sizeof(scalar_t) == 2) ? 8 : 4;
232 - CopyWithScaleOp<cache_t, scalar_t, kv_dt> k_op{k_scale_val};
233 - CopyWithScaleOp<cache_t, scalar_t, kv_dt> v_op{v_scale_val};
234 - if (is_contiguous_heads) {
235 - // NHD layout
231 +
232 + if (is_contiguous_heads && kv_scale_stride == 0) {
233 + // NHD layout and k/v_scales are [1] (i.e. single scale for all heads)
234 // kv cache: [num_blocks, block_size, num_heads, head_size]
235 + float k_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *k_scale;
236 + float v_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *v_scale;
237 +
238 + CopyWithScaleOp<cache_t, scalar_t, kv_dt> k_op{k_scale_val};
239 + CopyWithScaleOp<cache_t, scalar_t, kv_dt> v_op{v_scale_val};
240 +
241 vectorize_with_alignment<VEC_SIZE>(key_src, key_dst, n_elems, threadIdx.x,
242 blockDim.x, k_op);
239 -
243 vectorize_with_alignment<VEC_SIZE>(value_src, value_dst, n_elems,
244 threadIdx.x, blockDim.x, v_op);
242 -
245 } else {
246 + // HND layout OR k/v_scales are [num_heads] (i.e. per-attn-head)
247 // HND layout: heads are strided, but each head_size segment is contiguous
248 // kv cache: [num_blocks, num_heads, block_size, head_size]
249 const int lane = threadIdx.x & 31; // 0..31 within warp
@@ -256,6 +259,16 @@ __global__ void reshape_and_cache_flash_kernel(
259 cache_t* __restrict__ v_dst_h =
260 value_dst + static_cast<int64_t>(head) * head_stride;
261
262 + float k_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto)
263 + ? 0.f
264 + : k_scale[head * kv_scale_stride];
265 + float v_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto)
266 + ? 0.f
267 + : v_scale[head * kv_scale_stride];
268 +
269 + CopyWithScaleOp<cache_t, scalar_t, kv_dt> k_op{k_scale_val};
270 + CopyWithScaleOp<cache_t, scalar_t, kv_dt> v_op{v_scale_val};
271 +
272 // within each head, let the 32 threads of the warp perform the vector
273 // copy
274 vectorize_with_alignment<VEC_SIZE>(k_src_h, k_dst_h, head_size, lane, 32,
@@ -605,7 +618,8 @@ void reshape_and_cache(
618 slot_mapping.data_ptr<int64_t>(), block_stride, page_stride, \
619 head_stride, key_stride, value_stride, num_heads, head_size, \
620 block_size, reinterpret_cast<const float*>(k_scale.data_ptr()), \
608 - reinterpret_cast<const float*>(v_scale.data_ptr()));
621 + reinterpret_cast<const float*>(v_scale.data_ptr()), \
622 + kv_scale_stride);
623
624 void reshape_and_cache_flash(
625 torch::Tensor& key, // [num_tokens, num_heads, head_size]
@@ -614,8 +628,9 @@ void reshape_and_cache_flash(
628 torch::Tensor&
629 value_cache, // [num_blocks, block_size, num_heads, head_size]
630 torch::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens]
617 - const std::string& kv_cache_dtype, torch::Tensor& k_scale,
618 - torch::Tensor& v_scale) {
631 + const std::string& kv_cache_dtype,
632 + torch::Tensor& k_scale, // [1] or [num_heads]
633 + torch::Tensor& v_scale) { // [1] or [num_heads]
634 // NOTE(woosuk): In vLLM V1, key.size(0) can be different from
635 // slot_mapping.size(0) because of padding for CUDA graphs.
636 // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because
@@ -638,6 +653,12 @@ void reshape_and_cache_flash(
653 int64_t head_stride = key_cache.stride(2);
654 TORCH_CHECK(key_cache.stride(0) == value_cache.stride(0));
655
656 + TORCH_CHECK(k_scale.sizes() == v_scale.sizes(),
657 + "k_scale and v_scale must have the same shape");
658 + TORCH_CHECK(k_scale.numel() == 1 || k_scale.numel() == num_heads,
659 + "k_scale and v_scale must be of shape [1] or [num_heads]");
660 + int kv_scale_stride = (k_scale.numel() > 1) ? 1 : 0;
661 +
662 dim3 grid(num_tokens);
663 dim3 block(std::min(num_heads * head_size, 512));
664 const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
docs/features/quantization/quantized_kvcache.md
+139 -114
@@ -1,162 +1,187 @@
1 # Quantized KV Cache
2
3 -## FP8 KV Cache
3 +## FP8 KV Cache Overview
4
5 -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.
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 -### FP8 Formats
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 -[OCP (Open Compute Project)](https://www.opencompute.org) specifies two common 8-bit floating point data formats:
9 +### Supported FP8 KV-Cache Quantization Schemes
10
11 -- E5M2 (5 exponent bits and 2 mantissa bits)
12 -- E4M3FN (4 exponent bits and 3 mantissa bits, often shortened as E4M3)
11 +vLLM supports two main quantization strategies for the FP8 KV-cache:
12
14 -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.
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
16 -### Current Limitations
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
18 -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 +### Scale Calibration Approaches
22
20 -### How FP8 KV Cache Works
23 +You can configure how the quantization scales are computed in vLLM using three different approaches:
24
22 -The FP8 KV cache implementation follows this workflow:
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 + calculate_kv_scales=False
31 + ```
32
24 -1. **Storage**: Key and Value tensors are quantized to FP8 format using scaling factors before being stored in the KV cache
25 -2. **Retrieval**: When needed for attention computation, cached KV tensors are dequantized back to higher precision (FP16/BF16)
26 -3. **Attention**: The attention-value multiplication (softmax output × V) is performed using the dequantized higher-precision V tensor
33 +2. **Random token calibration (on-the-fly):**
34 + Scales are automatically estimated from a single batch of random tokens during warmup and then fixed.
35 + _Configure with:_
36 + ```python
37 + kv_cache_dtype="fp8"
38 + calculate_kv_scales=True
39 + ```
40
28 -This means the final attention computation operates on dequantized values, not FP8 tensors. The quantization reduces memory usage during storage but maintains computation accuracy by using higher precision during the actual attention operations.
41 +3. **[Recommended] Calibration with a dataset (via llm-compressor):**
42 + Scales are estimated using a curated calibration dataset for maximum accuracy.
43 + This requires the [llm-compressor](https://github.com/vllm-project/llm-compressor) library.
44 + _See example below!_
45
30 -### Performance Impact
46 +#### Additional `kv_cache_dtype` Options
47
32 -The current FP8 KV cache implementation primarily benefits throughput by allowing approximately double the amount of space for KV cache allocation. This enables either:
48 +- `kv_cache_dtype="auto"`: Use the model's default data type
49 +- `kv_cache_dtype="fp8_e4m3"`: Supported on CUDA 11.8+ and ROCm (AMD GPUs)
50 +- `kv_cache_dtype="fp8_e5m2"`: Supported on CUDA 11.8+
51
34 -- Processing longer context lengths for individual requests, or
35 -- Handling more concurrent request batches
52 +---
53
37 -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.
54 +## Examples
55
39 -Studies have shown that FP8 E4M3 quantization typically only minimally degrades inference accuracy, making it a practical choice for throughput optimization.
56 +### 1. No Calibration (`kv_cache_dtype="fp8"`, `calculate_kv_scales=False`)
57
41 -## Usage Example
58 +All quantization scales are set to 1.0.
59
43 -Here is an example of how to enable FP8 quantization:
60 +```python
61 +from vllm import LLM, SamplingParams
62
45 -??? code
63 +sampling_params = SamplingParams(temperature=0.7, top_p=0.8)
64 +llm = LLM(
65 + model="meta-llama/Llama-2-7b-chat-hf",
66 + kv_cache_dtype="fp8",
67 + calculate_kv_scales=False,
68 +)
69 +prompt = "London is the capital of"
70 +out = llm.generate(prompt, sampling_params)[0].outputs[0].text
71 +print(out)
72 +```
73
47 - ```python
48 - # To calculate kv cache scales on the fly enable the calculate_kv_scales
49 - # parameter
74 +---
75
51 - from vllm import LLM, SamplingParams
76 +### 2. Random Token Calibration (`kv_cache_dtype="fp8"`, `calculate_kv_scales=True`)
77
53 - sampling_params = SamplingParams(temperature=0.7, top_p=0.8)
54 - llm = LLM(
55 - model="meta-llama/Llama-2-7b-chat-hf",
56 - kv_cache_dtype="fp8",
57 - calculate_kv_scales=True,
58 - )
59 - prompt = "London is the capital of"
60 - out = llm.generate(prompt, sampling_params)[0].outputs[0].text
61 - print(out)
62 - ```
78 +Scales are automatically estimated from a single batch of tokens during warmup.
79
64 -The `kv_cache_dtype` argument specifies the data type for KV cache storage:
80 +```python
81 +from vllm import LLM, SamplingParams
82
66 -- `"auto"`: Uses the model's default "unquantized" data type
67 -- `"fp8"` or `"fp8_e4m3"`: Supported on CUDA 11.8+ and ROCm (AMD GPU)
68 -- `"fp8_e5m2"`: Supported on CUDA 11.8+
83 +sampling_params = SamplingParams(temperature=0.7, top_p=0.8)
84 +llm = LLM(
85 + model="meta-llama/Llama-2-7b-chat-hf",
86 + kv_cache_dtype="fp8",
87 + calculate_kv_scales=True,
88 +)
89 +prompt = "London is the capital of"
90 +out = llm.generate(prompt, sampling_params)[0].outputs[0].text
91 +print(out)
92 +```
93
70 -## Calibrated Scales for Better Accuracy
94 +---
95
72 -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.
96 +### 3. **[Recommended] Calibration Using a Dataset (with `llm-compressor`)**
97
74 -### Installation
98 +For the highest-quality quantization, we recommend calibrating against a dataset using `llm-compressor`. This enables advanced strategies such as per-attention-head quantization.
99
76 -First, install the required dependencies:
100 +#### Install the required package
101
102 ```bash
103 pip install llmcompressor
104 ```
105
82 -### Example Usage
106 +#### Example: Quantize Llama Attention & KV Cache to FP8
107
84 -Here's a complete example using `meta-llama/Llama-3.1-8B-Instruct` (most models can use this same pattern):
85 -
86 -??? code
108 +```python
109 +"""
110 +Quantize Llama attention + KV cache to FP8 (choose either 'tensor' or 'attn_head' strategy)
111 +using llm-compressor one-shot calibration.
112 +"""
113 +
114 +from datasets import load_dataset
115 +from transformers import AutoModelForCausalLM, AutoTokenizer
116 +
117 +from llmcompressor import oneshot
118 +from llmcompressor.modifiers.quantization import QuantizationModifier
119 +from compressed_tensors.quantization import QuantizationScheme, QuantizationArgs
120 +
121 +# -----------------------------
122 +# Config
123 +# -----------------------------
124 +MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"
125 +DATASET_ID = "HuggingFaceH4/ultrachat_200k"
126 +DATASET_SPLIT = "train_sft"
127 +STRATEGY = "tensor" # or "attn_head"
128 +NUM_CALIB_SAMPLES = 512 # Good starting value
129 +MAX_SEQ_LEN = 2048
130 +
131 +# -----------------------------
132 +# Helpers
133 +# -----------------------------
134 +def process_and_tokenize(example, tokenizer: AutoTokenizer):
135 + """Convert chat messages to tokens."""
136 + text = tokenizer.apply_chat_template(example["messages"], tokenize=False)
137 + return tokenizer(
138 + text,
139 + padding=False,
140 + max_length=MAX_SEQ_LEN,
141 + truncation=True,
142 + add_special_tokens=False,
143 + )
144
88 - ```python
89 - from datasets import load_dataset
90 - from transformers import AutoModelForCausalLM, AutoTokenizer
91 - from llmcompressor import oneshot
145 +def build_recipe(strategy: str) -> QuantizationModifier:
146 + fp8_args = QuantizationArgs(num_bits=8, type="float", strategy=strategy)
147 + return QuantizationModifier(
148 + config_groups={
149 + "attention": QuantizationScheme(
150 + targets=["LlamaAttention"], # Quantize queries: q_scale
151 + input_activations=fp8_args,
152 + )
153 + },
154 + kv_cache_scheme=fp8_args, # Quantize KV cache: k/v_scale
155 + )
156
93 - # Select model and load it
94 - MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"
95 - model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto", dtype="auto")
157 +# -----------------------------
158 +# Main
159 +# -----------------------------
160 +def main():
161 + model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto")
162 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
163 + ds = load_dataset(DATASET_ID, split=f"{DATASET_SPLIT}[:{NUM_CALIB_SAMPLES}]")
164 + ds = ds.shuffle(seed=42)
165 + ds = ds.map(
166 + lambda ex: process_and_tokenize(ex, tokenizer),
167 + remove_columns=ds.column_names,
168 + )
169
98 - # Select calibration dataset
99 - DATASET_ID = "HuggingFaceH4/ultrachat_200k"
100 - DATASET_SPLIT = "train_sft"
101 -
102 - # Configure calibration parameters
103 - NUM_CALIBRATION_SAMPLES = 512 # 512 samples is a good starting point
104 - MAX_SEQUENCE_LENGTH = 2048
105 -
106 - # Load and preprocess dataset
107 - ds = load_dataset(DATASET_ID, split=DATASET_SPLIT)
108 - ds = ds.shuffle(seed=42).select(range(NUM_CALIBRATION_SAMPLES))
109 -
110 - def process_and_tokenize(example):
111 - text = tokenizer.apply_chat_template(example["messages"], tokenize=False)
112 - return tokenizer(
113 - text,
114 - padding=False,
115 - max_length=MAX_SEQUENCE_LENGTH,
116 - truncation=True,
117 - add_special_tokens=False,
118 - )
119 -
120 - ds = ds.map(process_and_tokenize, remove_columns=ds.column_names)
121 -
122 - # Configure quantization settings
123 - recipe = """
124 - quant_stage:
125 - quant_modifiers:
126 - QuantizationModifier:
127 - kv_cache_scheme:
128 - num_bits: 8
129 - type: float
130 - strategy: tensor
131 - dynamic: false
132 - symmetric: true
133 - """
134 -
135 - # Apply quantization
170 + recipe = build_recipe(STRATEGY)
171 oneshot(
172 model=model,
173 dataset=ds,
174 recipe=recipe,
140 - max_seq_length=MAX_SEQUENCE_LENGTH,
141 - num_calibration_samples=NUM_CALIBRATION_SAMPLES,
175 + max_seq_length=MAX_SEQ_LEN,
176 + num_calibration_samples=NUM_CALIB_SAMPLES,
177 )
178
144 - # Save quantized model: Llama-3.1-8B-Instruct-FP8-KV
145 - SAVE_DIR = MODEL_ID.split("/")[1] + "-FP8-KV"
146 - model.save_pretrained(SAVE_DIR, save_compressed=True)
147 - tokenizer.save_pretrained(SAVE_DIR)
148 - ```
149 -
150 -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.
179 + save_dir = f"{MODEL_ID.rstrip('/').split('/')[-1]}-kvattn-fp8-{STRATEGY}"
180 + model.save_pretrained(save_dir, save_compressed=True)
181 + tokenizer.save_pretrained(save_dir)
182
152 -When running the model you must specify `kv_cache_dtype="fp8"` in order to enable the kv cache quantization and use the scales.
153 -
154 -```python
155 -from vllm import LLM, SamplingParams
156 -
157 -sampling_params = SamplingParams(temperature=0.7, top_p=0.8)
158 -llm = LLM(model="Llama-3.1-8B-Instruct-FP8-KV", kv_cache_dtype="fp8")
159 -prompt = "London is the capital of"
160 -out = llm.generate(prompt, sampling_params)[0].outputs[0].text
161 -print(out)
183 +if __name__ == "__main__":
184 + main()
185 ```
186 +
187 +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).
tests/kernels/attention/test_cache.py
+34 -13
@@ -8,6 +8,7 @@ import torch
8
9 from tests.kernels.utils import DEFAULT_OPCHECK_TEST_UTILS, opcheck
10 from vllm import _custom_ops as ops
11 +from vllm.model_executor.layers.quantization.utils.quant_utils import scaled_dequantize
12 from vllm.platforms import current_platform
13 from vllm.utils.torch_utils import set_random_seed
14
@@ -19,6 +20,7 @@ NUM_HEADS = [8] # Arbitrary values for testing
20 HEAD_SIZES = [64, 80, 256]
21 BLOCK_SIZES = [8, 16, 32]
22 CACHE_LAYOUTS = ["NHD", "HND"]
23 +KV_SCALE_TYPES = ["tensor", "attn_head"]
24
25 # Parameters for MLA tests.
26 KV_LORA_RANKS = [512]
@@ -170,6 +172,7 @@ def test_reshape_and_cache(
172 @pytest.mark.parametrize("device", CUDA_DEVICES)
173 @pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
174 @pytest.mark.parametrize("kv_cache_layout", CACHE_LAYOUTS)
175 +@pytest.mark.parametrize("kv_scale_type", KV_SCALE_TYPES)
176 @pytest.mark.parametrize("implementation", RESHAPE_FLASH_IMPLEMENTATIONS)
177 @torch.inference_mode()
178 def test_reshape_and_cache_flash(
@@ -184,6 +187,7 @@ def test_reshape_and_cache_flash(
187 device: str,
188 kv_cache_dtype: str,
189 kv_cache_layout: str,
190 + kv_scale_type: str,
191 implementation: str,
192 ) -> None:
193 set_random_seed(seed)
@@ -193,6 +197,9 @@ def test_reshape_and_cache_flash(
197 if implementation == "triton" and kv_cache_layout == "HND":
198 pytest.skip("Triton implementation only supports NHD layout.")
199
200 + if kv_scale_type == "attn_head" and implementation != "cuda":
201 + pytest.skip("Only CUDA implementation supports attn_head scaling.")
202 +
203 # fp8 conversion requires continugous memory buffer. Reduce the number of
204 # blocks and tokens to consume less memory.
205 num_tokens = num_tokens // 2
@@ -220,8 +227,12 @@ def test_reshape_and_cache_flash(
227 del key_caches
228 del value_caches
229
223 - k_scale = (key.amax() / 64.0).to(torch.float32)
224 - v_scale = (value.amax() / 64.0).to(torch.float32)
230 + if kv_scale_type == "tensor":
231 + k_scale = (key.amax() / 64.0).to(torch.float32)
232 + v_scale = (value.amax() / 64.0).to(torch.float32)
233 + else: # "attn_head"
234 + k_scale = (key.amax(dim=(0, 2)) / 64.0).to(torch.float32)
235 + v_scale = (value.amax(dim=(0, 2)) / 64.0).to(torch.float32)
236
237 def permute_and_compact(x):
238 y = x if kv_cache_layout == "NHD" else x.permute(0, 2, 1, 3)
@@ -230,15 +241,27 @@ def test_reshape_and_cache_flash(
241 key_cache_compact = permute_and_compact(key_cache)
242 value_cache_compact = permute_and_compact(value_cache)
243
244 + def convert_fp8_local(output, input, scale, kv_dtype):
245 + fp8_input = input.view(torch.float8_e4m3fn)
246 + if scale.numel() == 1: # per-tensor
247 + result = scaled_dequantize(
248 + fp8_input.flatten(0, 2), scale, group_shape=None, out_dtype=output.dtype
249 + ).reshape(*input.shape)
250 + else: # per-head: broadcast scale along the head dimension
251 + # Original code uses dim 2 for NHD, dim 1 for HND
252 + if kv_cache_layout == "NHD":
253 + result = fp8_input.to(output.dtype) * scale.view(1, 1, -1, 1)
254 + else:
255 + result = fp8_input.to(output.dtype) * scale.view(1, -1, 1, 1)
256 + output.copy_(result)
257 +
258 # Clone the KV caches.
259 if kv_cache_dtype == "fp8":
260 cloned_key_cache = torch.empty_like(key_cache_compact, dtype=torch.float16)
236 - ops.convert_fp8(
237 - cloned_key_cache, key_cache_compact, k_scale.item(), kv_cache_dtype
238 - )
261 + convert_fp8_local(cloned_key_cache, key_cache_compact, k_scale, kv_cache_dtype)
262 cloned_value_cache = torch.empty_like(value_cache_compact, dtype=torch.float16)
240 - ops.convert_fp8(
241 - cloned_value_cache, value_cache_compact, v_scale.item(), kv_cache_dtype
263 + convert_fp8_local(
264 + cloned_value_cache, value_cache_compact, v_scale, kv_cache_dtype
265 )
266 else:
267 cloned_key_cache = key_cache_compact.clone()
@@ -289,15 +312,13 @@ def test_reshape_and_cache_flash(
312
313 if kv_cache_dtype == "fp8":
314 result_key_cache = torch.empty_like(key_cache_compact, dtype=torch.float16)
292 - ops.convert_fp8(
293 - result_key_cache, key_cache_compact, k_scale.item(), kv_dtype=kv_cache_dtype
294 - )
315 + convert_fp8_local(result_key_cache, key_cache_compact, k_scale, kv_cache_dtype)
316 result_value_cache = torch.empty_like(value_cache_compact, dtype=torch.float16)
296 - ops.convert_fp8(
317 + convert_fp8_local(
318 result_value_cache,
319 value_cache_compact,
299 - v_scale.item(),
300 - kv_dtype=kv_cache_dtype,
320 + v_scale,
321 + kv_cache_dtype,
322 )
323
324 # Run the reference implementation.
tests/quantization/test_compressed_tensors.py
+21 -3
@@ -32,6 +32,7 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
32 sparse_cutlass_supported,
33 )
34 from vllm.platforms import current_platform
35 +from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
36
37 # AITER only supports per-channel-per-channel INT8 gemm
38 # and per-tensor-per-tensor INT8 GEMM.
@@ -360,9 +361,26 @@ def test_compressed_tensors_fp8(vllm_runner):
361 @pytest.mark.skipif(
362 not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
363 )
363 -def test_compressed_tensors_kv_cache(vllm_runner):
364 - model_path = "nm-testing/TinyLlama-1.1B-compressed-tensors-kv-cache-scheme"
365 - with vllm_runner(model_path, enforce_eager=True, kv_cache_dtype="fp8") as llm:
364 +def test_compressed_tensors_kv_cache_fp8_per_tensor(vllm_runner):
365 + model_path = "nm-testing/TinyLlama-1.1B-Chat-v1.0-kvcache-fp8-tensor"
366 + with vllm_runner(model_path) as llm:
367 + output = llm.generate_greedy("Hello world!", max_tokens=4)
368 + assert output
369 +
370 +
371 +@pytest.mark.skipif(
372 + not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
373 +)
374 +def test_compressed_tensors_kv_cache_fp8_per_attn_head(vllm_runner):
375 + model_path = "nm-testing/TinyLlama-1.1B-Chat-v1.0-kvcache-fp8-attn_head"
376 + try:
377 + fa_version = get_flash_attn_version()
378 + except Exception:
379 + pytest.skip("This test requires FlashAttention backend.")
380 + if fa_version is None or fa_version < 3:
381 + pytest.skip("This test requires FlashAttention version >= 3.")
382 +
383 + with vllm_runner(model_path, attention_config={"backend": "FLASH_ATTN"}) as llm:
384 output = llm.generate_greedy("Hello world!", max_tokens=4)
385 assert output
386
vllm/attention/layer.py
+37 -36
@@ -75,13 +75,16 @@ def set_default_quant_scales(layer: nn.Module, register_buffer: bool = False) ->
75 layer._v_scale_float = 1.0
76 layer._prob_scale_float = 1.0
77
78 + # Initialize q/k/v range constants used by calc_kv_scales
79 + layer.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32)
80 + layer.k_range = torch.tensor(envs.K_SCALE_CONSTANT, dtype=torch.float32)
81 + layer.v_range = torch.tensor(envs.V_SCALE_CONSTANT, dtype=torch.float32)
82 +
83
84 def _init_kv_cache_quant(
85 layer: nn.Module,
86 quant_config: QuantizationConfig | None,
87 prefix: str,
83 - kv_cache_dtype: str,
84 - calculate_kv_scales: bool,
88 ) -> None:
89 """Initializes KV cache scaling factors and quantization method.
90
@@ -94,16 +97,10 @@ def _init_kv_cache_quant(
97 layer: The attention layer instance to initialize.
98 quant_config: Optional quantization configuration.
99 prefix: Layer name prefix for quantization method lookup.
97 - kv_cache_dtype: The KV cache data type string.
98 - calculate_kv_scales: Whether to calculate KV scales dynamically.
100 """
100 - # The default k/v_scale is set to 1.0. This is ignored
101 - # when kv-cache is not fp8, and should be used with
102 - # kv-cache in fp8_e5m2. For kv-cache in fp8_e4m3, we
103 - # expect the pre-quantized k/v_scale to be loaded along
104 - # with the model weights.
105 - layer.kv_cache_dtype = kv_cache_dtype
106 - layer.calculate_kv_scales = calculate_kv_scales
101 + quant_method = (
102 + quant_config.get_quant_method(layer, prefix=prefix) if quant_config else None
103 + )
104
105 # Note [Register q/k/v/prob scales in state dict]
106 # When calling model.to(device), only parameters/buffers in state dict are
@@ -133,7 +130,7 @@ def _init_kv_cache_quant(
130 assert isinstance(quant_method, BaseKVCacheMethod)
131 # TODO (mgoin): kv cache dtype should be specified in the FP8
132 # checkpoint config and become the "auto" behavior
136 - if kv_cache_dtype == "fp8_e5m2":
133 + if layer.kv_cache_dtype == "fp8_e5m2":
134 raise ValueError("fp8_e5m2 kv-cache is not supported with fp8 checkpoints.")
135 # If quantization is enabled, we make "k_scale" and "v_scale"
136 # parameters so that it can be loaded from the model checkpoint.
@@ -197,9 +194,20 @@ class Attention(nn.Module, AttentionLayerBase):
194 kv_cache_dtype = "auto"
195 block_size = 16
196 calculate_kv_scales = False
197 +
198 + # llm-compressor mdls need to set cache_dtype to "fp8" manually.
199 + if getattr(quant_config, "kv_cache_scheme", None) is not None:
200 + kv_cache_dtype = "fp8"
201 + calculate_kv_scales = False
202 + if cache_config is not None:
203 + cache_config.cache_dtype = "fp8"
204 + cache_config.calculate_kv_scales = False
205 +
206 self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
207 kv_cache_dtype, vllm_config.model_config
208 )
209 + self.kv_cache_dtype = kv_cache_dtype
210 + self.calculate_kv_scales = calculate_kv_scales
211 if num_kv_heads is None:
212 num_kv_heads = num_heads
213 assert num_heads % num_kv_heads == 0, (
@@ -208,15 +216,6 @@ class Attention(nn.Module, AttentionLayerBase):
216 self.quant_config = quant_config
217 self.layer_name = prefix
218
211 - # Initialize KV cache quantization attributes
212 - _init_kv_cache_quant(
213 - self,
214 - self.quant_config,
215 - self.layer_name,
216 - kv_cache_dtype,
217 - calculate_kv_scales,
218 - )
219 -
219 self.num_heads = num_heads
220 self.head_size = head_size
221 self.head_size_v = self.head_size if head_size_v is None else head_size_v
@@ -318,18 +317,24 @@ class Attention(nn.Module, AttentionLayerBase):
317 for _ in range(vllm_config.parallel_config.pipeline_parallel_size)
318 ]
319
321 - # Initialize q/k/v range constants.
322 - self.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32)
323 - self.k_range = torch.tensor(envs.K_SCALE_CONSTANT, dtype=torch.float32)
324 - self.v_range = torch.tensor(envs.V_SCALE_CONSTANT, dtype=torch.float32)
320 + # Initialize KV cache quantization attributes
321 + _init_kv_cache_quant(self, quant_config, prefix)
322
323 # for attn backends supporting query quantization
324 self.query_quant = None
328 - if (
329 - self.kv_cache_dtype.startswith("fp8")
330 - and self.impl.supports_quant_query_input
325 + if self.impl.supports_quant_query_input and self.kv_cache_dtype.startswith(
326 + "fp8"
327 ):
332 - self.query_quant = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)
328 + is_per_head = (
329 + hasattr(self, "q_scale") and self.q_scale.numel() == self.num_kv_heads
330 + )
331 + block_size = self.head_size * self.num_heads // self.num_kv_heads
332 + self.query_quant = QuantFP8(
333 + static=True,
334 + group_shape=GroupShape(-1, block_size)
335 + if is_per_head
336 + else GroupShape.PER_TENSOR,
337 + )
338
339 def forward(
340 self,
@@ -524,13 +529,9 @@ class MLAAttention(nn.Module, AttentionLayerBase):
529 self.quant_config = quant_config
530
531 # Initialize KV cache quantization attributes
527 - _init_kv_cache_quant(
528 - self,
529 - self.quant_config,
530 - self.layer_name,
531 - kv_cache_dtype,
532 - calculate_kv_scales,
533 - )
532 + self.kv_cache_dtype = kv_cache_dtype
533 + self.calculate_kv_scales = calculate_kv_scales
534 + _init_kv_cache_quant(self, quant_config, prefix)
535
536 dtype = torch.get_default_dtype()
537 self.attn_backend = get_attn_backend(
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+168 -22
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from contextlib import suppress
5 +from functools import partial
6 from typing import TYPE_CHECKING, Any, Literal, Optional, cast
7
8 import torch
@@ -19,6 +20,10 @@ from compressed_tensors.transform import TransformConfig
20
21 import vllm.envs as envs
22 from vllm.attention.layer import Attention
23 +from vllm.distributed import (
24 + get_tensor_model_parallel_rank,
25 + get_tensor_model_parallel_world_size,
26 +)
27 from vllm.logger import init_logger
28 from vllm.model_executor.layers.fused_moe import FusedMoE
29 from vllm.model_executor.layers.linear import (
@@ -87,6 +92,8 @@ class CompressedTensorsConfig(QuantizationConfig):
92 kv_cache_scheme: dict[str, Any] | None = None,
93 config: dict[str, Any] | None = None,
94 transform_config: dict[str, Any] | None = None,
95 + total_num_heads: int | None = None,
96 + total_num_kv_heads: int | None = None,
97 ):
98 super().__init__()
99 self.ignore = ignore
@@ -97,6 +104,8 @@ class CompressedTensorsConfig(QuantizationConfig):
104 self.sparsity_scheme_map = sparsity_scheme_map
105 self.sparsity_ignore_list = sparsity_ignore_list
106 self.config = config
107 + self.total_num_heads = total_num_heads
108 + self.total_num_kv_heads = total_num_kv_heads
109
110 if transform_config:
111 self.transform_config = TransformConfig.model_validate(transform_config)
@@ -200,13 +209,29 @@ class CompressedTensorsConfig(QuantizationConfig):
209
210 @classmethod
211 def from_config(cls, config: dict[str, Any]) -> "CompressedTensorsConfig":
212 + # We keep only config groups which are not doing Attention quantization
213 + # because Attention quantization on its own is not supported by vLLM.
214 + # It is coupled with KV-cache quantization, and if scales are present in the
215 + # checkpoint, they will be used properly.
216 + grps_without_attn_quant = {}
217 + for k, v in config["config_groups"].items():
218 + # e.g. LlamaAttention, Qwen3Attention, etc.
219 + if len(v["targets"]) == 1 and v["targets"][0].endswith("Attention"):
220 + logger.warning(
221 + "Skipping CompressedTensors config group for %s. Attention quant "
222 + "is coupled with KV-cache quantization in vLLM.",
223 + v["targets"][0],
224 + )
225 + continue
226 + grps_without_attn_quant[k] = v
227 + config["config_groups"] = grps_without_attn_quant
228 +
229 ignore: list[str] = cast(list[str], config.get("ignore", []))
230 quant_format = cast(str, config.get("format"))
231 target_scheme_map = cls._quantization_scheme_map_from_config(config=config)
232 sparsity_scheme_map, sparsity_ignore_list = cls._parse_sparsity_config(
233 config=config
234 )
209 - transform_config = config.get("transform_config")
235
236 return cls(
237 target_scheme_map=target_scheme_map,
@@ -215,7 +240,10 @@ class CompressedTensorsConfig(QuantizationConfig):
240 sparsity_scheme_map=sparsity_scheme_map,
241 sparsity_ignore_list=sparsity_ignore_list,
242 config=config,
218 - transform_config=transform_config,
243 + transform_config=config.get("transform_config"),
244 + kv_cache_scheme=config.get("kv_cache_scheme"),
245 + total_num_heads=config.get("total_num_heads"),
246 + total_num_kv_heads=config.get("total_num_kv_heads"),
247 )
248
249 @classmethod
@@ -791,22 +819,6 @@ class CompressedTensorsConfig(QuantizationConfig):
819
820 return None
821
794 - def get_cache_scale(self, name: str) -> str | None:
795 - """
796 - Check whether the param name matches the format for k/v cache scales
797 - in compressed-tensors. If this is the case, return its equivalent
798 - param name expected by vLLM
799 -
800 - :param name: param name
801 - :return: matching param name for KV cache scale in vLLM
802 - """
803 - if name.endswith(".output_scale") and ".k_proj" in name:
804 - return name.replace(".k_proj.output_scale", ".attn.k_scale")
805 - if name.endswith(".output_scale") and ".v_proj" in name:
806 - return name.replace(".v_proj.output_scale", ".attn.v_scale")
807 - # If no matches, return None
808 - return None
809 -
822 def has_blocked_weights(self) -> bool:
823 for scheme in self.target_scheme_map.values():
824 weight_quant = scheme.get("weights")
@@ -965,12 +977,16 @@ class CompressedTensorsKVCacheMethod(BaseKVCacheMethod):
977 f"received num_bits={num_bits}, type={type_}"
978 )
979
980 + # TODO: delegate validation to compressed-tensors library so that we have a
981 + # single source of truth. Right now this is not possible until the next release
982 + # of compressed-tensors.
983 strategy = kv_cache_scheme.get("strategy")
969 - if strategy != "tensor":
984 + supported_strategies = ("tensor", "attn_head")
985 + if strategy not in supported_strategies:
986 raise NotImplementedError(
971 - "Only support per-tensor scaling factor "
972 - "for compressed-tensors KV cache. "
973 - f"Expected strategy: tensor, found strategy: {strategy}"
987 + "Invalid strategy for compressed-tensors KV cache. "
988 + f"Expected strategies: {supported_strategies}, found strategy:"
989 + f" {strategy}"
990 )
991
992 is_symmetric = kv_cache_scheme.get("symmetric")
@@ -980,3 +996,133 @@ class CompressedTensorsKVCacheMethod(BaseKVCacheMethod):
996 "for compressed-tensors KV cache. "
997 f"However found symmetric: {is_symmetric}"
998 )
999 +
1000 + def create_weights(self, layer: torch.nn.Module):
1001 + """
1002 + Initialize placeholder scales and zero points to enable loading of
1003 + quantized params from compressed-tensors checkpoints.
1004 + """
1005 + strategy = None # for backward compatibility
1006 + if (
1007 + hasattr(self.quant_config, "kv_cache_scheme")
1008 + and self.quant_config.kv_cache_scheme is not None
1009 + ):
1010 + strategy = self.quant_config.kv_cache_scheme["strategy"]
1011 +
1012 + if strategy == "attn_head":
1013 + assert layer.impl.supports_per_head_quant_scales, (
1014 + f"Layer {layer.__class__.__name__} with implementation "
1015 + f"{layer.impl.__class__.__name__} does not support per-head scales."
1016 + )
1017 + n_scales = int(layer.num_kv_heads)
1018 + else:
1019 + n_scales = 1
1020 +
1021 + layer.k_scale = torch.nn.Parameter(
1022 + torch.ones(n_scales, requires_grad=False, dtype=torch.float32)
1023 + )
1024 + layer.v_scale = torch.nn.Parameter(
1025 + torch.ones(n_scales, requires_grad=False, dtype=torch.float32)
1026 + )
1027 + layer.q_scale = torch.nn.Parameter(
1028 + torch.ones(n_scales, requires_grad=False, dtype=torch.float32)
1029 + )
1030 +
1031 + # Zero points are not used in vLLM as currently only symmetric quantization is
1032 + # supported. We need to create them here to enable loading of llm-compressor
1033 + # checkpoints which contain them irrespective of the symmetric/asymmetric
1034 + # scheme used during quantization.
1035 + layer.k_zero_point = torch.nn.Parameter(
1036 + torch.zeros(n_scales, requires_grad=False)
1037 + )
1038 + layer.v_zero_point = torch.nn.Parameter(
1039 + torch.zeros(n_scales, requires_grad=False)
1040 + )
1041 + layer.q_zero_point = torch.nn.Parameter(
1042 + torch.zeros(n_scales, requires_grad=False)
1043 + )
1044 +
1045 + # TP-aware loading for attn_head strategy follows attention head partitioning:
1046 + # - q_scale is partitioned over query heads.
1047 + # - k/v_scale is partitioned over kv heads when total_kv_heads >= tp_size,
1048 + # and replicated when total_kv_heads < tp_size.
1049 + if strategy == "attn_head":
1050 +
1051 + def _tp_aware_loader(
1052 + param: torch.Tensor,
1053 + loaded_weight: torch.Tensor,
1054 + kind: Literal["q", "k", "v"],
1055 + param_type: Literal["scale", "zero_point"],
1056 + ):
1057 + # Zero-points are not used as vLLM only supports symmetric quantization
1058 + if param_type == "zero_point":
1059 + return
1060 +
1061 + # LLM-Compressor stores scales as 3D tensors of shape [num_heads, 1, 1]
1062 + loaded_weight = loaded_weight.flatten()
1063 +
1064 + # FlashAttn expects [num_kv_heads] instead of [num_heads] for q_scale.
1065 + # We reduce by taking the max scale in each attention head group.
1066 + if kind == "q":
1067 + reduction_factor = (
1068 + self.quant_config.total_num_heads # type: ignore[attr-defined]
1069 + // self.quant_config.total_num_kv_heads # type: ignore[attr-defined]
1070 + )
1071 + loaded_weight = torch.amax(
1072 + loaded_weight.view(-1, reduction_factor), dim=1
1073 + )
1074 +
1075 + tp_rank = get_tensor_model_parallel_rank()
1076 + tp_size = get_tensor_model_parallel_world_size()
1077 +
1078 + if layer.num_kv_heads * tp_size == self.quant_config.total_num_kv_heads: # type: ignore[attr-defined]
1079 + # heads evenly distributed
1080 + loaded_weight = loaded_weight[
1081 + tp_rank * layer.num_kv_heads : (tp_rank + 1)
1082 + * layer.num_kv_heads
1083 + ]
1084 + else:
1085 + # heads replicated to match TP size
1086 + assert layer.num_kv_heads == 1
1087 + replicas = tp_size // self.quant_config.total_num_kv_heads # type: ignore[attr-defined]
1088 + shard_rank = tp_rank // replicas
1089 + loaded_weight = loaded_weight[shard_rank : shard_rank + 1]
1090 +
1091 + param.data.copy_(loaded_weight.to(dtype=param.dtype))
1092 +
1093 + layer.q_scale.weight_loader = partial(
1094 + _tp_aware_loader, kind="q", param_type="scale"
1095 + )
1096 + layer.k_scale.weight_loader = partial(
1097 + _tp_aware_loader, kind="k", param_type="scale"
1098 + )
1099 + layer.v_scale.weight_loader = partial(
1100 + _tp_aware_loader, kind="v", param_type="scale"
1101 + )
1102 +
1103 + layer.q_zero_point.weight_loader = partial(
1104 + _tp_aware_loader, kind="q", param_type="zero_point"
1105 + )
1106 + layer.k_zero_point.weight_loader = partial(
1107 + _tp_aware_loader, kind="k", param_type="zero_point"
1108 + )
1109 + layer.v_zero_point.weight_loader = partial(
1110 + _tp_aware_loader, kind="v", param_type="zero_point"
1111 + )
1112 +
1113 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
1114 + """
1115 + Override the default vLLM placeholder scales with the llm-compressor loaded
1116 + scales. Zero points are not used as only symmetric quantization is supported.
1117 + """
1118 + layer._k_scale = layer.k_scale
1119 + layer._v_scale = layer.v_scale
1120 + layer._q_scale = layer.q_scale
1121 +
1122 + # Discard all placeholders.
1123 + del layer.k_scale
1124 + del layer.v_scale
1125 + del layer.q_scale
1126 + del layer.k_zero_point
1127 + del layer.v_zero_point
1128 + del layer.q_zero_point
vllm/model_executor/layers/quantization/input_quant_fp8.py
+4 -1
@@ -11,6 +11,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
11 GroupShape,
12 get_fp8_min_max,
13 group_broadcast,
14 + prep_scale_for_group_broadcast,
15 )
16 from vllm.platforms import current_platform
17
@@ -40,7 +41,7 @@ class QuantFP8(CustomOp):
41 """
42 :param static: static or dynamic quantization
43 :param group_shape: quantization group shape (PER_TOKEN, PER_TENSOR,
43 - or arbitrary block size)
44 + PER_CHANNEL, or arbitrary block size)
45 :param num_token_padding: Pad the token dimension of output to this
46 size
47 :param column_major_scales: For group quantization, output scales in
@@ -157,6 +158,8 @@ class QuantFP8(CustomOp):
158 x_max = x.abs().max().unsqueeze(-1).to(torch.float32)
159
160 scale = (x_max / _FP8_MAX).clamp(min=_FP8_MIN_SCALING_FACTOR)
161 + else:
162 + scale = prep_scale_for_group_broadcast(scale, x, self.group_shape)
163
164 # Even for dynamic per-token scales,
165 # reciprocal performs slightly better than division
vllm/model_executor/layers/quantization/utils/quant_utils.py
+47 -24
@@ -191,6 +191,51 @@ def group_broadcast(t, shape):
191 return t
192
193
194 +def prep_scale_for_group_broadcast(
195 + scale: torch.Tensor,
196 + x: torch.Tensor,
197 + group_shape: GroupShape | None,
198 +) -> torch.Tensor:
199 + """
200 + Prepare the input quantization scale for group broadcasting.
201 +
202 + Args:
203 + scale: The scale tensor (scalar or 1D).
204 + x: Target tensor whose shape determines broadcast dimensions.
205 + group_shape: GroupShape to broadcast over.
206 +
207 + Returns:
208 + scale reshaped for correct broadcasting.
209 + """
210 + if scale.numel() == 1:
211 + # For per-tensor quant, keep the scale as a scalar (not reshaped to (1, 1)).
212 + # This avoids misclassifying it as channelwise quant in Fp8LinearOp.apply,
213 + # where the "per_tensor_activations" check relies on "x_scale.dim() < 2":
214 + # per_tensor_activations = (x_scale.numel() == 1) and x_scale.dim() < 2
215 + # For all other cases, reshape scalar scales to (1, 1) for broadcasting.
216 + return (
217 + scale
218 + if group_shape is not None and group_shape.is_per_tensor()
219 + else scale.reshape(1, 1)
220 + )
221 + if scale.ndim == 1:
222 + assert group_shape is not None, (
223 + "group_shape must be provided to correctly broadcast 1D scale"
224 + )
225 + rows, cols = _normalize_quant_group_shape(x, group_shape)
226 + # Determine broadcasting dimension: either rows or columns match group size
227 + if rows == x.shape[-2]:
228 + scale = scale.unsqueeze(-2)
229 + elif cols == x.shape[-1]:
230 + scale = scale.unsqueeze(-1)
231 + else:
232 + raise ValueError(
233 + f"1D scale with shape {scale.shape} cannot be broadcast to x with shape"
234 + f" {x.shape}, group_shape={(rows, cols)}"
235 + )
236 + return scale
237 +
238 +
239 # Quantize assuming once scale per group of elements with shape group_shape,
240 # example group shapes:
241 # * (-1, -1) for per-tensor quantization
@@ -241,7 +286,7 @@ def scaled_quantize(
286 _, fp8_max = get_fp8_min_max()
287 scale = fp8_max / amax
288
244 - # Apply scale and convert form:
289 + # Apply scale and convert from:
290 # (BLK_M, BLK_N, BLOCK_SIZE_M * BLOCK_SIZE_N) to (M, N)
291 x_scl_sat = (
292 (x_blkd_permd * scale.unsqueeze(-1))
@@ -261,29 +306,7 @@ def scaled_dequantize(
306 group_shape: GroupShape | None = None,
307 out_dtype: torch.dtype = torch.float32,
308 ) -> torch.Tensor:
264 - if group_shape is not None:
265 - group_shape = _normalize_quant_group_shape(x_q, group_shape)
266 -
267 - if x_s.numel() == 1: # scalar
268 - x_s = x_s.reshape(1, 1) # normalize all scalar-like tensors to (1, 1)
269 - if x_s.ndim == 1:
270 - if group_shape is None:
271 - raise AssertionError(
272 - "if x_s is 1D tensor, group_shape must be provided otherwise "
273 - "its ambiguous which dimension to broadcast x_s to"
274 - )
275 - # unsqueeze the scales for the dimension where we want to broadcast
276 - # across the full extent
277 - if group_shape[0] == x_q.shape[-2]:
278 - x_s = x_s.unsqueeze(-2)
279 - elif group_shape[1] == x_q.shape[-1]:
280 - x_s = x_s.unsqueeze(-1)
281 - else:
282 - raise AssertionError(
283 - "if x_s is a vector we should be broadcasting it to the full "
284 - "extent of one of the dimensions"
285 - )
286 -
309 + x_s = prep_scale_for_group_broadcast(x_s, x_q, group_shape)
310 if group_shape is not None:
311 assert x_s.shape[-1] == x_q.shape[-1] // group_shape[1]
312 assert x_s.shape[-2] == x_q.shape[-2] // group_shape[0]
vllm/model_executor/model_loader/weight_utils.py
+29 -2
@@ -246,6 +246,23 @@ def get_quant_config(
246 # compressed-tensors uses a compressions_config
247 hf_quant_config = getattr(model_config.hf_config, "compression_config", None)
248
249 + # Pipe information about heads to enable TP-aware loading of attn_head scales
250 + if (
251 + hf_quant_config is not None
252 + and hf_quant_config.get("quant_method") == "compressed-tensors"
253 + ):
254 + if hf_text_config is not None:
255 + n_heads = getattr(hf_text_config, "num_attention_heads", None)
256 + n_kv_heads = getattr(hf_text_config, "num_key_value_heads", None)
257 + else:
258 + n_heads = getattr(model_config.hf_config, "num_attention_heads", None)
259 + n_kv_heads = getattr(model_config.hf_config, "num_key_value_heads", None)
260 +
261 + hf_quant_config["total_num_heads"] = n_heads
262 + hf_quant_config["total_num_kv_heads"] = (
263 + n_kv_heads if n_kv_heads is not None else n_heads
264 + )
265 +
266 if hf_quant_config is not None:
267 return quant_cls.from_config(hf_quant_config)
268
@@ -1157,11 +1174,21 @@ def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None:
1174 # .mixer.attn.{k,v}_scale
1175 (r"\.mixer\.[kv]_proj\.([kv])_scale$", r".mixer.attn.\1_scale"),
1176 # Default format: .{k,v}_scale -> .attn.{k,v}_scale
1160 - (r"\.([kv])_scale$", r".attn.\1_scale"),
1177 + (r"\.([qkv])_scale$", r".attn.\1_scale"),
1178 + (r"\.([qkv])_zero_point$", r".attn.\1_zero_point"),
1179 ]
1180
1181 # Check if name ends with k_scale or v_scale
1164 - if name.endswith((".k_scale", ".v_scale")):
1182 + if name.endswith(
1183 + (
1184 + ".k_scale",
1185 + ".v_scale",
1186 + ".q_scale",
1187 + ".k_zero_point",
1188 + ".v_zero_point",
1189 + ".q_zero_point",
1190 + )
1191 + ):
1192 import regex as re
1193
1194 for pattern, replacement in scale_mapping_patterns:
vllm/model_executor/models/apertus.py
+1 -1
@@ -437,7 +437,7 @@ class ApertusModel(nn.Module):
437 weight_loader(param, loaded_weight)
438 loaded_params.add(scale_name)
439 continue
440 - if "scale" in name:
440 + if "scale" in name or "zero_point" in name:
441 # Remapping the name of FP8 kv-scale.
442 name = maybe_remap_kv_scale_name(name, params_dict)
443 if name is None:
vllm/model_executor/models/arcee.py
+1 -1
@@ -303,7 +303,7 @@ class ArceeModel(nn.Module):
303 loaded_params.add(scale_name)
304 continue
305
306 - if "scale" in name:
306 + if "scale" in name or "zero_point" in name:
307 remapped_name = maybe_remap_kv_scale_name(name, params_dict)
308 if remapped_name is None:
309 continue
vllm/model_executor/models/llama.py
+2 -2
@@ -465,8 +465,8 @@ class LlamaModel(nn.Module):
465 weight_loader(param, loaded_weight)
466 loaded_params.add(scale_name)
467 continue
468 - if "scale" in name:
469 - # Remapping the name of FP8 kv-scale.
468 + if "scale" in name or "zero_point" in name:
469 + # Remapping the name of FP8 kv-scale or zero point.
470 name = maybe_remap_kv_scale_name(name, params_dict)
471 if name is None:
472 continue
vllm/model_executor/models/llama_eagle.py
+2 -2
@@ -140,8 +140,8 @@ class LlamaModel(nn.Module):
140 weight_loader(param, loaded_weight)
141 loaded_params.add(scale_name)
142 continue
143 - # Remapping the name FP8 kv-scale
144 - if "scale" in name:
143 + # Remapping the name FP8 kv-scale or zero point.
144 + if "scale" in name or "zero_point" in name:
145 name = maybe_remap_kv_scale_name(name, params_dict)
146 if name is None:
147 continue
vllm/model_executor/models/llama_eagle3.py
+2 -2
@@ -238,8 +238,8 @@ class LlamaModel(nn.Module):
238 weight_loader(param, loaded_weight)
239 loaded_params.add(scale_name)
240 continue
241 - # Remapping the name FP8 kv-scale
242 - if "scale" in name:
241 + # Remapping the name FP8 kv-scale or zero point.
242 + if "scale" in name or "zero_point" in name:
243 name = maybe_remap_kv_scale_name(name, params_dict)
244 if name is None:
245 continue
vllm/model_executor/models/nemotron_h.py
+1 -1
@@ -661,7 +661,7 @@ class NemotronHModel(nn.Module):
661 params_dict = dict(self.named_parameters())
662 loaded_params: set[str] = set()
663 for name, loaded_weight in weights:
664 - if "scale" in name:
664 + if "scale" in name or "zero_point" in name:
665 # Remapping the name of FP8 kv-scale.
666 name = maybe_remap_kv_scale_name(name, params_dict)
667 if name is None:
vllm/model_executor/models/nemotron_nas.py
+1 -1
@@ -342,7 +342,7 @@ class DeciModel(nn.Module):
342 weight_loader(param, loaded_weight)
343 loaded_params.add(scale_name)
344 continue
345 - if "scale" in name:
345 + if "scale" in name or "zero_point" in name:
346 # Remapping the name of FP8 kv-scale.
347 name = maybe_remap_kv_scale_name(name, params_dict)
348 if name is None:
vllm/v1/attention/backend.py
+1
@@ -620,6 +620,7 @@ class AttentionImpl(ABC, Generic[T]):
620 # TODO add support to more backends:
621 # https://github.com/vllm-project/vllm/issues/25584
622 supports_quant_query_input: bool = False
623 + supports_per_head_quant_scales: bool = False
624
625 dcp_world_size: int
626 dcp_rank: int
vllm/v1/attention/backends/flash_attn.py
+15 -6
@@ -576,6 +576,11 @@ class FlashAttentionImpl(AttentionImpl):
576 )
577
578 self.supports_quant_query_input = True
579 + self.supports_per_head_quant_scales = (
580 + self.vllm_flash_attn_version >= 3
581 + if self.vllm_flash_attn_version is not None
582 + else False
583 + )
584
585 def forward(
586 self,
@@ -691,6 +696,10 @@ class FlashAttentionImpl(AttentionImpl):
696
697 descale_shape = (cu_seqlens_q.shape[0] - 1, self.num_kv_heads)
698
699 + q_descale = layer._q_scale.expand(descale_shape)
700 + k_descale = layer._k_scale.expand(descale_shape)
701 + v_descale = layer._v_scale.expand(descale_shape)
702 +
703 if self.dcp_world_size > 1:
704 self._forward_with_dcp(
705 query[:num_actual_tokens],
@@ -700,9 +709,9 @@ class FlashAttentionImpl(AttentionImpl):
709 value_cache,
710 output[:num_actual_tokens],
711 attn_metadata,
703 - q_descale=layer._q_scale.expand(descale_shape),
704 - k_descale=layer._k_scale.expand(descale_shape),
705 - v_descale=layer._v_scale.expand(descale_shape),
712 + q_descale=q_descale,
713 + k_descale=k_descale,
714 + v_descale=v_descale,
715 )
716 return output
717 else:
@@ -728,9 +737,9 @@ class FlashAttentionImpl(AttentionImpl):
737 softcap=self.logits_soft_cap,
738 scheduler_metadata=scheduler_metadata,
739 fa_version=self.vllm_flash_attn_version,
731 - q_descale=layer._q_scale.expand(descale_shape),
732 - k_descale=layer._k_scale.expand(descale_shape),
733 - v_descale=layer._v_scale.expand(descale_shape),
740 + q_descale=q_descale,
741 + k_descale=k_descale,
742 + v_descale=v_descale,
743 num_splits=attn_metadata.max_num_splits,
744 s_aux=self.sinks,
745 )