From fa4321de3d894c50c5ca0766dffa352d3fb07423 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Sun, 5 Jul 2026 04:20:48 -0400 Subject: [PATCH] [Bugfix][TurboQuant] Preserve KV cache dtype in backend shape (#47609) --- vllm/v1/worker/gpu/attn_utils.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 3906717b5b7..5fcc9053bf4 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -25,6 +25,7 @@ from vllm.v1.kv_cache_interface import ( KVCacheSpec, KVQuantMode, MambaSpec, + TQFullAttentionSpec, UniformTypeKVCacheSpecs, ) from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata @@ -307,6 +308,7 @@ 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( @@ -391,7 +393,10 @@ def _update_hybrid_attention_layout( # (quantization only changes the last dim), so this is a no-op today, # but it keeps both call sites consistent for skip layers. layer_cache_dtype = ( - "auto" if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE else cache_dtype + "auto" + if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE + and not isinstance(kv_cache_spec, TQFullAttentionSpec) + else cache_dtype ) block_dim = group.backend.get_kv_cache_block_dim( kernel_block_sizes[group.kv_cache_group_id],