From e6466cf0e621f51432f1c5afe7f23df908862763 Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Thu, 7 May 2026 11:23:59 -0400 Subject: [PATCH] [Refactor] Replace and ban a few torch.cuda functions in favor of torch.accelerator replacements. (#3365) Signed-off-by: Nick Cao --- benchmarks/diffusion/quantization_quality.py | 12 ++++++------ .../cross_node_mooncake_transfer_engine.py | 4 ++-- .../offline_inference/internvla_a1/end2end.py | 2 +- .../text_to_speech/voxtral_tts/end2end.py | 6 +++--- pyproject.toml | 6 ++++++ .../reliability/test_reliability_qwen3_omni.py | 2 +- .../quantization/test_quantization_quality.py | 12 ++++++------ tests/diffusion/test_worker_wrapper_base.py | 8 ++++---- tests/e2e/accuracy/conftest.py | 2 +- .../test_diffusers_backend_similarity.py | 4 ++-- .../e2e/accuracy/test_ltx2_3_video_similarity.py | 4 ++-- tests/e2e/accuracy/test_qwen_image.py | 2 +- tests/e2e/accuracy/test_qwen_image_edit.py | 2 +- tests/e2e/accuracy/test_qwen_image_layered.py | 2 +- .../wan22_i2v/test_wan22_i2v_video_similarity.py | 4 ++-- .../test_hunyuanimage3_text2img.py | 2 +- .../offline_inference/test_omni_sleep_mode.py | 6 +++--- .../offline_inference/test_quantization_fp8.py | 12 ++++++------ tests/entrypoints/test_omni_sleep_mode.py | 16 ++++++++-------- tests/helpers/env.py | 6 +++--- tests/helpers/fixtures/env.py | 4 ++-- .../models/ming_flash_omni/test_talker_cfm.py | 10 +++++----- tests/worker/test_process_gpu_memory.py | 2 +- tools/configure_stage_memory.py | 2 +- .../cache/teacache/coefficient_estimator.py | 2 +- .../diffusion/models/ltx2/pipeline_ltx2_3.py | 6 +++--- .../models/magi_human/pipeline_magi_human.py | 8 ++++---- vllm_omni/diffusion/worker/diffusion_worker.py | 2 +- .../mooncake_transfer_engine_connector.py | 2 +- .../models/fish_speech/fish_speech_fast_ar.py | 2 +- .../models/hunyuan_image3/hunyuan_image3.py | 2 +- .../models/mimo_audio/mimo_audio.py | 4 +++- .../models/mimo_audio/mimo_audio_llm.py | 4 ++-- .../qwen3_tts/cuda_graph_decoder_wrapper.py | 4 ++-- .../models/voxcpm2/voxcpm2_talker.py | 4 ++-- .../cuda_graph_acoustic_transformer_wrapper.py | 4 ++-- vllm_omni/platforms/cuda/platform.py | 4 ++-- vllm_omni/platforms/rocm/platform.py | 4 ++-- vllm_omni/worker/base.py | 2 +- vllm_omni/worker/gpu_ar_model_runner.py | 2 +- 40 files changed, 98 insertions(+), 90 deletions(-) diff --git a/benchmarks/diffusion/quantization_quality.py b/benchmarks/diffusion/quantization_quality.py index 4a916e7ea..81132e5aa 100644 --- a/benchmarks/diffusion/quantization_quality.py +++ b/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): diff --git a/benchmarks/distributed/omni_connectors/cross_node_mooncake_transfer_engine.py b/benchmarks/distributed/omni_connectors/cross_node_mooncake_transfer_engine.py index fd01739cc..53cd0fad5 100644 --- a/benchmarks/distributed/omni_connectors/cross_node_mooncake_transfer_engine.py +++ b/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)}") diff --git a/examples/offline_inference/internvla_a1/end2end.py b/examples/offline_inference/internvla_a1/end2end.py index a7799d7ce..6bc69f4e0 100755 --- a/examples/offline_inference/internvla_a1/end2end.py +++ b/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]: diff --git a/examples/offline_inference/text_to_speech/voxtral_tts/end2end.py b/examples/offline_inference/text_to_speech/voxtral_tts/end2end.py index a54277725..6852b2b36 100644 --- a/examples/offline_inference/text_to_speech/voxtral_tts/end2end.py +++ b/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: diff --git a/pyproject.toml b/pyproject.toml index 3c3f7a026..0eaf07d32 100644 --- a/pyproject.toml +++ b/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" diff --git a/tests/dfx/reliability/test_reliability_qwen3_omni.py b/tests/dfx/reliability/test_reliability_qwen3_omni.py index 18251946d..37f63b942 100644 --- a/tests/dfx/reliability/test_reliability_qwen3_omni.py +++ b/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)) diff --git a/tests/diffusion/quantization/test_quantization_quality.py b/tests/diffusion/quantization/test_quantization_quality.py index ba6a150c4..1e94b9cd0 100644 --- a/tests/diffusion/quantization/test_quantization_quality.py +++ b/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() # --------------------------------------------------------------------------- diff --git a/tests/diffusion/test_worker_wrapper_base.py b/tests/diffusion/test_worker_wrapper_base.py index 986bad9a8..502a923ba 100644 --- a/tests/diffusion/test_worker_wrapper_base.py +++ b/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) diff --git a/tests/e2e/accuracy/conftest.py b/tests/e2e/accuracy/conftest.py index 0108e86da..a4804fc19 100644 --- a/tests/e2e/accuracy/conftest.py +++ b/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: diff --git a/tests/e2e/accuracy/test_diffusers_backend_similarity.py b/tests/e2e/accuracy/test_diffusers_backend_similarity.py index b9d452bdc..6392d9bcf 100644 --- a/tests/e2e/accuracy/test_diffusers_backend_similarity.py +++ b/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) diff --git a/tests/e2e/accuracy/test_ltx2_3_video_similarity.py b/tests/e2e/accuracy/test_ltx2_3_video_similarity.py index dec533d58..ae8a3ce87 100644 --- a/tests/e2e/accuracy/test_ltx2_3_video_similarity.py +++ b/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 diff --git a/tests/e2e/accuracy/test_qwen_image.py b/tests/e2e/accuracy/test_qwen_image.py index 4b8215d54..7ddf0fe6c 100644 --- a/tests/e2e/accuracy/test_qwen_image.py +++ b/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) diff --git a/tests/e2e/accuracy/test_qwen_image_edit.py b/tests/e2e/accuracy/test_qwen_image_edit.py index 07deecca9..fd733b9de 100644 --- a/tests/e2e/accuracy/test_qwen_image_edit.py +++ b/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) diff --git a/tests/e2e/accuracy/test_qwen_image_layered.py b/tests/e2e/accuracy/test_qwen_image_layered.py index 30ad2966f..cd1c6fbc5 100644 --- a/tests/e2e/accuracy/test_qwen_image_layered.py +++ b/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) diff --git a/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py b/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py index 1caef3bff..17e4fc0f2 100644 --- a/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py +++ b/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") diff --git a/tests/e2e/offline_inference/test_hunyuanimage3_text2img.py b/tests/e2e/offline_inference/test_hunyuanimage3_text2img.py index bd0d132d0..5b34faa98 100644 --- a/tests/e2e/offline_inference/test_hunyuanimage3_text2img.py +++ b/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, diff --git a/tests/e2e/offline_inference/test_omni_sleep_mode.py b/tests/e2e/offline_inference/test_omni_sleep_mode.py index 5a3ae9ab7..6d4cbdb26 100644 --- a/tests/e2e/offline_inference/test_omni_sleep_mode.py +++ b/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") diff --git a/tests/e2e/offline_inference/test_quantization_fp8.py b/tests/e2e/offline_inference/test_quantization_fp8.py index 9801e0ae7..8db2d90e6 100644 --- a/tests/e2e/offline_inference/test_quantization_fp8.py +++ b/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", diff --git a/tests/entrypoints/test_omni_sleep_mode.py b/tests/entrypoints/test_omni_sleep_mode.py index aa7be1ba0..4a098e04d 100644 --- a/tests/entrypoints/test_omni_sleep_mode.py +++ b/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") diff --git a/tests/helpers/env.py b/tests/helpers/env.py index 8df78e09e..d6f682659 100644 --- a/tests/helpers/env.py +++ b/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() diff --git a/tests/helpers/fixtures/env.py b/tests/helpers/fixtures/env.py index 939bad02c..aa7a2fd04 100644 --- a/tests/helpers/fixtures/env.py +++ b/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) diff --git a/tests/model_executor/models/ming_flash_omni/test_talker_cfm.py b/tests/model_executor/models/ming_flash_omni/test_talker_cfm.py index 419ce00da..de6060ba4 100644 --- a/tests/model_executor/models/ming_flash_omni/test_talker_cfm.py +++ b/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) diff --git a/tests/worker/test_process_gpu_memory.py b/tests/worker/test_process_gpu_memory.py index c27da311d..d5cc2eda8 100644 --- a/tests/worker/test_process_gpu_memory.py +++ b/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 diff --git a/tools/configure_stage_memory.py b/tools/configure_stage_memory.py index cd91c5fce..9f396051b 100644 --- a/tools/configure_stage_memory.py +++ b/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( diff --git a/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py b/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py index 38c805c28..37ea454de 100644 --- a/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py +++ b/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. diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py index dd0a57179..c3bbf3aa0 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py +++ b/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)) diff --git a/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py b/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py index 881c72edc..7f073ab87 100644 --- a/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py +++ b/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)) diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 5c8cde3e2..b25addb46 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/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"] diff --git a/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py b/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py index bd4160f3e..4b21ed0ec 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py +++ b/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: diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py index df3017152..633652c92 100644 --- a/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py +++ b/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: diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py index d2552bddd..75e4adb7e 100644 --- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py +++ b/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, diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py index ee924dc46..9a09ab030 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py +++ b/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: diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py index de78773f2..0a886bd5f 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py +++ b/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, ): diff --git a/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py b/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py index 26d69fe22..fd4b62520 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py +++ b/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(): diff --git a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py index 0da965d99..635b570bc 100644 --- a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py +++ b/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 diff --git a/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py b/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py index d7407afe5..491f185a4 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py +++ b/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(): diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py index 6afdfcb66..849e7028a 100644 --- a/vllm_omni/platforms/cuda/platform.py +++ b/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: diff --git a/vllm_omni/platforms/rocm/platform.py b/vllm_omni/platforms/rocm/platform.py index 87cc08418..96ee5a498 100644 --- a/vllm_omni/platforms/rocm/platform.py +++ b/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: diff --git a/vllm_omni/worker/base.py b/vllm_omni/worker/base.py index 8bd9efc89..67b6f77b3 100644 --- a/vllm_omni/worker/base.py +++ b/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( diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 143a373fe..1decf13cb 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/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: