[Refactor] Replace and ban a few torch.cuda functions in favor of torch.accelerator replacements. (#3365)

Signed-off-by: Nick Cao <ncao@redhat.com>

benchmarks/diffusion/quantization_quality.py

@@ -151,7 +151,7 @@ def _generate_image(omni, args, prompt, seed):
     from vllm_omni.platforms import current_omni_platform
 
     generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
-    torch.cuda.reset_peak_memory_stats()
+    torch.accelerator.reset_peak_memory_stats()
     start = time.perf_counter()
     outputs = omni.generate(
         {"prompt": prompt},
@@ -163,7 +163,7 @@ def _generate_image(omni, args, prompt, seed):
         ),
     )
     elapsed = time.perf_counter() - start
-    peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
+    peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
 
     first = outputs[0]
     req_out = first.request_output[0] if hasattr(first, "request_output") else first
@@ -178,7 +178,7 @@ def _generate_video(omni, args, prompt, seed):
     from vllm_omni.platforms import current_omni_platform
 
     generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
-    torch.cuda.reset_peak_memory_stats()
+    torch.accelerator.reset_peak_memory_stats()
     start = time.perf_counter()
     outputs = omni.generate(
         {"prompt": prompt, "negative_prompt": ""},
@@ -192,7 +192,7 @@ def _generate_video(omni, args, prompt, seed):
         ),
     )
     elapsed = time.perf_counter() - start
-    peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
+    peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
 
     first = outputs[0]
     if hasattr(first, "request_output") and isinstance(first.request_output, list):
@@ -228,8 +228,8 @@ def _unload_omni(omni):
     del omni
     gc.collect()
     if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        torch.cuda.synchronize()
+        torch.accelerator.empty_cache()
+        torch.accelerator.synchronize()
 
 
 def run_benchmark(args):

benchmarks/distributed/omni_connectors/cross_node_mooncake_transfer_engine.py

@@ -614,8 +614,8 @@ Examples:
         if not torch.cuda.is_available():
             print("[ERROR] CUDA is not available but GPU mode was requested.")
             sys.exit(1)
-        if args.gpu_id >= torch.cuda.device_count():
-            print(f"[ERROR] GPU {args.gpu_id} not available. Found {torch.cuda.device_count()} GPUs.")
+        if args.gpu_id >= torch.accelerator.device_count():
+            print(f"[ERROR] GPU {args.gpu_id} not available. Found {torch.accelerator.device_count()} GPUs.")
             sys.exit(1)
         print(f"[INFO] Using GPU {args.gpu_id}: {torch.cuda.get_device_name(args.gpu_id)}")
 

examples/offline_inference/internvla_a1/end2end.py

@@ -172,7 +172,7 @@ def run_pipeline_forward(
 
 def _synchronize(device: str) -> None:
     if device.startswith("cuda") and torch.cuda.is_available():
-        torch.cuda.synchronize()
+        torch.accelerator.synchronize()
 
 
 def _latency_summary(values_ms: list[float]) -> dict[str, float]:

examples/offline_inference/text_to_speech/voxtral_tts/end2end.py

@@ -184,7 +184,7 @@ async def run_streaming(inputs, sampling_params_list, model_name, args, output_d
     )
 
     async_omni.shutdown()
-    torch.cuda.empty_cache()
+    torch.accelerator.empty_cache()
     gc.collect()
 
 
@@ -223,7 +223,7 @@ def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir
     print(f"RTF: {output_audio_dur / vllm_elapsed:.4f}")
 
     del llm
-    torch.cuda.empty_cache()
+    torch.accelerator.empty_cache()
     gc.collect()
 
 
@@ -379,7 +379,7 @@ def main(args: Any) -> None:
             f"--num-prompts ({args.num_prompts}) must be divisible by --concurrency ({args.concurrency})"
         )
 
-    torch.cuda.empty_cache()
+    torch.accelerator.empty_cache()
     gc.collect()
 
     if args.streaming:

pyproject.toml

