[Refactor] Replace hardcoded CUDA device selection with platform-agnostic API (#4496)

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

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py

@@ -49,6 +49,7 @@ from vllm_omni.model_executor.models.cosyvoice3.utils import (
     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 @@ class CosyVoice3MultiModalProcessor(BaseMultiModalProcessor[CosyVoice3MultiModal
             ) 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

vllm_omni/model_executor/models/fish_speech/dac_encoder.py

@@ -17,6 +17,7 @@ from vllm_omni.model_executor.models.fish_speech.dac_utils import (
     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

vllm_omni/model_executor/models/glm_tts/glm_tts.py

@@ -53,6 +53,7 @@ from vllm.v1.sample.sampler import Sampler
 
 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 @@ class GLMTTSMultiModalProcessor(BaseMultiModalProcessor[GLMTTSMultiModalProcessi
 
         # 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

vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py

@@ -41,6 +41,7 @@ from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_decoder import (
     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 @@ class HiggsAudioV2Code2Wav(nn.Module):
         2. Otherwise fall back to ``<model_dir>/<audio_tokenizer_subdir>``.
         """
         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:

vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py

@@ -29,6 +29,8 @@ from typing import Any
 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()
 

vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py

@@ -27,6 +27,7 @@ from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3
     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 @@ class HiggsAudioV3Code2Wav(nn.Module):
             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 @@ class HiggsAudioV3Code2Wav(nn.Module):
         )
 
         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)

vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py

@@ -23,6 +23,7 @@ from vllm_omni.model_executor.models.mimo_audio.config_mimo_audio import TALKER_
 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 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
 
         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)

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py

@@ -291,7 +291,7 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
                 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.

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py

@@ -25,6 +25,8 @@ import torch.nn.functional as F
 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 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
                 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)
 

vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py

@@ -38,6 +38,7 @@ from vllm_omni.model_executor.models.utils import (
     reinit_rotary_inv_freq,
     transformers_keys_to_ignore_compat,
 )
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
@@ -182,7 +183,7 @@ class MossTTSNanoForGeneration(nn.Module):
         # 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():

vllm_omni/model_executor/models/omnivoice/omnivoice.py

@@ -38,6 +38,7 @@ from vllm.multimodal.processing import (
 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__)
@@ -494,7 +495,7 @@ class OmniVoiceModel(
         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()
 

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -46,6 +46,7 @@ from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
     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
@@ -1222,7 +1223,7 @@ class Qwen3OmniMoeForConditionalGeneration(
             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()
             )
         )
 

vllm_omni/platforms/interface.py

@@ -248,6 +248,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

vllm_omni/worker/gpu_ar_worker.py

@@ -12,6 +12,7 @@ from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
 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 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase):
                     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)

vllm_omni/worker/gpu_generation_worker.py

@@ -11,6 +11,7 @@ from vllm.v1.utils import report_usage_stats
 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 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase):
                     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)