| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """TurboQuant configuration.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | import logging |
| 8 | import math |
| 9 | from dataclasses import dataclass |
| 10 | from typing import TYPE_CHECKING |
| 11 | |
| 12 | if TYPE_CHECKING: |
| 13 | from vllm.config import ModelConfig |
| 14 | |
| 15 | logger = logging.getLogger(__name__) |
| 16 | |
| 17 | # Named TQ presets: each maps to frozen config parameters. |
| 18 | # key_quant_bits: 8 = FP8 keys, 3-4 = MSE (Lloyd-Max) quantized keys. |
| 19 | # value_quant_bits: 3-4 = uniform quantized values. |
| 20 | TQ_PRESETS: dict[str, dict] = { |
| 21 | "turboquant_k8v4": { |
| 22 | "key_quant_bits": 8, |
| 23 | "value_quant_bits": 4, |
| 24 | "norm_correction": False, |
| 25 | }, |
| 26 | "turboquant_4bit_nc": { |
| 27 | "key_quant_bits": 4, |
| 28 | "value_quant_bits": 4, |
| 29 | "norm_correction": True, |
| 30 | }, |
| 31 | "turboquant_k3v4_nc": { |
| 32 | "key_quant_bits": 3, |
| 33 | "value_quant_bits": 4, |
| 34 | "norm_correction": True, |
| 35 | }, |
| 36 | "turboquant_3bit_nc": { |
| 37 | "key_quant_bits": 3, |
| 38 | "value_quant_bits": 3, |
| 39 | "norm_correction": True, |
| 40 | }, |
| 41 | } |
| 42 | |
| 43 | |
| 44 | @dataclass |
| 45 | class TurboQuantConfig: |
| 46 | """Configuration for TurboQuant KV-cache quantization. |
| 47 | |
| 48 | Applies Hadamard rotation followed by per-coordinate Lloyd-Max scalar |
| 49 | quantization for keys, and uniform quantization for values. |
| 50 | |
| 51 | Historical note: the core algorithmic pattern implemented for key |
| 52 | quantization (Hadamard rotation followed by deterministic scalar |
| 53 | quantization and re-normalization) was originally established in DRIVE |
| 54 | (Vargaftik et al., NeurIPS 2021) and EDEN (Vargaftik et al., ICML |
| 55 | 2022). This formulation is also mathematically equivalent to the |
| 56 | scalar case of the HIGGS quantization method (Malinovskii et al., |
| 57 | "Pushing the Limits of Large Language Model Quantization via the |
| 58 | Linearity Theorem", NAACL 2025; preprint arXiv:2411.17525), which |
| 59 | subsequently generalized these concepts. |
| 60 | |
| 61 | A first application of this approach to KV-cache compression is in |
| 62 | "Cache Me If You Must: Adaptive Key-Value Quantization for Large |
| 63 | Language Models" (Shutova et al., ICML 2025; preprint |
| 64 | arXiv:2501.19392). All of these foundational and application |
| 65 | references pre-date the TurboQuant paper (Zandieh et al., ICLR 2026). |
| 66 | |
| 67 | QJL is intentionally omitted: community consensus (5+ independent |
| 68 | groups) found it hurts attention quality by amplifying variance |
| 69 | through softmax. |
| 70 | |
| 71 | Named presets (use via --kv-cache-dtype): |
| 72 | turboquant_k8v4: FP8 keys + 4-bit values, 2.6x, +1.17% PPL |
| 73 | turboquant_4bit_nc: 4-bit MSE keys + 4-bit values + NC, 3.8x, +2.71% |
| 74 | turboquant_k3v4_nc: 3-bit MSE keys + 4-bit values + NC, ~3.5x, +10.63% |
| 75 | turboquant_3bit_nc: 3-bit MSE keys + 3-bit values + NC, 4.9x, +20.59% |
| 76 | |
| 77 | Args: |
| 78 | head_dim: Attention head dimension (e.g. 64, 96, 128). |
| 79 | key_quant_bits: Bits for key quantization. 8 = FP8 keys (no |
| 80 | rotation/MSE). 3-4 = Lloyd-Max MSE quantized keys. |
| 81 | value_quant_bits: Bits per value dimension for uniform quantization. |
| 82 | 3 = 8 levels, 4 = 16 levels (default). |
| 83 | norm_correction: Re-normalize centroid vectors to unit norm before |
| 84 | inverse rotation during dequant. Fixes quantization-induced norm |
| 85 | distortion, improving PPL by ~0.8% at 4-bit. |
| 86 | |
| 87 | """ |
| 88 | |
| 89 | head_dim: int = 128 |
| 90 | key_quant_bits: int = 3 # 3-4 = MSE keys, 8 = FP8 keys |
| 91 | value_quant_bits: int = 4 # 3-4 = uniform quantized values |
| 92 | seed: int = 42 # kept for backward compatibility; no longer used internally |
| 93 | norm_correction: bool = False |
| 94 | |
| 95 | @property |
| 96 | def key_fp8(self) -> bool: |
| 97 | """Whether keys are stored as FP8 — no rotation/quantization needed.""" |
| 98 | return self.key_quant_bits == 8 |
| 99 | |
| 100 | @property |
| 101 | def mse_bits(self) -> int: |
| 102 | """MSE quantizer bit-width (determines centroid count: 2^mse_bits). |
| 103 | |
| 104 | For MSE key modes, equals key_quant_bits. |
| 105 | For FP8 key mode, falls back to value_quant_bits (centroids are still |
| 106 | needed for continuation-prefill dequant and decode kernel params). |
| 107 | """ |
| 108 | if self.key_fp8: |
| 109 | return self.value_quant_bits |
| 110 | return self.key_quant_bits |
| 111 | |
| 112 | @property |
| 113 | def key_mse_bits(self) -> int: |
| 114 | """MSE bits actually used for key quantization (0 if FP8 keys).""" |
| 115 | if self.key_fp8: |
| 116 | return 0 |
| 117 | return self.key_quant_bits |
| 118 | |
| 119 | @property |
| 120 | def centroid_bits(self) -> int: |
| 121 | """Bits for centroid generation — always non-zero.""" |
| 122 | return self.mse_bits |
| 123 | |
| 124 | @property |
| 125 | def n_centroids(self) -> int: |
| 126 | return 2**self.mse_bits |
| 127 | |
| 128 | @property |
| 129 | def key_packed_size(self) -> int: |
| 130 | """Packed bytes for a single KEY vector. |
| 131 | |
| 132 | FP8 mode (key_quant_bits=8): |
| 133 | head_dim bytes (1 byte per element, no overhead). |
| 134 | |
| 135 | TQ mode: |
| 136 | - MSE indices: ceil(head_dim * key_mse_bits / 8) bytes |
| 137 | - vec_norm: 2 bytes (float16) |
| 138 | """ |
| 139 | if self.key_fp8: |
| 140 | return self.head_dim # 1 byte per element |
| 141 | mse_bytes = math.ceil(self.head_dim * self.key_mse_bits / 8) |
| 142 | norm_bytes = 2 # vec_norm fp16 |
| 143 | return mse_bytes + norm_bytes |
| 144 | |
| 145 | @property |
| 146 | def effective_value_quant_bits(self) -> int: |
| 147 | """Actual bits used for value storage.""" |
| 148 | return self.value_quant_bits |
| 149 | |
| 150 | @property |
| 151 | def value_packed_size(self) -> int: |
| 152 | """Packed bytes for a single VALUE vector. |
| 153 | |
| 154 | Uniform quantization: ceil(head_dim * bits / 8) + 4 bytes (scale + zero fp16). |
| 155 | """ |
| 156 | data_bytes = math.ceil(self.head_dim * self.value_quant_bits / 8) |
| 157 | return data_bytes + 4 # +2 scale(fp16) +2 zero(fp16) |
| 158 | |
| 159 | @property |
| 160 | def slot_size(self) -> int: |
| 161 | """Total packed bytes per head per position (key + value combined). |
| 162 | |
| 163 | Layout: [key_packed | value_packed] |
| 164 | """ |
| 165 | return self.key_packed_size + self.value_packed_size |
| 166 | |
| 167 | @property |
| 168 | def slot_size_aligned(self) -> int: |
| 169 | """Slot size rounded up to next even number. |
| 170 | |
| 171 | Even-number is required so effective_head_size = slot_size_aligned // 2 |
| 172 | is integral. |
| 173 | """ |
| 174 | s = self.slot_size |
| 175 | return s + (s % 2) # round up to even |
| 176 | |
| 177 | @staticmethod |
| 178 | def get_boundary_skip_layers( |
| 179 | model_config: ModelConfig, |
| 180 | n: int = 2, |
| 181 | ) -> list[str]: |
| 182 | """Layer indices to skip TQ compression (boundary protection). |
| 183 | |
| 184 | For hybrid models (attention + Mamba/linear-attention), boundary |
| 185 | protection is disabled — hybrids typically have only 8-12 |
| 186 | full-attention layers and a hard n=2 on each side would cover |
| 187 | ~40 % of them. The dense GSM8K baselines that motivate n=2 |
| 188 | don't apply to hybrids. |
| 189 | |
| 190 | For dense models, skips first N and last N attention layers. |
| 191 | Empirically required for aggressive presets (k3v4_nc, 3bit_nc) |
| 192 | — without it GSM8K drops ~30 points on Qwen3-4B. |
| 193 | """ |
| 194 | if model_config.is_hybrid: |
| 195 | attn_indices = _get_full_attention_layer_indices(model_config) |
| 196 | if not attn_indices: |
| 197 | raise NotImplementedError( |
| 198 | "TurboQuant KV cache requires identifiable " |
| 199 | "full-attention layers, but none were found in " |
| 200 | "the hybrid model config." |
| 201 | ) |
| 202 | logger.info("TQ hybrid: full-attention layers %s", attn_indices) |
| 203 | return [] |
| 204 | |
| 205 | num_layers = model_config.hf_text_config.num_hidden_layers |
| 206 | if n <= 0 or num_layers <= 0: |
| 207 | return [] |
| 208 | n = min(n, num_layers // 2) # don't skip more than half |
| 209 | first = list(range(n)) |
| 210 | last = list(range(num_layers - n, num_layers)) |
| 211 | # Deduplicate (if num_layers <= 2*n) |
| 212 | indices = sorted(set(first + last)) |
| 213 | return [str(i) for i in indices] |
| 214 | |
| 215 | @staticmethod |
| 216 | def from_cache_dtype(cache_dtype: str, head_dim: int) -> TurboQuantConfig: |
| 217 | """Create config from a named preset. |
| 218 | |
| 219 | Valid presets: turboquant_k8v4, turboquant_4bit_nc, etc. |
| 220 | """ |
| 221 | if cache_dtype not in TQ_PRESETS: |
| 222 | valid = ", ".join(TQ_PRESETS.keys()) |
| 223 | raise ValueError( |
| 224 | f"Unknown TurboQuant cache dtype: {cache_dtype!r}. " |
| 225 | f"Valid presets: {valid}" |
| 226 | ) |
| 227 | preset = TQ_PRESETS[cache_dtype] |
| 228 | return TurboQuantConfig( |
| 229 | head_dim=head_dim, |
| 230 | key_quant_bits=preset["key_quant_bits"], |
| 231 | value_quant_bits=preset["value_quant_bits"], |
| 232 | norm_correction=preset["norm_correction"], |
| 233 | ) |
| 234 | |
| 235 | |
| 236 | def _get_full_attention_layer_indices(model_config: ModelConfig) -> list[int]: |
| 237 | """Global indices of full-attention layers in a hybrid model. |
| 238 | |
| 239 | Covers the conventions used across vLLM: ``layer_types`` (Qwen3.5/Next), |
| 240 | ``layers_block_type`` (Jamba/Zamba2), ``attn_type_list`` (Minimax). |
| 241 | """ |
| 242 | text_cfg = model_config.hf_text_config |
| 243 | hf_cfg = model_config.hf_config |
| 244 | |
| 245 | layer_types = getattr(text_cfg, "layer_types", None) |
| 246 | if layer_types is not None: |
| 247 | return [ |
| 248 | i for i, t in enumerate(layer_types) if t in ("full_attention", "attention") |
| 249 | ] |
| 250 | |
| 251 | layers_block_type = getattr(text_cfg, "layers_block_type", None) |
| 252 | if layers_block_type is not None: |
| 253 | return [ |
| 254 | i for i, t in enumerate(layers_block_type) if t in ("attention", "hybrid") |
| 255 | ] |
| 256 | |
| 257 | attn_type_list = getattr(hf_cfg, "attn_type_list", None) |
| 258 | if attn_type_list is not None: |
| 259 | return [i for i, t in enumerate(attn_type_list) if t == 1] |
| 260 | |
| 261 | return [] |