[XPU][UT] Fix OOM and skip graph case (vllm-project#49287)

Signed-off-by: mayuyuace <qiming1.zhang@intel.com>
Signed-off-by: Qiming Zhang <qiming1.zhang@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>

Author: 2 people

tests/conftest.py

@@ -914,13 +914,13 @@ def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
-        from tests.utils import wait_for_rocm_memory_to_settle
+        from tests.utils import wait_for_memory_to_settle
 
         del self.model
         cleanup_dist_env_and_memory()
         # ROCm frees VRAM lazily; wait so a runner started right after this HF
         # model exits does not OOM on its startup memory guard.
-        wait_for_rocm_memory_to_settle(
+        wait_for_memory_to_settle(
             threshold_ratio=getattr(self, "threshold_ratios", None)
         )
         if hasattr(self, "threshold_ratios"):
@@ -998,9 +998,24 @@ def __init__(
             # V1 startup requires free_memory >= total * gpu_memory_utilization.
             # ROCm CI can hand a test a device that is still lazily releasing
             # VRAM from a previous process, so wait before constructing LLM.
-            from tests.utils import wait_for_rocm_memory_to_settle
-
-            wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)
+            from tests.utils import wait_for_memory_to_settle
+
+            wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)
+        elif current_platform.is_xpu():
+            # The XPU/oneAPI runtime keeps ~1 GiB of context resident in the
+            # parent pytest process for its whole lifetime (grown by in-process
+            # HfRunner models), and distributed tests additionally allocate a
+            # CCL context in the engine subprocess. The default utilization of
+            # 0.92 leaves too little headroom for both, so lower it on XPU when
+            # the caller did not request an explicit value.
+            if "gpu_memory_utilization" not in kwargs:
+                kwargs["gpu_memory_utilization"] = 0.9
+            gpu_memory_utilization = kwargs["gpu_memory_utilization"]
+            # XPU (Level Zero) can also release device memory lazily after a
+            # previous engine shuts down, so wait before constructing LLM.
+            from tests.utils import wait_for_memory_to_settle
+
+            wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)
 
         with init_ctx:
             self.llm = LLM(
@@ -1326,14 +1341,14 @@ def collective_rpc(self, *args, **kwargs):
     def __enter__(self):
         return self
 
-    def _wait_for_rocm_memory_release(self, gpu_memory_utilization: float) -> None:
-        from tests.utils import wait_for_rocm_memory_to_settle
+    def _wait_for_memory_release(self, gpu_memory_utilization: float) -> None:
+        from tests.utils import wait_for_memory_to_settle
 
         # V1 startup requires free_memory >= total * gpu_memory_utilization.
         # Wait for the complementary used-memory ratio so the next runner does
         # not fail the startup guard immediately after this runner exits. The
         # wait is bounded so cleanup failures fail this test instead of hanging.
-        wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)
+        wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)
 
     def __exit__(self, exc_type, exc_value, traceback):
         # Explicitly shutdown the engine core to release GPU resources
@@ -1359,7 +1374,7 @@ def __exit__(self, exc_type, exc_value, traceback):
         del self.llm
         torch._dynamo.reset()
         cleanup_dist_env_and_memory()
-        self._wait_for_rocm_memory_release(gpu_memory_utilization)
+        self._wait_for_memory_release(gpu_memory_utilization)
 
 
 @pytest.fixture(scope="session")
@@ -1746,7 +1761,7 @@ def clean_gpu_memory_between_tests():
 
     import gc
 
-    from tests.utils import wait_for_gpu_memory_to_clear, wait_for_rocm_memory_to_settle
+    from tests.utils import wait_for_gpu_memory_to_clear, wait_for_memory_to_settle
 
     num_gpus = torch.accelerator.device_count()
 
@@ -1755,7 +1770,7 @@ def _wait_for_settled_gpu_memory() -> None:
             return
         try:
             if current_platform.is_rocm():
