[Feature][TTS] MOSS-TTS: place the reference-audio encoder on the code2wav stage's GPU (#7724)

Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>

tests/entrypoints/openai_api/test_tts_adapter.py

@@ -255,6 +255,115 @@ def test_moss_reference_transcript_mode(variant, ref_text, mode, mocker):
     assert request.input == "Target."
 
 
+def _moss_adapter_with_stage_configs(stage_configs, mocker):
+    engine_client = SimpleNamespace(
+        model_config=SimpleNamespace(model="OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5"),
+        stage_configs=stage_configs,
+    )
+    return MossTTSAdapter(SpeechServingContext(server=mocker.Mock(), engine_client=engine_client))
+
+
+def _moss_cuda_available(mocker, device_count=8):
+    mocker.patch("torch.cuda.is_available", return_value=True)
+    mocker.patch("torch.accelerator.device_count", return_value=device_count)
+
+
+@pytest.mark.parametrize(
+    "stage_configs,expected",
+    [
+        # Follows the code2wav stage's first device.
+        (
+            [
+                SimpleNamespace(model_stage="moss_tts_local", runtime_config=SimpleNamespace(devices="0")),
+                SimpleNamespace(model_stage="moss_tts_local_codec", runtime_config=SimpleNamespace(devices="3")),
+            ],
+            "cuda:3",
+        ),
+        # No codec-named stage: anchors on the last stage of the pipeline.
+        (
+            [
+                SimpleNamespace(model_stage="talker", runtime_config=SimpleNamespace(devices="1")),
+                SimpleNamespace(model_stage="decoder", runtime_config=SimpleNamespace(devices="2")),
+            ],
+            "cuda:2",
+        ),
+        # Multi-device codec stage pins replica 0 on the first device.
+        (
+            [
+                SimpleNamespace(model_stage="moss_tts_local_codec", runtime_config=SimpleNamespace(devices="5,6")),
+            ],
+            "cuda:5",
+        ),
+        # No explicit pinning anywhere: stage worker defaults to cuda:0.
+        (
+            [SimpleNamespace(model_stage="moss_tts_local_codec", runtime_config=SimpleNamespace(devices=None))],
+            "cuda:0",
+        ),
+    ],
+)
+def test_moss_ref_encoder_device_follows_codec_stage(stage_configs, expected, mocker):
+    import torch
+
+    _moss_cuda_available(mocker)
+    adapter = _moss_adapter_with_stage_configs(stage_configs, mocker)
+    assert adapter._resolve_ref_encoder_device() == torch.device(expected)
+
+
+def test_moss_ref_encoder_device_accepts_dict_runtime(mocker):
+    import torch
+
+    _moss_cuda_available(mocker)
+    adapter = _moss_adapter_with_stage_configs(
+        [SimpleNamespace(model_stage="moss_tts_local_codec", runtime={"devices": "2"})],
+        mocker,
+    )
+    assert adapter._resolve_ref_encoder_device() == torch.device("cuda:2")
+
+
+def test_moss_ref_encoder_device_cpu_when_cuda_unavailable(mocker):
+    import torch
+
+    mocker.patch("torch.cuda.is_available", return_value=False)
+    adapter = _moss_adapter_with_stage_configs(
+        [SimpleNamespace(model_stage="moss_tts_local_codec", runtime_config=SimpleNamespace(devices="3"))],
+        mocker,
+    )
+    assert adapter._resolve_ref_encoder_device() == torch.device("cpu")
+
+
+def test_moss_ref_encoder_device_cpu_when_index_out_of_range(mocker):
+    import torch
+
+    _moss_cuda_available(mocker, device_count=2)
+    adapter = _moss_adapter_with_stage_configs(
+        [SimpleNamespace(model_stage="moss_tts_local_codec", runtime_config=SimpleNamespace(devices="5"))],
+        mocker,
+    )
+    assert adapter._resolve_ref_encoder_device() == torch.device("cpu")
+
+
+def test_moss_get_processor_places_audio_tokenizer_on_codec_stage_device(mocker):
+    import torch
+
+    _moss_cuda_available(mocker)
+    adapter = _moss_adapter_with_stage_configs(
+        [SimpleNamespace(model_stage="moss_tts_local_codec", runtime_config=SimpleNamespace(devices="3"))],
+        mocker,
+    )
+    audio_tokenizer = mocker.Mock()
+    audio_tokenizer.to.return_value = audio_tokenizer
+    processor = SimpleNamespace(audio_tokenizer=audio_tokenizer)
+    from_pretrained = mocker.patch("transformers.AutoProcessor.from_pretrained", return_value=processor)
+
+    assert adapter._get_moss_processor() is processor
+    audio_tokenizer.to.assert_called_once_with(torch.device("cuda", 3))
+    audio_tokenizer.eval.assert_called_once_with()
+    # Lazy cache: no re-load / re-placement on the second call.
+    assert adapter._get_moss_processor() is processor
+    from_pretrained.assert_called_once()
+    audio_tokenizer.to.assert_called_once()
+
+
 def test_qwen3_tts_metadata():
     assert Qwen3TTSAdapter.backend == "ar"
     assert issubclass(Qwen3TTSAdapter, ARTTSAdapter)

vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py

@@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, cast
 from transformers import AutoModel, AutoTokenizer
 from transformers.dynamic_module_utils import get_class_from_dynamic_module
 from vllm.inputs import tokens_input
