A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 07:15 UTC
[TPU] fix kv cache dtype in model runner (#19244)
Signed-off-by: Chengji Yao <chengjiyao@google.com>
Chengji Yao committed
Jun 6, 2025 at 01:20 UTC
0d49483ea97705f531dd42383ecbb2476d7dfa2b
1 file changed
+9
-3
vllm/v1/worker/tpu_model_runner.py
+9
-3
@@ -29,7 +29,8 @@ from vllm.multimodal.inputs import (BatchedTensorInputs, MultiModalKwargs,
29
PlaceholderRange)
30
from vllm.multimodal.utils import group_mm_inputs_by_modality
31
from vllm.sequence import IntermediateTensors
32
-from vllm.utils import LayerBlockType, cdiv, is_pin_memory_available
32
+from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, LayerBlockType, cdiv,
33
+ is_pin_memory_available)
34
from vllm.v1.attention.backends.pallas import (PallasAttentionBackend,
35
PallasMetadata)
36
from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
@@ -138,6 +139,11 @@ class TPUModelRunner(LoRAModelRunnerMixin):
139
140
self.pin_memory = is_pin_memory_available()
141
self.dtype = self.model_config.dtype
142
+ if cache_config.cache_dtype == "auto":
143
+ self.kv_cache_dtype = self.dtype
144
+ else:
145
+ self.kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[
146
+ cache_config.cache_dtype]
147
self._hidden_states_dtype = self.dtype
148
149
self.is_multimodal_model = model_config.is_multimodal_model
@@ -480,7 +486,7 @@ class TPUModelRunner(LoRAModelRunnerMixin):
486
block_size=block_size,
487
num_kv_heads=attn_module.num_kv_heads,
488
head_size=attn_module.head_size,
483
- dtype=attn_module.dtype,
489
+ dtype=self.kv_cache_dtype,
490
sliding_window=attn_module.sliding_window,
491
use_mla=False,
492
)
@@ -489,7 +495,7 @@ class TPUModelRunner(LoRAModelRunnerMixin):
495
block_size=block_size,
496
num_kv_heads=attn_module.num_kv_heads,
497
head_size=attn_module.head_size,
492
- dtype=attn_module.dtype,
498
+ dtype=self.kv_cache_dtype,
499
use_mla=False,
500
)
501
elif attn_module.attn_type in (AttentionType.ENCODER,