From 3045863dbf1461699bad564d622bfc7689e51c15 Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Tue, 16 Jun 2026 15:13:55 -0400 Subject: [PATCH 1/2] [Refactor] Replace hardcoded CUDA device selection with platform-agnostic API Replace all `torch.device("cuda" if torch.cuda.is_available() else "cpu")` and similar patterns with `current_omni_platform.get_torch_device()` This enables the same code to work on non-CUDA platforms (XPU, NPU) without per-model changes, using the existing platform abstraction layer. Co-Authored-By: Claude Signed-off-by: Nick Cao --- vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py | 3 ++- vllm_omni/model_executor/models/fish_speech/dac_encoder.py | 3 ++- vllm_omni/model_executor/models/glm_tts/glm_tts.py | 3 ++- .../models/higgs_audio_v2/higgs_audio_v2_code2wav.py | 3 ++- .../models/higgs_audio_v2/higgs_audio_v2_tokenizer.py | 4 +++- .../models/higgs_audio_v3/higgs_audio_v3_code2wav.py | 5 +++-- .../model_executor/models/mimo_audio/mimo_audio_code2wav.py | 4 ++-- .../model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py | 2 +- .../models/minicpmo_4_5/minicpmo_4_5_omni_tts.py | 4 +++- .../models/moss_tts_nano/modeling_moss_tts_nano.py | 3 ++- vllm_omni/model_executor/models/omnivoice/omnivoice.py | 3 ++- vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py | 3 ++- vllm_omni/worker/gpu_ar_worker.py | 3 ++- vllm_omni/worker/gpu_generation_worker.py | 3 ++- 14 files changed, 30 insertions(+), 16 deletions(-) diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index 1a387cdcc39..c67ff811f52 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -49,6 +49,7 @@ unpad_prompt_conditioning, ) from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.platforms import current_omni_platform from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config from vllm_omni.utils.speaker_cache import get_speaker_cache @@ -198,7 +199,7 @@ def _ensure_s3_model(cls): ) from e model = _s3.load_model("speech_tokenizer_v3_25hz") - device = "cuda" if torch.cuda.is_available() else "cpu" + device = current_omni_platform.get_torch_device() model = model.to(device).eval() cls._s3_model = (model, _s3, device) return cls._s3_model diff --git a/vllm_omni/model_executor/models/fish_speech/dac_encoder.py b/vllm_omni/model_executor/models/fish_speech/dac_encoder.py index cdf0da992fc..69ebd29eeb2 100644 --- a/vllm_omni/model_executor/models/fish_speech/dac_encoder.py +++ b/vllm_omni/model_executor/models/fish_speech/dac_encoder.py @@ -17,6 +17,7 @@ DAC_SAMPLE_RATE, build_dac_codec, ) +from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) @@ -129,7 +130,7 @@ def encode_reference_audio_codes( (dtype=torch.long). """ if device is None: - device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") + device = current_omni_platform.get_torch_device() else: device = torch.device(device) dtype = torch.float32 diff --git a/vllm_omni/model_executor/models/glm_tts/glm_tts.py b/vllm_omni/model_executor/models/glm_tts/glm_tts.py index 6e14840e24b..664a794c4c0 100644 --- a/vllm_omni/model_executor/models/glm_tts/glm_tts.py +++ b/vllm_omni/model_executor/models/glm_tts/glm_tts.py @@ -53,6 +53,7 @@ from vllm_omni.model_executor.models.common.nucleus_ras_sampling import ras_sample_one as _ras_sample_one from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.platforms import current_omni_platform from vllm_omni.transformers_utils.configs.glm_tts import GLMTTSConfig from vllm_omni.utils.speaker_cache import get_speaker_cache @@ -411,7 +412,7 @@ def _ensure_cached_runtime_components( # Load voice clone components (speech tokenizer + CampPlus) if load_voice_clone: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = current_omni_platform.get_torch_device() self.processor_device = device campplus_path = getattr(self, "_campplus_path", None) or resolve_glm_tts_campplus_path(model_dir) self._campplus_path = campplus_path diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py index 6a14707ee5b..91609e1a4ee 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py +++ b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py @@ -41,6 +41,7 @@ load_higgs_audio_codec, ) from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.platforms import current_omni_platform __all__ = [ "HiggsAudioV2Code2Wav", @@ -173,7 +174,7 @@ def load_codec_from_disk(self, model_dir: str, device: torch.device | None = Non 2. Otherwise fall back to ``/``. """ if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = current_omni_platform.get_torch_device() audio_tokenizer_path = self._resolve_audio_tokenizer_path(model_dir) quantizer, fc2, acoustic_decoder, _tokenizer_config = load_higgs_audio_codec(audio_tokenizer_path, device) if len(quantizer.quantizers) != self.num_codebooks: diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py index cacf82ee56c..960f6cfef02 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py +++ b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py @@ -29,6 +29,8 @@ import numpy as np import torch +from vllm_omni.platforms import current_omni_platform + __all__ = [ "UnsupportedInputError", "MULTI_SPEAKER_TAG_PATTERN", @@ -239,7 +241,7 @@ def _load_audio_tokenizer(): allow_patterns=[f"{_K2_OMNIVOICE_SUBDIR}/*"], ) audio_tokenizer_dir = os.path.join(repo_path, _K2_OMNIVOICE_SUBDIR) - device = "cuda" if torch.cuda.is_available() else "cpu" + device = current_omni_platform.get_torch_device() model = HiggsAudioV2TokenizerModel.from_pretrained(audio_tokenizer_dir).to(device) return model.eval() diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py index 81c18a32e92..bc82c474b2d 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py +++ b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py @@ -27,6 +27,7 @@ HiggsAudioV3Config, ) from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.platforms import current_omni_platform __all__ = [ "HiggsAudioV3Code2Wav", @@ -126,7 +127,7 @@ def _load_from_tokenizer_repo(self) -> None: else: raise FileNotFoundError(f"No cached snapshot for {tokenizer_id}") - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = current_omni_platform.get_torch_device() quantizer, fc2, acoustic_decoder, _cfg = load_higgs_audio_codec(tokenizer_dir, device) self.quantizer = quantizer self.fc2 = fc2 @@ -278,7 +279,7 @@ def _load_from_bundled_state( ) if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = current_omni_platform.get_torch_device() # The bundled codec may use boson-ai key naming. Try remapping. needs_remap = any(k.startswith("quantizer.vq.layers.") for k in codec_state) diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py index bd178afe667..26482c6bce0 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py @@ -23,6 +23,7 @@ from vllm_omni.model_executor.models.mimo_audio.cuda_graph_decoder_wrapper import CUDAGraphMiMoDecoderWrapper from vllm_omni.model_executor.models.mimo_audio.modeling_audio_tokenizer import MiMoAudioTokenizer from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.platforms import current_omni_platform logger = logging.getLogger(__name__) @@ -442,8 +443,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = ""): self.logits_processor = LogitsProcessor(config.vocab_size) - device_str = "cuda" if torch.cuda.is_available() else "cpu" - self.device = torch.device(device_str) + self.device = current_omni_platform.get_torch_device() self.sample_rate = getattr(config, "audio_sample_rate", 24000) self.tokenizer = AutoTokenizer.from_pretrained(self.config.name_or_path, trust_remote_code=True) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py index 7746b423e9e..d9b48807ae4 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py @@ -291,7 +291,7 @@ def forward( device = inputs_embeds.device else: num_tokens = 1 - device = torch.device("cuda") + device = current_omni_platform.get_torch_device() hidden_dim = self.config.hidden_size if hasattr(self.config, "hidden_size") else 2560 # Profile/dummy run: both input_ids and inputs_embeds are None. diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index c7208752ba8..083f54afe3d 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -25,6 +25,8 @@ from vllm.config import VllmConfig from vllm.model_executor.models.interfaces import SupportsPP +from vllm_omni.platforms import current_omni_platform + try: from stepaudio2 import Token2wav as _Token2wav @@ -379,7 +381,7 @@ def _dummy_hidden_states( break else: num_tokens = 1 - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = current_omni_platform.get_torch_device() hidden_size = int(getattr(self, "_hidden_size", 768) or 768) return torch.zeros((num_tokens, hidden_size), device=device, dtype=torch.bfloat16) diff --git a/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py b/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py index 87b93288684..1a86c2b82fd 100644 --- a/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py +++ b/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py @@ -38,6 +38,7 @@ reinit_rotary_inv_freq, transformers_keys_to_ignore_compat, ) +from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) @@ -183,7 +184,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: # Eager construction (not in load_weights) -- see module docstring. _patch_torchaudio_load() - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = current_omni_platform.get_torch_device() self._device: torch.device = device if device.type == "cuda" and torch.cuda.is_bf16_supported(): diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice.py b/vllm_omni/model_executor/models/omnivoice/omnivoice.py index c5026086e23..c37caeef540 100644 --- a/vllm_omni/model_executor/models/omnivoice/omnivoice.py +++ b/vllm_omni/model_executor/models/omnivoice/omnivoice.py @@ -38,6 +38,7 @@ from vllm.sequence import IntermediateTensors from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.platforms import current_omni_platform from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig logger = init_logger(__name__) @@ -495,7 +496,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: try: device = next(self.parameters()).device except StopIteration: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = current_omni_platform.get_torch_device() model_dir = self._resolve_model_dir() diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index 4a144bff027..b800d8d143f 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -46,6 +46,7 @@ Qwen3OmniMoeThinkerProcessingInfo, ) from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, safe_tensor_reshape +from vllm_omni.platforms import current_omni_platform # Special token IDs for Qwen3 Omni MoE # Reference: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct/blob/main/tokenizer_config.json @@ -1213,7 +1214,7 @@ def _talker_to_code_predictor( else ( talker_hidden_states.device if isinstance(talker_hidden_states, torch.Tensor) - else torch.device("cuda" if torch.cuda.is_available() else "cpu") + else current_omni_platform.get_torch_device() ) ) diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py index 08551ad6948..8eafd531ef8 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/vllm_omni/worker/gpu_ar_worker.py @@ -12,6 +12,7 @@ from vllm.v1.worker.workspace import init_workspace_manager from vllm_omni.diffusion.data import OmniACK, OmniSleepTask, OmniWakeTask +from vllm_omni.platforms import current_omni_platform from vllm_omni.worker.base import OmniGPUWorkerBase from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner from vllm_omni.worker.memory_utils import request_memory_tolerant @@ -58,7 +59,7 @@ def init_device(self): f"be less than or equal to the number of visible devices " f"({visible_device_count})." ) - self.device = torch.device(f"cuda:{self.local_rank}") + self.device = current_omni_platform.get_torch_device(self.local_rank) torch.accelerator.set_device_index(self.device) current_platform.check_if_supports_dtype(self.model_config.dtype) diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py index 2a85436a007..38a88c84c05 100644 --- a/vllm_omni/worker/gpu_generation_worker.py +++ b/vllm_omni/worker/gpu_generation_worker.py @@ -11,6 +11,7 @@ from vllm.v1.worker.gpu_worker import init_worker_distributed_environment from vllm.v1.worker.workspace import init_workspace_manager +from vllm_omni.platforms import current_omni_platform from vllm_omni.worker.base import OmniGPUWorkerBase from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner from vllm_omni.worker.memory_utils import request_memory_tolerant @@ -57,7 +58,7 @@ def init_device(self): f"be less than or equal to the number of visible devices " f"({visible_device_count})." ) - self.device = torch.device(f"cuda:{self.local_rank}") + self.device = current_omni_platform.get_torch_device(self.local_rank) torch.accelerator.set_device_index(self.device) current_platform.check_if_supports_dtype(self.model_config.dtype) From 08d21d62cc9814088821be66ac73fee96509ce84 Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Mon, 22 Jun 2026 10:52:54 -0400 Subject: [PATCH 2/2] [Refactor] Return cpu device for UnspecifiedOmniPlatform Fall back to torch.device("cpu") instead of raising NotImplementedError, so helper paths that used to run on CPU keep working when no accelerator platform is resolved. Co-Authored-By: Claude Signed-off-by: Nick Cao --- vllm_omni/platforms/interface.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index 14aa93b20fb..19c569e996c 100644 --- a/vllm_omni/platforms/interface.py +++ b/vllm_omni/platforms/interface.py @@ -243,6 +243,10 @@ class UnspecifiedOmniPlatform(OmniPlatform): _enum = PlatformEnum.UNSPECIFIED device_type = "cpu" + @classmethod + def get_torch_device(cls, local_rank: int | None = None) -> torch.device: + return torch.device("cpu") + @classmethod def get_device_count(cls) -> int: return 0