diff --git a/tests/entrypoints/openai_api/test_tts_adapter.py b/tests/entrypoints/openai_api/test_tts_adapter.py index 500feab35..a2507a056 100644 --- a/tests/entrypoints/openai_api/test_tts_adapter.py +++ b/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) diff --git a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py index 555fec42d..5f5d8ecb9 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py +++ b/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