[Bugfix][XPU] Fix free-memory abort and num_speculative_steps crash on XPU (#5330)

Signed-off-by: Joshna Medisetty <joshna.medisetty@intel.com>
Co-authored-by: Chendi.Xue <chendi.xue@intel.com>

.buildkite/intel/scripts/run-xpu-test.sh

@@ -74,6 +74,6 @@ time timeout -k 30 30m docker run \
     cd /workspace/vllm-omni
     pytest -v -s -m "core_model and xpu and B60"
     pytest -v -s tests/diffusion/quantization/test_mxfp8_config.py
-    export VLLM_XPU_USE_SAMPLER_KERNEL=0    # NOTE: Remove this after vLLM v0.21.1 is merged. Fixes Qwen2-5 omni-expansion tests.
     pytest -v -s -m "advanced_model and xpu and B60"
+    pytest -v -s -m "omni and xpu and B60"
 '

tests/e2e/offline_inference/test_qwen3_tts_base.py

@@ -68,7 +68,7 @@ def get_prompt():
 
 @pytest.mark.advanced_model
 @pytest.mark.tts
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@hardware_test(res={"cuda": "L4", "xpu": "B60"}, num_cards=1)
 @pytest.mark.parametrize("omni_runner", tts_server_params, indirect=True)
 def test_text_to_audio_001(omni_runner, omni_runner_handler) -> None:
     """

vllm_omni/platforms/xpu/platform.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import os
+
 import torch
 from vllm.config import VllmConfig
 from vllm.config.kernel import IrOpPriorityConfig
@@ -13,6 +15,13 @@ from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
 logger = init_logger(__name__)
 
 
+# Use the native query; vLLM's XPU custom op reports free=0 in spawned workers.
+torch.accelerator.get_memory_info = lambda device=None: torch.xpu.mem_get_info(device)
+
+# The XPU sampler kernel is broken for omni on the current vLLM; use the native path.
+os.environ.setdefault("VLLM_XPU_USE_SAMPLER_KERNEL", "0")
+
+
 class XPUOmniPlatform(OmniPlatform, XPUPlatform):
     """XPU/Intel GPU implementation of OmniPlatform.
 

vllm_omni/platforms/xpu/worker/xpu_ar_worker.py

@@ -1,15 +1,22 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from vllm.logger import init_logger
 from vllm.v1.worker.xpu_worker import XPUWorker
 
 from vllm_omni.platforms.xpu.worker.xpu_ar_model_runner import XPUARModelRunner
 from vllm_omni.worker.mixins import OmniWorkerMixin
 
+logger = init_logger(__name__)
+
 
 class XPUARWorker(OmniWorkerMixin, XPUWorker):
     """XPU AR worker for thinker/talker stages in Omni model."""
 
     def init_device(self):
         super().init_device()
+        if self.use_v2_model_runner:
+            # OMNI: v2 model runner does not yet include omni hooks.
+            logger.warning("OMNI XPUARWorker forces v1 model runner for omni hooks.")
+            self.use_v2_model_runner = False
         self.model_runner: XPUARModelRunner = XPUARModelRunner(self.vllm_config, self.device)

vllm_omni/platforms/xpu/worker/xpu_generation_worker.py

@@ -1,15 +1,22 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from vllm.logger import init_logger
 from vllm.v1.worker.xpu_worker import XPUWorker
 
 from vllm_omni.platforms.xpu.worker.xpu_generation_model_runner import XPUGenerationModelRunner
 from vllm_omni.worker.mixins import OmniWorkerMixin
 
+logger = init_logger(__name__)
+
 
 class XPUGenerationWorker(OmniWorkerMixin, XPUWorker):
     """XPU generation worker for the code2wav (non-AR waveform generation) stage in the Omni model."""
 
     def init_device(self):
         super().init_device()
+        if self.use_v2_model_runner:
+            # OMNI: v2 model runner does not yet include omni hooks.
+            logger.warning("OMNI XPUGenerationWorker forces v1 model runner for omni hooks.")
+            self.use_v2_model_runner = False
         self.model_runner: XPUGenerationModelRunner = XPUGenerationModelRunner(self.vllm_config, self.device)