@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2ce3d0ce36
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 5, 16:00 UTC

[Feature] KV cache per-token-head INT8/FP8 quantization (#38378)

Signed-off-by: JartX <sagformas@epdcenter.es> Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: yangyang4991 <yangyang4991@gmail.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Isotr0py <2037008807@qq.com>

JartX committed Apr 2, 2026 at 14:13 UTC 2ce3d0ce360bf1e31e7ea3b2b0971a37bf5707c6
16 files changed +1308 -66
docs/design/attention_backends.md
+1 -1
@@ -177,7 +177,7 @@ Priority is **1 = highest** (tried first).
177 | `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ✅ | ❌ | All | N/A |
178 | `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
179 | `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
180 -| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
180 +| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
181
182 > **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
183 >
tests/models/quantization/test_per_token_kv_cache.py new
+94
@@ -0,0 +1,94 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""End-to-end accuracy tests for per-token-head KV cache quantization.
5 +
6 +Compares logprobs between a baseline bf16 model and the same model with
7 +per-token-head quantized KV cache (int8 or fp8) using the Triton attention
8 +backend.
9 +
10 +Run: pytest tests/models/quantization/test_per_token_kv_cache.py -v -s
11 +"""
12 +
13 +import pytest
14 +
15 +from vllm.platforms import current_platform
16 +
17 +from ..utils import check_logprobs_close
18 +
19 +
20 +@pytest.mark.skipif(
21 + not current_platform.is_cuda_alike(),
22 + reason="Per-token-head KV cache requires CUDA or ROCm GPU.",
23 +)
24 +@pytest.mark.parametrize(
25 + "base_model,test_model",
26 + [
27 + (
28 + "meta-llama/Llama-3.2-1B-Instruct",
29 + "meta-llama/Llama-3.2-1B-Instruct",
30 + ),
31 + ],
32 +)
33 +@pytest.mark.parametrize(
34 + "kv_cache_dtype", ["int8_per_token_head", "fp8_per_token_head"]
35 +)
36 +@pytest.mark.parametrize("max_tokens", [4])
37 +@pytest.mark.parametrize("enforce_eager", [True])
38 +@pytest.mark.parametrize("backend", ["TRITON_ATTN"])
39 +@pytest.mark.parametrize("tensor_parallel_size", [1])
40 +def test_per_token_head_kv_cache_accuracy(
41 + vllm_runner,
42 + example_prompts,
43 + base_model: str,
44 + test_model: str,
45 + kv_cache_dtype: str,
46 + max_tokens: int,
47 + enforce_eager: bool,
48 + backend: str,
49 + tensor_parallel_size: int,
50 + monkeypatch: pytest.MonkeyPatch,
51 +) -> None:
52 + """Compare logprobs between bf16 baseline and per-token-head quantized KV
53 + cache.
54 +
55 + Uses calculate_kv_scales (dynamic scale computation) since there are
56 + no per-token-head calibrated checkpoints available yet.
57 + """
58 + with monkeypatch.context() as m:
59 + m.setenv("TOKENIZERS_PARALLELISM", "true")
60 +
61 + MAX_MODEL_LEN = 1024
62 + NUM_LOG_PROBS = 8
63 +
64 + with vllm_runner(
65 + base_model,
66 + max_model_len=MAX_MODEL_LEN,
67 + tensor_parallel_size=tensor_parallel_size,
68 + enforce_eager=enforce_eager,
69 + kv_cache_dtype="auto",
70 + attention_config={"backend": backend},
71 + ) as vllm_model:
72 + baseline_outputs = vllm_model.generate_greedy_logprobs(
73 + example_prompts, max_tokens, NUM_LOG_PROBS
74 + )
75 +
76 + with vllm_runner(
77 + test_model,
78 + max_model_len=MAX_MODEL_LEN,
79 + tensor_parallel_size=tensor_parallel_size,
80 + enforce_eager=enforce_eager,
81 + kv_cache_dtype=kv_cache_dtype,
82 + calculate_kv_scales=True,
83 + attention_config={"backend": backend},
84 + ) as vllm_model:
85 + test_outputs = vllm_model.generate_greedy_logprobs(
86 + example_prompts, max_tokens, NUM_LOG_PROBS
87 + )
88 +
89 + check_logprobs_close(
90 + outputs_0_lst=baseline_outputs,
91 + outputs_1_lst=test_outputs,
92 + name_0="bf16_kv_cache",
93 + name_1=f"{kv_cache_dtype}_kv_cache",
94 + )
tests/quantization/test_per_token_kv_cache.py new
+560
@@ -0,0 +1,560 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for per-token-head KV cache quantization (INT8 and FP8).
4 +
5 +Covers:
6 +- Per-token-head Triton reshape-and-cache kernel
7 +- Round-trip quantize/dequantize accuracy
8 +- process_weights_after_loading early-return path
9 +- End-to-end integration with Triton unified attention kernel
10 +
11 +Run: pytest tests/quantization/test_per_token_kv_cache.py -v -s
12 +"""
13 +
14 +import random
15 +from dataclasses import dataclass
16 +from unittest.mock import MagicMock
17 +
18 +import pytest
19 +import torch
20 +
21 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
22 + get_fp8_min_max,
23 +)
24 +from vllm.platforms import current_platform
25 +from vllm.utils.torch_utils import set_random_seed
26 +from vllm.v1.kv_cache_interface import KVQuantMode, is_quantized_kv_cache
27 +
28 +# Skip entire module if no CUDA/ROCm GPU available
29 +pytestmark = [
30 + pytest.mark.skipif(
31 + not current_platform.is_cuda_alike(),
32 + reason="Per-token-head KV cache tests require CUDA or ROCm GPU.",
33 + ),
34 +]
35 +
36 +# ---------------------------------------------------------------------------
37 +# Test parameters
38 +# ---------------------------------------------------------------------------
39 +NUM_TOKENS = [1, 7, 42]
40 +NUM_KV_HEADS = [1, 4, 8]
41 +HEAD_SIZES = [64, 128]
42 +BLOCK_SIZES = [16]
43 +SEEDS = [0]
44 +
45 +# Platform-dependent FP8 dtype and range
46 +FP8_DTYPE = current_platform.fp8_dtype()
47 +FP8_MIN, FP8_MAX = get_fp8_min_max()
48 +
49 +
50 +# ---------------------------------------------------------------------------
51 +# Per-dtype quantization config
52 +# ---------------------------------------------------------------------------
53 +@dataclass(frozen=True)
54 +class QuantConfig:
55 + """Quantization parameters for a given cache dtype."""
56 +
57 + cache_dtype: torch.dtype # torch.int8 or FP8_DTYPE
58 + kv_cache_dtype_str: str # "int8_per_token_head" or "fp8_per_token_head"
59 + quant_max: float
60 + quant_min: float
61 + kv_quant_mode: KVQuantMode
62 + # INT8 Triton stores truncate; FP8 hardware casts round.
63 + uses_trunc: bool
64 +
65 +
66 +INT8_CONFIG = QuantConfig(
67 + cache_dtype=torch.int8,
68 + kv_cache_dtype_str="int8_per_token_head",
69 + quant_max=127.0,
70 + quant_min=-128.0,
71 + kv_quant_mode=KVQuantMode.INT8_PER_TOKEN_HEAD,
72 + uses_trunc=True,
73 +)
74 +FP8_CONFIG = QuantConfig(
75 + cache_dtype=FP8_DTYPE,
76 + kv_cache_dtype_str="fp8_per_token_head",
77 + quant_max=FP8_MAX,
78 + quant_min=FP8_MIN,
79 + kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD,
80 + uses_trunc=False,
81 +)
82 +
83 +QUANT_CONFIGS = [INT8_CONFIG, FP8_CONFIG]
84 +
85 +
86 +@pytest.fixture(params=QUANT_CONFIGS, ids=["int8", "fp8"])
87 +def qcfg(request) -> QuantConfig:
88 + return request.param
89 +
90 +
91 +# ---------------------------------------------------------------------------
92 +# Helpers
93 +# ---------------------------------------------------------------------------
94 +def _quantize_per_token_head_ref(
95 + data: torch.Tensor, # [num_tokens, num_heads, head_size]
96 + cfg: QuantConfig,
97 +) -> tuple[torch.Tensor, torch.Tensor]:
98 + """Reference per-token-head quantization (one scale per token per head).
99 +
100 + Returns (quantized, scales) where scales is [num_tokens, num_heads].
101 + """
102 + absmax = data.float().abs().amax(dim=2) # [num_tokens, num_heads]
103 + scales = (absmax / cfg.quant_max).clamp(min=1e-6)
104 + scaled = data.float() * (1.0 / scales[:, :, None])
105 + if cfg.uses_trunc:
106 + q = scaled.round().clamp(cfg.quant_min, cfg.quant_max).to(cfg.cache_dtype)
107 + else:
108 + q = scaled.clamp(cfg.quant_min, cfg.quant_max).to(cfg.cache_dtype)
109 + return q, scales
110 +
111 +
112 +# ===========================================================================
113 +# 1. is_quantized_kv_cache / get_kv_quant_mode
114 +# ===========================================================================
115 +class TestIsQuantizedKvCache:
116 + def test_fp8_variants(self):
117 + assert is_quantized_kv_cache("fp8")
118 + assert is_quantized_kv_cache("fp8_e4m3")
119 + assert is_quantized_kv_cache("fp8_e5m2")
120 +
121 + def test_int8_per_token_head(self):
122 + assert is_quantized_kv_cache("int8_per_token_head")
123 +
124 + def test_fp8_per_token_head(self):
125 + assert is_quantized_kv_cache("fp8_per_token_head")
126 +
127 + def test_auto(self):
128 + assert not is_quantized_kv_cache("auto")
129 +
130 + def test_bfloat16(self):
131 + assert not is_quantized_kv_cache("bfloat16")
132 +
133 + def test_kv_quant_mode_int8(self):
134 + from vllm.v1.kv_cache_interface import get_kv_quant_mode
135 +
136 + assert (
137 + get_kv_quant_mode("int8_per_token_head") == KVQuantMode.INT8_PER_TOKEN_HEAD
138 + )
139 +
140 + def test_kv_quant_mode_fp8(self):
141 + from vllm.v1.kv_cache_interface import get_kv_quant_mode
142 +
143 + assert get_kv_quant_mode("fp8_per_token_head") == KVQuantMode.FP8_PER_TOKEN_HEAD
144 +
145 +
146 +# ===========================================================================
147 +# 2. Triton per-token-head kernel (reshape-and-cache)
148 +# ===========================================================================
149 +@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
150 +@pytest.mark.parametrize("num_heads", NUM_KV_HEADS)
151 +@pytest.mark.parametrize("head_size", HEAD_SIZES)
152 +@pytest.mark.parametrize("block_size", BLOCK_SIZES)
153 +@pytest.mark.parametrize("seed", SEEDS)
154 +@torch.inference_mode()
155 +def test_reshape_and_cache_per_token_head(
156 + qcfg: QuantConfig,
157 + num_tokens: int,
158 + num_heads: int,
159 + head_size: int,
160 + block_size: int,
161 + seed: int,
162 +):
163 + """Test triton_reshape_and_cache_flash_per_token_head_quant kernel."""
164 + from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
165 + triton_reshape_and_cache_flash_per_token_head_quant,
166 + )
167 +
168 + set_random_seed(seed)
169 + torch.set_default_device("cuda")
170 +
171 + num_blocks = (num_tokens + block_size - 1) // block_size + 4
172 +
173 + key = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16)
174 + value = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16)
175 +
176 + key_cache = torch.zeros(
177 + num_blocks, block_size, num_heads, head_size, dtype=qcfg.cache_dtype
178 + )
179 + value_cache = torch.zeros(
180 + num_blocks, block_size, num_heads, head_size, dtype=qcfg.cache_dtype
181 + )
182 + k_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)
183 + v_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)
184 +
185 + num_slots = block_size * num_blocks
186 + slot_mapping = torch.tensor(
187 + random.sample(range(num_slots), num_tokens), dtype=torch.long
188 + )
189 +
190 + triton_reshape_and_cache_flash_per_token_head_quant(
191 + key,
192 + value,
193 + key_cache,
194 + value_cache,
195 + k_scale_cache,
196 + v_scale_cache,
197 + slot_mapping,
198 + )
199 +
200 + # Reference
201 + ref_k_quant, ref_k_scales = _quantize_per_token_head_ref(key, qcfg)
202 + ref_v_quant, ref_v_scales = _quantize_per_token_head_ref(value, qcfg)
203 +
204 + # Compare dequantized values rather than raw quantized values.
205 + # Triton and PyTorch reductions can differ at FP8 rounding boundaries
206 + # (up to 32 in quantized domain for fp8_e4m3), but the dequantized
207 + # error is bounded by the scale.
208 + for i, slot in enumerate(slot_mapping.tolist()):
209 + blk = slot // block_size
210 + off = slot % block_size
211 +
212 + actual_k_scale = k_scale_cache[blk, off] # [num_heads]
213 + k_deq = key_cache[blk, off].float() * actual_k_scale[:, None]
214 + k_ref_deq = key[i].float()
215 + torch.testing.assert_close(
216 + k_deq,
217 + k_ref_deq,
218 + atol=0.1,
219 + rtol=0.1,
220 + )
221 + actual_v_scale = v_scale_cache[blk, off] # [num_heads]
222 + v_deq = value_cache[blk, off].float() * actual_v_scale[:, None]
223 + v_ref_deq = value[i].float()
224 + torch.testing.assert_close(
225 + v_deq,
226 + v_ref_deq,
227 + atol=0.1,
228 + rtol=0.1,
229 + )
230 + # Per-head scales: [num_heads]
231 + torch.testing.assert_close(
232 + k_scale_cache[blk, off], ref_k_scales[i], atol=1e-4, rtol=1e-3
233 + )
234 + torch.testing.assert_close(
235 + v_scale_cache[blk, off], ref_v_scales[i], atol=1e-4, rtol=1e-3
236 + )
237 +
238 +
239 +# ===========================================================================
240 +# 3. Per-token-head round-trip accuracy (quantize -> dequantize)
241 +# ===========================================================================
242 +@pytest.mark.parametrize("num_tokens", [1, 16])
243 +@pytest.mark.parametrize("num_heads", [4])
244 +@pytest.mark.parametrize("head_size", [128])
245 +@pytest.mark.parametrize("block_size", [16])
246 +@torch.inference_mode()
247 +def test_per_token_head_round_trip_accuracy(
248 + qcfg: QuantConfig,
249 + num_tokens: int,
250 + num_heads: int,
251 + head_size: int,
252 + block_size: int,
253 +):
254 + """Verify per-token-head round-trip: kernel dequant matches reference.
255 +
256 + INT8: Triton truncates on float->int8 store.
257 + FP8: hardware cast (clamp then cast).
258 + """
259 + from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
260 + triton_reshape_and_cache_flash_per_token_head_quant,
261 + )
262 +
263 + torch.set_default_device("cuda")
264 + set_random_seed(42)
265 +
266 + num_blocks = (num_tokens + block_size - 1) // block_size + 2
267 +
268 + key = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16) * 0.5
269 + value = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16) * 0.5
270 +
271 + key_cache = torch.zeros(
272 + num_blocks, block_size, num_heads, head_size, dtype=qcfg.cache_dtype
273 + )
274 + value_cache = torch.zeros(
275 + num_blocks, block_size, num_heads, head_size, dtype=qcfg.cache_dtype
276 + )
277 + k_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)
278 + v_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)
279 +
280 + slot_mapping = torch.arange(num_tokens, dtype=torch.long)
281 +
282 + triton_reshape_and_cache_flash_per_token_head_quant(
283 + key,
284 + value,
285 + key_cache,
286 + value_cache,
287 + k_scale_cache,
288 + v_scale_cache,
289 + slot_mapping,
290 + )
291 +
292 + for i in range(num_tokens):
293 + blk = i // block_size
294 + off = i % block_size
295 +
296 + for label, data, cache, sc in [
297 + ("key", key, key_cache, k_scale_cache),
298 + ("val", value, value_cache, v_scale_cache),
299 + ]:
300 + for h in range(num_heads):
301 + orig = data[i, h].float() # [head_size]
302 +
303 + actual_q = cache[blk, off, h]
304 + actual_sc = sc[blk, off, h]
305 + actual_deq = actual_q.float() * actual_sc
306 +
307 + # Round-trip: dequantized should be close to original
308 + torch.testing.assert_close(
309 + actual_deq,
310 + orig,
311 + atol=0.1,
312 + rtol=0.1,
313 + )
314 +
315 +
316 +# ===========================================================================
317 +# 4. Negative slot mapping (padding tokens should be skipped)
318 +# ===========================================================================
319 +@torch.inference_mode()
320 +def test_per_token_head_negative_slot_skipped(qcfg: QuantConfig):
321 + """Tokens with slot_mapping=-1 should leave the cache unchanged."""
322 + from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
323 + triton_reshape_and_cache_flash_per_token_head_quant,
324 + )
325 +
326 + torch.set_default_device("cuda")
327 + num_tokens = 4
328 + num_heads = 2
329 + head_size = 64
330 + block_size = 16
331 + num_blocks = 2
332 +
333 + key = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16)
334 + value = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16)
335 +
336 + key_cache = torch.zeros(
337 + num_blocks, block_size, num_heads, head_size, dtype=qcfg.cache_dtype
338 + )
339 + value_cache = torch.zeros(
340 + num_blocks, block_size, num_heads, head_size, dtype=qcfg.cache_dtype
341 + )
342 + k_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)
343 + v_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)
344 +
345 + slot_mapping = torch.tensor([0, -1, 1, -1], dtype=torch.long)
346 +
347 + key_cache_before = key_cache.clone()
348 + val_cache_before = value_cache.clone()
349 +
350 + triton_reshape_and_cache_flash_per_token_head_quant(
351 + key,
352 + value,
353 + key_cache,
354 + value_cache,
355 + k_scale_cache,
356 + v_scale_cache,
357 + slot_mapping,
358 + )
359 +
360 + # Slots 0 and 1 should have been written (tokens 0 and 2)
361 + assert not torch.equal(key_cache[0, 0], key_cache_before[0, 0])
362 + assert not torch.equal(key_cache[0, 1], key_cache_before[0, 1])
363 + assert not torch.equal(value_cache[0, 0], val_cache_before[0, 0])
364 +
365 + # All other slots should be unchanged
366 + assert torch.equal(key_cache[0, 2:], key_cache_before[0, 2:])
367 + assert torch.equal(key_cache[1], key_cache_before[1])
368 + assert torch.equal(value_cache[0, 2:], val_cache_before[0, 2:])
369 +
370 +
371 +# ===========================================================================
372 +# 5. process_weights_after_loading -- per-token-head early return
373 +# ===========================================================================
374 +@pytest.mark.parametrize(
375 + "kv_cache_dtype", ["int8_per_token_head", "fp8_per_token_head"]
376 +)
377 +def test_process_weights_sets_placeholder_scales(kv_cache_dtype: str):
378 + """Per-token-head should set _k_scale=1.0, _v_scale=1.0
379 + and delete checkpoint attrs."""
380 + from vllm.model_executor.layers.quantization.kv_cache import (
381 + BaseKVCacheMethod,
382 + )
383 +
384 + layer = MagicMock()
385 + layer.kv_cache_dtype = kv_cache_dtype
386 + layer.calculate_kv_scales = False
387 + layer.k_scale = torch.nn.Parameter(torch.tensor(-1.0), requires_grad=False)
388 + layer.v_scale = torch.nn.Parameter(torch.tensor(-1.0), requires_grad=False)
389 + layer.q_scale = torch.nn.Parameter(torch.tensor(-1.0), requires_grad=False)
390 + layer.prob_scale = torch.nn.Parameter(torch.tensor(-1.0), requires_grad=False)
391 + layer._k_scale = torch.tensor(0.0)
392 + layer._v_scale = torch.tensor(0.0)
393 + layer._k_scale_float = 0.0
394 + layer._v_scale_float = 0.0
395 +
396 + method = BaseKVCacheMethod.__new__(BaseKVCacheMethod)
397 + method.quant_config = MagicMock()
398 + method.process_weights_after_loading(layer)
399 +
400 + assert layer._k_scale_float == 1.0
401 + assert layer._v_scale_float == 1.0
402 + assert not hasattr(layer, "k_scale")
403 + assert not hasattr(layer, "v_scale")
404 + assert not hasattr(layer, "q_scale")
405 + assert not hasattr(layer, "prob_scale")
406 +
407 +
408 +# ===========================================================================
409 +# 6. Triton unified_attention -- per-token-head scale cache (INT8 and FP8)
410 +# ===========================================================================
411 +@pytest.mark.parametrize(
412 + "seq_lens",
413 + [
414 + [(1, 128)],
415 + [(1, 64), (1, 32)],
416 + ],
417 +)
418 +@pytest.mark.parametrize("num_heads", [(4, 4)])
419 +@pytest.mark.parametrize("head_size", [128])
420 +@pytest.mark.parametrize("block_size", [16])
421 +@torch.inference_mode()
422 +def test_triton_unified_attention_per_token_head_scale(
423 + qcfg: QuantConfig,
424 + seq_lens: list[tuple[int, int]],
425 + num_heads: tuple[int, int],
426 + head_size: int,
427 + block_size: int,
428 +):
429 + """End-to-end: quantized KV with per-token-head scale caches."""
430 + from vllm.utils.math_utils import next_power_of_2
431 + from vllm.v1.attention.ops.triton_unified_attention import unified_attention
432 +
433 + torch.set_default_device("cuda")
434 + set_random_seed(0)
435 +
436 + num_seqs = len(seq_lens)
437 + query_lens = [s[0] for s in seq_lens]
438 + kv_lens = [s[1] for s in seq_lens]
439 + num_query_heads, num_kv_heads = num_heads
440 + max_query_len = max(query_lens)
441 + max_kv_len = max(kv_lens)
442 + scale = head_size**-0.5
443 + num_blocks = 2048
444 +
445 + query = torch.randn(
446 + sum(query_lens), num_query_heads, head_size, dtype=torch.bfloat16
447 + )
448 +
449 + key_cache_bf16 = torch.randn(
450 + num_blocks, block_size, num_kv_heads, head_size, dtype=torch.bfloat16
451 + )
452 + value_cache_bf16 = torch.randn_like(key_cache_bf16)
453 +
454 + # Per-token-head quantization: one scale per (block, slot, head)
455 + k_absmax = key_cache_bf16.float().abs().amax(dim=-1) # [..., num_kv_heads]
456 + v_absmax = value_cache_bf16.float().abs().amax(dim=-1)
457 + k_scale_cache = (k_absmax / qcfg.quant_max).clamp(min=1e-6).to(torch.float32)
458 + v_scale_cache = (v_absmax / qcfg.quant_max).clamp(min=1e-6).to(torch.float32)
459 +
460 + scaled_k = key_cache_bf16.float() / k_scale_cache[:, :, :, None]
461 + scaled_v = value_cache_bf16.float() / v_scale_cache[:, :, :, None]
462 + if qcfg.uses_trunc:
463 + key_cache_q = (
464 + scaled_k.round().clamp(qcfg.quant_min, qcfg.quant_max).to(qcfg.cache_dtype)
465 + )
466 + value_cache_q = (
467 + scaled_v.round().clamp(qcfg.quant_min, qcfg.quant_max).to(qcfg.cache_dtype)
468 + )
469 + else:
470 + key_cache_q = scaled_k.clamp(qcfg.quant_min, qcfg.quant_max).to(
471 + qcfg.cache_dtype
472 + )
473 + value_cache_q = scaled_v.clamp(qcfg.quant_min, qcfg.quant_max).to(
474 + qcfg.cache_dtype
475 + )
476 +
477 + # Dequantized reference
478 + key_cache_deq = key_cache_q.float() * k_scale_cache[:, :, :, None]
479 + value_cache_deq = value_cache_q.float() * v_scale_cache[:, :, :, None]
480 +
481 + cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
482 + dim=0, dtype=torch.int32
483 + )
484 + kv_lens_t = torch.tensor(kv_lens, dtype=torch.int32)
485 +
486 + max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
487 + block_tables = torch.randint(
488 + 0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
489 + )
490 +
491 + head_size_padded = next_power_of_2(head_size)
492 + seq_threshold_3D = 0
493 + num_par_softmax_segments = 16
494 + softmax_segm_output = torch.empty(
495 + (seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded),
496 + dtype=torch.float32,
497 + )
498 + softmax_segm_max = torch.empty(
499 + (seq_threshold_3D, num_query_heads, num_par_softmax_segments),
500 + dtype=torch.float32,
501 + )
502 + softmax_segm_expsum = torch.empty(
503 + (seq_threshold_3D, num_query_heads, num_par_softmax_segments),
504 + dtype=torch.float32,
505 + )
506 +
507 + output_q = torch.empty_like(query)
508 + unified_attention(
509 + q=query,
510 + k=key_cache_q,
511 + v=value_cache_q,
512 + out=output_q,
513 + cu_seqlens_q=cu_query_lens,
514 + seqused_k=kv_lens_t,
515 + max_seqlen_q=max_query_len,
516 + max_seqlen_k=max_kv_len,
517 + softmax_scale=scale,
518 + causal=True,
519 + window_size=(-1, -1),
520 + block_table=block_tables,
521 + softcap=0,
522 + q_descale=None,
523 + k_descale=None,
524 + v_descale=None,
525 + seq_threshold_3D=seq_threshold_3D,
526 + num_par_softmax_segments=num_par_softmax_segments,
527 + softmax_segm_output=softmax_segm_output,
528 + softmax_segm_max=softmax_segm_max,
529 + softmax_segm_expsum=softmax_segm_expsum,
530 + kv_quant_mode=qcfg.kv_quant_mode,
531 + k_scale_cache=k_scale_cache,
532 + v_scale_cache=v_scale_cache,
533 + )
534 +
535 + output_ref = torch.empty_like(query)
536 + unified_attention(
537 + q=query,
538 + k=key_cache_deq.to(torch.bfloat16),
539 + v=value_cache_deq.to(torch.bfloat16),
540 + out=output_ref,
541 + cu_seqlens_q=cu_query_lens,
542 + seqused_k=kv_lens_t,
543 + max_seqlen_q=max_query_len,
544 + max_seqlen_k=max_kv_len,
545 + softmax_scale=scale,
546 + causal=True,
547 + window_size=(-1, -1),
548 + block_table=block_tables,
549 + softcap=0,
550 + q_descale=None,
551 + k_descale=None,
552 + v_descale=None,
553 + seq_threshold_3D=seq_threshold_3D,
554 + num_par_softmax_segments=num_par_softmax_segments,
555 + softmax_segm_output=softmax_segm_output,
556 + softmax_segm_max=softmax_segm_max,
557 + softmax_segm_expsum=softmax_segm_expsum,
558 + )
559 +
560 + torch.testing.assert_close(output_q, output_ref, atol=5e-2, rtol=5e-2)
vllm/config/cache.py
+17 -4
@@ -8,7 +8,10 @@ from pydantic import Field, SkipValidation, field_validator, model_validator
8
9 from vllm.config.utils import config
10 from vllm.logger import init_logger
11 -from vllm.utils.torch_utils import is_quantized_kv_cache
11 +from vllm.utils.torch_utils import (
12 + is_quantized_kv_cache,
13 + kv_cache_uses_per_token_head_scales,
14 +)
15
16 logger = init_logger(__name__)
17
@@ -21,6 +24,8 @@ CacheDType = Literal[
24 "fp8_e5m2",
25 "fp8_inc",
26 "fp8_ds_mla",
27 + "int8_per_token_head",
28 + "fp8_per_token_head",
29 ]
30 MambaDType = Literal["auto", "float32", "float16"]
31 MambaCacheMode = Literal["all", "align", "none"]
@@ -237,12 +242,20 @@ class CacheConfig:
242 @field_validator("cache_dtype", mode="after")
243 @classmethod
244 def _validate_cache_dtype(cls, cache_dtype: CacheDType) -> CacheDType:
240 - if is_quantized_kv_cache(cache_dtype):
245 + if kv_cache_uses_per_token_head_scales(cache_dtype):
246 logger.info(
242 - "Using fp8 data type to store kv cache. It reduces the GPU "
247 + "Using %s data type to store kv cache. It reduces the GPU "
248 + "memory footprint and boosts the performance. "
249 + "Dynamic per-token-head scales will be computed at runtime.",
250 + str(cache_dtype),
251 + )
252 + elif is_quantized_kv_cache(cache_dtype):
253 + logger.info(
254 + "Using %s data type to store kv cache. It reduces the GPU "
255 "memory footprint and boosts the performance. "
256 "Meanwhile, it may cause accuracy drop without a proper "
245 - "scaling factor."
257 + "scaling factor",
258 + str(cache_dtype),
259 )
260 return cache_dtype
261
vllm/model_executor/layers/attention/attention.py
+8 -2
@@ -38,6 +38,7 @@ from vllm.v1.kv_cache_interface import (
38 FullAttentionSpec,
39 KVCacheSpec,
40 SlidingWindowSpec,
41 + get_kv_quant_mode,
42 )
43
44 if TYPE_CHECKING:
@@ -381,8 +382,10 @@ class Attention(nn.Module, AttentionLayerBase):
382
383 # for attn backends supporting query quantization
384 self.query_quant = None
384 - if self.impl.supports_quant_query_input and self.kv_cache_dtype.startswith(
385 - "fp8"
385 + if (
386 + self.impl.supports_quant_query_input
387 + and self.kv_cache_dtype.startswith("fp8")
388 + and not self.kv_cache_dtype.endswith("per_token_head")
389 ):
390 is_per_head = (
391 hasattr(self, "q_scale") and self.q_scale.numel() == self.num_kv_heads
@@ -539,6 +542,7 @@ class Attention(nn.Module, AttentionLayerBase):
542 block_size = vllm_config.cache_config.block_size
543 # Should not be called for enc-dec or encoder-only attention.
544 assert self.attn_type == AttentionType.DECODER
545 + quant_mode = get_kv_quant_mode(self.kv_cache_dtype)
546 if self.sliding_window is not None:
547 assert not vllm_config.model_config.use_mla, (
548 "MLA is not supported for slidingwindow"
@@ -548,6 +552,7 @@ class Attention(nn.Module, AttentionLayerBase):
552 num_kv_heads=self.num_kv_heads,
553 head_size=self.head_size,
554 dtype=self.kv_cache_torch_dtype,
555 + kv_quant_mode=quant_mode,
556 sliding_window=self.sliding_window,
557 )
558 else:
@@ -557,6 +562,7 @@ class Attention(nn.Module, AttentionLayerBase):
562 head_size=self.head_size,
563 head_size_v=self.head_size_v,
564 dtype=self.kv_cache_torch_dtype,
565 + kv_quant_mode=quant_mode,
566 )
567
568
vllm/model_executor/layers/attention/chunked_local_attention.py
+2
@@ -23,6 +23,7 @@ from vllm.v1.kv_cache_interface import (
23 AttentionSpec,
24 ChunkedLocalAttentionSpec,
25 KVCacheSpec,
26 + get_kv_quant_mode,
27 )
28
29
@@ -123,5 +124,6 @@ class ChunkedLocalAttention(Attention):
124 num_kv_heads=self.num_kv_heads,
125 head_size=self.head_size,
126 dtype=self.kv_cache_torch_dtype,
127 + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
128 attention_chunk_size=self.attention_chunk_size,
129 )
vllm/model_executor/layers/attention/cross_attention.py
+6 -1
@@ -18,7 +18,11 @@ from vllm.v1.attention.backend import (
18 subclass_attention_backend_with_overrides,
19 )
20 from vllm.v1.attention.selector import get_attn_backend
21 -from vllm.v1.kv_cache_interface import CrossAttentionSpec, KVCacheSpec
21 +from vllm.v1.kv_cache_interface import (
22 + CrossAttentionSpec,
23 + KVCacheSpec,
24 + get_kv_quant_mode,
25 +)
26
27 logger = init_logger(__name__)
28
@@ -220,4 +224,5 @@ class CrossAttention(Attention):
224 num_kv_heads=self.num_kv_heads,
225 head_size=self.head_size,
226 dtype=self.kv_cache_torch_dtype,
227 + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
228 )
vllm/model_executor/layers/attention/static_sink_attention.py
+2
@@ -26,6 +26,7 @@ from vllm.v1.kv_cache_interface import (
26 AttentionSpec,
27 KVCacheSpec,
28 SinkFullAttentionSpec,
29 + get_kv_quant_mode,
30 )
31
32 logger = init_logger(__name__)
@@ -217,6 +218,7 @@ class StaticSinkAttention(Attention, CustomOp):
218 head_size_v=self.head_size_v,
219 sink_len=self.sink_len,
220 dtype=self.kv_cache_torch_dtype,
221 + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
222 )
223
224
vllm/model_executor/layers/quantization/kv_cache.py
+15
@@ -10,6 +10,7 @@ from vllm.model_executor.layers.quantization.base_config import (
10 )
11 from vllm.platforms import current_platform
12 from vllm.utils.torch_utils import is_quantized_kv_cache
13 +from vllm.v1.kv_cache_interface import kv_cache_uses_per_token_head_scales
14
15 logger = init_logger(__name__)
16
@@ -53,6 +54,20 @@ class BaseKVCacheMethod(QuantizeMethodBase):
54 assert not hasattr(layer, "prob_scale")
55 return
56
57 + # Per-token-head quantized KV cache: scales are computed dynamically
58 + # per (token, head) in the kernel at cache-write time. Checkpoint
59 + # scales are never used regardless of calculate_kv_scales.
60 + if kv_cache_uses_per_token_head_scales(layer.kv_cache_dtype):
61 + layer._k_scale.copy_(1.0)
62 + layer._v_scale.copy_(1.0)
63 + layer._k_scale_float = 1.0
64 + layer._v_scale_float = 1.0
65 + del layer.k_scale
66 + del layer.v_scale
67 + del layer.q_scale
68 + del layer.prob_scale
69 + return
70 +
71 # If the kv-cache is not quantized, we enforce the k/v_scale to be 1.0
72 # regardless whether the kv-scale is available in the checkpoint.
73 # No need to process kv scales after loading if we are going to
vllm/platforms/interface.py
+5
@@ -505,6 +505,7 @@ class Platform:
505 FullAttentionSpec,
506 MambaSpec,
507 MLAAttentionSpec,
508 + get_kv_quant_mode,
509 )
510
511 cache_config = vllm_config.cache_config
@@ -516,6 +517,8 @@ class Platform:
517 else:
518 kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype]
519
520 + kv_quant_mode = get_kv_quant_mode(cache_config.cache_dtype)
521 +
522 # Compute attention page size for 1 token
523 if model_config.use_mla:
524 attn_page_size_1_token = MLAAttentionSpec(
@@ -523,6 +526,7 @@ class Platform:
526 num_kv_heads=model_config.get_num_kv_heads(parallel_config),
527 head_size=model_config.get_head_size(),
528 dtype=kv_cache_dtype,
529 + kv_quant_mode=kv_quant_mode,
530 ).page_size_bytes
531 else:
532 attn_page_size_1_token = FullAttentionSpec(
@@ -530,6 +534,7 @@ class Platform:
534 num_kv_heads=model_config.get_num_kv_heads(parallel_config),
535 head_size=model_config.get_head_size(),
536 dtype=kv_cache_dtype,
537 + kv_quant_mode=kv_quant_mode,
538 ).page_size_bytes
539
540 # Compute mamba page size
vllm/utils/torch_utils.py
+8 -1
@@ -37,6 +37,8 @@ STR_DTYPE_TO_TORCH_DTYPE = {
37 "fp8_e4m3": torch.uint8,
38 "fp8_e5m2": torch.uint8,
39 "int8": torch.int8,
40 + "int8_per_token_head": torch.int8,
41 + "fp8_per_token_head": torch.uint8,
42 "fp8_inc": torch.float8_e4m3fn,
43 "fp8_ds_mla": torch.uint8,
44 }
@@ -62,7 +64,12 @@ T = TypeVar("T")
64
65
66 def is_quantized_kv_cache(kv_cache_dtype: str) -> bool:
65 - return kv_cache_dtype.startswith("fp8")
67 + return kv_cache_dtype.startswith("fp8") or kv_cache_dtype.endswith("per_token_head")
68 +
69 +
70 +def kv_cache_uses_per_token_head_scales(kv_cache_dtype: str) -> bool:
71 + """Return True if *kv_cache_dtype* needs per-token-head scales."""
72 + return kv_cache_dtype.endswith("per_token_head")
73
74
75 def is_strictly_contiguous(t: torch.Tensor) -> bool:
vllm/v1/attention/backend.py
+10 -1
@@ -17,7 +17,9 @@ if TYPE_CHECKING:
17 from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
18 from vllm.platforms.interface import DeviceCapability
19 from vllm.v1.attention.backends.utils import KVCacheLayoutType
20 - from vllm.v1.kv_cache_interface import AttentionSpec
20 + from vllm.v1.kv_cache_interface import AttentionSpec, KVQuantMode
21 +
22 +from vllm.v1.kv_cache_interface import get_kv_quant_mode
23
24
25 class AttentionType(str, Enum):
@@ -740,6 +742,13 @@ class AttentionImplBase(ABC, Generic[T]):
742 class AttentionImpl(AttentionImplBase[T], Generic[T]):
743 """Standard attention implementation with forward method."""
744
745 + kv_cache_dtype: str
746 +
747 + @property
748 + def kv_quant_mode(self) -> "KVQuantMode":
749 + """Return the KV cache quantization mode for this layer."""
750 + return get_kv_quant_mode(self.kv_cache_dtype)
751 +
752 @abstractmethod
753 def __init__(
754 self,
vllm/v1/attention/backends/triton_attn.py
+132 -18
@@ -33,9 +33,14 @@ from vllm.v1.attention.backends.utils import get_kv_cache_layout
33 from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd
34 from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
35 triton_reshape_and_cache_flash,
36 + triton_reshape_and_cache_flash_per_token_head_quant,
37 )
38 from vllm.v1.attention.ops.triton_unified_attention import unified_attention
38 -from vllm.v1.kv_cache_interface import AttentionSpec
39 +from vllm.v1.kv_cache_interface import (
40 + AttentionSpec,
41 + get_kv_quant_mode,
42 + kv_cache_uses_per_token_head_scales,
43 +)
44
45 logger = init_logger(__name__)
46
@@ -270,6 +275,8 @@ class TritonAttentionBackend(AttentionBackend):
275 "fp8",
276 "fp8_e4m3",
277 "fp8_e5m2",
278 + "int8_per_token_head",
279 + "fp8_per_token_head",
280 ]
281
282 @staticmethod
@@ -302,6 +309,18 @@ class TritonAttentionBackend(AttentionBackend):
309 ) -> tuple[int, ...]:
310 if block_size % 16 != 0:
311 raise ValueError("Block size must be a multiple of 16.")
312 + if kv_cache_uses_per_token_head_scales(cache_dtype_str):
313 + # Pad head_size by sizeof(float32)/sizeof(cache_dtype) so
314 + # the per-head scale fits inline. The backend extracts
315 + # data[:head_size] and scale[head_size:] via typed views.
316 + from vllm.utils.torch_utils import (
317 + STR_DTYPE_TO_TORCH_DTYPE,
318 + get_dtype_size,
319 + )
320 +
321 + cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_dtype_str]
322 + scale_pad = get_dtype_size(torch.float32) // get_dtype_size(cache_dtype)
323 + return (num_blocks, 2, block_size, num_kv_heads, head_size + scale_pad)
324 return (num_blocks, 2, block_size, num_kv_heads, head_size)
325
326 @staticmethod
@@ -365,6 +384,62 @@ class TritonAttentionBackend(AttentionBackend):
384
385
386 class TritonAttentionImpl(AttentionImpl):
387 + # Per-token-head quant: scale views carved from inline head padding.
388 + _k_scale_cache: torch.Tensor | None = None
389 + _v_scale_cache: torch.Tensor | None = None
390 +
391 + def _ensure_scale_caches(self, kv_cache: torch.Tensor) -> None:
392 + """Extract per-head scale views from the padded head dimension.
393 +
394 + The KV cache shape is ``(num_blocks, 2, block_size, nkv, hs+pad)``
395 + where ``pad = sizeof(float32) / sizeof(cache_dtype)``. The last
396 + ``pad`` elements of each head hold one float32 scale. We create
397 + strided float32 views over those bytes.
398 +
399 + Scale shape: ``(num_blocks, block_size, num_kv_heads)``
400 + """
401 + if self._k_scale_cache is not None:
402 + return
403 + from vllm.utils.torch_utils import get_dtype_size
404 +
405 + num_blocks, _, block_size, nkv, padded_hs = kv_cache.shape
406 + dtype_sz = kv_cache.element_size()
407 + scale_pad = get_dtype_size(torch.float32) // dtype_sz # e.g. 4
408 + hs = padded_hs - scale_pad
409 +
410 + raw = kv_cache.untyped_storage()
411 + base_f32 = torch.tensor([], dtype=torch.float32, device=kv_cache.device).set_(
412 + raw
413 + )
414 +
415 + # In the raw bytes, each (block, kv_half, slot, head) occupies
416 + # padded_hs * dtype_sz bytes. The scale float32 sits at byte
417 + # offset hs * dtype_sz within that region.
418 + kv_half_bytes = block_size * nkv * padded_hs * dtype_sz
419 + full_block_f32 = 2 * kv_half_bytes // 4 # stride between blocks
420 + slot_f32 = nkv * padded_hs * dtype_sz // 4 # stride between slots
421 + head_f32 = padded_hs * dtype_sz // 4 # stride between heads
422 + scale_off_f32 = hs * dtype_sz // 4 # offset to scale within head
423 +
424 + # K scales: kv_half=0
425 + self._k_scale_cache = torch.as_strided(
426 + base_f32,
427 + size=(num_blocks, block_size, nkv),
428 + stride=(full_block_f32, slot_f32, head_f32),
429 + storage_offset=scale_off_f32,
430 + )
431 + self._k_scale_cache.fill_(1.0)
432 +
433 + # V scales: kv_half=1, offset by kv_half_bytes
434 + v_base_f32 = kv_half_bytes // 4
435 + self._v_scale_cache = torch.as_strided(
436 + base_f32,
437 + size=(num_blocks, block_size, nkv),
438 + stride=(full_block_f32, slot_f32, head_f32),
439 + storage_offset=v_base_f32 + scale_off_f32,
440 + )
441 + self._v_scale_cache.fill_(1.0)
442 +
443 def fused_output_quant_supported(self, quant_key: QuantKey):
444 return quant_key == kFp8StaticTensorSym
445
@@ -418,6 +493,9 @@ class TritonAttentionImpl(AttentionImpl):
493 self.use_alibi_sqrt = use_alibi_sqrt
494 self.supports_quant_query_input = current_platform.is_cuda()
495
496 + self._kv_quant_mode = get_kv_quant_mode(kv_cache_dtype)
497 + self._is_per_token_head_quant = self._kv_quant_mode.is_per_token_head
498 +
499 def forward(
500 self,
501 layer: torch.nn.Module,
@@ -480,15 +558,35 @@ class TritonAttentionImpl(AttentionImpl):
558 layer,
559 )
560
483 - # For decoder and cross-attention, use KV cache as before
484 - key_cache, value_cache = kv_cache.unbind(1)
485 - if is_quantized_kv_cache(self.kv_cache_dtype):
486 - if key_cache.dtype != self.fp8_dtype:
561 + # Per-token-head quantized KV cache: use separate scale caches.
562 + if self._is_per_token_head_quant:
563 + self._ensure_scale_caches(kv_cache)
564 + key_cache, value_cache = kv_cache.unbind(1)
565 + if key_cache.dtype == torch.uint8:
566 key_cache = key_cache.view(self.fp8_dtype)
567 value_cache = value_cache.view(self.fp8_dtype)
489 - assert layer._q_scale_float == 1.0, (
490 - "A non 1.0 q_scale is not currently supported."
568 + k_descale = None
569 + v_descale = None
570 + k_scale_cache = self._k_scale_cache
571 + v_scale_cache = self._v_scale_cache
572 + # FP8 per-tensor / auto path (original flow).
573 + else:
574 + key_cache, value_cache = kv_cache.unbind(1)
575 + if is_quantized_kv_cache(self.kv_cache_dtype):
576 + if key_cache.dtype != self.fp8_dtype:
577 + key_cache = key_cache.view(self.fp8_dtype)
578 + value_cache = value_cache.view(self.fp8_dtype)
579 + assert layer._q_scale_float == 1.0, (
580 + "A non 1.0 q_scale is not currently supported."
581 + )
582 + descale_shape = (
583 + attn_metadata.query_start_loc.shape[0] - 1,
584 + key_cache.shape[2],
585 )
586 + k_descale = layer._k_scale.expand(descale_shape)
587 + v_descale = layer._v_scale.expand(descale_shape)
588 + k_scale_cache = None
589 + v_scale_cache = None
590
591 cu_seqlens_q = attn_metadata.query_start_loc
592 seqused_k = attn_metadata.seq_lens
@@ -502,7 +600,6 @@ class TritonAttentionImpl(AttentionImpl):
600 softmax_segm_max = attn_metadata.softmax_segm_max
601 softmax_segm_expsum = attn_metadata.softmax_segm_expsum
602
505 - descale_shape = (cu_seqlens_q.shape[0] - 1, key_cache.shape[2])
603 mm_prefix_range_tensor = attn_metadata.mm_prefix_range_tensor
604
605 unified_attention(
@@ -522,8 +619,8 @@ class TritonAttentionImpl(AttentionImpl):
619 block_table=block_table,
620 softcap=self.logits_soft_cap,
621 q_descale=None, # Not supported
525 - k_descale=layer._k_scale.expand(descale_shape),
526 - v_descale=layer._v_scale.expand(descale_shape),
622 + k_descale=k_descale,
623 + v_descale=v_descale,
624 seq_threshold_3D=seq_threshold_3D,
625 num_par_softmax_segments=num_par_softmax_segments,
626 softmax_segm_output=softmax_segm_output,
@@ -532,6 +629,9 @@ class TritonAttentionImpl(AttentionImpl):
629 sinks=self.sinks,
630 output_scale=output_scale,
631 mm_prefix_range=mm_prefix_range_tensor,
632 + kv_quant_mode=self._kv_quant_mode,
633 + k_scale_cache=k_scale_cache,
634 + v_scale_cache=v_scale_cache,
635 )
636
637 return output
@@ -555,10 +655,10 @@ class TritonAttentionImpl(AttentionImpl):
655 attn_metadata: Encoder attention metadata
656 layer: The attention layer
657 """
558 - # For encoder attention, process FP8 quantization if needed
658 + # Quantized KV cache is not supported for encoder attention.
659 if is_quantized_kv_cache(self.kv_cache_dtype):
660 raise NotImplementedError(
561 - "quantization is not supported for encoder attention"
661 + "quantized KV cache is not supported for encoder attention"
662 )
663
664 # Use encoder-specific metadata for sequence information
@@ -594,16 +694,28 @@ class TritonAttentionImpl(AttentionImpl):
694 # For encoder attention,
695 # we use direct Q, K, V tensors without caching
696 return
597 - # For decoder and cross-attention, use KV cache as before
598 - key_cache, value_cache = kv_cache.unbind(1)
599 -
697 # Reshape the input keys and values and store them in the cache.
698 + if self._is_per_token_head_quant:
699 + self._ensure_scale_caches(kv_cache)
700 + key_cache, value_cache = kv_cache.unbind(1)
701 + if key_cache.dtype == torch.uint8:
702 + key_cache = key_cache.view(self.fp8_dtype)
703 + value_cache = value_cache.view(self.fp8_dtype)
704 + triton_reshape_and_cache_flash_per_token_head_quant(
705 + key,
706 + value,
707 + key_cache,
708 + value_cache,
709 + self._k_scale_cache,
710 + self._v_scale_cache,
711 + slot_mapping,
712 + )
713 + return
714 + # For decoder and cross-attention, use KV cache as before.
715 + key_cache, value_cache = kv_cache.unbind(1)
716 if is_quantized_kv_cache(self.kv_cache_dtype):
717 key_cache = key_cache.view(self.fp8_dtype)
718 value_cache = value_cache.view(self.fp8_dtype)
604 - # triton kernel does not support uint8 kv_cache
605 - # (because some explicit casts (e.g. float8_e4m3fnuz)
606 - # are not supported)
719 triton_reshape_and_cache_flash(
720 key,
721 value,
@@ -616,6 +728,8 @@ class TritonAttentionImpl(AttentionImpl):
728 )
729
730 def fused_rope_kvcache_supported(self):
731 + if self._is_per_token_head_quant:
732 + return False
733 return rocm_aiter_ops.is_enabled()
734
735 def do_rope_and_kv_cache_update(
vllm/v1/attention/ops/triton_reshape_and_cache_flash.py
+198 -1
@@ -3,10 +3,16 @@
3
4 import torch
5
6 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
7 + FP8_DTYPE,
8 + get_fp8_min_max,
9 +)
10 from vllm.platforms import current_platform
11 from vllm.triton_utils import tl, triton
12 from vllm.utils.torch_utils import is_quantized_kv_cache
13
14 +FP8_MIN, FP8_MAX = get_fp8_min_max()
15 +
16
17 @triton.jit
18 def reshape_and_cache_kernel_flash(
@@ -118,6 +124,198 @@ def reshape_and_cache_kernel_flash(
124 return
125
126
127 +# ---------------------------------------------------------------------------
128 +# Per-token-head dynamic quantization kernel
129 +# Grid: (num_tokens, NUM_KV_HEADS)
130 +# Each program handles one (token, head) pair:
131 +# 1. Loads K (or V) for that single head
132 +# 2. Computes absmax across head_size → scale = absmax / QUANT_MAX
133 +# 3. Quantizes and stores the data + per-head scale
134 +#
135 +# Parametrised by QUANT_MAX / QUANT_MIN so the same code path works
136 +# for int8 (±127/128), fp8_e4m3 (±448), and other formats.
137 +# ---------------------------------------------------------------------------
138 +@triton.jit
139 +def _reshape_cache_per_token_head(
140 + key_ptr, # [num_tokens, num_kv_heads, head_size]
141 + value_ptr, # [num_tokens, num_kv_heads, head_size_v]
142 + key_cache_ptr, # [num_blocks, block_size, num_kv_heads, head_size]
143 + value_cache_ptr, # [num_blocks, block_size, num_kv_heads, head_size_v]
144 + k_scale_cache_ptr, # [num_blocks, block_size, num_kv_heads] float32
145 + v_scale_cache_ptr, # [num_blocks, block_size, num_kv_heads] float32
146 + slot_mapping_ptr, # [num_tokens]
147 + stride_key_tok: tl.int64,
148 + stride_key_head: tl.int64,
149 + stride_val_tok: tl.int64,
150 + stride_val_head: tl.int64,
151 + stride_kc_blk: tl.int64, # key_cache stride over blocks
152 + stride_kc_slot: tl.int64, # key_cache stride over slots
153 + stride_kc_head: tl.int64, # key_cache stride over heads
154 + stride_vc_blk: tl.int64,
155 + stride_vc_slot: tl.int64,
156 + stride_vc_head: tl.int64,
157 + stride_ks_blk: tl.int64, # k_scale_cache stride[0] (blocks)
158 + stride_ks_slot: tl.int64, # k_scale_cache stride[1] (slots)
159 + stride_ks_head: tl.int64, # k_scale_cache stride[2] (heads)
160 + stride_vs_blk: tl.int64, # v_scale_cache stride[0] (blocks)
161 + stride_vs_slot: tl.int64, # v_scale_cache stride[1] (slots)
162 + stride_vs_head: tl.int64, # v_scale_cache stride[2] (heads)
163 + block_size: tl.constexpr,
164 + head_size: tl.constexpr,
165 + head_size_v: tl.constexpr,
166 + HEAD_SIZE_PADDED: tl.constexpr, # next_power_of_2(max(head_size, head_size_v))
167 + QUANT_MAX: tl.constexpr = 127.0,
168 + QUANT_MIN: tl.constexpr = -128.0,
169 +):
170 + tok = tl.program_id(0)
171 + head = tl.program_id(1)
172 +
173 + slot = tl.load(slot_mapping_ptr + tok).to(tl.int64)
174 + if slot < 0:
175 + return
176 +
177 + blk = slot // block_size
178 + slot_in_blk = slot % block_size
179 +
180 + dim_offs = tl.arange(0, HEAD_SIZE_PADDED)
181 +
182 + # ---- Key: load one head → absmax → quantize → store -------------------
183 + k_mask = dim_offs < head_size
184 + k_h = tl.load(
185 + key_ptr + tok * stride_key_tok + head * stride_key_head + dim_offs,
186 + mask=k_mask,
187 + other=0.0,
188 + ).to(tl.float32)
189 +
190 + k_scale = tl.maximum(tl.max(tl.abs(k_h)) / QUANT_MAX, 1e-6)
191 + tl.store(
192 + k_scale_cache_ptr
193 + + blk * stride_ks_blk
194 + + slot_in_blk * stride_ks_slot
195 + + head * stride_ks_head,
196 + k_scale,
197 + )
198 +
199 + k_q = tl.clamp(k_h * (1.0 / k_scale), QUANT_MIN, QUANT_MAX)
200 + tl.store(
201 + key_cache_ptr
202 + + blk * stride_kc_blk
203 + + slot_in_blk * stride_kc_slot
204 + + head * stride_kc_head
205 + + dim_offs,
206 + k_q,
207 + mask=k_mask,
208 + )
209 +
210 + # ---- Value: same per-head approach ------------------------------------
211 + v_mask = dim_offs < head_size_v
212 + v_h = tl.load(
213 + value_ptr + tok * stride_val_tok + head * stride_val_head + dim_offs,
214 + mask=v_mask,
215 + other=0.0,
216 + ).to(tl.float32)
217 +
218 + v_scale = tl.maximum(tl.max(tl.abs(v_h)) / QUANT_MAX, 1e-6)
219 + tl.store(
220 + v_scale_cache_ptr
221 + + blk * stride_vs_blk
222 + + slot_in_blk * stride_vs_slot
223 + + head * stride_vs_head,
224 + v_scale,
225 + )
226 +
227 + v_q = tl.clamp(v_h * (1.0 / v_scale), QUANT_MIN, QUANT_MAX)
228 + tl.store(
229 + value_cache_ptr
230 + + blk * stride_vc_blk
231 + + slot_in_blk * stride_vc_slot
232 + + head * stride_vc_head
233 + + dim_offs,
234 + v_q,
235 + mask=v_mask,
236 + )
237 +
238 +
239 +# Mapping from cache torch dtype to (QUANT_MAX, QUANT_MIN) for the
240 +# per-token-head quantization kernel.
241 +_PER_TOKEN_HEAD_QUANT_PARAMS: dict[torch.dtype, tuple[float, float]] = {
242 + torch.int8: (127.0, -128.0),
243 + FP8_DTYPE: (FP8_MAX, FP8_MIN),
244 +}
245 +
246 +
247 +def triton_reshape_and_cache_flash_per_token_head_quant(
248 + key: torch.Tensor, # [num_tokens, num_kv_heads, head_size]
249 + value: torch.Tensor, # [num_tokens, num_kv_heads, head_size_v]
250 + key_cache: torch.Tensor, # [num_blocks, block_size, num_kv_heads, head_size]
251 + value_cache: torch.Tensor, # [num_blocks, block_size, num_kv_heads, head_size_v]
252 + k_scale_cache: torch.Tensor, # [num_blocks, block_size, num_kv_heads] float32
253 + v_scale_cache: torch.Tensor, # [num_blocks, block_size, num_kv_heads] float32
254 + slot_mapping: torch.Tensor, # [num_tokens]
255 +):
256 + """Quantize key/value per (token, head) and write to paged cache.
257 +
258 + Computes one scale = absmax / QUANT_MAX per (token, head), stores
259 + quantized data in key_cache/value_cache, and stores the float32
260 + scale in k_scale_cache/v_scale_cache.
261 +
262 + The quantization range (QUANT_MAX, QUANT_MIN) is derived from the
263 + cache tensor dtype so the same code path works for int8 and fp8.
264 + """
265 + cache_dtype = key_cache.dtype
266 + quant_params = _PER_TOKEN_HEAD_QUANT_PARAMS.get(cache_dtype)
267 + if quant_params is None:
268 + raise ValueError(
269 + f"Per-token-head quantization not supported for cache dtype "
270 + f"{cache_dtype}. Supported: {list(_PER_TOKEN_HEAD_QUANT_PARAMS)}"
271 + )
272 + quant_max, quant_min = quant_params
273 +
274 + num_tokens, num_kv_heads, head_size = key.shape
275 + head_size_v = value.shape[2]
276 + head_size_padded = triton.next_power_of_2(max(head_size, head_size_v))
277 +
278 + block_size = key_cache.shape[1]
279 +
280 + if current_platform.is_rocm() or current_platform.is_xpu():
281 + num_warps = 4
282 + else:
283 + num_warps = min(16, max(1, head_size_padded // 32))
284 +
285 + _reshape_cache_per_token_head[(num_tokens, num_kv_heads)](
286 + key_ptr=key,
287 + value_ptr=value,
288 + key_cache_ptr=key_cache,
289 + value_cache_ptr=value_cache,
290 + k_scale_cache_ptr=k_scale_cache,
291 + v_scale_cache_ptr=v_scale_cache,
292 + slot_mapping_ptr=slot_mapping,
293 + stride_key_tok=key.stride(0),
294 + stride_key_head=key.stride(1),
295 + stride_val_tok=value.stride(0),
296 + stride_val_head=value.stride(1),
297 + stride_kc_blk=key_cache.stride(0),
298 + stride_kc_slot=key_cache.stride(1),
299 + stride_kc_head=key_cache.stride(2),
300 + stride_vc_blk=value_cache.stride(0),
301 + stride_vc_slot=value_cache.stride(1),
302 + stride_vc_head=value_cache.stride(2),
303 + stride_ks_blk=k_scale_cache.stride(0),
304 + stride_ks_slot=k_scale_cache.stride(1),
305 + stride_ks_head=k_scale_cache.stride(2),
306 + stride_vs_blk=v_scale_cache.stride(0),
307 + stride_vs_slot=v_scale_cache.stride(1),
308 + stride_vs_head=v_scale_cache.stride(2),
309 + block_size=block_size,
310 + head_size=head_size,
311 + head_size_v=head_size_v,
312 + HEAD_SIZE_PADDED=head_size_padded,
313 + QUANT_MAX=quant_max,
314 + QUANT_MIN=quant_min,
315 + num_warps=num_warps,
316 + )
317 +
318 +
319 def triton_reshape_and_cache_flash(
320 key: torch.Tensor, # [num_tokens, num_heads, head_size]
321 value: torch.Tensor, # [num_tokens, num_heads, head_size]
@@ -224,7 +422,6 @@ def triton_reshape_and_cache_flash(
422 block_size=block_size,
423 x=x,
424 USE_HEAD_MAJOR_LAYOUT=use_head_major_layout,
227 - # FP8 flags
425 FP8_KV_CACHE=FP8_KV_CACHE,
426 # autotune parameters
427 TILE_SIZE=TILE_SIZE,
vllm/v1/attention/ops/triton_unified_attention.py
+189 -36
@@ -13,6 +13,7 @@ import vllm.envs as envs
13 from vllm.logger import init_logger
14 from vllm.platforms import current_platform
15 from vllm.triton_utils import tl, triton
16 +from vllm.v1.kv_cache_interface import KVQuantMode
17
18 logger = init_logger(__name__)
19 is_batch_invariant = envs.VLLM_BATCH_INVARIANT
@@ -32,6 +33,63 @@ def apply_softcap(S, x):
33 return x * (p1 - p2) / (p1 + p2)
34
35
36 +@triton.jit
37 +def _prepare_kv_tile(
38 + data,
39 + Q,
40 + tensor_scale,
41 + scale_cache_ptr,
42 + physical_block_idx,
43 + seq_offset,
44 + kv_head_idx,
45 + stride_s_blk,
46 + stride_s_slot,
47 + stride_s_head,
48 + tile_mask,
49 + BLOCK_SIZE: tl.constexpr,
50 + KV_QUANT_MODE: tl.constexpr,
51 +):
52 + """Prepare a loaded KV tile for attention computation.
53 +
54 + Casts the raw KV data to Q's dtype and loads per-token-head scales
55 + when applicable:
56 +
57 + - ``KV_QUANT_MODE == 0``: cast only (no-op for bf16/fp16).
58 + - ``KV_QUANT_MODE == 1`` (FP8 per-tensor): dequantize inline
59 + using the tensor-wide scale.
60 + - ``KV_QUANT_MODE >= 2`` (per-token-head int8/fp8): cast to Q's
61 + dtype and return per-head scales separately — the caller applies
62 + them after the dot product for better numerical efficiency.
63 +
64 + Returns ``(data, token_head_scales)``. *token_head_scales* is only
65 + meaningful when ``KV_QUANT_MODE >= 2``; callers gate its use on
66 + the same constexpr so the compiler eliminates dead code.
67 + """
68 + # KV_QUANT_MODE values: 0=none, 1=fp8 per-tensor,
69 + # 2=int8 per-token-head, 3=fp8 per-token-head
70 +
71 + # Placeholder scales (float32) — never read when KV_QUANT_MODE < 2.
72 + unused_scales = tile_mask.to(tl.float32)
73 +
74 + if KV_QUANT_MODE == 1: # FP8 per-tensor
75 + if Q.dtype.is_fp8():
76 + return data.to(Q.dtype), unused_scales
77 + return (data.to(tl.float32) * tl.load(tensor_scale)).to(Q.dtype), unused_scales
78 + if KV_QUANT_MODE >= 2: # per-token-head (int8 or fp8)
79 + scale_idx = (
80 + physical_block_idx * stride_s_blk
81 + + (seq_offset % BLOCK_SIZE) * stride_s_slot
82 + + kv_head_idx * stride_s_head
83 + )
84 + token_head_scales = tl.load(
85 + scale_cache_ptr + scale_idx, mask=tile_mask, other=1.0
86 + )
87 + return data.to(Q.dtype), token_head_scales
88 + # .to(Q.dtype) is a no-op when data is already Q's type (bf16/fp16),
89 + # but required so Triton sees consistent return types across branches.
90 + return data.to(Q.dtype), unused_scales
91 +
92 +
93 @triton.jit
94 def find_seq_idx(
95 query_start_len_ptr,
@@ -105,8 +163,20 @@ def kernel_unified_attention_2d(
163 num_seqs: tl.int32,
164 BLOCK_M: tl.constexpr, # int
165 USE_FP8: tl.constexpr, # bool
166 + # KV cache quantization: 0=none, 1=fp8, 2=per-token-head
167 + KV_QUANT_MODE: tl.constexpr = 0,
168 FP8_MIN: tl.constexpr = float8_info.min,
169 FP8_MAX: tl.constexpr = float8_info.max,
170 + # Per-token-head scale caches (KV_QUANT_MODE >= 2)
171 + # Shape: [num_blocks, block_size, num_kv_heads]
172 + k_scale_cache_ptr=None,
173 + v_scale_cache_ptr=None,
174 + stride_ks_blk=0,
175 + stride_ks_slot=0,
176 + stride_ks_head=0,
177 + stride_vs_blk=0,
178 + stride_vs_slot=0,
179 + stride_vs_head=0,
180 ):
181 q_block_global_idx = tl.program_id(0)
182 kv_head_idx = tl.program_id(1)
@@ -258,14 +328,21 @@ def kernel_unified_attention_2d(
328 mask=dim_mask[:, None] & tile_mask[None, :],
329 other=0.0,
330 )
261 -
262 - if K_load.dtype.is_fp8():
263 - if Q.dtype.is_fp8():
264 - K = K_load
265 - else:
266 - K = (K_load.to(tl.float32) * tl.load(k_scale)).to(Q.dtype)
267 - else:
268 - K = K_load
331 + K, k_token_head_scales = _prepare_kv_tile(
332 + K_load,
333 + Q,
334 + k_scale,
335 + k_scale_cache_ptr,
336 + physical_block_idx,
337 + seq_offset,
338 + kv_head_idx,
339 + stride_ks_blk,
340 + stride_ks_slot,
341 + stride_ks_head,
342 + tile_mask,
343 + BLOCK_SIZE,
344 + KV_QUANT_MODE,
345 + )
346
347 # V : (TILE_SIZE, HEAD_SIZE)
348 V_load = tl.load(
@@ -273,14 +350,21 @@ def kernel_unified_attention_2d(
350 mask=dim_mask[None, :] & tile_mask[:, None],
351 other=0.0,
352 )
276 -
277 - if V_load.dtype.is_fp8():
278 - if Q.dtype.is_fp8():
279 - V = V_load
280 - else:
281 - V = (V_load.to(tl.float32) * tl.load(v_scale)).to(Q.dtype)
282 - else:
283 - V = V_load
353 + V, v_token_head_scales = _prepare_kv_tile(
354 + V_load,
355 + Q,
356 + v_scale,
357 + v_scale_cache_ptr,
358 + physical_block_idx,
359 + seq_offset,
360 + kv_head_idx,
361 + stride_vs_blk,
362 + stride_vs_slot,
363 + stride_vs_head,
364 + tile_mask,
365 + BLOCK_SIZE,
366 + KV_QUANT_MODE,
367 + )
368
369 # Compute attention mask: causal by default (key <= query)
370 query_abs_pos = context_len + query_pos[:, None]
@@ -318,7 +402,12 @@ def kernel_unified_attention_2d(
402 # S : (BLOCK_M, TILE_SIZE)
403 S = tl.zeros(shape=(BLOCK_M, TILE_SIZE), dtype=tl.float32)
404
321 - S += scale * tl.dot(Q, K)
405 + # Per-token-head quant: fuse softmax_scale with per-head k_scale
406 + # to avoid a separate BLOCK_M × TILE_SIZE multiply on S.
407 + if KV_QUANT_MODE >= 2:
408 + S += tl.dot(Q, K) * (scale * k_token_head_scales[None, :])
409 + else:
410 + S += scale * tl.dot(Q, K)
411
412 if USE_SOFTCAP:
413 S = apply_softcap(S, softcap)
@@ -382,7 +471,12 @@ def kernel_unified_attention_2d(
471 )
472
473 # acc : (BLOCK_M, HEAD_SIZE_PADDED)
385 - acc += tl.dot(P.to(V.dtype), V)
474 + # Per-token-head quant: apply v_scale to P instead of V.
475 + if KV_QUANT_MODE >= 2:
476 + P_v = (P * v_token_head_scales[None, :]).to(V.dtype)
477 + acc += tl.dot(P_v, V)
478 + else:
479 + acc += tl.dot(P.to(V.dtype), V)
480
481 # epilogue
482 acc = acc / L[:, None]
@@ -453,6 +547,18 @@ def kernel_unified_attention_3d(
547 USE_MM_PREFIX: tl.constexpr, # bool
548 MAX_MM_RANGES: tl.constexpr, # int
549 mm_prefix_range_ptr, # [num_seqs] - prefix length for each sequence
550 + # KV cache quantization: 0=none, 1=fp8, 2=per-token-head
551 + KV_QUANT_MODE: tl.constexpr = 0,
552 + # Per-token-head scale caches (KV_QUANT_MODE >= 2)
553 + # Shape: [num_blocks, block_size, num_kv_heads]
554 + k_scale_cache_ptr=None,
555 + v_scale_cache_ptr=None,
556 + stride_ks_blk=0,
557 + stride_ks_slot=0,
558 + stride_ks_head=0,
559 + stride_vs_blk=0,
560 + stride_vs_slot=0,
561 + stride_vs_head=0,
562 ):
563 q_block_global_idx = tl.program_id(0)
564 kv_head_idx = tl.program_id(1)
@@ -613,14 +719,21 @@ def kernel_unified_attention_3d(
719 mask=dim_mask[:, None] & tile_mask[None, :],
720 other=0.0,
721 )
616 -
617 - if K_load.dtype.is_fp8():
618 - if Q.dtype.is_fp8():
619 - K = K_load
620 - else:
621 - K = (K_load.to(tl.float32) * tl.load(k_scale)).to(Q.dtype)
622 - else:
623 - K = K_load
722 + K, k_token_head_scales = _prepare_kv_tile(
723 + K_load,
724 + Q,
725 + k_scale,
726 + k_scale_cache_ptr,
727 + physical_block_idx,
728 + seq_offset,
729 + kv_head_idx,
730 + stride_ks_blk,
731 + stride_ks_slot,
732 + stride_ks_head,
733 + tile_mask,
734 + BLOCK_SIZE,
735 + KV_QUANT_MODE,
736 + )
737
738 # V : (TILE_SIZE, HEAD_SIZE)
739 V_load = tl.load(
@@ -628,14 +741,21 @@ def kernel_unified_attention_3d(
741 mask=dim_mask[None, :] & tile_mask[:, None],
742 other=0.0,
743 )
631 -
632 - if V_load.dtype.is_fp8():
633 - if Q.dtype.is_fp8():
634 - V = V_load
635 - else:
636 - V = (V_load.to(tl.float32) * tl.load(v_scale)).to(Q.dtype)
637 - else:
638 - V = V_load
744 + V, v_token_head_scales = _prepare_kv_tile(
745 + V_load,
746 + Q,
747 + v_scale,
748 + v_scale_cache_ptr,
749 + physical_block_idx,
750 + seq_offset,
751 + kv_head_idx,
752 + stride_vs_blk,
753 + stride_vs_slot,
754 + stride_vs_head,
755 + tile_mask,
756 + BLOCK_SIZE,
757 + KV_QUANT_MODE,
758 + )
759
760 # Compute attention mask: causal by default (key <= query)
761 query_abs_pos = context_len + query_pos[:, None]
@@ -672,7 +792,13 @@ def kernel_unified_attention_3d(
792
793 # S : (BLOCK_M, TILE_SIZE)
794 S = tl.zeros(shape=(BLOCK_M, TILE_SIZE), dtype=tl.float32)
675 - S += scale * tl.dot(Q, K)
795 +
796 + # Per-token-head quant: fuse softmax_scale with per-head k_scale
797 + # to avoid a separate BLOCK_M × TILE_SIZE multiply on S.
798 + if KV_QUANT_MODE >= 2:
799 + S += tl.dot(Q, K) * (scale * k_token_head_scales[None, :])
800 + else:
801 + S += scale * tl.dot(Q, K)
802
803 if USE_SOFTCAP:
804 S = apply_softcap(S, softcap)
@@ -736,7 +862,12 @@ def kernel_unified_attention_3d(
862 )
863
864 # acc : (BLOCK_M, HEAD_SIZE_PADDED)
739 - acc += tl.dot(P.to(V.dtype), V)
865 + # Per-token-head quant: apply v_scale to P instead of V.
866 + if KV_QUANT_MODE >= 2:
867 + P_v = (P * v_token_head_scales[None, :]).to(V.dtype)
868 + acc += tl.dot(P_v, V)
869 + else:
870 + acc += tl.dot(P.to(V.dtype), V)
871
872 segm_output_offset = (
873 query_offset_0[:, None].to(tl.int64)
@@ -911,6 +1042,10 @@ def unified_attention(
1042 # Optional tensor for prefix lengths (PrefixLM support)
1043 mm_prefix_range=None,
1044 use_alibi_sqrt=False,
1045 + # KV cache quantization mode and per-token-head scale caches.
1046 + kv_quant_mode: KVQuantMode = KVQuantMode.NONE,
1047 + k_scale_cache=None, # [num_blocks, block_size, num_kv_heads] float32
1048 + v_scale_cache=None, # [num_blocks, block_size, num_kv_heads] float32
1049 ):
1050 assert causal, "Only causal attention is supported"
1051 assert q_descale is None, "Q scales not supported"
@@ -1040,6 +1175,15 @@ def unified_attention(
1175 num_seqs=num_seqs,
1176 BLOCK_M=BLOCK_M,
1177 USE_FP8=output_scale is not None,
1178 + KV_QUANT_MODE=kv_quant_mode,
1179 + k_scale_cache_ptr=k_scale_cache,
1180 + v_scale_cache_ptr=v_scale_cache,
1181 + stride_ks_blk=k_scale_cache.stride(0) if k_scale_cache is not None else 0,
1182 + stride_ks_slot=k_scale_cache.stride(1) if k_scale_cache is not None else 0,
1183 + stride_ks_head=k_scale_cache.stride(2) if k_scale_cache is not None else 0,
1184 + stride_vs_blk=v_scale_cache.stride(0) if v_scale_cache is not None else 0,
1185 + stride_vs_slot=v_scale_cache.stride(1) if v_scale_cache is not None else 0,
1186 + stride_vs_head=v_scale_cache.stride(2) if v_scale_cache is not None else 0,
1187 )
1188 else:
1189 kernel_unified_attention_3d[
@@ -1092,6 +1236,15 @@ def unified_attention(
1236 num_seqs=num_seqs,
1237 BLOCK_M=BLOCK_M,
1238 NUM_SEGMENTS_PER_SEQ=num_par_softmax_segments,
1239 + KV_QUANT_MODE=kv_quant_mode,
1240 + k_scale_cache_ptr=k_scale_cache,
1241 + v_scale_cache_ptr=v_scale_cache,
1242 + stride_ks_blk=k_scale_cache.stride(0) if k_scale_cache is not None else 0,
1243 + stride_ks_slot=k_scale_cache.stride(1) if k_scale_cache is not None else 0,
1244 + stride_ks_head=k_scale_cache.stride(2) if k_scale_cache is not None else 0,
1245 + stride_vs_blk=v_scale_cache.stride(0) if v_scale_cache is not None else 0,
1246 + stride_vs_slot=v_scale_cache.stride(1) if v_scale_cache is not None else 0,
1247 + stride_vs_head=v_scale_cache.stride(2) if v_scale_cache is not None else 0,
1248 )
1249 reduce_segments[(q.shape[0], num_query_heads)](
1250 output_ptr=out,
vllm/v1/kv_cache_interface.py
+61 -1
@@ -1,21 +1,70 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from __future__ import annotations
5 +
6 import copy
7 from dataclasses import dataclass, fields, replace
8 +from enum import IntEnum
9 from math import prod
10 +from typing import TYPE_CHECKING
11
12 import torch
13 from typing_extensions import Self
14
11 -from vllm.config import VllmConfig
15 from vllm.logger import init_logger
16 +
17 +if TYPE_CHECKING:
18 + from vllm.config import VllmConfig
19 from vllm.utils.math_utils import cdiv
20 from vllm.utils.torch_utils import get_dtype_size
21
22 logger = init_logger(__name__)
23
24
25 +# ---------------------------------------------------------------------------
26 +# KV cache quantization mode
27 +# ---------------------------------------------------------------------------
28 +
29 +
30 +class KVQuantMode(IntEnum):
31 + """KV cache quantization mode.
32 +
33 + Used by attention backends and kernels to dispatch quantization logic
34 + without string matching on ``kv_cache_dtype``.
35 + """
36 +
37 + NONE = 0
38 + FP8_PER_TENSOR = 1 # per-tensor scales (current fp8 path)
39 + INT8_PER_TOKEN_HEAD = 2 # per-token-head dynamic scales for int8
40 + FP8_PER_TOKEN_HEAD = 3 # per-token-head dynamic scales for fp8
41 +
42 + @property
43 + def is_per_token_head(self) -> bool:
44 + """True for any per-token-head quantization mode."""
45 + return self >= 2
46 +
47 +
48 +def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode:
49 + """Map a ``kv_cache_dtype`` string to a :class:`KVQuantMode`."""
50 + if kv_cache_dtype == "int8_per_token_head":
51 + return KVQuantMode.INT8_PER_TOKEN_HEAD
52 + if kv_cache_dtype == "fp8_per_token_head":
53 + return KVQuantMode.FP8_PER_TOKEN_HEAD
54 + if kv_cache_dtype.startswith("fp8"):
55 + return KVQuantMode.FP8_PER_TENSOR
56 + return KVQuantMode.NONE
57 +
58 +
59 +def is_quantized_kv_cache(kv_cache_dtype: str) -> bool:
60 + return get_kv_quant_mode(kv_cache_dtype) != KVQuantMode.NONE
61 +
62 +
63 +def kv_cache_uses_per_token_head_scales(kv_cache_dtype: str) -> bool:
64 + """Return True if *kv_cache_dtype* needs per-token-head scales."""
65 + return get_kv_quant_mode(kv_cache_dtype).is_per_token_head
66 +
67 +
68 @dataclass(frozen=True)
69 class KVCacheSpec:
70 """
@@ -66,11 +115,19 @@ class AttentionSpec(KVCacheSpec):
115 num_kv_heads: int
116 head_size: int
117 dtype: torch.dtype
118 + kv_quant_mode: KVQuantMode = KVQuantMode.NONE
119 page_size_padded: int | None = None
120
121 @property
122 def page_size_bytes(self) -> int:
123 real_page_size = self.real_page_size_bytes
124 + # Per-token-head scales are stored in separate tensors managed
125 + # by the attention backend, but the memory is carved from the
126 + # raw KV cache allocation so it must be budgeted here.
127 + if self.kv_quant_mode.is_per_token_head:
128 + real_page_size += (
129 + 2 * self.block_size * self.num_kv_heads * get_dtype_size(torch.float32)
130 + )
131 if self.page_size_padded is not None:
132 assert self.page_size_padded >= real_page_size
133 return self.page_size_padded
@@ -159,6 +216,7 @@ class FullAttentionSpec(AttentionSpec):
216 head_size=specs[0].head_size,
217 head_size_v=specs[0].head_size_v,
218 dtype=specs[0].dtype,
219 + kv_quant_mode=specs[0].kv_quant_mode,
220 page_size_padded=specs[0].page_size_padded,
221 sliding_window=cls.merge_window_sizes(sliding_window),
222 attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),
@@ -220,6 +278,7 @@ class MLAAttentionSpec(FullAttentionSpec):
278 num_kv_heads=specs[0].num_kv_heads,
279 head_size=specs[0].head_size,
280 dtype=specs[0].dtype,
281 + kv_quant_mode=specs[0].kv_quant_mode,
282 page_size_padded=specs[0].page_size_padded,
283 cache_dtype_str=cache_dtype_str_set.pop(),
284 )
@@ -352,6 +411,7 @@ class SinkFullAttentionSpec(FullAttentionSpec):
411 head_size_v=specs[0].head_size_v,
412 sink_len=specs[0].sink_len,
413 dtype=specs[0].dtype,
414 + kv_quant_mode=specs[0].kv_quant_mode,
415 page_size_padded=specs[0].page_size_padded,
416 sliding_window=cls.merge_window_sizes(sliding_window),
417 attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),