default unit test kv cache dtype

Signed-off-by: NickLucche <[email protected]>

Author: NickLucche

vllm/attention/layer.py

@@ -174,7 +174,7 @@ def __init__(
             calculate_kv_scales = False
         self.block_size = block_size
         self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
-            kv_cache_dtype, vllm_config.model_config.dtype
+            kv_cache_dtype, vllm_config.model_config
         )
         if num_kv_heads is None:
             num_kv_heads = num_heads

vllm/utils/__init__.py

@@ -157,9 +157,12 @@ def set_default_torch_num_threads(num_threads: int):
     torch.set_num_threads(old_num_threads)
 
 
-def kv_cache_dtype_str_to_dtype(kv_cache_dtype: str, model_dtype: str) -> torch.dtype:
+def kv_cache_dtype_str_to_dtype(
+    kv_cache_dtype: str, model_config: ModelConfig
+) -> torch.dtype:
     if kv_cache_dtype == "auto":
-        return model_dtype
+        # Model config may not be specified for unit tests, default to float16
+        return model_config.dtype if model_config else torch.half
     return STR_DTYPE_TO_TORCH_DTYPE[kv_cache_dtype]
 
 

vllm/v1/worker/gpu_model_runner.py

@@ -235,7 +235,7 @@ def __init__(
         self.pin_memory = is_pin_memory_available()
         self.dtype = self.model_config.dtype
         self.kv_cache_dtype = kv_cache_dtype_str_to_dtype(
-            cache_config.cache_dtype, self.dtype
+            cache_config.cache_dtype, self.model_config
         )
 
         self.is_pooling_model = model_config.runner_type == "pooling"