+from vllm.logger import init_logger
 
 from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
 from vllm_omni.entrypoints.openai.tts_adapters.base import (
@@ -27,6 +28,8 @@ from vllm_omni.model_executor.models.moss_tts.realtime_prompt import build_realt
 if TYPE_CHECKING:
     from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
 
+logger = init_logger(__name__)
+
 
 class _MossTTSAdapterBase(ARTTSAdapter):
     accumulate_nonstreaming: bool = False
@@ -143,13 +146,81 @@ class _MossTTSAdapterBase(ARTTSAdapter):
             return "voice_generator"
         return "tts"
 
+    def _codec_stage_devices(self) -> str | None:
+        """Read the runtime ``devices`` entry of this pipeline's codec stage.
+
+        The codec (code2wav) stage is the final stage of every MOSS pipeline
+        (``moss_tts_local_codec`` / ``moss_tts_codec``); prefer a stage whose
+        ``model_stage`` mentions "codec", else anchor on the last stage. The
+        runtime mapping accepts the resolved ``runtime_config`` object, the
+        omegaconf ``runtime`` mapping, and the legacy ``yaml_runtime`` dict.
+        """
+        stages = list(getattr(self.engine_client, "stage_configs", None) or ())
+        if not stages:
+            return None
+        codec_stages = [s for s in stages if "codec" in str(getattr(s, "model_stage", "") or "")]
+        anchor = codec_stages[-1] if codec_stages else stages[-1]
+        for attr in ("runtime_config", "runtime", "yaml_runtime"):
+            runtime_cfg = getattr(anchor, attr, None)
+            if runtime_cfg is None:
+                continue
+            devices = (
+                runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None)
+            )
+            if devices is not None and str(devices).strip():
+                return str(devices)
+        devices = getattr(anchor, "devices", None)
+        return str(devices) if devices is not None and str(devices).strip() else None
+
+    def _resolve_ref_encoder_device(self):
+        """Pick the device of the API-process reference-audio encoder.
+
+        The encoder follows the code2wav stage onto its first GPU: the stage
+        worker remaps its ``devices`` entry through ``CUDA_VISIBLE_DEVICES``
+        and sees it as logical index 0, while replica launch restores the
+        parent environment, so the same index names the same physical GPU in
+        this process. When the stage pins no devices the stage worker and the
+        encoder both default to the first visible GPU. Falls back to CPU when
+        CUDA is unavailable or the stage device is not parseable.
+        """
+        import torch  # local to avoid pulling torch at module import time
+
+        if not torch.cuda.is_available():
+            return torch.device("cpu")
+        devices_str = self._codec_stage_devices()
+        if devices_str is None:
+            # No explicit pinning: the stage worker lands on the first visible
+            # GPU (vLLM picks cuda:0 of the un-remapped environment).
+            return torch.device("cuda", 0)
+        first = devices_str.split(",")[0].strip()
+        if not first:
+            return torch.device("cuda", 0)
+        if first == "cpu":
+            return torch.device("cpu")
+        if not first.isdigit():
+            logger.warning("MOSS ref encoder: codec stage devices %r not parseable; keeping CPU", devices_str)
+            return torch.device("cpu")
+        index = int(first)
+        device_count = torch.accelerator.device_count()
+        if index >= device_count:
+            logger.warning(
+                "MOSS ref encoder: codec stage device cuda:%d out of range (%d visible); keeping CPU",
+                index,
+                device_count,
+            )
+            return torch.device("cpu")
+        return torch.device("cuda", index)
+
     def _get_moss_processor(self):
         """Lazily load the upstream MOSS-TTS processor once per server.
 
         Cached on ``self._moss_processor_cache``. The processor owns its own
-        audio_tokenizer (~1.6 B params); we keep it on CPU so it doesn't
-        compete with the talker (~8 GiB) and codec (~7 GiB) for our 96 GiB
-        GPU — per-request ref-audio encoding is fast enough on CPU.
+        audio_tokenizer (~1.6 B params) used for per-request reference-audio
+        encoding; it is placed on the code2wav stage's GPU (see
+        ``_resolve_ref_encoder_device``) so cold encodes run on device instead
+        of paying a CPU forward per request. Note this memory is charged to
+        the API process, outside each stage's ``gpu_memory_utilization``
+        budget — keep some headroom on the codec GPU.
         """
         cached = getattr(self, "_moss_processor_cache", None)
         if cached is not None:
@@ -159,7 +230,9 @@ class _MossTTSAdapterBase(ARTTSAdapter):
         model_id = self.engine_client.model_config.model
         proc = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
         if hasattr(proc, "audio_tokenizer"):
-            proc.audio_tokenizer = proc.audio_tokenizer.to("cpu").eval()
+            device = self._resolve_ref_encoder_device()
+            proc.audio_tokenizer = proc.audio_tokenizer.to(device).eval()
+            logger.info("MOSS reference-audio encoder (audio_tokenizer) placed on %s", device)
         self._moss_processor_cache = proc
         return proc
 
@@ -188,16 +261,13 @@ class _MossTTSAdapterBase(ARTTSAdapter):
                     ),
                 )
             )
-            codec = (
-                AutoModel.from_pretrained(
-                    codec_path,
-                    trust_remote_code=True,
-                )
-                # Reference encoding runs in the API process. Keep the codec on CPU
-                # so it does not reserve accelerator memory outside the model workers.
-                .to("cpu")
-                .eval()
+            codec = AutoModel.from_pretrained(
+                codec_path,
+                trust_remote_code=True,
             )
+            # Reference encoding runs in the API process; follow the codec
+            # (code2wav) stage's GPU like the processor path above.
+            codec = codec.to(self._resolve_ref_encoder_device()).eval()
             cached = (tokenizer, processor, codec)
             self._moss_realtime_components = cached
             return cached