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 4, 23:45 UTC
[Attention] TurboQuant: remove redundant random signs, add prior art attribution (#40194)
Signed-off-by: Dan Alistarh <d.alistarh@gmail.com>
Dan Alistarh committed
Apr 18, 2026 at 20:31 UTC
ed0622e3a809fe399b81009c9dbb9cf7299414e6
6 files changed
+53
-103
tests/quantization/test_turboquant.py
+14
-42
@@ -18,9 +18,6 @@ from vllm.model_executor.layers.quantization.turboquant.config import (
18
TQ_PRESETS,
19
TurboQuantConfig,
20
)
21
-from vllm.model_executor.layers.quantization.turboquant.quantizer import (
22
- generate_wht_signs,
23
-)
21
from vllm.platforms import current_platform
22
from vllm.utils.math_utils import next_power_of_2
23
@@ -393,7 +390,7 @@ class TestRotationMatrix:
390
391
392
# ============================================================================
396
-# WHT rotation tests (serving path: generate_wht_signs + _build_hadamard)
393
+# Hadamard rotation tests (serving path: _build_hadamard)
394
# ============================================================================
395
396
@@ -406,50 +403,26 @@ def _build_hadamard(d: int, device: str = "cpu") -> torch.Tensor:
403
404
405
@pytest.mark.skipif(not GPGPU_AVAILABLE, reason="GPGPU not available")
409
-class TestWHTRotation:
410
- """Tests for the WHT rotation actually used in serving."""
406
+class TestHadamardRotation:
407
+ """Tests for the Hadamard rotation used in serving."""
408
409
@pytest.mark.parametrize("dim", [64, 128, 256])
413
- def test_wht_orthonormal(self, dim):
414
- """signs * H must be orthonormal: (signs*H) @ (signs*H)^T = I."""
415
- signs = generate_wht_signs(dim, seed=42, device=DEVICE_TYPE)
410
+ def test_hadamard_orthonormal(self, dim):
411
+ """H must be orthonormal: H @ H^T = I."""
412
H = _build_hadamard(dim, DEVICE_TYPE)
417
- PiT = (signs.unsqueeze(1) * H).contiguous()
418
- eye = PiT @ PiT.T
413
+ eye = H @ H.T
414
assert torch.allclose(eye, torch.eye(dim, device=DEVICE_TYPE), atol=1e-5), (
420
- f"WHT rotation not orthonormal for dim={dim}"
415
+ f"Hadamard not orthonormal for dim={dim}"
416
)
417
418
@pytest.mark.parametrize("dim", [64, 128, 256])
424
- def test_wht_self_inverse(self, dim):
425
- """PiT should be self-inverse: PiT @ PiT = I (up to sign flip)."""
426
- signs = generate_wht_signs(dim, seed=42, device=DEVICE_TYPE)
419
+ def test_hadamard_symmetric(self, dim):
420
+ """Sylvester Hadamard must be symmetric: H = H^T."""
421
H = _build_hadamard(dim, DEVICE_TYPE)
428
- PiT = (signs.unsqueeze(1) * H).contiguous()
429
- Pi = PiT.T.contiguous()
430
- # Pi @ PiT should be identity (rotation then inverse)
431
- result = Pi @ PiT
432
- assert torch.allclose(result, torch.eye(dim, device=DEVICE_TYPE), atol=1e-5), (
433
- f"WHT rotation not self-inverse for dim={dim}"
422
+ assert torch.allclose(H, H.T, atol=1e-6), (
423
+ f"Hadamard not symmetric for dim={dim}"
424
)
425
436
- def test_wht_signs_deterministic(self):
437
- """Same seed must produce identical signs."""
438
- s1 = generate_wht_signs(128, seed=42)
439
- s2 = generate_wht_signs(128, seed=42)
440
- assert torch.equal(s1, s2)
441
-
442
- def test_wht_signs_different_seeds(self):
443
- """Different seeds must produce different signs."""
444
- s1 = generate_wht_signs(128, seed=42)
445
- s2 = generate_wht_signs(128, seed=99)
446
- assert not torch.equal(s1, s2)
447
-
448
- def test_wht_signs_are_pm1(self):
449
- """All sign values must be exactly +1 or -1."""
450
- signs = generate_wht_signs(128, seed=42)
451
- assert torch.all(signs.abs() == 1.0)
452
-
426
427
# ============================================================================
428
# Store → Decode round-trip test (GPU + Triton required)
@@ -491,11 +464,10 @@ class TestStoreDecodeRoundTrip:
464
465
device = torch.device(DEVICE_TYPE)
466
494
- # Generate rotation
495
- signs = generate_wht_signs(D, seed=42, device=device)
467
+ # Pure Hadamard rotation (symmetric: H = H^T, so Pi = PiT = H)
468
H = _build_hadamard(D, DEVICE_TYPE)
497
- PiT = (signs.unsqueeze(1) * H).contiguous().float()
498
- Pi = PiT.T.contiguous()
469
+ PiT = H
470
+ Pi = H
471
472
# Generate centroids
473
centroids, _ = solve_lloyd_max(D, cfg.centroid_bits)
vllm/model_executor/layers/attention/attention.py
+1
-18
@@ -406,33 +406,16 @@ class Attention(nn.Module, AttentionLayerBase):
406
def _init_turboquant_buffers(
407
self, cache_dtype: str, head_size: int, prefix: str
408
) -> None:
409
- """Initialize TurboQuant rotation/projection matrices and centroids."""
409
+ """Initialize TurboQuant centroids for Lloyd-Max quantization."""
410
from vllm.model_executor.layers.quantization.turboquant.centroids import (
411
get_centroids,
412
)
413
from vllm.model_executor.layers.quantization.turboquant.config import (
414
TurboQuantConfig,
415
)
416
- from vllm.model_executor.layers.quantization.turboquant.quantizer import (
417
- generate_wht_signs,
418
- )
416
417
tq_config = TurboQuantConfig.from_cache_dtype(cache_dtype, head_size)
418
422
- # Each layer needs a unique rotation matrix so quantization errors
423
- # don't correlate across layers. Stride must exceed max head_dim to
424
- # ensure non-overlapping RNG streams between adjacent layers.
425
- _TQ_LAYER_SEED_STRIDE = 1337
426
-
427
- from vllm.model_executor.models.utils import extract_layer_index
428
-
429
- layer_idx = extract_layer_index(prefix)
430
- seed = tq_config.seed + layer_idx * _TQ_LAYER_SEED_STRIDE
431
-
432
- self.register_buffer(
433
- "_tq_signs",
434
- generate_wht_signs(head_size, seed=seed),
435
- )
419
self.register_buffer(
420
"_tq_centroids",
421
get_centroids(head_size, tq_config.centroid_bits),
vllm/model_executor/layers/quantization/turboquant/__init__.py
+12
-5
@@ -1,12 +1,19 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""TurboQuant: Near-optimal KV-cache quantization for vLLM.
3
+"""TurboQuant: KV-cache quantization for vLLM.
4
5
-PolarQuant compression: random rotation + per-coordinate Lloyd-Max
6
-scalar quantization for keys, uniform quantization for values.
5
+Hadamard rotation + per-coordinate Lloyd-Max scalar quantization for
6
+keys, uniform quantization for values.
7
8
-Reference: "TurboQuant: Online Vector Quantization with Near-optimal
9
-Distortion Rate" (ICLR 2026), Zandieh et al.
8
+The technique implemented here consists of the scalar case of the HIGGS
9
+quantization method (Malinovskii et al., "Pushing the Limits of Large
10
+Language Model Quantization via the Linearity Theorem", NAACL 2025;
11
+preprint arXiv:2411.17525): rotation + optimized grid + optional
12
+re-normalization, applied to KV cache compression. A first application
13
+of this approach to KV-cache compression is in "Cache Me If You Must:
14
+Adaptive Key-Value Quantization for Large Language Models" (Shutova
15
+et al., ICML 2025; preprint arXiv:2501.19392). Both these references
16
+pre-date the TurboQuant paper (Zandieh et al., ICLR 2026).
17
"""
18
19
from vllm.model_executor.layers.quantization.turboquant.config import TurboQuantConfig
vllm/model_executor/layers/quantization/turboquant/config.py
+17
-7
@@ -36,10 +36,22 @@ TQ_PRESETS: dict[str, dict] = {
36
class TurboQuantConfig:
37
"""Configuration for TurboQuant KV-cache quantization.
38
39
- Uses PolarQuant (WHT rotation + Lloyd-Max scalar quantization) for keys
40
- and uniform quantization for values. QJL is intentionally omitted —
41
- community consensus (5+ independent groups) found it hurts attention
42
- quality by amplifying variance through softmax.
39
+ Applies Hadamard rotation followed by per-coordinate Lloyd-Max scalar
40
+ quantization for keys, and uniform quantization for values.
41
+
42
+ Historical note: this is the scalar case of the HIGGS quantization
43
+ method (Malinovskii et al., "Pushing the Limits of Large Language Model
44
+ Quantization via the Linearity Theorem", NAACL 2025; preprint
45
+ arXiv:2411.17525): rotation + optimized grid + optional re-normalization,
46
+ applied to KV cache compression. A first application of this approach to
47
+ KV-cache compression is in "Cache Me If You Must: Adaptive Key-Value
48
+ Quantization for Large Language Models" (Shutova et al., ICML 2025;
49
+ preprint arXiv:2501.19392). Both these references pre-date the
50
+ TurboQuant paper.
51
+
52
+ QJL is intentionally omitted — community consensus (5+ independent
53
+ groups) found it hurts attention quality by amplifying variance through
54
+ softmax.
55
56
Named presets (use via --kv-cache-dtype):
57
turboquant_k8v4: FP8 keys + 4-bit values, 2.6x, +1.17% PPL
@@ -53,8 +65,6 @@ class TurboQuantConfig:
65
rotation/MSE). 3-4 = Lloyd-Max MSE quantized keys.
66
value_quant_bits: Bits per value dimension for uniform quantization.
67
3 = 8 levels, 4 = 16 levels (default).
56
- seed: Base seed for deterministic random matrix generation.
57
- Actual seed per layer = seed + layer_idx * 1337.
68
norm_correction: Re-normalize centroid vectors to unit norm before
69
inverse rotation during dequant. Fixes quantization-induced norm
70
distortion, improving PPL by ~0.8% at 4-bit.
@@ -63,7 +73,7 @@ class TurboQuantConfig:
73
head_dim: int = 128
74
key_quant_bits: int = 3 # 3-4 = MSE keys, 8 = FP8 keys
75
value_quant_bits: int = 4 # 3-4 = uniform quantized values
66
- seed: int = 42
76
+ seed: int = 42 # kept for backward compatibility; no longer used internally
77
norm_correction: bool = False
78
79
@property
vllm/model_executor/layers/quantization/turboquant/quantizer.py
-18
@@ -2,23 +2,5 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""TurboQuant quantizer utilities.
4
5
-Serving path uses generate_wht_signs() for WHT rotation sign buffers.
5
Triton kernels handle all quantization, packing, and dequantization on GPU.
6
"""
8
-
9
-import torch
10
-
11
-_CPU = torch.device("cpu")
12
-
13
-
14
-def generate_wht_signs(d: int, seed: int, device: torch.device = _CPU) -> torch.Tensor:
15
- """Generate deterministic random ±1 signs for WHT rotation.
16
-
17
- Used with Walsh-Hadamard Transform for per-layer rotation randomization.
18
- Same seed derivation as QR (per-layer via seed + layer_idx * stride).
19
- """
20
- gen = torch.Generator(device="cpu")
21
- gen.manual_seed(seed)
22
- bits = torch.randint(0, 2, (d,), generator=gen, device="cpu")
23
- signs = bits.float() * 2 - 1
24
- return signs.to(device)
vllm/v1/attention/backends/turboquant_attn.py
+9
-13
@@ -279,29 +279,25 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
279
)
280
281
def _ensure_on_device(self, layer, device):
282
- """One-time derivation of TQ buffers (rotation matrices, midpoints).
282
+ """One-time derivation of TQ buffers (rotation matrix, midpoints).
283
284
- Registered buffers (_tq_signs, _tq_centroids) are already on the
285
- correct device via register_buffer + model.to(device).
284
+ The Hadamard rotation is shared across all layers: random sign
285
+ flips do not improve Lloyd-Max quantization quality because the
286
+ quantizer is symmetric around zero (sign-flipping a coordinate
287
+ maps it to the mirror centroid with identical distortion).
288
"""
289
if not hasattr(layer, "_tq_cached"):
288
- D = layer._tq_signs.shape[0]
289
- signs = layer._tq_signs.to(device=device, dtype=torch.float32)
290
+ D = self.head_size
291
291
- # WHT rotation: orthonormal + self-inverse, enabling future
292
+ # Pure Hadamard: orthonormal + symmetric (H = H^T), enabling
293
# in-kernel butterfly fusion and trivial inverse for continuation.
294
H = _build_hadamard(D, str(device))
294
- layer._tq_PiT = (signs.unsqueeze(1) * H).contiguous()
295
- layer._tq_Pi = layer._tq_PiT.T.contiguous()
295
+ layer._tq_PiT = H
296
+ layer._tq_Pi = H
297
298
c = layer._tq_centroids.to(device=device, dtype=torch.float32)
298
- # Precompute midpoints for threshold-based quantization
299
c_sorted, _ = c.sort()
300
layer._tq_midpoints = (c_sorted[:-1] + c_sorted[1:]) / 2
301
- # Decode buffers (_tq_mid_o_buf, _tq_output_buf, _tq_lse_buf)
302
- # are pre-allocated via register_buffer in Attention.__init__
303
- # and moved to GPU by model.to(device) — no allocation needed
304
- # here. The memory profiler sees them before KV cache sizing.
301
layer._tq_cached = True
302
303
def do_kv_cache_update(