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, 08:30 UTC
optimize get_kv_cache_torch_dtype (#18531)
Signed-off-by: idellzheng <idellzheng@tencent.com>
chunxiaozheng committed
May 27, 2025 at 21:08 UTC
6b6d4961147220fb80f9cc7dcb74db478f9c9a23
1 file changed
+3
-4
vllm/utils.py
+3
-4
@@ -759,16 +759,15 @@ def get_kv_cache_torch_dtype(
759
model_dtype: Optional[Union[str, torch.dtype]] = None) -> torch.dtype:
760
if isinstance(cache_dtype, str):
761
if cache_dtype == "auto":
762
- if isinstance(model_dtype, str):
762
+ if isinstance(model_dtype,
763
+ str) and model_dtype in STR_DTYPE_TO_TORCH_DTYPE:
764
torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[model_dtype]
765
elif isinstance(model_dtype, torch.dtype):
766
torch_dtype = model_dtype
767
else:
768
raise ValueError(f"Invalid model dtype: {model_dtype}")
768
- elif cache_dtype in ["half", "bfloat16", "float"]:
769
+ elif cache_dtype in STR_DTYPE_TO_TORCH_DTYPE:
770
torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_dtype]
770
- elif cache_dtype == "fp8":
771
- torch_dtype = torch.uint8
771
else:
772
raise ValueError(f"Invalid kv cache dtype: {cache_dtype}")
773
elif isinstance(cache_dtype, torch.dtype):