[V1] V1 Enablement Oracle (#13726)

Signed-off-by: rshaw@neuralmagic.com <rshaw@neuralmagic.com>
Co-authored-by: rshaw@neuralmagic.com <rshaw@neuralmagic.com>
Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>
Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>
Co-authored-by: Michael Goin <michael@neuralmagic.com>

tests/models/test_initialization.py

@@ -6,6 +6,8 @@ import pytest
 from transformers import PretrainedConfig
 
 from vllm import LLM
+from vllm.engine.llm_engine import LLMEngine as V0LLMEngine
+from vllm.v1.engine.core import EngineCore as V1EngineCore
 
 from .registry import HF_EXAMPLE_MODELS
 
@@ -36,12 +38,18 @@ def test_can_initialize(model_arch):
         return hf_config
 
     # Avoid calling model.forward()
-    def _initialize_kv_caches(self) -> None:
+    def _initialize_kv_caches_v0(self) -> None:
         self.cache_config.num_gpu_blocks = 0
         self.cache_config.num_cpu_blocks = 0
 
-    with patch.object(LLM.get_engine_class(), "_initialize_kv_caches",
-                      _initialize_kv_caches):
+    def _initalize_kv_caches_v1(self, vllm_config):
+        # gpu_blocks (> 0), cpu_blocks
+        return 1, 0
+
+    with (patch.object(V0LLMEngine, "_initialize_kv_caches",
+                       _initialize_kv_caches_v0),
+          patch.object(V1EngineCore, "_initialize_kv_caches",
+                       _initalize_kv_caches_v1)):
         LLM(
             model_info.default,
             tokenizer=model_info.tokenizer,