-                wait_for_rocm_memory_to_settle()
+                wait_for_memory_to_settle()
             else:
                 wait_for_gpu_memory_to_clear(
                     devices=list(range(num_gpus)),

tests/models/language/generation/test_hybrid.py

@@ -9,6 +9,7 @@
 from tests.models.registry import HF_EXAMPLE_MODELS
 from tests.utils import multi_gpu_test
 from vllm import LLM
+from vllm.config import CUDAGraphMode
 from vllm.engine.arg_utils import EngineArgs
 from vllm.platforms import current_platform
 from vllm.sampling_params import SamplingParams
@@ -210,6 +211,8 @@ def test_mamba_cache_cg_padding(
     cudagraph_dispatcher.initialize_cudagraph_keys(
         vllm_config.compilation_config.cudagraph_mode
     )
+    if cudagraph_dispatcher.cudagraph_mode == CUDAGraphMode.NONE:
+        pytest.skip("CUDA/XPU graph is disabled.Please enable it to run this test. ")
     while (
         len(example_prompts)
         == cudagraph_dispatcher.dispatch(len(example_prompts))[1].num_tokens

tests/models/language/pooling/test_colbert.py

@@ -11,7 +11,7 @@
 import pytest
 import torch
 
-from tests.utils import wait_for_rocm_memory_to_settle
+from tests.utils import wait_for_memory_to_settle
 from vllm.distributed import cleanup_dist_env_and_memory
 from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score
 
@@ -165,7 +165,7 @@ def _hf_colbert_model(model_name: str, hf_spec: dict, device: torch.device):
     finally:
         del hf_model, linear_weight
         cleanup_dist_env_and_memory()
-        wait_for_rocm_memory_to_settle()
+        wait_for_memory_to_settle()
 
 
 def _assert_embeddings_close(vllm_outputs, hf_embeddings):

tests/models/multimodal/generation/test_voxtral_realtime.py

@@ -104,9 +104,9 @@ async def async_engine():
     from vllm.platforms import current_platform
 
     if current_platform.is_rocm():
-        from tests.utils import wait_for_rocm_memory_to_settle
+        from tests.utils import wait_for_memory_to_settle
 
-        wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)
+        wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)
 
     engine_args = AsyncEngineArgs(**ENGINE_CONFIG)
     llm = AsyncLLM.from_engine_args(engine_args)
@@ -122,9 +122,9 @@ async def async_engine():
         from vllm.distributed import cleanup_dist_env_and_memory
 
         cleanup_dist_env_and_memory()
-        from tests.utils import wait_for_rocm_memory_to_settle
+        from tests.utils import wait_for_memory_to_settle
 
-        wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)
+        wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)
 
 
 def test_voxtral_realtime_forward(audio_assets, tokenizer, vllm_runner, monkeypatch):

tests/utils.py

@@ -1540,6 +1540,13 @@ def record_gpu_memory_usage_stats(
             mem_info = amdsmi_get_gpu_vram_usage(dev_handle)
             gb_used = mem_info["vram_used"] / 2**10
             gb_total = mem_info["vram_total"] / 2**10
+        elif current_platform.is_xpu():
+            # nvml/amdsmi are unavailable on XPU. Query device memory through
+            # torch.accelerator.get_memory_info, which the XPU platform patches
+            # to return (free, total) bytes via Level Zero.
+            free_b, total_b = torch.accelerator.get_memory_info(device)
+            gb_used = (total_b - free_b) / 2**30
+            gb_total = total_b / 2**30
         else:
             dev_handle = get_nvml_device_handle(device)
             mem_info = nvmlDeviceGetMemoryInfo(dev_handle)
@@ -1680,19 +1687,19 @@ def wait_for_gpu_memory_to_clear(
         time.sleep(poll_interval_s)
 
 
-def wait_for_rocm_memory_to_settle(
+def wait_for_memory_to_settle(
     *,
     threshold_ratio: float | dict[int, float] | None = 0.1,
     timeout_s: float = 240,
 ) -> None:
-    """Block until ROCm device VRAM usage drops below ``threshold_ratio``.
+    """Block until ROCm or XPU device VRAM usage drops below ``threshold_ratio``.
 
-    ROCm reclaims GPU memory more lazily than CUDA, so back-to-back model
+    ROCm and XPU reclaims GPU memory more lazily than CUDA, so back-to-back model
     loads in a single test process can OOM the *next* engine/model startup
     even after ``cleanup_dist_env_and_memory``. This gives the driver time to
     actually release VRAM before the next allocation. No-op off ROCm.
     """
-    if not current_platform.is_rocm():
+    if not current_platform.is_rocm() and not current_platform.is_xpu():
         return
 
     num_gpus = current_platform.device_count()

tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh

@@ -112,7 +112,7 @@ wait_for_server() {
 # Function to clean up previous instances
 wait_for_gpu_memory_release() {
   if [[ "$SMI_BIN" == *"rocm"* ]]; then
-    PYTHONPATH="${GIT_ROOT}" python3 -c "from tests.utils import wait_for_rocm_memory_to_settle; wait_for_rocm_memory_to_settle()"
+    PYTHONPATH="${GIT_ROOT}" python3 -c "from tests.utils import wait_for_memory_to_settle; wait_for_memory_to_settle()"
   fi
 }