@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 782 lines 28.7 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Unit tests for TurboQuant KV-cache quantization.
4
5 Run: .venv/bin/python -m pytest tests/quantization/test_turboquant.py -v
6 """
7
8 import math
9 from types import SimpleNamespace
10
11 import pytest
12 import torch
13
14 from vllm.model_executor.layers.quantization.turboquant.centroids import (
15 get_centroids,
16 solve_lloyd_max,
17 )
18 from vllm.model_executor.layers.quantization.turboquant.config import (
19 TQ_PRESETS,
20 TurboQuantConfig,
21 )
22 from vllm.platforms import current_platform
23 from vllm.utils.math_utils import next_power_of_2
24
25 # ============================================================================
26 # Helpers
27 # ============================================================================
28
29 ALL_PRESETS = list(TQ_PRESETS.keys())
30
31
32 def _assert_strictly_sorted(seq, name="sequence"):
33 for i in range(len(seq) - 1):
34 assert seq[i] < seq[i + 1], f"{name} not sorted at index {i}"
35
36
37 def _is_power_of_2(n: int) -> bool:
38 return n > 0 and next_power_of_2(n) == n
39
40
41 # Expected concrete values for each preset at head_dim=128.
42 # fmt: off
43 PRESET_EXPECTED = {
44 "turboquant_k8v4": dict(
45 key_fp8=True, key_quant_bits=8,
46 key_mse_bits=0, value_quant_bits=4,
47 mse_bits=4, n_centroids=16, centroid_bits=4,
48 norm_correction=False,
49 key_packed_size=128, value_packed_size=68,
50 slot_size=196, slot_size_aligned=196,
51 ),
52 "turboquant_4bit_nc": dict(
53 key_fp8=False, key_quant_bits=4,
54 key_mse_bits=4, value_quant_bits=4,
55 mse_bits=4, n_centroids=16, centroid_bits=4,
56 norm_correction=True,
57 key_packed_size=66, value_packed_size=68,
58 slot_size=134, slot_size_aligned=134,
59 ),
60 "turboquant_k3v4_nc": dict(
61 key_fp8=False, key_quant_bits=3,
62 key_mse_bits=3, value_quant_bits=4,
63 mse_bits=3, n_centroids=8, centroid_bits=3,
64 norm_correction=True,
65 key_packed_size=50, value_packed_size=68,
66 slot_size=118, slot_size_aligned=118,
67 ),
68 "turboquant_3bit_nc": dict(
69 key_fp8=False, key_quant_bits=3,
70 key_mse_bits=3, value_quant_bits=3,
71 mse_bits=3, n_centroids=8, centroid_bits=3,
72 norm_correction=True,
73 key_packed_size=50, value_packed_size=52,
74 slot_size=102, slot_size_aligned=102,
75 ),
76 }
77 # fmt: on
78
79
80 # ============================================================================
81 # Config tests (CPU-only, no dependencies beyond config.py)
82 # ============================================================================
83
84
85 class TestTurboQuantConfig:
86 @pytest.mark.parametrize("preset", ALL_PRESETS)
87 def test_preset_parses(self, preset):
88 cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
89 assert isinstance(cfg, TurboQuantConfig)
90
91 def test_invalid_preset_raises(self):
92 with pytest.raises(ValueError, match="Unknown TurboQuant"):
93 TurboQuantConfig.from_cache_dtype("turboquant_invalid", head_dim=128)
94
95 # ---- Per-preset concrete value checks (table-driven) ----
96
97 @pytest.mark.parametrize("preset", ALL_PRESETS)
98 def test_key_mode(self, preset):
99 cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
100 exp = PRESET_EXPECTED[preset]
101 assert cfg.key_fp8 is exp["key_fp8"]
102 assert cfg.key_quant_bits == exp["key_quant_bits"]
103 assert cfg.key_mse_bits == exp["key_mse_bits"]
104
105 @pytest.mark.parametrize("preset", ALL_PRESETS)
106 def test_value_mode(self, preset):
107 cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
108 exp = PRESET_EXPECTED[preset]
109 assert cfg.value_quant_bits == exp["value_quant_bits"]
110
111 @pytest.mark.parametrize("preset", ALL_PRESETS)
112 def test_bits_and_centroids(self, preset):
113 cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
114 exp = PRESET_EXPECTED[preset]
115 assert cfg.mse_bits == exp["mse_bits"]
116 assert cfg.n_centroids == exp["n_centroids"]
117 assert cfg.centroid_bits == exp["centroid_bits"]
118
119 @pytest.mark.parametrize("preset", ALL_PRESETS)
120 def test_norm_correction(self, preset):
121 cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
122 assert cfg.norm_correction is PRESET_EXPECTED[preset]["norm_correction"]
123
124 @pytest.mark.parametrize("preset", ALL_PRESETS)
125 def test_packed_sizes(self, preset):
126 cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
127 exp = PRESET_EXPECTED[preset]
128 assert cfg.key_packed_size == exp["key_packed_size"]
129 assert cfg.value_packed_size == exp["value_packed_size"]
130 assert cfg.slot_size == exp["slot_size"]
131 assert cfg.slot_size_aligned == exp["slot_size_aligned"]
132
133 # ---- Cross-preset structural invariants ----
134
135 @pytest.mark.parametrize("preset", ALL_PRESETS)
136 def test_slot_equals_key_plus_value(self, preset):
137 cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
138 assert cfg.slot_size == cfg.key_packed_size + cfg.value_packed_size
139
140 @pytest.mark.parametrize("preset", ALL_PRESETS)
141 def test_padded_slot_is_even(self, preset):
142 cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
143 assert cfg.slot_size_aligned >= cfg.slot_size
144 assert cfg.slot_size_aligned % 2 == 0, (
145 f"slot_size_aligned={cfg.slot_size_aligned} is not even"
146 )
147
148 @pytest.mark.parametrize("preset", ALL_PRESETS)
149 def test_key_value_packed_sizes_positive(self, preset):
150 cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
151 assert cfg.key_packed_size > 0
152 assert cfg.value_packed_size > 0
153
154 @pytest.mark.parametrize("preset", ALL_PRESETS)
155 def test_n_centroids_is_2_to_mse_bits(self, preset):
156 cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
157 assert cfg.n_centroids == 2**cfg.mse_bits
158
159 @pytest.mark.parametrize("preset", ALL_PRESETS)
160 def test_centroid_bits_always_positive(self, preset):
161 cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
162 assert cfg.centroid_bits > 0
163
164 @pytest.mark.parametrize("preset", ALL_PRESETS)
165 def test_mse_key_or_fp8_exclusive(self, preset):
166 """Each preset is either FP8 keys or MSE keys, never both."""
167 cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
168 if cfg.key_fp8:
169 assert cfg.key_mse_bits == 0
170 assert cfg.key_quant_bits == 8
171 else:
172 assert cfg.key_mse_bits > 0
173 assert cfg.key_quant_bits in (3, 4)
174
175 @pytest.mark.parametrize("preset", ALL_PRESETS)
176 @pytest.mark.parametrize("head_dim", [64, 96, 128, 256])
177 def test_all_presets_all_head_dims(self, preset, head_dim):
178 cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=head_dim)
179 assert cfg.head_dim == head_dim
180 assert cfg.slot_size == cfg.key_packed_size + cfg.value_packed_size
181 assert cfg.slot_size_aligned >= cfg.slot_size
182 assert cfg.slot_size_aligned % 2 == 0
183
184 # ---- Boundary skip layers ----
185
186 @staticmethod
187 def _dense_model_config(num_layers):
188 from types import SimpleNamespace
189
190 return SimpleNamespace(
191 is_hybrid=False,
192 hf_text_config=SimpleNamespace(num_hidden_layers=num_layers),
193 )
194
195 def test_boundary_skip_layers_basic(self):
196 mc = self._dense_model_config(32)
197 layers = TurboQuantConfig.get_boundary_skip_layers(mc)
198 assert layers == ["0", "1", "30", "31"]
199
200 def test_boundary_skip_layers_zero(self):
201 mc = self._dense_model_config(32)
202 assert TurboQuantConfig.get_boundary_skip_layers(mc, 0) == []
203
204 def test_boundary_skip_layers_small_model(self):
205 mc = self._dense_model_config(4)
206 layers = TurboQuantConfig.get_boundary_skip_layers(mc)
207 assert layers == ["0", "1", "2", "3"]
208
209 def test_boundary_skip_layers_cap_at_half(self):
210 mc = self._dense_model_config(8)
211 layers = TurboQuantConfig.get_boundary_skip_layers(mc, 10)
212 assert len(layers) == 8
213
214
215 class TestHybridAttentionIndices:
216 """Regression tests for boundary protection on hybrid models.
217
218 Hybrid models (attention + Mamba / linear-attention) identify KV-carrying
219 layers via layer_types / layers_block_type / attn_type_list. The helper
220 must return the *global* layer indices of the full-attention layers so
221 that kv_cache_dtype_skip_layers matches what extract_layer_index(prefix)
222 reports on the Attention layers at runtime.
223 """
224
225 @staticmethod
226 def _fake_model_config(text_cfg=None, hf_cfg=None):
227 from types import SimpleNamespace
228
229 return SimpleNamespace(
230 hf_text_config=text_cfg if text_cfg is not None else SimpleNamespace(),
231 hf_config=hf_cfg if hf_cfg is not None else SimpleNamespace(),
232 )
233
234 def test_layer_types_full_attention(self):
235 from vllm.model_executor.layers.quantization.turboquant.config import (
236 _get_full_attention_layer_indices,
237 )
238
239 cfg = type("C", (), {})()
240 cfg.layer_types = [
241 "linear_attention",
242 "linear_attention",
243 "full_attention",
244 "linear_attention",
245 "full_attention",
246 "full_attention",
247 ]
248 mc = self._fake_model_config(text_cfg=cfg)
249 assert _get_full_attention_layer_indices(mc) == [2, 4, 5]
250
251 def test_layers_block_type_jamba(self):
252 from vllm.model_executor.layers.quantization.turboquant.config import (
253 _get_full_attention_layer_indices,
254 )
255
256 cfg = type("C", (), {})()
257 cfg.layers_block_type = ["mamba", "attention", "mamba", "attention"]
258 mc = self._fake_model_config(text_cfg=cfg)
259 assert _get_full_attention_layer_indices(mc) == [1, 3]
260
261 def test_attn_type_list_minimax(self):
262 from vllm.model_executor.layers.quantization.turboquant.config import (
263 _get_full_attention_layer_indices,
264 )
265
266 hf = type("C", (), {})()
267 hf.attn_type_list = [0, 1, 0, 1, 1]
268 mc = self._fake_model_config(hf_cfg=hf)
269 assert _get_full_attention_layer_indices(mc) == [1, 3, 4]
270
271 def test_no_hybrid_hints_returns_empty(self):
272 from vllm.model_executor.layers.quantization.turboquant.config import (
273 _get_full_attention_layer_indices,
274 )
275
276 mc = self._fake_model_config()
277 assert _get_full_attention_layer_indices(mc) == []
278
279
280 class TestTurboQuantKVCacheSpec:
281 @pytest.mark.parametrize("preset", ALL_PRESETS)
282 def test_kv_cache_spec_sets_kv_quant_mode(self, preset):
283 from vllm.model_executor.layers.attention.attention import Attention
284 from vllm.v1.attention.backends.turboquant_attn import (
285 TurboQuantAttentionBackend,
286 )
287 from vllm.v1.kv_cache_interface import FullAttentionSpec
288
289 layer = SimpleNamespace(
290 attn_type="decoder",
291 kv_cache_dtype=preset,
292 kv_cache_torch_dtype=torch.uint8,
293 head_size=128,
294 head_size_v=128,
295 num_kv_heads=4,
296 sliding_window=None,
297 get_attn_backend=lambda: TurboQuantAttentionBackend,
298 )
299 vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=32))
300
301 # The layer builds an unpacked spec; the worker's spec-collection
302 # loop applies TQ slot packing via the backend's customize_spec hook.
303 spec = Attention.get_kv_cache_spec(layer, vllm_config)
304 assert isinstance(spec, FullAttentionSpec)
305 assert spec.kv_quant_mode.is_turboquant
306 assert spec.state_content_bytes is None
307
308 spec = TurboQuantAttentionBackend.customize_spec(spec)
309 expected_slot = TurboQuantConfig.from_cache_dtype(preset, 128).slot_size_aligned
310 assert spec.state_content_bytes == expected_slot
311
312
313 class TestTurboQuantWorkspaceReservation:
314 @staticmethod
315 def _fake_vllm_config(
316 *,
317 max_num_seqs: int = 16,
318 max_num_batched_tokens: int = 4096,
319 enable_chunked_prefill: bool = True,
320 max_model_len: int = 8192,
321 dtype: torch.dtype = torch.float16,
322 max_num_kv_splits: int = 4,
323 ):
324 return SimpleNamespace(
325 scheduler_config=SimpleNamespace(
326 max_num_seqs=max_num_seqs,
327 max_num_batched_tokens=max_num_batched_tokens,
328 enable_chunked_prefill=enable_chunked_prefill,
329 ),
330 model_config=SimpleNamespace(
331 max_model_len=max_model_len,
332 dtype=dtype,
333 get_num_attention_heads=lambda parallel_config: 8,
334 ),
335 parallel_config=SimpleNamespace(
336 tensor_parallel_size=2,
337 decode_context_parallel_size=1,
338 ),
339 attention_config=SimpleNamespace(
340 tq_max_kv_splits_for_cuda_graph=max_num_kv_splits
341 ),
342 )
343
344 @staticmethod
345 def _fake_kv_cache_spec():
346 from vllm.v1.kv_cache_interface import FullAttentionSpec
347
348 return FullAttentionSpec(
349 block_size=32,
350 num_kv_heads=4,
351 head_size=128,
352 head_size_v=128,
353 dtype=torch.uint8,
354 state_content_bytes=102,
355 )
356
357 def test_metadata_builder_reserves_decode_and_continuation_prefill_workspace(
358 self, monkeypatch
359 ):
360 from vllm.v1.attention.backends import turboquant_attn
361
362 calls = []
363
364 class FakeWorkspaceManager:
365 def get_simultaneous(self, *shapes_and_dtypes):
366 calls.append(shapes_and_dtypes)
367
368 monkeypatch.setattr(
369 turboquant_attn,
370 "current_workspace_manager",
371 lambda: FakeWorkspaceManager(),
372 )
373 monkeypatch.setattr(
374 turboquant_attn,
375 "is_workspace_manager_initialized",
376 lambda: True,
377 )
378
379 turboquant_attn.TurboQuantMetadataBuilder(
380 kv_cache_spec=self._fake_kv_cache_spec(),
381 layer_names=["layers.0.self_attn.attn"],
382 vllm_config=self._fake_vllm_config(),
383 device=torch.device("cuda"),
384 )
385
386 assert calls == [
387 (
388 ((16, 8, 4, 129), torch.float32),
389 ((16, 8, 128), torch.float16),
390 ((16, 8), torch.float32),
391 ),
392 (
393 ((1, 4, 8192, 128), torch.float16),
394 ((1, 4, 8192, 128), torch.float16),
395 ),
396 ]
397
398 def test_metadata_builder_skips_continuation_prefill_when_disabled(
399 self, monkeypatch
400 ):
401 from vllm.v1.attention.backends import turboquant_attn
402
403 calls = []
404
405 class FakeWorkspaceManager:
406 def get_simultaneous(self, *shapes_and_dtypes):
407 calls.append(shapes_and_dtypes)
408
409 monkeypatch.setattr(
410 turboquant_attn,
411 "current_workspace_manager",
412 lambda: FakeWorkspaceManager(),
413 )
414 monkeypatch.setattr(
415 turboquant_attn,
416 "is_workspace_manager_initialized",
417 lambda: True,
418 )
419
420 turboquant_attn.TurboQuantMetadataBuilder(
421 kv_cache_spec=self._fake_kv_cache_spec(),
422 layer_names=["layers.0.self_attn.attn"],
423 vllm_config=self._fake_vllm_config(enable_chunked_prefill=False),
424 device=torch.device("cuda"),
425 )
426
427 assert calls == [
428 (
429 ((16, 8, 4, 129), torch.float32),
430 ((16, 8, 128), torch.float16),
431 ((16, 8), torch.float32),
432 )
433 ]
434
435
436 # ============================================================================
437 # Centroids tests (CPU-only)
438 # ============================================================================
439
440
441 class TestCentroids:
442 @pytest.mark.parametrize("bits,expected_n", [(2, 4), (3, 8), (4, 16)])
443 def test_centroids_shape(self, bits, expected_n):
444 c = get_centroids(128, bits)
445 assert c.shape == (expected_n,)
446
447 @pytest.mark.parametrize("bits", [2, 3, 4])
448 def test_centroids_sorted(self, bits):
449 _assert_strictly_sorted(get_centroids(128, bits), "centroids")
450
451 def test_centroids_cached(self):
452 c1 = get_centroids(128, 3)
453 c2 = get_centroids(128, 3)
454 assert c1 is c2, "get_centroids should return cached object"
455
456 def test_centroids_different_dims_not_identical(self):
457 c64 = get_centroids(64, 3)
458 c128 = get_centroids(128, 3)
459 assert not torch.equal(c64, c128)
460
461 @pytest.mark.parametrize("bits", [2, 3, 4])
462 def test_centroids_symmetric_around_zero(self, bits):
463 """N(0, 1/d) is symmetric, so centroids should be ~symmetric."""
464 c = get_centroids(128, bits)
465 assert abs(c.mean().item()) < 0.01, "Centroids not centered near 0"
466 assert abs(c[0].item() + c[-1].item()) < 0.01
467
468 @pytest.mark.parametrize("bits", [2, 3, 4])
469 def test_centroids_within_4sigma(self, bits):
470 """All centroids should be within ~4 sigma of N(0, 1/d)."""
471 sigma = math.sqrt(1.0 / 128)
472 c = get_centroids(128, bits)
473 for i, val in enumerate(c):
474 assert abs(val.item()) < 4 * sigma, (
475 f"Centroid {i}={val:.6f} outside 4*sigma={4 * sigma:.6f}"
476 )
477
478
479 class TestLloydMax:
480 @pytest.mark.parametrize("bits,expected_n", [(2, 4), (3, 8), (4, 16)])
481 def test_solve_shapes(self, bits, expected_n):
482 centroids, boundaries = solve_lloyd_max(128, bits)
483 assert centroids.shape == (expected_n,)
484 assert boundaries.shape == (expected_n - 1,)
485
486 @pytest.mark.parametrize("bits", [2, 3, 4])
487 def test_centroids_sorted(self, bits):
488 centroids, _ = solve_lloyd_max(128, bits)
489 _assert_strictly_sorted(centroids, "centroids")
490
491 @pytest.mark.parametrize("bits", [2, 3, 4])
492 def test_boundaries_sorted(self, bits):
493 _, boundaries = solve_lloyd_max(128, bits)
494 _assert_strictly_sorted(boundaries, "boundaries")
495
496 @pytest.mark.parametrize("bits", [2, 3, 4])
497 def test_boundaries_between_centroids(self, bits):
498 """Each boundary must lie between its adjacent centroids."""
499 centroids, boundaries = solve_lloyd_max(128, bits)
500 for i in range(len(boundaries)):
501 assert centroids[i] < boundaries[i] < centroids[i + 1], (
502 f"Boundary {i}={boundaries[i]:.6f} not between "
503 f"c[{i}]={centroids[i]:.6f} and c[{i + 1}]={centroids[i + 1]:.6f}"
504 )
505
506 @pytest.mark.parametrize("bits", [2, 3, 4])
507 def test_boundaries_are_midpoints(self, bits):
508 """Lloyd-Max boundaries are midpoints of adjacent centroids."""
509 centroids, boundaries = solve_lloyd_max(128, bits)
510 for i in range(len(boundaries)):
511 expected = (centroids[i] + centroids[i + 1]) / 2.0
512 assert abs(boundaries[i].item() - expected.item()) < 1e-6
513
514 def test_solve_deterministic(self):
515 c1, b1 = solve_lloyd_max(128, 3)
516 c2, b2 = solve_lloyd_max(128, 3)
517 assert torch.equal(c1, c2)
518 assert torch.equal(b1, b2)
519
520 def test_solve_dtype_float32(self):
521 centroids, boundaries = solve_lloyd_max(128, 3)
522 assert centroids.dtype == torch.float32
523 assert boundaries.dtype == torch.float32
524
525 @pytest.mark.parametrize("bits", [3, 4])
526 def test_centroids_match_scipy_reference(self, bits):
527 """Verify _trapz(n=200) centroids match scipy.integrate.quad reference.
528
529 This ensures our scipy-free trapezoid integration doesn't silently
530 drift from the published Lloyd-Max quality.
531 """
532 pytest.importorskip("scipy")
533 from scipy.integrate import quad
534
535 d = 128
536 sigma2 = 1.0 / d
537 sigma = math.sqrt(sigma2)
538
539 def pdf(x):
540 return (1.0 / math.sqrt(2 * math.pi * sigma2)) * math.exp(
541 -x * x / (2 * sigma2)
542 )
543
544 n_levels = 2**bits
545 lo, hi = -3.5 * sigma, 3.5 * sigma
546 ref_centroids = [lo + (hi - lo) * (i + 0.5) / n_levels for i in range(n_levels)]
547 for _ in range(200):
548 boundaries = [
549 (ref_centroids[i] + ref_centroids[i + 1]) / 2.0
550 for i in range(n_levels - 1)
551 ]
552 edges = [lo * 3] + boundaries + [hi * 3]
553 new_centroids = []
554 for i in range(n_levels):
555 a, b = edges[i], edges[i + 1]
556 num, _ = quad(lambda x: x * pdf(x), a, b)
557 den, _ = quad(pdf, a, b)
558 new_centroids.append(num / den if den > 1e-15 else ref_centroids[i])
559 if (
560 max(abs(new_centroids[i] - ref_centroids[i]) for i in range(n_levels))
561 < 1e-10
562 ):
563 break
564 ref_centroids = new_centroids
565
566 # Compare our _trapz centroids against scipy reference
567 our_centroids, _ = solve_lloyd_max(d, bits)
568 ref_t = torch.tensor(ref_centroids, dtype=torch.float32)
569 max_err = (our_centroids - ref_t).abs().max().item()
570 # _trapz(n=200) has ~O(h^2) error vs adaptive quad; 1e-3 is tight
571 # enough to catch regression while allowing trapezoid approximation.
572 assert max_err < 1e-3, (
573 f"d={d}, bits={bits}: max centroid error vs scipy = {max_err:.2e}"
574 )
575
576
577 # ============================================================================
578 # Rotation matrix tests (GPU required)
579 # ============================================================================
580
581 GPGPU_AVAILABLE = torch.cuda.is_available() or torch.xpu.is_available()
582 DEVICE_TYPE = current_platform.device_type
583
584
585 def generate_rotation_matrix(d: int, seed: int, device: str = "cpu") -> torch.Tensor:
586 """Haar-distributed random orthogonal matrix via QR (test/benchmark only)."""
587 gen = torch.Generator(device="cpu")
588 gen.manual_seed(seed)
589 G = torch.randn(d, d, generator=gen, device="cpu", dtype=torch.float32)
590 # torch.linalg.qr on CPU requires LAPACK, which some torch wheels
591 # (ROCm) ship without. Run QR on accelerator instead
592 qr_device = "cuda" if torch.cuda.is_available() else "cpu"
593 Q, R = torch.linalg.qr(G.to(qr_device))
594 diag_sign = torch.sign(torch.diag(R))
595 diag_sign[diag_sign == 0] = 1.0
596 Q = Q * diag_sign.unsqueeze(0)
597 return Q.to(device)
598
599
600 @pytest.mark.skipif(not GPGPU_AVAILABLE, reason="GPGPU not available")
601 class TestRotationMatrix:
602 """Tests for the QR-based rotation (standalone benchmarks only)."""
603
604 @pytest.mark.parametrize("dim", [64, 96, 128, 256])
605 def test_rotation_matrix_shape_and_orthogonal(self, dim):
606 Pi = generate_rotation_matrix(dim, seed=42, device=DEVICE_TYPE)
607 assert Pi.shape == (dim, dim)
608 eye = Pi @ Pi.T
609 assert torch.allclose(eye, torch.eye(dim, device=DEVICE_TYPE), atol=1e-5), (
610 f"Pi not orthogonal for dim={dim}"
611 )
612
613 def test_rotation_matrix_deterministic(self):
614 Pi1 = generate_rotation_matrix(128, seed=42)
615 Pi2 = generate_rotation_matrix(128, seed=42)
616 assert torch.equal(Pi1, Pi2)
617
618 def test_rotation_matrix_different_seeds(self):
619 Pi1 = generate_rotation_matrix(128, seed=42)
620 Pi2 = generate_rotation_matrix(128, seed=99)
621 assert not torch.equal(Pi1, Pi2)
622
623 def test_rotation_matrix_det_is_pm1(self):
624 """Orthogonal matrix determinant must be +1 or -1."""
625 Pi = generate_rotation_matrix(128, seed=42, device=DEVICE_TYPE)
626 det = torch.linalg.det(Pi)
627 assert abs(abs(det.item()) - 1.0) < 1e-4
628
629
630 # ============================================================================
631 # Hadamard rotation tests (serving path: _build_hadamard)
632 # ============================================================================
633
634
635 def _build_hadamard(d: int, device: str = "cpu") -> torch.Tensor:
636 """Reproduce the serving-path Hadamard construction."""
637 H = torch.tensor([[1.0]])
638 while H.shape[0] < d:
639 H = torch.cat([torch.cat([H, H], 1), torch.cat([H, -H], 1)], 0)
640 return (H / math.sqrt(d)).to(torch.device(device))
641
642
643 @pytest.mark.skipif(not GPGPU_AVAILABLE, reason="GPGPU not available")
644 class TestHadamardRotation:
645 """Tests for the Hadamard rotation used in serving."""
646
647 @pytest.mark.parametrize("dim", [64, 128, 256])
648 def test_hadamard_orthonormal(self, dim):
649 """H must be orthonormal: H @ H^T = I."""
650 H = _build_hadamard(dim, DEVICE_TYPE)
651 eye = H @ H.T
652 assert torch.allclose(eye, torch.eye(dim, device=DEVICE_TYPE), atol=1e-5), (
653 f"Hadamard not orthonormal for dim={dim}"
654 )
655
656 @pytest.mark.parametrize("dim", [64, 128, 256])
657 def test_hadamard_symmetric(self, dim):
658 """Sylvester Hadamard must be symmetric: H = H^T."""
659 H = _build_hadamard(dim, DEVICE_TYPE)
660 assert torch.allclose(H, H.T, atol=1e-6), (
661 f"Hadamard not symmetric for dim={dim}"
662 )
663
664
665 # ============================================================================
666 # Store → Decode round-trip test (GPU + Triton required)
667 # ============================================================================
668
669
670 @pytest.mark.skipif(not GPGPU_AVAILABLE, reason="GPGPU not available")
671 class TestStoreDecodeRoundTrip:
672 """End-to-end: store KV into TQ cache, decode, compare vs fp16 ref."""
673
674 @pytest.mark.parametrize(
675 "preset",
676 ["turboquant_k8v4", "turboquant_4bit_nc"],
677 )
678 def test_single_token_roundtrip(self, preset):
679 """Store 1 token, decode with query=key, check attention output.
680
681 For a single token with query=key, attention output should equal
682 the value (softmax over single key = 1.0). Quantization error
683 means we check cosine similarity rather than exact equality.
684 """
685 from vllm.model_executor.layers.quantization.turboquant.centroids import (
686 solve_lloyd_max,
687 )
688 from vllm.v1.attention.ops.triton_turboquant_decode import (
689 triton_turboquant_decode_attention,
690 )
691 from vllm.v1.attention.ops.triton_turboquant_store import (
692 triton_turboquant_store,
693 )
694
695 cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128)
696 D = 128
697 Hk = 4 # num_kv_heads
698 Hq = 4 # num_q_heads (no GQA for simplicity)
699 B = 1 # single token
700 block_size = 16
701 num_blocks = 1
702
703 device = torch.device(DEVICE_TYPE)
704
705 # Pure Hadamard rotation (symmetric: H = H^T, so Pi = PiT = H)
706 H = _build_hadamard(D, DEVICE_TYPE)
707 PiT = H
708 Pi = H
709
710 # Generate centroids
711 centroids, _ = solve_lloyd_max(D, cfg.centroid_bits)
712 centroids = centroids.float().to(device)
713 c_sorted, _ = centroids.sort()
714 midpoints = ((c_sorted[:-1] + c_sorted[1:]) / 2).to(device)
715
716 # Random K, V
717 torch.manual_seed(123)
718 key = torch.randn(B, Hk, D, device=device, dtype=torch.float16)
719 value = torch.randn(B, Hk, D, device=device, dtype=torch.float16)
720
721 # Allocate KV cache
722 padded_slot = cfg.slot_size_aligned
723 kv_cache = torch.zeros(
724 num_blocks,
725 block_size,
726 Hk,
727 padded_slot,
728 device=device,
729 dtype=torch.uint8,
730 )
731 slot_mapping = torch.tensor([0], device=device, dtype=torch.int32)
732
733 # Store
734 triton_turboquant_store(
735 key,
736 value,
737 kv_cache,
738 slot_mapping,
739 PiT,
740 midpoints,
741 mse_bits=cfg.key_mse_bits,
742 key_packed_size=cfg.key_packed_size,
743 value_quant_bits=cfg.effective_value_quant_bits,
744 key_fp8=cfg.key_fp8,
745 )
746
747 # Decode: use key as query so attention = softmax([1]) * V = V
748 query = key.expand(B, Hq, D).contiguous().to(torch.float16)
749 block_table = torch.tensor([[0]], device=device, dtype=torch.int32)
750 seq_lens = torch.tensor([1], device=device, dtype=torch.int32)
751
752 output = triton_turboquant_decode_attention(
753 query=query,
754 kv_cache=kv_cache,
755 block_table=block_table,
756 seq_lens=seq_lens,
757 Pi=Pi,
758 centroids=centroids,
759 scale=1.0 / math.sqrt(D),
760 mse_bits=cfg.key_mse_bits,
761 key_packed_size=cfg.key_packed_size,
762 value_quant_bits=cfg.effective_value_quant_bits,
763 key_fp8=cfg.key_fp8,
764 norm_correction=cfg.norm_correction,
765 PiT=PiT,
766 max_num_kv_splits=4,
767 )
768
769 # With single KV, output should approximate the stored value.
770 # Check per-head cosine similarity > threshold.
771 out_fp32 = output.float()
772 val_fp32 = value.expand(B, Hq, D).float()
773 for h in range(Hq):
774 cos_sim = torch.nn.functional.cosine_similarity(
775 out_fp32[0, h].unsqueeze(0),
776 val_fp32[0, h].unsqueeze(0),
777 ).item()
778 # FP8 keys should be very accurate; MSE keys have more error
779 threshold = 0.95 if cfg.key_fp8 else 0.85
780 assert cos_sim > threshold, (
781 f"Preset {preset} head {h}: cosine_sim={cos_sim:.4f} < {threshold}"
782 )