[Bugfix][TurboQuant] Add KV quant mode for turboquant (#50533)
Signed-off-by: Soila Kavulya <soila.p.kavulya@intel.com> Co-authored-by: Claude <noreply@anthropic.com>
Soila Kavulya committed
Jul 31, 2026 at 10:42 UTC
aef85aed5d8970052fa14407b716cdac1a8164ac
4 files changed
+32
-2
tests/quantization/test_turboquant.py
+23
@@ -277,6 +277,29 @@ class TestHybridAttentionIndices:
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.kv_cache_interface import KVQuantMode, TQFullAttentionSpec
285
+
286
+ layer = SimpleNamespace(
287
+ attn_type="decoder",
288
+ kv_cache_dtype=preset,
289
+ kv_cache_torch_dtype=torch.uint8,
290
+ head_size=128,
291
+ head_size_v=128,
292
+ num_kv_heads=4,
293
+ sliding_window=None,
294
+ )
295
+ vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=32))
296
+
297
+ spec = Attention.get_kv_cache_spec(layer, vllm_config)
298
+
299
+ assert isinstance(spec, TQFullAttentionSpec)
300
+ assert spec.kv_quant_mode == KVQuantMode.TURBOQUANT
301
+
302
+
303
class TestTurboQuantWorkspaceReservation:
304
@staticmethod
305
def _fake_vllm_config(
vllm/model_executor/layers/attention/attention.py
+1
@@ -680,6 +680,7 @@ class Attention(nn.Module, AttentionLayerBase):
680
head_size=self.head_size,
681
head_size_v=self.head_size,
682
dtype=self.kv_cache_torch_dtype,
683
+ kv_quant_mode=quant_mode,
684
tq_slot_size=tq_config.slot_size_aligned,
685
)
686
else:
vllm/v1/kv_cache_interface.py
+8
@@ -43,6 +43,7 @@ class KVQuantMode(IntEnum):
43
FP8_PER_TOKEN_HEAD = 3 # per-token-head dynamic scales for fp8
44
INT4_PER_TOKEN_HEAD = 4 # packed 2×int4/byte, RHT + asymmetric zp
45
NVFP4 = 5 # packed fp4 data + fp8 block scales
46
+ TURBOQUANT = 6 # Hadamard-rotated Lloyd-Max quant, packed K+V per slot
47
48
@property
49
def is_per_token_head(self) -> bool:
@@ -58,6 +59,11 @@ class KVQuantMode(IntEnum):
59
"""True for NVFP4 packed quantization mode."""
60
return self == KVQuantMode.NVFP4
61
62
+ @property
63
+ def is_turboquant(self) -> bool:
64
+ """True for turboquant quantization mode."""
65
+ return self == KVQuantMode.TURBOQUANT
66
+
67
68
def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode:
69
"""Map a ``kv_cache_dtype`` string to a :class:`KVQuantMode`."""
@@ -69,6 +75,8 @@ def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode:
75
return KVQuantMode.FP8_PER_TOKEN_HEAD
76
if kv_cache_dtype == "nvfp4":
77
return KVQuantMode.NVFP4
78
+ if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("turboquant_"):
79
+ return KVQuantMode.TURBOQUANT
80
if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("fp8"):
81
return KVQuantMode.FP8_PER_TENSOR
82
return KVQuantMode.NONE
vllm/v1/worker/gpu/attn_utils.py
-2
@@ -27,7 +27,6 @@ from vllm.v1.kv_cache_interface import (
27
KVCacheSpec,
28
KVQuantMode,
29
MambaSpec,
30
- TQFullAttentionSpec,
30
UniformTypeKVCacheSpecs,
31
)
32
from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata
@@ -324,7 +323,6 @@ def _reshape_kv_cache(
323
layer_cache_dtype = (
324
"auto"
325
if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE
327
- and not isinstance(kv_cache_spec, TQFullAttentionSpec)
326
else cache_dtype
327
)
328
kv_cache_shape = group.backend.get_kv_cache_shape(