@tduncan / Nvidia-Optimization-Stack-VLLM / commits / f4b161d7fc

[XPU][UT] Fix OOM and skip graph case (#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>

Qiming Zhang committed Aug 17, 2026 at 18:47 UTC f4b161d7fca438bfe29509984759be1943a5aa88

6 files changed +47 -22

tests/conftest.py

+26 -11

@@ -914,13 +914,13 @@ class HfRunner:

914 return self

915

916 def __exit__(self, exc_type, exc_value, traceback):

917 - from tests.utils import wait_for_rocm_memory_to_settle

917 + from tests.utils import wait_for_memory_to_settle

918

919 del self.model

920 cleanup_dist_env_and_memory()

921 # ROCm frees VRAM lazily; wait so a runner started right after this HF

922 # model exits does not OOM on its startup memory guard.

923 - wait_for_rocm_memory_to_settle(

923 + wait_for_memory_to_settle(

924 threshold_ratio=getattr(self, "threshold_ratios", None)

925 )

926 if hasattr(self, "threshold_ratios"):

@@ -998,9 +998,24 @@ class VllmRunner:

998 # V1 startup requires free_memory >= total * gpu_memory_utilization.

999 # ROCm CI can hand a test a device that is still lazily releasing

1000 # VRAM from a previous process, so wait before constructing LLM.

1001 - from tests.utils import wait_for_rocm_memory_to_settle

1002 -

1003 - wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)

1001 + from tests.utils import wait_for_memory_to_settle

1002 +

1003 + wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)

1004 + elif current_platform.is_xpu():

1005 + # The XPU/oneAPI runtime keeps ~1 GiB of context resident in the

1006 + # parent pytest process for its whole lifetime (grown by in-process

1007 + # HfRunner models), and distributed tests additionally allocate a

1008 + # CCL context in the engine subprocess. The default utilization of

1009 + # 0.92 leaves too little headroom for both, so lower it on XPU when

1010 + # the caller did not request an explicit value.

1011 + if "gpu_memory_utilization" not in kwargs:

1012 + kwargs["gpu_memory_utilization"] = 0.9

1013 + gpu_memory_utilization = kwargs["gpu_memory_utilization"]

1014 + # XPU (Level Zero) can also release device memory lazily after a

1015 + # previous engine shuts down, so wait before constructing LLM.

1016 + from tests.utils import wait_for_memory_to_settle

1017 +

1018 + wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)

1019

1020 with init_ctx:

1021 self.llm = LLM(

@@ -1326,14 +1341,14 @@ class VllmRunner:

1341 def __enter__(self):

1342 return self

1343

1329 - def _wait_for_rocm_memory_release(self, gpu_memory_utilization: float) -> None:

1330 - from tests.utils import wait_for_rocm_memory_to_settle

1344 + def _wait_for_memory_release(self, gpu_memory_utilization: float) -> None:

1345 + from tests.utils import wait_for_memory_to_settle

1346

1347 # V1 startup requires free_memory >= total * gpu_memory_utilization.

1348 # Wait for the complementary used-memory ratio so the next runner does

1349 # not fail the startup guard immediately after this runner exits. The

1350 # wait is bounded so cleanup failures fail this test instead of hanging.

1336 - wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)

1351 + wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)

1352

1353 def __exit__(self, exc_type, exc_value, traceback):

1354 # Explicitly shutdown the engine core to release GPU resources

@@ -1359,7 +1374,7 @@ class VllmRunner:

1374 del self.llm

1375 torch._dynamo.reset()

1376 cleanup_dist_env_and_memory()

1362 - self._wait_for_rocm_memory_release(gpu_memory_utilization)

1377 + self._wait_for_memory_release(gpu_memory_utilization)

1378

1379

1380 @pytest.fixture(scope="session")

@@ -1746,7 +1761,7 @@ def clean_gpu_memory_between_tests():

1761

1762 import gc

1763

1749 - from tests.utils import wait_for_gpu_memory_to_clear, wait_for_rocm_memory_to_settle

1764 + from tests.utils import wait_for_gpu_memory_to_clear, wait_for_memory_to_settle

1765

1766 num_gpus = torch.accelerator.device_count()

1767

@@ -1755,7 +1770,7 @@ def clean_gpu_memory_between_tests():

1770 return

1771 try:

1772 if current_platform.is_rocm():

1758 - wait_for_rocm_memory_to_settle()

1773 + wait_for_memory_to_settle()

1774 else:

