[Bugfix][TurboQuant] Add KV quant mode for turboquant (#50533)

Signed-off-by: Soila Kavulya <[email protected]>
Co-authored-by: Claude <[email protected]>

tests/quantization/test_turboquant.py

@@ -277,6 +277,29 @@ class TestHybridAttentionIndices:
         assert _get_full_attention_layer_indices(mc) == []
 
 
+class TestTurboQuantKVCacheSpec:
+    @pytest.mark.parametrize("preset", ALL_PRESETS)
+    def test_kv_cache_spec_sets_kv_quant_mode(self, preset):
+        from vllm.model_executor.layers.attention.attention import Attention
+        from vllm.v1.kv_cache_interface import KVQuantMode, TQFullAttentionSpec
+
+        layer = SimpleNamespace(
+            attn_type="decoder",
+            kv_cache_dtype=preset,
+            kv_cache_torch_dtype=torch.uint8,
+            head_size=128,
+            head_size_v=128,
+            num_kv_heads=4,
+            sliding_window=None,
+        )
+        vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=32))
+
+        spec = Attention.get_kv_cache_spec(layer, vllm_config)
+
+        assert isinstance(spec, TQFullAttentionSpec)
+        assert spec.kv_quant_mode == KVQuantMode.TURBOQUANT
+
+
 class TestTurboQuantWorkspaceReservation:
     @staticmethod
     def _fake_vllm_config(

vllm/model_executor/layers/attention/attention.py

@@ -680,6 +680,7 @@ class Attention(nn.Module, AttentionLayerBase):
                 head_size=self.head_size,
                 head_size_v=self.head_size,
                 dtype=self.kv_cache_torch_dtype,
+                kv_quant_mode=quant_mode,
                 tq_slot_size=tq_config.slot_size_aligned,
             )
         else:

vllm/v1/kv_cache_interface.py

@@ -43,6 +43,7 @@ class KVQuantMode(IntEnum):
     FP8_PER_TOKEN_HEAD = 3  # per-token-head dynamic scales for fp8
     INT4_PER_TOKEN_HEAD = 4  # packed 2×int4/byte, RHT + asymmetric zp
     NVFP4 = 5  # packed fp4 data + fp8 block scales
+    TURBOQUANT = 6  # Hadamard-rotated Lloyd-Max quant, packed K+V per slot
 
     @property
     def is_per_token_head(self) -> bool:
@@ -58,6 +59,11 @@ class KVQuantMode(IntEnum):
         """True for NVFP4 packed quantization mode."""
         return self == KVQuantMode.NVFP4
 
+    @property
+    def is_turboquant(self) -> bool:
+        """True for turboquant quantization mode."""
+        return self == KVQuantMode.TURBOQUANT
+
 
 def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode:
     """Map a ``kv_cache_dtype`` string to a :class:`KVQuantMode`."""
@@ -69,6 +75,8 @@ def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode:
         return KVQuantMode.FP8_PER_TOKEN_HEAD
     if kv_cache_dtype == "nvfp4":
         return KVQuantMode.NVFP4
+    if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("turboquant_"):
+        return KVQuantMode.TURBOQUANT
     if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("fp8"):
         return KVQuantMode.FP8_PER_TENSOR
     return KVQuantMode.NONE

vllm/v1/worker/gpu/attn_utils.py

@@ -27,7 +27,6 @@ from vllm.v1.kv_cache_interface import (
     KVCacheSpec,
     KVQuantMode,
     MambaSpec,
-    TQFullAttentionSpec,
     UniformTypeKVCacheSpecs,
 )
 from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata
@@ -324,7 +323,6 @@ def _reshape_kv_cache(
                 layer_cache_dtype = (
                     "auto"
                     if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE
-                    and not isinstance(kv_cache_spec, TQFullAttentionSpec)
                     else cache_dtype
                 )
                 kv_cache_shape = group.backend.get_kv_cache_shape(