[Bugfix][TurboQuant] Preserve KV cache dtype in backend shape (#47609)
Lucas Wilkinson committed
Jul 5, 2026 at 04:20 UTC
fa4321de3d894c50c5ca0766dffa352d3fb07423
1 file changed
+6
-1
vllm/v1/worker/gpu/attn_utils.py
+6
-1
@@ -25,6 +25,7 @@ from vllm.v1.kv_cache_interface import (
25
KVCacheSpec,
26
KVQuantMode,
27
MambaSpec,
28
+ TQFullAttentionSpec,
29
UniformTypeKVCacheSpecs,
30
)
31
from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata
@@ -307,6 +308,7 @@ def _reshape_kv_cache(
308
layer_cache_dtype = (
309
"auto"
310
if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE
311
+ and not isinstance(kv_cache_spec, TQFullAttentionSpec)
312
else cache_dtype
313
)
314
kv_cache_shape = group.backend.get_kv_cache_shape(
@@ -391,7 +393,10 @@ def _update_hybrid_attention_layout(
393
# (quantization only changes the last dim), so this is a no-op today,
394
# but it keeps both call sites consistent for skip layers.
395
layer_cache_dtype = (
394
- "auto" if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE else cache_dtype
396
+ "auto"
397
+ if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE
398
+ and not isinstance(kv_cache_spec, TQFullAttentionSpec)
399
+ else cache_dtype
400
)
401
block_dim = group.backend.get_kv_cache_block_dim(
402
kernel_block_sizes[group.kv_cache_group_id],