From ca3ccfc9a3054e65e4267bb3e3cc95db48a98f09 Mon Sep 17 00:00:00 2001 From: Joshna-Medisetty Date: Fri, 18 Sep 2026 12:41:41 -0700 Subject: [PATCH] [XPU][Bugfix] Make AR-Diffusion KV preallocation device-agnostic (#7777) Signed-off-by: Joshna Medisetty Co-authored-by: Cursor --- .../diffusion/models/dreamzero/pipeline_dreamzero.py | 8 ++++---- vllm_omni/experimental/ar_diffusion/runner.py | 6 +++--- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py index 939544a9f..3ebda6b6e 100644 --- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py +++ b/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 # ----------------------------------------------------------------------- diff --git a/vllm_omni/experimental/ar_diffusion/runner.py b/vllm_omni/experimental/ar_diffusion/runner.py index 23da43054..740587f38 100644 --- a/vllm_omni/experimental/ar_diffusion/runner.py +++ b/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."""