1775 wait_for_gpu_memory_to_clear(

1776 devices=list(range(num_gpus)),

tests/models/language/generation/test_hybrid.py

@@ -9,6 +9,7 @@ import pytest

9 from tests.models.registry import HF_EXAMPLE_MODELS

10 from tests.utils import multi_gpu_test

11 from vllm import LLM

12 +from vllm.config import CUDAGraphMode

13 from vllm.engine.arg_utils import EngineArgs

14 from vllm.platforms import current_platform

15 from vllm.sampling_params import SamplingParams

@@ -210,6 +211,8 @@ def test_mamba_cache_cg_padding(

211 cudagraph_dispatcher.initialize_cudagraph_keys(

212 vllm_config.compilation_config.cudagraph_mode

213 )

214 + if cudagraph_dispatcher.cudagraph_mode == CUDAGraphMode.NONE:

215 + pytest.skip("CUDA/XPU graph is disabled.Please enable it to run this test. ")

216 while (

217 len(example_prompts)

218 == cudagraph_dispatcher.dispatch(len(example_prompts))[1].num_tokens

tests/models/language/pooling/test_colbert.py

+2 -2

@@ -11,7 +11,7 @@ from contextlib import contextmanager

11 import pytest

12 import torch

14 -from tests.utils import wait_for_rocm_memory_to_settle

14 +from tests.utils import wait_for_memory_to_settle

15 from vllm.distributed import cleanup_dist_env_and_memory

16 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):

165 finally:

166 del hf_model, linear_weight

167 cleanup_dist_env_and_memory()

168 - wait_for_rocm_memory_to_settle()

168 + wait_for_memory_to_settle()

169

170

171 def _assert_embeddings_close(vllm_outputs, hf_embeddings):

tests/models/multimodal/generation/test_voxtral_realtime.py

+4 -4

@@ -104,9 +104,9 @@ async def async_engine():

104 from vllm.platforms import current_platform

105

106 if current_platform.is_rocm():

107 - from tests.utils import wait_for_rocm_memory_to_settle

107 + from tests.utils import wait_for_memory_to_settle

108

109 - wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)

109 + wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)

110

111 engine_args = AsyncEngineArgs(**ENGINE_CONFIG)

112 llm = AsyncLLM.from_engine_args(engine_args)

@@ -122,9 +122,9 @@ async def async_engine():

122 from vllm.distributed import cleanup_dist_env_and_memory

123

124 cleanup_dist_env_and_memory()

125 - from tests.utils import wait_for_rocm_memory_to_settle

125 + from tests.utils import wait_for_memory_to_settle

126

127 - wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)

127 + wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)

128

129

130 def test_voxtral_realtime_forward(audio_assets, tokenizer, vllm_runner, monkeypatch):

tests/utils.py

+11 -4

@@ -1540,6 +1540,13 @@ def record_gpu_memory_usage_stats(

1540 mem_info = amdsmi_get_gpu_vram_usage(dev_handle)

1541 gb_used = mem_info["vram_used"] / 2**10

1542 gb_total = mem_info["vram_total"] / 2**10

1543 + elif current_platform.is_xpu():

1544 + # nvml/amdsmi are unavailable on XPU. Query device memory through

1545 + # torch.accelerator.get_memory_info, which the XPU platform patches

1546 + # to return (free, total) bytes via Level Zero.

1547 + free_b, total_b = torch.accelerator.get_memory_info(device)

1548 + gb_used = (total_b - free_b) / 2**30

1549 + gb_total = total_b / 2**30

1550 else:

1551 dev_handle = get_nvml_device_handle(device)

1552 mem_info = nvmlDeviceGetMemoryInfo(dev_handle)

@@ -1680,19 +1687,19 @@ def wait_for_gpu_memory_to_clear(

1687 time.sleep(poll_interval_s)

1688

1689

1683 -def wait_for_rocm_memory_to_settle(

1690 +def wait_for_memory_to_settle(

1691 *,

1692 threshold_ratio: float | dict[int, float] | None = 0.1,

1693 timeout_s: float = 240,

1694 ) -> None:

1688 - """Block until ROCm device VRAM usage drops below ``threshold_ratio``.

1695 + """Block until ROCm or XPU device VRAM usage drops below ``threshold_ratio``.

1696

1690 - ROCm reclaims GPU memory more lazily than CUDA, so back-to-back model

1697 + ROCm and XPU reclaims GPU memory more lazily than CUDA, so back-to-back model

1698 loads in a single test process can OOM the *next* engine/model startup

1699 even after ``cleanup_dist_env_and_memory``. This gives the driver time to

1700 actually release VRAM before the next allocation. No-op off ROCm.

1701 """

1695 - if not current_platform.is_rocm():

1702 + if not current_platform.is_rocm() and not current_platform.is_xpu():

1703 return

1704

1705 num_gpus = current_platform.device_count()

tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh

+1 -1

@@ -112,7 +112,7 @@ wait_for_server() {

112 # Function to clean up previous instances

113 wait_for_gpu_memory_release() {

114 if [[ "$SMI_BIN" == *"rocm"* ]]; then

115 - PYTHONPATH="${GIT_ROOT}" python3 -c "from tests.utils import wait_for_rocm_memory_to_settle; wait_for_rocm_memory_to_settle()"

115 + PYTHONPATH="${GIT_ROOT}" python3 -c "from tests.utils import wait_for_memory_to_settle; wait_for_memory_to_settle()"

116 fi

117 }

118