[Bugfix] Record device-agnostic torch.Event on XPU to fix async D2H image corruption [XPU] (#5571)

Signed-off-by: Thanaji Rao Thakkalapelli <thanaji.rao.thakkalapelli@intel.com>
Co-authored-by: Chendi.Xue <chendi.xue@intel.com>

tests/platforms/test_record_device_event.py

@@ -4,6 +4,8 @@
 
 import pytest
 
+from tests.helpers.mark import hardware_test
+
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
@@ -38,6 +40,71 @@ class TestCudaOmniPlatformRecordDeviceEvent:
         assert result is None
 
 
+class TestXPUOmniPlatformRecordDeviceEvent:
+    """XPU must record a device-agnostic ``torch.Event``, not ``torch.xpu.Event``.
+
+    The async diffusion output thread gates its D2H copy with
+    ``torch.Stream.wait_event`` on a generic ``torch.Stream``, and that binding
+    silently drops the dependency for a ``torch.xpu.Event`` — the copy then races
+    the still-running forward and the host reads a partially written tensor
+    (garbage rows at the bottom of the image).
+    """
+
+    def test_records_device_agnostic_torch_event(self, mocker):
+        # vLLM's XPUPlatform imports vllm_xpu_kernels at module scope, and that
+        # package is absent on non-XPU images.
+        pytest.importorskip("vllm_xpu_kernels")
+
+        from vllm_omni.platforms.xpu.platform import XPUOmniPlatform
+
+        mock_event = mocker.MagicMock()
+        mocker.patch("torch.Event", return_value=mock_event)
+        xpu_event = mocker.patch("torch.xpu.Event")
+
+        result = XPUOmniPlatform.record_device_event()
+
+        assert result is mock_event
+        mock_event.record.assert_called_once()
+        xpu_event.assert_not_called()
+
+    @hardware_test(res={"xpu": "B60"})
+    def test_recorded_event_orders_side_stream_d2h(self):
+        """On-device: the event must make a side-stream D2H wait for compute."""
+        import torch
+
+        if not torch.xpu.is_available():
+            pytest.skip("XPU not available")
+
+        from vllm_omni.platforms.xpu.platform import XPUOmniPlatform
+
+        device = torch.device("xpu", 0)
+        # Normalized so the matmul chain stays finite; the ``* 0.0`` term keeps
+        # every element exactly 7.0 while still depending on the whole chain.
+        base = torch.randn(4096, 4096, device=device) / 64.0
+        chained = base.clone()
+        for _ in range(40):
+            chained = chained @ base
+        tensor = torch.full_like(chained, 7.0) + chained * 0.0
+
+        event = XPUOmniPlatform.record_device_event()
+        assert event is not None
+
+        # Mirrors WorkerProc._async_output_loop.
+        d2h_stream = torch.Stream(device=device)
+        d2h_stream.wait_event(event)
+        previous = torch.accelerator.current_stream()
+        torch.accelerator.set_stream(d2h_stream)
+        try:
+            host = torch.empty(tensor.shape, dtype=tensor.dtype, pin_memory=True)
+            host.copy_(tensor, non_blocking=True)
+        finally:
+            torch.accelerator.set_stream(previous)
+        d2h_stream.synchronize()
+
+        copied_early = int((host != 7.0).sum())
+        assert copied_early == 0, f"{copied_early} elements copied before compute finished"
+
+
 class TestNPUOmniPlatformRecordDeviceEvent:
     """Test NPUOmniPlatform.record_device_event with mocked torch.npu."""
 

vllm_omni/platforms/xpu/platform.py

@@ -102,8 +102,20 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform):
 
     @classmethod
     def record_device_event(cls) -> torch.Event | None:
+        """Record an XPU event on the current stream to mark tensor readiness.
+
+        Deliberately a device-agnostic ``torch.Event`` rather than a
+        ``torch.xpu.Event``. The consumer (the async diffusion output thread)
+        waits with ``torch.Stream.wait_event`` on a generic ``torch.Stream``,
+        and that C-level binding silently no-ops for a ``torch.xpu.Event``
+        instead of enqueuing the dependency — the side stream then starts its
+        D2H copy while the compute stream is still writing the tensor, so the
+        host reads a partially-written image (garbage rows at the bottom of the
+        output). ``torch.Event`` dispatches through the accelerator hooks and
+        the wait is honored, which is the actual fix.
+        """
         try:
-            event = torch.xpu.Event()
+            event = torch.Event()
             event.record()
             return event
         except Exception: