[XPU][Bugfix] Make AR-Diffusion KV preallocation device-agnostic (#7777)

Signed-off-by: Joshna Medisetty <joshna.medisetty@intel.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py

@@ -720,7 +720,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
     def warmup_compile(self) -> None:
         """Warm up compiled text/image/VAE paths before timed inference."""
-        if not torch.cuda.is_available():
+        if not torch.accelerator.is_available():
             return
 
         state = self.state
@@ -801,12 +801,12 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         latents: tuple[torch.Tensor, torch.Tensor],
         do_true_cfg: bool,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Post-step sync: .contiguous() + cuda.synchronize()"""
+        """Post-step sync: .contiguous() + accelerator stream synchronize"""
         latents = tuple(t.contiguous() for t in latents)
         if do_true_cfg and get_classifier_free_guidance_world_size() > 1:
-            device = next((t.device for t in latents if t.is_cuda), None)
+            device = next((t.device for t in latents if t.device.type != "cpu"), None)
             if device is not None:
-                torch.cuda.current_stream(device).synchronize()
+                torch.accelerator.current_stream(device).synchronize()
         return latents
 
     # -----------------------------------------------------------------------

vllm_omni/experimental/ar_diffusion/runner.py

@@ -93,9 +93,9 @@ class ARDiffusionModelRunner(DiffusionModelRunner):
             self._warmup_ar_rollout()
 
     def _available_memory_bytes(self) -> int:
-        if self.device is None or torch.device(self.device).type != "cuda":
-            raise RuntimeError("AR-Diffusion KV preallocation currently requires a CUDA device")
-        return int(torch.cuda.mem_get_info(self.device)[0])
+        if self.device is None:
+            raise RuntimeError("AR-Diffusion KV preallocation requires an initialized device")
+        return current_omni_platform.get_free_memory(torch.device(self.device))
 
     def _preallocate_kv_cache(self, *, available_bytes: int | None = None) -> None:
         """Build pools solely from the pipeline capability and runner config."""