@@ -142,6 +142,12 @@ ignore = [
 
 [tool.ruff.lint.flake8-tidy-imports.banned-api]
 "librosa".msg = "The librosa module is banned, use vllm.multimodal helpers instead"
+"torch.cuda.max_memory_allocated".msg = "Use torch.accelerator.max_memory_allocated"
+"torch.cuda.reset_peak_memory_stats".msg = "Use torch.accelerator.reset_peak_memory_stats"
+"torch.cuda.device_count".msg = "Use torch.accelerator.device_count"
+"torch.cuda.current_device".msg = "Use torch.accelerator.current_device_index"
+"torch.cuda.empty_cache".msg  = "Use torch.accelerator.empty_cache"
+"torch.cuda.synchronize".msg = "Use torch.accelerator.synchronize"
 
 [tool.mypy]
 python_version = "3.12, 3.13"

tests/dfx/reliability/test_reliability_qwen3_omni.py

@@ -57,7 +57,7 @@ DEPLOY_CONFIGS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "vll
 
 def _default_oom_device_spec() -> str:
     """Use currently visible CUDA ordinals to avoid invalid device index in sidecar."""
-    count = torch.cuda.device_count()
+    count = torch.accelerator.device_count()
     if count <= 0:
         return "0"
     return ",".join(str(i) for i in range(count))

tests/diffusion/quantization/test_quantization_quality.py

@@ -104,7 +104,7 @@ def _generate_image(omni, config: QualityTestConfig):
     generator = torch.Generator(
         device=current_omni_platform.device_type,
     ).manual_seed(config.seed)
-    torch.cuda.reset_peak_memory_stats()
+    torch.accelerator.reset_peak_memory_stats()
 
     outputs = omni.generate(
         {"prompt": config.prompt},
@@ -116,7 +116,7 @@ def _generate_image(omni, config: QualityTestConfig):
         ),
     )
 
-    peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
+    peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
     first = outputs[0]
     if hasattr(first, "images") and first.images:
         return first.images[0], peak_mem
@@ -135,7 +135,7 @@ def _generate_video(omni, config: QualityTestConfig):
     generator = torch.Generator(
         device=current_omni_platform.device_type,
     ).manual_seed(config.seed)
-    torch.cuda.reset_peak_memory_stats()
+    torch.accelerator.reset_peak_memory_stats()
 
     outputs = omni.generate(
         {"prompt": config.prompt, "negative_prompt": ""},
@@ -148,7 +148,7 @@ def _generate_video(omni, config: QualityTestConfig):
         ),
     )
 
-    peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
+    peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
     first = outputs[0]
     if hasattr(first, "request_output") and isinstance(first.request_output, list):
         inner = first.request_output[0]
@@ -190,8 +190,8 @@ def _unload(omni):
     del omni
     gc.collect()
     if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        torch.cuda.synchronize()
+        torch.accelerator.empty_cache()
+        torch.accelerator.synchronize()
 
 
 # ---------------------------------------------------------------------------

tests/diffusion/test_worker_wrapper_base.py

@@ -371,7 +371,7 @@ class TestCustomPipelineWorkerExtension:
 
     def test_re_init_pipeline_basic(self, mocker: MockerFixture, mock_od_config):
         """Test basic re_init_pipeline functionality."""
-        mocker.patch("torch.cuda.empty_cache")
+        mocker.patch("torch.accelerator.empty_cache")
         mocker.patch("gc.collect")
         mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
 
@@ -405,7 +405,7 @@ class TestCustomPipelineWorkerExtension:
     def test_re_init_pipeline_cleanup(self, mocker: MockerFixture, mock_od_config):
         """Test that re_init_pipeline properly cleans up old pipeline."""
         mock_gc_collect = mocker.patch("gc.collect")
