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, 07:30 UTC
[Feature] TurboQuant: support hybrid models and uniform quantization (#39931)
Signed-off-by: JartX <sagformas@epdcenter.es> Signed-off-by: Jim Smith <jhsmith0@me.com> Co-authored-by: Jim Smith <jhsmith0@me.com> Co-authored-by: Sandermage <sandermage@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com>
JartX committed
May 5, 2026 at 02:14 UTC
4f2af1a7c03aae2b3227dd7e69d726104d44a711
4 files changed
+186
-27
tests/quantization/test_turboquant.py
+82
-4
@@ -182,22 +182,100 @@ class TestTurboQuantConfig:
182
183
# ---- Boundary skip layers ----
184
185
+ @staticmethod
186
+ def _dense_model_config(num_layers):
187
+ from types import SimpleNamespace
188
+
189
+ return SimpleNamespace(
190
+ is_hybrid=False,
191
+ hf_text_config=SimpleNamespace(num_hidden_layers=num_layers),
192
+ )
193
+
194
def test_boundary_skip_layers_basic(self):
186
- layers = TurboQuantConfig.get_boundary_skip_layers(32)
195
+ mc = self._dense_model_config(32)
196
+ layers = TurboQuantConfig.get_boundary_skip_layers(mc)
197
assert layers == ["0", "1", "30", "31"]
198
199
def test_boundary_skip_layers_zero(self):
190
- assert TurboQuantConfig.get_boundary_skip_layers(32, 0) == []
200
+ mc = self._dense_model_config(32)
201
+ assert TurboQuantConfig.get_boundary_skip_layers(mc, 0) == []
202
203
def test_boundary_skip_layers_small_model(self):
193
- layers = TurboQuantConfig.get_boundary_skip_layers(4)
204
+ mc = self._dense_model_config(4)
205
+ layers = TurboQuantConfig.get_boundary_skip_layers(mc)
206
assert layers == ["0", "1", "2", "3"]
207
208
def test_boundary_skip_layers_cap_at_half(self):
197
- layers = TurboQuantConfig.get_boundary_skip_layers(8, 10)
209
+ mc = self._dense_model_config(8)
210
+ layers = TurboQuantConfig.get_boundary_skip_layers(mc, 10)
211
assert len(layers) == 8
212
213
214
+class TestHybridAttentionIndices:
215
+ """Regression tests for boundary protection on hybrid models.
216
+
217
+ Hybrid models (attention + Mamba / linear-attention) identify KV-carrying
218
+ layers via layer_types / layers_block_type / attn_type_list. The helper
219
+ must return the *global* layer indices of the full-attention layers so
220
+ that kv_cache_dtype_skip_layers matches what extract_layer_index(prefix)
221
+ reports on the Attention layers at runtime.
222
+ """
223
+
224
+ @staticmethod
225
+ def _fake_model_config(text_cfg=None, hf_cfg=None):
226
+ from types import SimpleNamespace
227
+
228
+ return SimpleNamespace(
229
+ hf_text_config=text_cfg if text_cfg is not None else SimpleNamespace(),
230
+ hf_config=hf_cfg if hf_cfg is not None else SimpleNamespace(),
231
+ )
232
+
233
+ def test_layer_types_full_attention(self):
234
+ from vllm.model_executor.layers.quantization.turboquant.config import (
235
+ _get_full_attention_layer_indices,
236
+ )
237
+
238
+ cfg = type("C", (), {})()
239
+ cfg.layer_types = [
240
+ "linear_attention",
241
+ "linear_attention",
242
+ "full_attention",
243
+ "linear_attention",
244
+ "full_attention",
245
+ "full_attention",
246
+ ]
247
+ mc = self._fake_model_config(text_cfg=cfg)
248
+ assert _get_full_attention_layer_indices(mc) == [2, 4, 5]
249
+
250
+ def test_layers_block_type_jamba(self):
251
+ from vllm.model_executor.layers.quantization.turboquant.config import (
252
+ _get_full_attention_layer_indices,
253
+ )
254
+
255
+ cfg = type("C", (), {})()
256
+ cfg.layers_block_type = ["mamba", "attention", "mamba", "attention"]
257
+ mc = self._fake_model_config(text_cfg=cfg)
258
+ assert _get_full_attention_layer_indices(mc) == [1, 3]
259
+
260
+ def test_attn_type_list_minimax(self):
261
+ from vllm.model_executor.layers.quantization.turboquant.config import (
262
+ _get_full_attention_layer_indices,
263
+ )
264
+
265
+ hf = type("C", (), {})()
266
+ hf.attn_type_list = [0, 1, 0, 1, 1]
267
+ mc = self._fake_model_config(hf_cfg=hf)
268
+ assert _get_full_attention_layer_indices(mc) == [1, 3, 4]
269
+
270
+ def test_no_hybrid_hints_returns_empty(self):
271
+ from vllm.model_executor.layers.quantization.turboquant.config import (
272
+ _get_full_attention_layer_indices,
273
+ )
274
+
275
+ mc = self._fake_model_config()
276
+ assert _get_full_attention_layer_indices(mc) == []
277
+
278
+
279
# ============================================================================
280
# Centroids tests (CPU-only)
281
# ============================================================================
vllm/engine/arg_utils.py
+3
-17
@@ -1699,29 +1699,15 @@ class EngineArgs:
1699
kv_offloading_backend=self.kv_offloading_backend,
1700
)
1701
1702
- # TurboQuant: auto-skip first/last 2 layers (boundary protection).
1703
- # These layers are most sensitive to quantization error.
1704
- # Users can add extra layers via --kv-cache-dtype-skip-layers.
1702
if resolved_cache_dtype.startswith("turboquant_"):
1706
- if model_config.is_hybrid:
1707
- raise NotImplementedError(
1708
- "TurboQuant KV cache is not supported for hybrid "
1709
- "(attention + Mamba) models. Boundary layer protection "
1710
- "requires uniform attention layers."
1711
- )
1703
from vllm.model_executor.layers.quantization.turboquant.config import (
1704
TurboQuantConfig,
1705
)
1706
1716
- num_layers = model_config.hf_text_config.num_hidden_layers
1717
- boundary = TurboQuantConfig.get_boundary_skip_layers(num_layers)
1707
+ boundary = TurboQuantConfig.get_boundary_skip_layers(model_config)
1708
existing = set(cache_config.kv_cache_dtype_skip_layers)
1719
- merged = sorted(existing | set(boundary), key=lambda x: int(x))
1720
- cache_config.kv_cache_dtype_skip_layers = merged
1721
- logger.info(
1722
- "TQ: skipping layers %s for boundary protection (num_layers=%d)",
1723
- merged,
1724
- num_layers,
1709
+ cache_config.kv_cache_dtype_skip_layers = sorted(
1710
+ existing | set(boundary), key=int
1711
)
1712
1713
ray_runtime_env = None
vllm/model_executor/layers/quantization/turboquant/config.py
+65
-6
@@ -2,8 +2,17 @@
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.
@@ -159,12 +168,34 @@ class TurboQuantConfig:
168
return s + (s % 2) # round up to even
169
170
@staticmethod
162
- def get_boundary_skip_layers(num_layers: int, n: int = 2) -> list[str]:
163
- """Get layer indices to skip TQ compression (boundary protection).
164
-
165
- Returns first N and last N layer indices as strings, suitable for
166
- kv_cache_dtype_skip_layers.
171
+ def get_boundary_skip_layers(
172
+ model_config: ModelConfig,
173
+ n: int = 2,
174
+ ) -> list[str]:
175
+ """Layer indices to skip TQ compression (boundary protection).
176
+
177
+ For hybrid models (attention + Mamba/linear-attention), boundary
178
+ protection is disabled — hybrids typically have only 8-12
179
+ full-attention layers and a hard n=2 on each side would cover
180
+ ~40 % of them. The dense GSM8K baselines that motivate n=2
181
+ don't apply to hybrids.
182
+
183
+ For dense models, skips first N and last N attention layers.
184
+ Empirically required for aggressive presets (k3v4_nc, 3bit_nc)
185
+ — without it GSM8K drops ~30 points on Qwen3-4B.
186
"""
187
+ if model_config.is_hybrid:
188
+ attn_indices = _get_full_attention_layer_indices(model_config)
189
+ if not attn_indices:
190
+ raise NotImplementedError(
191
+ "TurboQuant KV cache requires identifiable "
192
+ "full-attention layers, but none were found in "
193
+ "the hybrid model config."
194
+ )
195
+ logger.info("TQ hybrid: full-attention layers %s", attn_indices)
196
+ return []
197
+
198
+ num_layers = model_config.hf_text_config.num_hidden_layers
199
if n <= 0 or num_layers <= 0:
200
return []
201
n = min(n, num_layers // 2) # don't skip more than half
@@ -175,7 +206,7 @@ class TurboQuantConfig:
206
return [str(i) for i in indices]
207
208
@staticmethod
178
- def from_cache_dtype(cache_dtype: str, head_dim: int) -> "TurboQuantConfig":
209
+ def from_cache_dtype(cache_dtype: str, head_dim: int) -> TurboQuantConfig:
210
"""Create config from a named preset.
211
212
Valid presets: turboquant_k8v4, turboquant_4bit_nc, etc.
@@ -193,3 +224,31 @@ class TurboQuantConfig:
224
value_quant_bits=preset["value_quant_bits"],
225
norm_correction=preset["norm_correction"],
226
)
227
+
228
+
229
+def _get_full_attention_layer_indices(model_config: ModelConfig) -> list[int]:
230
+ """Global indices of full-attention layers in a hybrid model.
231
+
232
+ Covers the conventions used across vLLM: ``layer_types`` (Qwen3.5/Next),
233
+ ``layers_block_type`` (Jamba/Zamba2), ``attn_type_list`` (Minimax).
234
+ """
235
+ text_cfg = model_config.hf_text_config
236
+ hf_cfg = model_config.hf_config
237
+
238
+ layer_types = getattr(text_cfg, "layer_types", None)
239
+ if layer_types is not None:
240
+ return [
241
+ i for i, t in enumerate(layer_types) if t in ("full_attention", "attention")
242
+ ]
243
+
244
+ layers_block_type = getattr(text_cfg, "layers_block_type", None)
245
+ if layers_block_type is not None:
246
+ return [
247
+ i for i, t in enumerate(layers_block_type) if t in ("attention", "hybrid")
248
+ ]
249
+
250
+ attn_type_list = getattr(hf_cfg, "attn_type_list", None)
251
+ if attn_type_list is not None:
252
+ return [i for i, t in enumerate(attn_type_list) if t == 1]
253
+
254
+ return []
vllm/platforms/interface.py
+36
@@ -545,6 +545,42 @@ class Platform:
545
dtype=kv_cache_dtype,
546
kv_quant_mode=kv_quant_mode,
547
).page_size_bytes
548
+ elif cache_config.cache_dtype.startswith("turboquant_"):
549
+ # TQ has a packed K|V layout; the standard FullAttentionSpec
550
+ # formula over-sizes it and trips unify_kv_cache_spec_page_size
551
+ # when all attention layers are TQ. With mixed skip+TQ the skip
552
+ # layers still use the standard layout — take max so mamba
553
+ # padding covers the largest actual page.
554
+ from vllm.model_executor.layers.quantization.turboquant.config import (
555
+ TurboQuantConfig,
556
+ )
557
+ from vllm.v1.kv_cache_interface import TQFullAttentionSpec
558
+
559
+ tq_cfg = TurboQuantConfig.from_cache_dtype(
560
+ cache_config.cache_dtype, model_config.get_head_size()
561
+ )
562
+ tq_page = TQFullAttentionSpec(
563
+ block_size=1,
564
+ num_kv_heads=model_config.get_num_kv_heads(parallel_config),
565
+ head_size=model_config.get_head_size(),
566
+ head_size_v=model_config.get_head_size(),
567
+ dtype=kv_cache_dtype,
568
+ kv_quant_mode=kv_quant_mode,
569
+ tq_slot_size=tq_cfg.slot_size_aligned,
570
+ ).page_size_bytes
571
+ if cache_config.kv_cache_dtype_skip_layers:
572
+ skip_page = FullAttentionSpec(
573
+ block_size=1,
574
+ num_kv_heads=model_config.get_num_kv_heads(parallel_config),
575
+ head_size=model_config.get_head_size(),
576
+ dtype=model_config.dtype,
577
+ ).page_size_bytes
578
+ # lcm, not max: skip_page is often not a multiple of
579
+ # tq_page, so max would leave per-layer page sizes
580
+ # un-unifiable downstream.
581
+ attn_page_size_1_token = lcm(tq_page, skip_page)
582
+ else:
583
+ attn_page_size_1_token = tq_page
584
else:
585
attn_page_size_1_token = FullAttentionSpec(
586
block_size=1,