@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 725 lines 26 KB
Raw
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 (INT4, 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.attention.ops.int4_per_token_head import single_rht
27 from vllm.v1.kv_cache_interface import KVQuantMode, is_quantized_kv_cache
28
29 DEVICE_TYPE = current_platform.device_type
30
31 # Skip entire module if no CUDA/ROCm GPU available
32 pytestmark = [
33 pytest.mark.skipif(
34 current_platform.is_cpu(),
35 reason="Per-token-head KV cache tests require GPU.",
36 ),
37 ]
38
39 # ---------------------------------------------------------------------------
40 # Test parameters
41 # ---------------------------------------------------------------------------
42 NUM_TOKENS = [1, 7, 42]
43 NUM_KV_HEADS = [1, 4, 8]
44 HEAD_SIZES = [64, 128]
45 BLOCK_SIZES = [16]
46 SEEDS = [0]
47
48 # Platform-dependent FP8 dtype and range
49 FP8_DTYPE = current_platform.fp8_dtype()
50 FP8_MIN, FP8_MAX = get_fp8_min_max()
51
52
53 # ---------------------------------------------------------------------------
54 # Per-dtype quantization config
55 # ---------------------------------------------------------------------------
56 @dataclass(frozen=True)
57 class QuantConfig:
58 """Quantization parameters for a given cache dtype."""
59
60 cache_dtype: torch.dtype # torch.int8 or FP8_DTYPE
61 kv_cache_dtype_str: str # "int8_per_token_head" or "fp8_per_token_head"
62 quant_max: float
63 quant_min: float
64 kv_quant_mode: KVQuantMode
65 # INT8 rounds explicitly; FP8 relies on dtype cast rounding.
66 rounds_before_store: bool
67
68
69 INT8_CONFIG = QuantConfig(
70 cache_dtype=torch.int8,
71 kv_cache_dtype_str="int8_per_token_head",
72 quant_max=127.0,
73 quant_min=-128.0,
74 kv_quant_mode=KVQuantMode.INT8_PER_TOKEN_HEAD,
75 rounds_before_store=True,
76 )
77 FP8_CONFIG = QuantConfig(
78 cache_dtype=FP8_DTYPE,
79 kv_cache_dtype_str="fp8_per_token_head",
80 quant_max=FP8_MAX,
81 quant_min=FP8_MIN,
82 kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD,
83 rounds_before_store=False,
84 )
85 INT4_CONFIG = QuantConfig(
86 cache_dtype=torch.uint8,
87 kv_cache_dtype_str="int4_per_token_head",
88 quant_max=7.0,
89 quant_min=-8.0,
90 kv_quant_mode=KVQuantMode.INT4_PER_TOKEN_HEAD,
91 # Unused for int4 (handled by its own rint path); kept for the dataclass.
92 rounds_before_store=False,
93 )
94 QUANT_CONFIGS = [INT4_CONFIG, INT8_CONFIG, FP8_CONFIG]
95
96
97 @pytest.fixture(params=QUANT_CONFIGS, ids=["int4", "int8", "fp8"])
98 def qcfg(request) -> QuantConfig:
99 return request.param
100
101
102 # ---------------------------------------------------------------------------
103 # Helpers
104 # ---------------------------------------------------------------------------
105 def _quantize_per_token_head_ref(
106 data: torch.Tensor, # [num_tokens, num_heads, head_size]
107 cfg: QuantConfig,
108 ) -> tuple[torch.Tensor, torch.Tensor]:
109 """Reference per-token-head quantization (one scale per token per head).
110
111 Returns (quantized, scales) where scales is [num_tokens, num_heads].
112 """
113 absmax = data.float().abs().amax(dim=2) # [num_tokens, num_heads]
114 scales = (absmax / cfg.quant_max).clamp(min=1e-6)
115 scaled = data.float() * (1.0 / scales[:, :, None])
116 if cfg.rounds_before_store:
117 q = scaled.round().clamp(cfg.quant_min, cfg.quant_max).to(cfg.cache_dtype)
118 else:
119 q = scaled.clamp(cfg.quant_min, cfg.quant_max).to(cfg.cache_dtype)
120 return q, scales
121
122
123 # ===========================================================================
124 # 1. is_quantized_kv_cache / get_kv_quant_mode
125 # ===========================================================================
126 class TestIsQuantizedKvCache:
127 def test_fp8_variants(self):
128 assert is_quantized_kv_cache("fp8")
129 assert is_quantized_kv_cache("fp8_e4m3")
130 assert is_quantized_kv_cache("fp8_e5m2")
131
132 def test_int4_per_token_head(self):
133 assert is_quantized_kv_cache("int4_per_token_head")
134
135 def test_int8_per_token_head(self):
136 assert is_quantized_kv_cache("int8_per_token_head")
137
138 def test_fp8_per_token_head(self):
139 assert is_quantized_kv_cache("fp8_per_token_head")
140
141 def test_auto(self):
142 assert not is_quantized_kv_cache("auto")
143
144 def test_bfloat16(self):
145 assert not is_quantized_kv_cache("bfloat16")
146
147 def test_kv_quant_mode_int4(self):
148 from vllm.v1.kv_cache_interface import get_kv_quant_mode
149
150 assert (
151 get_kv_quant_mode("int4_per_token_head") == KVQuantMode.INT4_PER_TOKEN_HEAD
152 )
153
154 def test_kv_quant_mode_int8(self):
155 from vllm.v1.kv_cache_interface import get_kv_quant_mode
156
157 assert (
158 get_kv_quant_mode("int8_per_token_head") == KVQuantMode.INT8_PER_TOKEN_HEAD
159 )
160
161 def test_kv_quant_mode_fp8(self):
162 from vllm.v1.kv_cache_interface import get_kv_quant_mode
163
164 assert get_kv_quant_mode("fp8_per_token_head") == KVQuantMode.FP8_PER_TOKEN_HEAD
165
166
167 # ===========================================================================
168 # 2. Triton per-token-head kernel (reshape-and-cache)
169 # ===========================================================================
170 @pytest.mark.parametrize("num_tokens", NUM_TOKENS)
171 @pytest.mark.parametrize("num_heads", NUM_KV_HEADS)
172 @pytest.mark.parametrize("head_size", HEAD_SIZES)
173 @pytest.mark.parametrize("block_size", BLOCK_SIZES)
174 @pytest.mark.parametrize("seed", SEEDS)
175 @torch.inference_mode()
176 def test_reshape_and_cache_per_token_head(
177 qcfg: QuantConfig,
178 num_tokens: int,
179 num_heads: int,
180 head_size: int,
181 block_size: int,
182 seed: int,
183 ):
184 """Test triton_reshape_and_cache_flash_per_token_head_quant kernel."""
185 from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
186 triton_reshape_and_cache_flash_per_token_head_quant,
187 )
188
189 set_random_seed(seed)
190 torch.set_default_device(DEVICE_TYPE)
191
192 num_blocks = (num_tokens + block_size - 1) // block_size + 4
193 is_int4 = qcfg.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD
194 cache_head_size = head_size // 2 if is_int4 else head_size
195
196 key = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16)
197 value = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16)
198
199 key_cache = torch.zeros(
200 num_blocks, block_size, num_heads, cache_head_size, dtype=qcfg.cache_dtype
201 )
202 value_cache = torch.zeros(
203 num_blocks, block_size, num_heads, cache_head_size, dtype=qcfg.cache_dtype
204 )
205 k_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)
206 v_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)
207
208 num_slots = block_size * num_blocks
209 slot_mapping = torch.tensor(
210 random.sample(range(num_slots), num_tokens), dtype=torch.long
211 )
212
213 triton_reshape_and_cache_flash_per_token_head_quant(
214 key,
215 value,
216 key_cache,
217 value_cache,
218 k_scale_cache,
219 v_scale_cache,
220 slot_mapping,
221 kv_quant_mode=qcfg.kv_quant_mode,
222 )
223
224 # INT4 (RHT + asymmetric), INT8/FP8 have different dequant paths. Only
225 # INT8/FP8 can be compared to a PyTorch reference.
226 if not is_int4:
227 ref_k_quant, ref_k_scales = _quantize_per_token_head_ref(key, qcfg)
228 ref_v_quant, ref_v_scales = _quantize_per_token_head_ref(value, qcfg)
229
230 for i, slot in enumerate(slot_mapping.tolist()):
231 blk = slot // block_size
232 off = slot % block_size
233
234 if is_int4:
235 # Coarser quantization → wider tolerance.
236 deq_atol = deq_rtol = 0.5
237 for label, data, cache, sc in [
238 ("key", key, key_cache, k_scale_cache),
239 ("val", value, value_cache, v_scale_cache),
240 ]:
241 packed_scale = sc[blk, off] # [num_heads] float32
242 scale_bits = packed_scale.view(torch.int32)
243 zp = (scale_bits & 0xF).to(torch.float32)
244 clean_scale = (scale_bits & -16).view(torch.float32)
245
246 packed = cache[blk, off]
247 lo = (packed & 0xF).to(torch.float32)
248 hi = ((packed >> 4) & 0xF).to(torch.float32)
249 full = torch.zeros(num_heads, head_size, dtype=torch.float32)
250 full[:, 0::2] = lo
251 full[:, 1::2] = hi
252 # Asymmetric dequant in RHT domain, then IRHT/d → original
253 deq_rht = (full - zp[:, None]) * clean_scale[:, None]
254 deq = single_rht(deq_rht, inverse=True) / head_size
255 ref_deq = data[i].float()
256 torch.testing.assert_close(deq, ref_deq, atol=deq_atol, rtol=deq_rtol)
257 else:
258 actual_k_scale = k_scale_cache[blk, off] # [num_heads]
259 k_deq = key_cache[blk, off].float() * actual_k_scale[:, None]
260 k_ref_deq = key[i].float()
261 torch.testing.assert_close(
262 k_deq,
263 k_ref_deq,
264 atol=0.1,
265 rtol=0.1,
266 )
267 actual_v_scale = v_scale_cache[blk, off] # [num_heads]
268 v_deq = value_cache[blk, off].float() * actual_v_scale[:, None]
269 v_ref_deq = value[i].float()
270 torch.testing.assert_close(
271 v_deq,
272 v_ref_deq,
273 atol=0.1,
274 rtol=0.1,
275 )
276 # Per-head scales: [num_heads]
277 torch.testing.assert_close(
278 k_scale_cache[blk, off], ref_k_scales[i], atol=1e-4, rtol=1e-3
279 )
280 torch.testing.assert_close(
281 v_scale_cache[blk, off], ref_v_scales[i], atol=1e-4, rtol=1e-3
282 )
283
284
285 # ===========================================================================
286 # 3. Per-token-head round-trip accuracy (quantize -> dequantize)
287 # ===========================================================================
288 @pytest.mark.parametrize("num_tokens", [1, 16])
289 @pytest.mark.parametrize("num_heads", [4])
290 @pytest.mark.parametrize("head_size", [128])
291 @pytest.mark.parametrize("block_size", [16])
292 @torch.inference_mode()
293 def test_per_token_head_round_trip_accuracy(
294 qcfg: QuantConfig,
295 num_tokens: int,
296 num_heads: int,
297 head_size: int,
298 block_size: int,
299 ):
300 """Verify per-token-head round-trip: kernel dequant matches reference.
301
302 INT8: round-to-nearest before int8 store.
303 FP8: hardware cast (clamp then cast).
304 """
305 from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
306 triton_reshape_and_cache_flash_per_token_head_quant,
307 )
308
309 torch.set_default_device(DEVICE_TYPE)
310 set_random_seed(42)
311
312 is_int4 = qcfg.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD
313 num_blocks = (num_tokens + block_size - 1) // block_size + 2
314 cache_head_size = head_size // 2 if is_int4 else head_size
315
316 key = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16) * 0.5
317 value = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16) * 0.5
318
319 key_cache = torch.zeros(
320 num_blocks, block_size, num_heads, cache_head_size, dtype=qcfg.cache_dtype
321 )
322 value_cache = torch.zeros(
323 num_blocks, block_size, num_heads, cache_head_size, dtype=qcfg.cache_dtype
324 )
325 k_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)
326 v_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)
327
328 slot_mapping = torch.arange(num_tokens, dtype=torch.long)
329
330 triton_reshape_and_cache_flash_per_token_head_quant(
331 key,
332 value,
333 key_cache,
334 value_cache,
335 k_scale_cache,
336 v_scale_cache,
337 slot_mapping,
338 kv_quant_mode=qcfg.kv_quant_mode,
339 )
340
341 rt_atol = 0.5 if is_int4 else 0.1
342
343 for i in range(num_tokens):
344 blk = i // block_size
345 off = i % block_size
346
347 for label, data, cache, sc in [
348 ("key", key, key_cache, k_scale_cache),
349 ("val", value, value_cache, v_scale_cache),
350 ]:
351 for h in range(num_heads):
352 orig = data[i, h].float()
353 actual_sc = sc[blk, off, h]
354 if is_int4:
355 sc_bits = actual_sc.view(torch.int32)
356 zp = (sc_bits & 0xF).to(torch.float32)
357 clean_sc = (sc_bits & -16).view(torch.float32)
358 packed = cache[blk, off, h]
359 lo = (packed & 0xF).to(torch.float32)
360 hi = ((packed >> 4) & 0xF).to(torch.float32)
361 full = torch.zeros(head_size)
362 full[0::2] = lo
363 full[1::2] = hi
364 deq_rht = (full - zp) * clean_sc
365 actual_deq = (
366 single_rht(deq_rht.unsqueeze(0), inverse=True).squeeze(0)
367 / head_size
368 )
369 else:
370 actual_deq = cache[blk, off, h].float() * actual_sc
371 torch.testing.assert_close(
372 actual_deq,
373 orig,
374 atol=rt_atol,
375 rtol=rt_atol,
376 )
377
378
379 @torch.inference_mode()
380 def test_int8_per_token_head_raw_cache_matches_round_reference():
381 """INT8 cache writes should match round-to-nearest quantization exactly."""
382 from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
383 triton_reshape_and_cache_flash_per_token_head_quant,
384 )
385
386 torch.set_default_device(DEVICE_TYPE)
387
388 head_size = 8
389 block_size = 4
390
391 key = torch.tensor(
392 [[[-127.0, -2.6, -2.4, -1.6, -1.4, -0.6, -0.4, 127.0]]],
393 dtype=torch.bfloat16,
394 )
395 value = -key
396
397 key_cache = torch.zeros(1, block_size, 1, head_size, dtype=torch.int8)
398 value_cache = torch.zeros_like(key_cache)
399 k_scale_cache = torch.ones(1, block_size, 1, dtype=torch.float32)
400 v_scale_cache = torch.ones_like(k_scale_cache)
401 slot_mapping = torch.tensor([2], dtype=torch.long)
402
403 triton_reshape_and_cache_flash_per_token_head_quant(
404 key,
405 value,
406 key_cache,
407 value_cache,
408 k_scale_cache,
409 v_scale_cache,
410 slot_mapping,
411 kv_quant_mode=INT8_CONFIG.kv_quant_mode,
412 )
413
414 ref_k_quant, ref_k_scales = _quantize_per_token_head_ref(key, INT8_CONFIG)
415 ref_v_quant, ref_v_scales = _quantize_per_token_head_ref(value, INT8_CONFIG)
416
417 slot = slot_mapping.item()
418 blk = slot // block_size
419 off = slot % block_size
420 assert torch.equal(key_cache[blk, off], ref_k_quant[0])
421 assert torch.equal(value_cache[blk, off], ref_v_quant[0])
422 torch.testing.assert_close(k_scale_cache[blk, off], ref_k_scales[0])
423 torch.testing.assert_close(v_scale_cache[blk, off], ref_v_scales[0])
424
425
426 # ===========================================================================
427 # 4. Negative slot mapping (padding tokens should be skipped)
428 # ===========================================================================
429 @torch.inference_mode()
430 def test_per_token_head_negative_slot_skipped(qcfg: QuantConfig):
431 """Tokens with slot_mapping=-1 should leave the cache unchanged."""
432 from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
433 triton_reshape_and_cache_flash_per_token_head_quant,
434 )
435
436 torch.set_default_device(DEVICE_TYPE)
437 num_tokens = 4
438 num_heads = 2
439 head_size = 64
440 block_size = 16
441 num_blocks = 2
442 is_int4 = qcfg.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD
443 cache_head_size = head_size // 2 if is_int4 else head_size
444
445 key = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16)
446 value = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16)
447
448 key_cache = torch.zeros(
449 num_blocks, block_size, num_heads, cache_head_size, dtype=qcfg.cache_dtype
450 )
451 value_cache = torch.zeros(
452 num_blocks, block_size, num_heads, cache_head_size, dtype=qcfg.cache_dtype
453 )
454 k_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)
455 v_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32)
456
457 slot_mapping = torch.tensor([0, -1, 1, -1], dtype=torch.long)
458
459 key_cache_before = key_cache.clone()
460 val_cache_before = value_cache.clone()
461
462 triton_reshape_and_cache_flash_per_token_head_quant(
463 key,
464 value,
465 key_cache,
466 value_cache,
467 k_scale_cache,
468 v_scale_cache,
469 slot_mapping,
470 kv_quant_mode=qcfg.kv_quant_mode,
471 )
472
473 # Slots 0 and 1 should have been written (tokens 0 and 2)
474 assert not torch.equal(key_cache[0, 0], key_cache_before[0, 0])
475 assert not torch.equal(key_cache[0, 1], key_cache_before[0, 1])
476 assert not torch.equal(value_cache[0, 0], val_cache_before[0, 0])
477
478 # All other slots should be unchanged
479 assert torch.equal(key_cache[0, 2:], key_cache_before[0, 2:])
480 assert torch.equal(key_cache[1], key_cache_before[1])
481 assert torch.equal(value_cache[0, 2:], val_cache_before[0, 2:])
482
483
484 # ===========================================================================
485 # 5. process_weights_after_loading -- per-token-head early return
486 # ===========================================================================
487 @pytest.mark.parametrize(
488 "kv_cache_dtype",
489 ["int4_per_token_head", "int8_per_token_head", "fp8_per_token_head"],
490 )
491 def test_process_weights_sets_placeholder_scales(kv_cache_dtype: str):
492 """Per-token-head should set _k_scale=1.0, _v_scale=1.0
493 and delete checkpoint attrs."""
494 from vllm.model_executor.layers.quantization.kv_cache import (
495 BaseKVCacheMethod,
496 )
497
498 layer = MagicMock()
499 layer.kv_cache_dtype = kv_cache_dtype
500 layer.k_scale = torch.nn.Parameter(torch.tensor(-1.0), requires_grad=False)
501 layer.v_scale = torch.nn.Parameter(torch.tensor(-1.0), requires_grad=False)
502 layer.q_scale = torch.nn.Parameter(torch.tensor(-1.0), requires_grad=False)
503 layer.prob_scale = torch.nn.Parameter(torch.tensor(-1.0), requires_grad=False)
504 layer._k_scale = torch.tensor(0.0)
505 layer._v_scale = torch.tensor(0.0)
506 layer._k_scale_float = 0.0
507 layer._v_scale_float = 0.0
508
509 method = BaseKVCacheMethod.__new__(BaseKVCacheMethod)
510 method.quant_config = MagicMock()
511 method.process_weights_after_loading(layer)
512
513 assert layer._k_scale_float == 1.0
514 assert layer._v_scale_float == 1.0
515 assert not hasattr(layer, "k_scale")
516 assert not hasattr(layer, "v_scale")
517 assert not hasattr(layer, "q_scale")
518 assert not hasattr(layer, "prob_scale")
519
520
521 # ===========================================================================
522 # 6. Triton unified_attention -- per-token-head scale cache (INT4/INT8/FP8)
523 # ===========================================================================
524 @pytest.mark.parametrize(
525 "seq_lens",
526 [
527 [(1, 128)],
528 [(1, 64), (1, 32)],
529 ],
530 )
531 @pytest.mark.parametrize("num_heads", [(4, 4)])
532 @pytest.mark.parametrize("head_size", [128])
533 @pytest.mark.parametrize("block_size", [16])
534 @torch.inference_mode()
535 def test_triton_unified_attention_per_token_head_scale(
536 qcfg: QuantConfig,
537 seq_lens: list[tuple[int, int]],
538 num_heads: tuple[int, int],
539 head_size: int,
540 block_size: int,
541 ):
542 """End-to-end: quantized KV with per-token-head scale caches."""
543 from vllm.utils.math_utils import next_power_of_2
544 from vllm.v1.attention.ops.triton_unified_attention import unified_attention
545
546 torch.set_default_device(DEVICE_TYPE)
547 set_random_seed(0)
548
549 is_int4 = qcfg.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD
550
551 num_seqs = len(seq_lens)
552 query_lens = [s[0] for s in seq_lens]
553 kv_lens = [s[1] for s in seq_lens]
554 num_query_heads, num_kv_heads = num_heads
555 max_query_len = max(query_lens)
556 max_kv_len = max(kv_lens)
557 scale = head_size**-0.5
558 num_blocks = 2048
559
560 query = torch.randn(
561 sum(query_lens), num_query_heads, head_size, dtype=torch.bfloat16
562 )
563
564 key_cache_bf16 = torch.randn(
565 num_blocks, block_size, num_kv_heads, head_size, dtype=torch.bfloat16
566 )
567 value_cache_bf16 = torch.randn_like(key_cache_bf16)
568
569 if is_int4:
570 # Asymmetric quantization reference (matches the Triton kernel).
571 kf = key_cache_bf16.float()
572 vf = value_cache_bf16.float()
573 k_min = kf.amin(dim=-1)
574 k_max = kf.amax(dim=-1)
575 v_min = vf.amin(dim=-1)
576 v_max = vf.amax(dim=-1)
577 k_scale_cache = ((k_max - k_min) / 15.0).clamp(min=1e-6).to(torch.float32)
578 v_scale_cache = ((v_max - v_min) / 15.0).clamp(min=1e-6).to(torch.float32)
579 k_zp = (-k_min / k_scale_cache).round().clamp(0, 15)
580 v_zp = (-v_min / v_scale_cache).round().clamp(0, 15)
581
582 key_cache_q_full = (
583 (kf / k_scale_cache[..., None] + k_zp[..., None]).round().clamp(0, 15)
584 )
585 value_cache_q_full = (
586 (vf / v_scale_cache[..., None] + v_zp[..., None]).round().clamp(0, 15)
587 )
588
589 # Dequantized reference: x_hat = (q - zp) * scale
590 key_cache_deq = (key_cache_q_full - k_zp[..., None]) * k_scale_cache[..., None]
591 value_cache_deq = (value_cache_q_full - v_zp[..., None]) * v_scale_cache[
592 ..., None
593 ]
594
595 # Pack two uint4 values into one byte
596 def _pack_int4(data_float):
597 u = data_float.to(torch.uint8)
598 lo = u[..., 0::2]
599 hi = u[..., 1::2]
600 return (lo & 0xF) | ((hi & 0xF) << 4)
601
602 key_cache_q = _pack_int4(key_cache_q_full)
603 value_cache_q = _pack_int4(value_cache_q_full)
604
605 # Steganography: pack zp into low 4 bits of scale
606 k_zp_int = k_zp.to(torch.int32)
607 k_bits = k_scale_cache.view(torch.int32)
608 k_scale_cache = ((k_bits & -16) | (k_zp_int & 0xF)).view(torch.float32)
609 v_zp_int = v_zp.to(torch.int32)
610 v_bits = v_scale_cache.view(torch.int32)
611 v_scale_cache = ((v_bits & -16) | (v_zp_int & 0xF)).view(torch.float32)
612 else:
613 # Symmetric quantization for int8/fp8.
614 k_absmax = key_cache_bf16.float().abs().amax(dim=-1)
615 v_absmax = value_cache_bf16.float().abs().amax(dim=-1)
616 k_scale_cache = (k_absmax / qcfg.quant_max).clamp(min=1e-6).to(torch.float32)
617 v_scale_cache = (v_absmax / qcfg.quant_max).clamp(min=1e-6).to(torch.float32)
618 scaled_k = key_cache_bf16.float() / k_scale_cache[:, :, :, None]
619 scaled_v = value_cache_bf16.float() / v_scale_cache[:, :, :, None]
620
621 key_cache_q_full = scaled_k.round().clamp(qcfg.quant_min, qcfg.quant_max)
622 value_cache_q_full = scaled_v.round().clamp(qcfg.quant_min, qcfg.quant_max)
623
624 key_cache_deq = key_cache_q_full * k_scale_cache[:, :, :, None]
625 value_cache_deq = value_cache_q_full * v_scale_cache[:, :, :, None]
626
627 if not is_int4 and qcfg.rounds_before_store:
628 key_cache_q = key_cache_q_full.to(qcfg.cache_dtype)
629 value_cache_q = value_cache_q_full.to(qcfg.cache_dtype)
630 elif not is_int4:
631 key_cache_q = scaled_k.clamp(qcfg.quant_min, qcfg.quant_max).to(
632 qcfg.cache_dtype
633 )
634 value_cache_q = scaled_v.clamp(qcfg.quant_min, qcfg.quant_max).to(
635 qcfg.cache_dtype
636 )
637
638 cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
639 dim=0, dtype=torch.int32
640 )
641 kv_lens_t = torch.tensor(kv_lens, dtype=torch.int32)
642
643 max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
644 block_tables = torch.randint(
645 0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
646 )
647
648 head_size_padded = next_power_of_2(head_size)
649 seq_threshold_3D = 0
650 num_par_softmax_segments = 16
651 softmax_segm_output = torch.empty(
652 (seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded),
653 dtype=torch.float32,
654 )
655 softmax_segm_max = torch.empty(
656 (seq_threshold_3D, num_query_heads, num_par_softmax_segments),
657 dtype=torch.float32,
658 )
659 softmax_segm_expsum = torch.empty(
660 (seq_threshold_3D, num_query_heads, num_par_softmax_segments),
661 dtype=torch.float32,
662 )
663
664 output_q = torch.empty_like(query)
665 unified_attention(
666 q=query,
667 k=key_cache_q,
668 v=value_cache_q,
669 out=output_q,
670 cu_seqlens_q=cu_query_lens,
671 seqused_k=kv_lens_t,
672 max_seqlen_q=max_query_len,
673 max_seqlen_k=max_kv_len,
674 softmax_scale=scale,
675 causal=True,
676 window_size=(-1, -1),
677 block_table=block_tables,
678 softcap=0,
679 q_descale=None,
680 k_descale=None,
681 v_descale=None,
682 seq_threshold_3D=seq_threshold_3D,
683 num_par_softmax_segments=num_par_softmax_segments,
684 softmax_segm_output=softmax_segm_output,
685 softmax_segm_max=softmax_segm_max,
686 softmax_segm_expsum=softmax_segm_expsum,
687 kv_quant_mode=qcfg.kv_quant_mode,
688 k_scale_cache=k_scale_cache,
689 v_scale_cache=v_scale_cache,
690 )
691
692 output_ref = torch.empty_like(query)
693 unified_attention(
694 q=query,
695 k=key_cache_deq.to(torch.bfloat16),
696 v=value_cache_deq.to(torch.bfloat16),
697 out=output_ref,
698 cu_seqlens_q=cu_query_lens,
699 seqused_k=kv_lens_t,
700 max_seqlen_q=max_query_len,
701 max_seqlen_k=max_kv_len,
702 softmax_scale=scale,
703 causal=True,
704 window_size=(-1, -1),
705 block_table=block_tables,
706 softcap=0,
707 q_descale=None,
708 k_descale=None,
709 v_descale=None,
710 seq_threshold_3D=seq_threshold_3D,
711 num_par_softmax_segments=num_par_softmax_segments,
712 softmax_segm_output=softmax_segm_output,
713 softmax_segm_max=softmax_segm_max,
714 softmax_segm_expsum=softmax_segm_expsum,
715 )
716
717 # Coarser quantization → wider tolerance.
718 if is_int4:
719 # Hopper's attention reduction order can move a few BF16 elements by
720 # just over 1.0 after INT4 quantization.
721 atol = 1.1 if current_platform.is_device_capability_family(90) else 0.5
722 rtol = 0.5
723 else:
724 atol, rtol = 5e-2, 5e-2
725 torch.testing.assert_close(output_q, output_ref, atol=atol, rtol=rtol)