-        mock_empty_cache = mocker.patch("torch.cuda.empty_cache")
+        mock_empty_cache = mocker.patch("torch.accelerator.empty_cache")
         mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
 
         wrapper = WorkerWrapperBase(
@@ -434,7 +434,7 @@ class TestCustomPipelineWorkerExtension:
 
     def test_re_init_pipeline_none_pipeline(self, mocker: MockerFixture, mock_od_config):
         """Test re_init_pipeline when pipeline is None."""
-        mocker.patch("torch.cuda.empty_cache")
+        mocker.patch("torch.accelerator.empty_cache")
         mocker.patch("gc.collect")
         mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
 
@@ -510,7 +510,7 @@ class TestCustomPipelineWorkerExtension:
 
     def test_re_init_pipeline_multiple_calls(self, mocker: MockerFixture, mock_od_config):
         """Test calling re_init_pipeline multiple times."""
-        mocker.patch("torch.cuda.empty_cache")
+        mocker.patch("torch.accelerator.empty_cache")
         mocker.patch("gc.collect")
         mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
 

tests/e2e/accuracy/conftest.py

@@ -230,7 +230,7 @@ def _build_accuracy_server_config(
     run_level: str,
     model_prefix: str,
 ) -> AccuracyServerConfig:
-    if torch.cuda.device_count() < 1:
+    if torch.accelerator.device_count() < 1:
         pytest.skip("Need at least 1 CUDA GPU for accuracy benchmark smoke tests.")
 
     if not generate_model:

tests/e2e/accuracy/test_diffusers_backend_similarity.py

@@ -160,7 +160,7 @@ def _run_diffusers_wan22_i2v(*, model: str, output_path: Path, conditioning_imag
         del pipe
         gc.collect()
         if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+            torch.accelerator.empty_cache()
         run_post_test_cleanup(enable_force=True)
 
 
@@ -240,7 +240,7 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> tuple[Image.I
         del pipe
         gc.collect()
         if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+            torch.accelerator.empty_cache()
         run_post_test_cleanup(enable_force=True)
 
 

tests/e2e/accuracy/test_ltx2_3_video_similarity.py

@@ -199,7 +199,7 @@ def _run_diffusers_baseline(model: str, output_dir: Path) -> list[Image.Image]:
         del pipe
         gc.collect()
         if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+            torch.accelerator.empty_cache()
         run_post_test_cleanup(enable_force=True)
 
 
@@ -278,7 +278,7 @@ def _run_with_custom_transformer(model: str, output_dir: Path) -> list[Image.Ima
     del pipe, result, our_transformer
     gc.collect()
     if torch.cuda.is_available():
-        torch.cuda.empty_cache()
+        torch.accelerator.empty_cache()
     return frames
 
 

tests/e2e/accuracy/test_qwen_image.py

@@ -98,7 +98,7 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image:
         del pipe
         gc.collect()
         if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+            torch.accelerator.empty_cache()
         run_post_test_cleanup(enable_force=True)
 
 

tests/e2e/accuracy/test_qwen_image_edit.py

@@ -109,7 +109,7 @@ def _run_diffusers_image_edit(
         del pipe
         gc.collect()
         if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+            torch.accelerator.empty_cache()
         run_post_test_cleanup(enable_force=True)
 
 

tests/e2e/accuracy/test_qwen_image_layered.py

@@ -125,7 +125,7 @@ def _run_diffusers_qwen_image_layered(*, model: str, input_image: Image.Image, o
         del pipe
         gc.collect()
         if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+            torch.accelerator.empty_cache()
         run_post_test_cleanup(enable_force=True)
 
 

tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py

@@ -572,7 +572,7 @@ def test_wan22_i2v_online_serving_generates_video(
     wan22_i2v_image_source: str | None,
     wan22_i2v_online_timeout_seconds: int,
 ) -> None:
-    if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
+    if not torch.cuda.is_available() or torch.accelerator.device_count() < 2:
         pytest.skip("Wan2.2 I2V similarity e2e test requires >= 2 CUDA GPUs.")
 
     _probe_binary("ffprobe")
@@ -597,7 +597,7 @@ def test_wan22_i2v_online_serving_generates_video(
 def test_wan22_i2v_serving_matches_diffusers_video_similarity(
     wan22_i2v_image_source: str | None,
 ) -> None:
-    if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
+    if not torch.cuda.is_available() or torch.accelerator.device_count() < 2:
         pytest.skip("Wan2.2 I2V similarity e2e test requires >= 2 CUDA GPUs.")
 
     _probe_binary("ffmpeg")

tests/e2e/offline_inference/test_hunyuanimage3_text2img.py

@@ -327,7 +327,7 @@ def _generate_image(omni: Omni, use_system_prompt: str | None) -> Image.Image:
     return _extract_generated_image(outputs)
 
 
-@pytest.mark.skipif(torch.cuda.device_count() < 8, reason="Need at least 8 CUDA GPUs for this test.")
+@pytest.mark.skipif(torch.accelerator.device_count() < 8, reason="Need at least 8 CUDA GPUs for this test.")
 @pytest.mark.parametrize("system_prompt_name,use_system_prompt,expected_embedding", SYSTEM_PROMPT_CASES)
 def test_system_prompt_scores(
     omni: Omni,

tests/e2e/offline_inference/test_omni_sleep_mode.py

@@ -26,7 +26,7 @@ def get_ack_info(ack, key, default=None):
 
 
 def get_dynamic_devices(stage_idx, num_stages, tp_size):
-    total_gpus = torch.cuda.device_count()
+    total_gpus = torch.accelerator.device_count()
     gpus_per_stage = tp_size
     start_idx = stage_idx * gpus_per_stage
     if start_idx + gpus_per_stage > total_gpus:
@@ -42,7 +42,7 @@ def get_dynamic_devices(stage_idx, num_stages, tp_size):
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.asyncio
 async def test_diffusion_model_sleep_tp(tp_size):
-    num_gpus = torch.cuda.device_count()
+    num_gpus = torch.accelerator.device_count()
     if num_gpus < tp_size:
         pytest.skip(f"Skipping TP={tp_size}")
 
@@ -88,7 +88,7 @@ async def test_diffusion_model_sleep_tp(tp_size):
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.asyncio
 async def test_multistage_sleep_h100(tp_size):
-    num_gpus = torch.cuda.device_count()
+    num_gpus = torch.accelerator.device_count()
     if num_gpus < tp_size * 2:
         pytest.skip("Not enough GPUs")
 

tests/e2e/offline_inference/test_quantization_fp8.py

@@ -67,7 +67,7 @@ def _generate_single_stage_image(
         omni_kwargs["quantization"] = quantization
 
     with OmniRunner(model, **omni_kwargs) as runner:
-        torch.cuda.reset_peak_memory_stats()
+        torch.accelerator.reset_peak_memory_stats()
 
         generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
         outputs = runner.omni.generate(
@@ -81,7 +81,7 @@ def _generate_single_stage_image(
             ),
         )
 
-        peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
+        peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
 
         first_output = outputs[0]
         assert first_output.final_output_type == "image"
@@ -116,7 +116,7 @@ def _generate_bagel_image(
     model_name = omni_kwargs.pop("model")
     with OmniRunner(model_name, **omni_kwargs) as runner:
         omni = runner.omni
-        torch.cuda.reset_peak_memory_stats()
+        torch.accelerator.reset_peak_memory_stats()
 
         params_list = omni.default_sampling_params_list
         if len(params_list) > 1:
@@ -134,7 +134,7 @@ def _generate_bagel_image(
             )
         )
 
-        peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
+        peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
 
         # Extract image
         generated_image = None
@@ -190,7 +190,7 @@ def test_single_stage_zimage_fp8_uses_less_memory():
         model="Tongyi-MAI/Z-Image-Turbo",
         quantization=None,
     )
-    torch.cuda.empty_cache()
+    torch.accelerator.empty_cache()
 
     _, mem_fp8 = _generate_single_stage_image(
         model="Tongyi-MAI/Z-Image-Turbo",
@@ -241,7 +241,7 @@ def test_single_stage_flux_fp8_uses_less_memory():
         width=512,
         num_inference_steps=4,
     )
-    torch.cuda.empty_cache()
+    torch.accelerator.empty_cache()
 
     _, mem_fp8 = _generate_single_stage_image(
         model="black-forest-labs/FLUX.1-dev",

tests/entrypoints/test_omni_sleep_mode.py

@@ -33,16 +33,16 @@ def get_vram_info(device_id: int) -> dict:
     """Obtain a snapshot of the specified GPU's memory (GiB)."""
     try:
         if current_omni_platform.is_rocm():
-            num_gpus = torch.cuda.device_count()
+            num_gpus = torch.accelerator.device_count()
             safe_id = device_id if device_id < num_gpus else 0
-            torch.cuda.synchronize(safe_id)
+            torch.accelerator.synchronize(safe_id)
             return {
                 "reserved": torch.cuda.memory_reserved(safe_id) / 1024**3,
                 "allocated": torch.cuda.memory_allocated(safe_id) / 1024**3,
             }
         else:
             with torch.cuda.device(device_id):
-                torch.cuda.synchronize()
+                torch.accelerator.synchronize()
                 return {
                     "reserved": torch.cuda.memory_reserved() / 1024**3,
                     "allocated": torch.cuda.memory_allocated() / 1024**3,
@@ -240,7 +240,7 @@ class TestOmniSleepMode:
             await diffusion_engine.wake_up(stage_ids=[0])
 
             get_vram_info(device_id)
-            torch.cuda.empty_cache()
+            torch.accelerator.empty_cache()
             await asyncio.sleep(2)
 
             initial_vram = get_vram_info(device_id)["reserved"]
@@ -253,7 +253,7 @@ class TestOmniSleepMode:
             await diffusion_engine.sleep(stage_ids=[0], level=2)
 
             await asyncio.sleep(3.0)
-            torch.cuda.empty_cache()
+            torch.accelerator.empty_cache()
 
             final_vram = get_vram_info(device_id)["reserved"]
             logger.info(f"GPU {device_id} Final VRAM after coordinated sleep: {final_vram:.2f} GiB")
@@ -276,7 +276,7 @@ class TestOmniSleepMode:
         device_id = 1
         try:
             get_vram_info(device_id)
-            torch.cuda.empty_cache()
+            torch.accelerator.empty_cache()
             vram_initial = get_vram_info(device_id)["reserved"]
             logger.info(f"Diffusion Initial VRAM: {vram_initial:.2f} GiB")
 
@@ -290,7 +290,7 @@ class TestOmniSleepMode:
 
             await asyncio.sleep(2)
             get_vram_info(device_id)
-            torch.cuda.empty_cache()
+            torch.accelerator.empty_cache()
 
             vram_sleeping = get_vram_info(device_id)["reserved"]
             logger.info(f"External VRAM measurement during Sleep: {vram_sleeping:.2f} GiB")
@@ -305,7 +305,7 @@ class TestOmniSleepMode:
 
             await asyncio.sleep(2)
             get_vram_info(device_id)
-            torch.cuda.empty_cache()
+            torch.accelerator.empty_cache()
             vram_restored = get_vram_info(device_id)["reserved"]
             logger.info(f"VRAM after Wake-up: {vram_restored:.2f} GiB")
 

tests/helpers/env.py

@@ -130,7 +130,7 @@ def wait_for_gpu_memory_to_clear(
                 )
 
             gc.collect()
-            torch.cuda.empty_cache()
+            torch.accelerator.empty_cache()
             time.sleep(5)
 
 
@@ -244,7 +244,7 @@ def run_pre_test_cleanup(enable_force: bool = False) -> None:
 
     print("Pre-test GPU status:")
 
-    num_gpus = torch.cuda.device_count()
+    num_gpus = torch.accelerator.device_count()
     if num_gpus > 0:
         try:
             wait_for_gpu_memory_to_clear(
@@ -262,7 +262,7 @@ def run_post_test_cleanup(enable_force: bool = False) -> None:
 
     if torch.cuda.is_available():
         gc.collect()
-        torch.cuda.empty_cache()
+        torch.accelerator.empty_cache()
 
         print("Post-test GPU status:")
         _print_gpu_processes()

tests/helpers/fixtures/env.py

@@ -11,7 +11,7 @@ def default_env():
     previous = {key: os.environ.get(key) for key in keys}
     os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = previous["VLLM_WORKER_MULTIPROC_METHOD"] or "spawn"
     os.environ["VLLM_TARGET_DEVICE"] = previous["VLLM_TARGET_DEVICE"] or (
-        "cuda" if torch.cuda.is_available() and torch.cuda.device_count() > 0 else "cpu"
+        "cuda" if torch.cuda.is_available() and torch.accelerator.device_count() > 0 else "cpu"
     )
     yield
     for key, value in previous.items():
@@ -51,7 +51,7 @@ def default_vllm_config():
     from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
 
     # Use CPU device if no GPU is available (e.g., in CI environments)
-    has_gpu = torch.cuda.is_available() and torch.cuda.device_count() > 0
+    has_gpu = torch.cuda.is_available() and torch.accelerator.device_count() > 0
     device = "cuda" if has_gpu else "cpu"
     device_config = DeviceConfig(device=device)
 

tests/model_executor/models/ming_flash_omni/test_talker_cfm.py

@@ -46,7 +46,7 @@ def _warmup_pipeline(cfm: CFM, aggregator: Aggregator, stop_head: torch.nn.Linea
         gen_lat = cfm.sample(llm_cond, lat_cond, y0, t, sde_args, sde_rnd)
         aggregator(gen_lat)
         stop_head(llm_cond[:, -1, :]).softmax(dim=-1)
-    torch.cuda.synchronize(device)
+    torch.accelerator.synchronize(device)
 
 
 def _build_pipeline():
@@ -95,7 +95,7 @@ class TestCFMGraphExecutor:
         his_lat = torch.randn(bsz, _HIS_PATCH_SIZE, _LATENT_DIM, device=device, dtype=_DTYPE)
 
         gen_lat, inputs_embeds, stop_out = executor.execute(input_tensor, his_lat)
-        torch.cuda.synchronize()
+        torch.accelerator.synchronize()
 
         assert gen_lat.shape == (bsz, _PATCH_SIZE, _LATENT_DIM)
         assert inputs_embeds.shape == (bsz, 1, _LLM_HIDDEN)
@@ -109,7 +109,7 @@ class TestCFMGraphExecutor:
         new_input = torch.randn_like(input_tensor)
         new_his = torch.randn_like(his_lat)
         gen_lat2, inputs_embeds2, stop_out2 = executor.execute(new_input, new_his)
-        torch.cuda.synchronize()
+        torch.accelerator.synchronize()
         assert gen_lat2.shape == gen_lat.shape
         assert inputs_embeds2.shape == inputs_embeds.shape
         assert stop_out2.shape == stop_out.shape
@@ -125,7 +125,7 @@ class TestCFMGraphExecutor:
         snapshot_his = his_lat.clone()
 
         executor.execute(input_tensor, his_lat)
-        torch.cuda.synchronize()
+        torch.accelerator.synchronize()
         assert torch.equal(input_tensor, snapshot_input)
         assert torch.equal(his_lat, snapshot_his)
 
@@ -139,7 +139,7 @@ class TestCFMGraphExecutorPool:
         his_lat = torch.randn(1, _HIS_PATCH_SIZE, _LATENT_DIM, device=device, dtype=_DTYPE)
 
         gen_lat, inputs_embeds, stop_out = pool.execute(input_tensor, his_lat)
-        torch.cuda.synchronize()
+        torch.accelerator.synchronize()
         assert gen_lat.shape == (1, _PATCH_SIZE, _LATENT_DIM)
         assert inputs_embeds.shape == (1, 1, _LLM_HIDDEN)
         assert stop_out.shape == (1, 2)

tests/worker/test_process_gpu_memory.py

@@ -68,7 +68,7 @@ class TestGetProcessGpuMemory:
         assert memory >= 0
 
         del tensor
-        torch.cuda.empty_cache()
+        torch.accelerator.empty_cache()
 
     def test_raises_on_invalid_device(self, mocker: MockerFixture):
         from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory

tools/configure_stage_memory.py

@@ -49,7 +49,7 @@ def get_gpu_info() -> list[dict]:
 
         if not torch.cuda.is_available():
             return []
-        for i in range(torch.cuda.device_count()):
+        for i in range(torch.accelerator.device_count()):
             free, total = torch.cuda.mem_get_info(i)
             props = torch.cuda.get_device_properties(i)
             gpus.append(

vllm_omni/diffusion/cache/teacache/coefficient_estimator.py

@@ -212,7 +212,7 @@ class TeaCacheCoefficientEstimator:
         trajectory = self.hook.stop_collection()
         if trajectory:
             self.collected_data.append(trajectory)
-        torch.cuda.empty_cache()
+        torch.accelerator.empty_cache()
 
     def estimate(self, poly_order: int = 4) -> list[float]:
         """Estimate polynomial coefficients from collected data.

vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py

@@ -267,7 +267,7 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin):
         )
         # Move text encoder back to CPU immediately
         self.text_encoder.to("cpu")
-        torch.cuda.empty_cache()
+        torch.accelerator.empty_cache()
 
         hidden_states = text_encoder_outputs.hidden_states
 
@@ -754,7 +754,7 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin):
         )
         self.connectors.to("cpu")
         if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+            torch.accelerator.empty_cache()
 
         # ---- Prepare latents ----
         latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1
@@ -978,7 +978,7 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin):
             self.vocoder.to(device)
             audio = self.vocoder(generated_mel_spectrograms)
             self.vocoder.to("cpu")
-            torch.cuda.empty_cache()
+            torch.accelerator.empty_cache()
 
         return DiffusionOutput(output=(video, audio))
 

vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py

@@ -1684,7 +1684,7 @@ class MagiHumanPipeline(nn.Module, ProgressBarMixin, DiffusionPipelineProfilerMi
         super().__init__()
         model_path = od_config.model
         local_files_only = os.path.exists(model_path)
-        device = f"cuda:{torch.cuda.current_device()}"
+        device = f"cuda:{torch.accelerator.current_device_index()}"
         self.device_str = device
         self.dtype = od_config.dtype or torch.bfloat16
 
@@ -2255,7 +2255,7 @@ class MagiHumanPipeline(nn.Module, ProgressBarMixin, DiffusionPipelineProfilerMi
                 1 - self.sr_audio_noise_scale
             )
 
-            torch.cuda.empty_cache()
+            torch.accelerator.empty_cache()
             sr_steps = sr_num_steps or self.sr_num_inference_steps_default
             final_latent_video, _ = self._evaluate_with_latent(
                 context,
@@ -2274,9 +2274,9 @@ class MagiHumanPipeline(nn.Module, ProgressBarMixin, DiffusionPipelineProfilerMi
             final_latent_video = br_latent_video
             final_latent_audio = br_latent_audio
 
-        torch.cuda.empty_cache()
+        torch.accelerator.empty_cache()
         videos_np = self._decode_video(final_latent_video)
-        torch.cuda.empty_cache()
+        torch.accelerator.empty_cache()
         audio_np = self._decode_audio(final_latent_audio)
 
         return DiffusionOutput(output=(videos_np, audio_np))

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -503,7 +503,7 @@ class CustomPipelineWorkerExtension:
         if self.model_runner.pipeline is not None:
             del self.model_runner.pipeline
             gc.collect()
-            torch.cuda.empty_cache()
+            torch.accelerator.empty_cache()
 
         # Get custom pipeline class name
         custom_pipeline_name = custom_pipeline_args["pipeline_class"]

vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py

@@ -878,7 +878,7 @@ class MooncakeTransferEngineConnector(OmniConnectorBase):
                 # Success
                 # Ensure data is visible on GPU
                 # Note: RDMA write visibility on GPU usually requires some form of fence/sync.
-                # torch.cuda.synchronize() is a heavy hammer but safe.
+                # torch.accelerator.synchronize() is a heavy hammer but safe.
                 # Ideally Mooncake engine provides a way to poll for completion that guarantees visibility.
                 # TODO(wzliu): Replace synchronize with cuda event in the future for better performance.
                 if self.pool.is_cuda:

vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py

@@ -372,7 +372,7 @@ class FishSpeechFastAR(nn.Module):
                 dtype=torch.long,
             )
             self(hidden, semantic, do_sample=False)
-        torch.cuda.synchronize(device)
+        torch.accelerator.synchronize(device)
 
     @torch.inference_mode()
     def _run_model(self, step_input: torch.Tensor, step_pos_ids: torch.Tensor, bsz: int) -> torch.Tensor:

vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py

@@ -1603,7 +1603,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
                 layer.mlp = None
                 del mlp
                 gc.collect()
-                torch.cuda.empty_cache()
+                torch.accelerator.empty_cache()
 
                 layer.mlp = HunyuanImage3SparseMoeBlock(
                     config=self.config,

vllm_omni/model_executor/models/mimo_audio/mimo_audio.py

@@ -276,7 +276,9 @@ class MiMoAudioDataParser(MultiModalDataParser):
                 self.device = torch.device(tokenizer_device)
         else:
             # Default to cuda (will use current GPU)
-            self.device = torch.device(f"cuda:{torch.cuda.current_device()}" if torch.cuda.is_available() else "cpu")
+            self.device = torch.device(
+                f"cuda:{torch.accelerator.current_device_index()}" if torch.cuda.is_available() else "cpu"
+            )
 
         self.audio_tokenizer_path = os.environ.get("MIMO_AUDIO_TOKENIZER_PATH", None)
         if not self.audio_tokenizer_path:

vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py

@@ -793,7 +793,7 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor
         local_embeds: torch.FloatTensor,  # [1, 1, hidden_size]
         tokens_dtype: torch.dtype = torch.int64,
         tokens_device: torch.device = torch.device(
-            f"cuda:{torch.cuda.current_device()}" if torch.cuda.is_available() else "cpu"
+            f"cuda:{torch.accelerator.current_device_index()}" if torch.cuda.is_available() else "cpu"
         ),
         local_sampler: MiMoSampler | MiMoLocalSamplerTensor | None = None,
     ):
@@ -847,7 +847,7 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor
         local_embeds: torch.FloatTensor,  # [1, 1, hidden_size]
         tokens_dtype: torch.dtype = torch.int64,
         tokens_device: torch.device = torch.device(
-            f"cuda:{torch.cuda.current_device()}" if torch.cuda.is_available() else "cpu"
+            f"cuda:{torch.accelerator.current_device_index()}" if torch.cuda.is_available() else "cpu"
         ),
         local_sampler: MiMoSampler | None = None,
     ):

vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py

@@ -110,7 +110,7 @@ class CUDAGraphDecoderWrapper:
             with torch.no_grad():
                 _ = self.decoder(dummy)
 
-        torch.cuda.synchronize(device)
+        torch.accelerator.synchronize(device)
 
         for size in self.capture_sizes:
             try:
@@ -126,7 +126,7 @@ class CUDAGraphDecoderWrapper:
         static_input = torch.zeros(1, self.num_quantizers, size, dtype=dtype, device=device)
         with torch.no_grad():
             _ = self.decoder(static_input)
-        torch.cuda.synchronize(device)
+        torch.accelerator.synchronize(device)
 
         graph = CUDAGraph()
         with torch.no_grad():

vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py

@@ -229,7 +229,7 @@ class _PerfTimer:
     def _resolve(self) -> None:
         if not self._pairs:
             return
-        torch.cuda.synchronize()
+        torch.accelerator.synchronize()
         for name, s, e in self._pairs:
             self._timers[name] = self._timers.get(name, 0.0) + s.elapsed_time(e)
             self._counts[name] = self._counts.get(name, 0) + 1
@@ -433,7 +433,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
         self._tts.base_lm = None
         del self._tts.residual_lm
         self._tts.residual_lm = None
-        torch.cuda.empty_cache()
+        torch.accelerator.empty_cache()
 
         self._inference_timesteps = 10
         self._cfg_value = 2.0

vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py

@@ -85,7 +85,7 @@ class CUDAGraphAcousticTransformerWrapper:
             with torch.no_grad():
                 self._forward_cudagraph_compatible(dummy, cfg_alpha=dummy_cfg_alpha, noise=dummy_noise)
 
-        torch.cuda.synchronize(device)
+        torch.accelerator.synchronize(device)
 
         # Phase 2: Capture graphs
         for size in self.capture_sizes:
@@ -199,7 +199,7 @@ class CUDAGraphAcousticTransformerWrapper:
         with torch.no_grad():
             _ = self._forward_cudagraph_compatible(static_input, cfg_alpha=static_cfg_alpha, noise=static_noise)
 
-        torch.cuda.synchronize(device)
+        torch.accelerator.synchronize(device)
 
         graph = CUDAGraph()
         with torch.no_grad():

vllm_omni/platforms/cuda/platform.py

@@ -114,7 +114,7 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
 
     @classmethod
     def get_device_count(cls) -> int:
-        return torch.cuda.device_count()
+        return torch.accelerator.device_count()
 
     @classmethod
     def get_device_version(cls) -> str | None:
@@ -122,7 +122,7 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
 
     @classmethod
     def synchronize(cls) -> None:
-        torch.cuda.synchronize()
+        torch.accelerator.synchronize()
 
     @classmethod
     def get_free_memory(cls, device: torch.device | None = None) -> int:

vllm_omni/platforms/rocm/platform.py

@@ -119,7 +119,7 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform):
 
     @classmethod
     def get_device_count(cls) -> int:
-        return torch.cuda.device_count()
+        return torch.accelerator.device_count()
 
     @classmethod
     def get_device_version(cls) -> str | None:
@@ -130,7 +130,7 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform):
 
     @classmethod
     def synchronize(cls) -> None:
-        torch.cuda.synchronize()
+        torch.accelerator.synchronize()
 
     @classmethod
     def get_free_memory(cls, device: torch.device | None = None) -> int:

vllm_omni/worker/base.py

@@ -110,7 +110,7 @@ class OmniGPUWorkerBase(GPUWorker):
         if kv_cache_memory_bytes := self.cache_config.kv_cache_memory_bytes:
             self.model_runner.profile_run()
             if current_omni_platform.is_rocm():
-                torch.cuda.synchronize()
+                torch.accelerator.synchronize()
             return kv_cache_memory_bytes
 
         with memory_profiling(

vllm_omni/worker/gpu_ar_model_runner.py

@@ -227,7 +227,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                         batch_descriptor=batch_desc,
                     ):
                         self.talker_mtp(ids, emb, hid, ts)
-                    torch.cuda.synchronize()
+                    torch.accelerator.synchronize()
 
             logger.info("Captured talker_mtp graphs for %d sizes", len(capture_sizes))
         except RuntimeError as e: