diff --git a/tests/quantization/test_turboquant.py b/tests/quantization/test_turboquant.py index ccdc69074c7..f4880abcb5f 100644 --- a/tests/quantization/test_turboquant.py +++ b/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( diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index 99be12012cf..c8e5a9c6ac3 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/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: diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index b7199bddcb6..cc61f1f5acc 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/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 diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index d346a421277..598bddf47ab 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/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(