From 7be49aa0ce95810a3b78fd6317f3c374bdd4422d Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Tue, 6 Oct 2026 08:20:46 -0700 Subject: [PATCH] [Model] Default MOSS Local 1.5 to adaptive CUDA MRV2 with MPS (#8438) Signed-off-by: Sy03 <1370724210@qq.com> --- docs/configuration/stage_configs.md | 19 ++- recipes/OpenMOSS/MOSS-TTS.md | 52 ++++++- tests/config/test_config_factory.py | 140 +++++++++++++++++- tests/config/test_omni_config.py | 4 +- .../sched/test_generation_scheduler_init.py | 14 +- .../openai_api/test_serving_speech.py | 37 +++++ .../openai_api/test_serving_speech_metrics.py | 8 +- .../moss_tts/test_first_chunk_fast_path.py | 79 ++++++++-- .../moss_tts/test_local_compile_config.py | 48 ++++++ .../models/moss_tts/test_local_model_state.py | 34 ++++- .../moss_tts/test_local_prefill_batch.py | 70 ++++++++- tests/worker_v2/test_omni_data_plane.py | 55 +++++++ vllm_omni/config/pipeline_registry.py | 4 +- vllm_omni/config/stage_config.py | 5 + vllm_omni/deploy/moss_tts_local.yaml | 93 +----------- vllm_omni/deploy/moss_tts_local_mrv2.yaml | 9 +- ...s_tts_local_mrv2_high_concurrency_mps.yaml | 8 + .../deploy/moss_tts_local_mrv2_optimized.yaml | 43 ++++++ vllm_omni/deploy/moss_tts_local_v1.yaml | 88 +++++++++++ .../entrypoints/openai/serving_speech.py | 4 +- .../entrypoints/openai/tts_adapters/base.py | 3 + .../openai/tts_adapters/moss_tts.py | 6 + .../models/moss_tts/first_chunk_fast_path.py | 3 + .../models/moss_tts/local_model_state.py | 17 ++- .../moss_tts/modeling_moss_tts_codec.py | 1 + .../moss_tts/modeling_moss_tts_local_depth.py | 10 +- .../moss_tts/modeling_moss_tts_talker.py | 18 ++- .../models/moss_tts/pipeline.py | 45 +++++- .../tokenizer_12hz/streaming_decoder.py | 9 +- vllm_omni/worker_v2/omni_data_plane.py | 34 ++++- 30 files changed, 808 insertions(+), 152 deletions(-) create mode 100644 tests/model_executor/models/moss_tts/test_local_compile_config.py create mode 100644 vllm_omni/deploy/moss_tts_local_mrv2_high_concurrency_mps.yaml create mode 100644 vllm_omni/deploy/moss_tts_local_mrv2_optimized.yaml create mode 100644 vllm_omni/deploy/moss_tts_local_v1.yaml diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index c2c0eb733..10d467554 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -598,8 +598,9 @@ the intended GPU, and warm the complete pipeline before measuring performance. ## MOSS-TTS Local 1.5 with Model Runner V2 -`OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5` can opt into CUDA MRV2 using -the shared runtime introduced for Qwen3-TTS: +`OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5` defaults to CUDA MRV2 using +the shared runtime introduced for Qwen3-TTS. To select the C64 profile without +MPS explicitly: ```bash vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \ @@ -607,7 +608,7 @@ vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \ ``` This profile inherits the batching, codec graph buckets and 1-frame/15-frame -chunk geometry from `moss_tts_local.yaml`, and selects V2 for both the Local +chunk geometry from `moss_tts_local_v1.yaml`, and selects V2 for both the Local Talker and codec stages. The Local depth predictor exposes the MRV2 `mtp` capabilities while retaining its V1 `talker_mtp` implementation, sampling defaults and explicit request-seed handling. Codec chunks use the native data @@ -631,8 +632,12 @@ guarantee that 128 maximum-length prompts fit simultaneously. See the [MOSS recipe](gh-file:recipes/OpenMOSS/MOSS-TTS.md#local-15-mrv2-and-slot-attention) for activation, backend comparisons, memory requirements and benchmark commands. -Omitting `--deploy-config`, or selecting `moss_tts_local.yaml`, retains V1. -NPU, XPU, ROCm and MUSA overrides also retain V1. This profile does not enable +Omitting `--deploy-config` selects the C128 system profile on CUDA GPUs with +at least 140 GiB memory and MPS available. Smaller GPUs or a failed memory +query use C64 with utilization-based memory budgets; missing MPS selects +C64 without MPS. Selecting `moss_tts_local.yaml` explicitly uses C64 with MPS. +Use `moss_tts_local_v1.yaml` for the previous V1 deployment. +NPU, XPU, ROCm and MUSA overrides retain V1. These profiles do not enable MRV2 for MOSS Delay, Realtime or Nano. Local 1.5 outputs 48 kHz stereo audio; set `VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=48000` and `VLLM_OMNI_BENCH_AUDIO_CHANNELS=2` when benchmarking raw PCM. @@ -640,5 +645,5 @@ set `VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=48000` and Event-driven orchestration remains independently selectable with `VLLM_OMNI_EVENT_DRIVEN_ORCH=0` or `1`. Keep the runner and deployment identical when comparing these modes. Model-runner selection does not change the -orchestration default or enable experimental reference encoding, chunk ramps, -generation-output draining or MPS. +orchestration default. MPS, codec dispatch and graph settings belong to the +deployment profile; the automatic CUDA default selects MPS when available. diff --git a/recipes/OpenMOSS/MOSS-TTS.md b/recipes/OpenMOSS/MOSS-TTS.md index c9ec7670b..f25aac978 100644 --- a/recipes/OpenMOSS/MOSS-TTS.md +++ b/recipes/OpenMOSS/MOSS-TTS.md @@ -35,9 +35,9 @@ the requested `input` speech. Without a nonblank `ref_text`, Local uses audio-reference generation. Reference transcripts must match the reference audio; they are not style instructions. -For the optional CUDA MRV2 runner, see the +For the CUDA MRV2 runner and platform fallbacks, see the [Local 1.5 deployment profile](../../docs/configuration/stage_configs.md#moss-tts-local-15-with-model-runner-v2). -The default Local deployment continues to use V1. +Local 1.5 defaults to MRV2 on CUDA and retains V1 on other platforms. ## References @@ -159,8 +159,52 @@ curl -X POST http://localhost:8091/v1/audio/speech \ ## Local 1.5 MRV2 and slot attention -`MOSS-TTS-Local-Transformer-v1.5` supports the native CUDA MRV2 pipeline with -an explicit deploy profile. The default Local profile continues to use V1. +`MOSS-TTS-Local-Transformer-v1.5` defaults to the native MRV2 pipeline on +CUDA, with the original Local projection and sampler. When +`nvidia-cuda-mps-control` is on `PATH`, it starts private full-quota MPS. +GPUs with at least 140 GiB total memory use the C128 system profile with +prefix caching, Triton backbone attention, bounded codec first-chunk decode, +native Torch sampler compilation and dense codec graph buckets through 128. Its Talker KV budget is 32 GiB; smaller GPUs or a failed memory +query use C64 with utilization-based memory budgets. When the MPS executable +is unavailable, the automatic default is C64 MRV2 without MPS. NPU, XPU, +ROCm and MUSA retain V1. Explicit deploy configs override automatic selection. + +```bash +CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=1 \ +vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \ + --stage-init-timeout 1200 --init-timeout 1500 +``` + +Use `--deploy-config vllm_omni/deploy/moss_tts_local_v1.yaml` to select the +previous V1 profile explicitly. C128's batch-prefill and direct-token switches +remain confined to the throughput profile; they are not enabled in C64. + +The same `moss_tts_local_mrv2_optimized.yaml` system profile can be selected +explicitly. It retains the original Local depth projection and sampling algorithm. +The codec admits fast first chunks below 32 active streams; crowded streams use +the regular decoder with a dispatch target of 16 and a maximum wait of 6 ms. +This preserves the low-load first-audio path while coalescing high-load work. +`local_compile_audio_sampler: true` in the Talker HF overrides compiles the native +Torch sampler; explicit request generators use the original helper. Set the +override to `false` in a deployment file to disable that compilation. The setting +is specific to this C128 system profile. +Both stages share one GPU, with a 32 GiB Talker KV budget; this profile requires H200-class memory and +`nvidia-cuda-mps-control` on `PATH`. + +```bash +CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=1 \ +vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \ + --deploy-config vllm_omni/deploy/moss_tts_local_mrv2_optimized.yaml \ + --stage-init-timeout 1200 --init-timeout 1500 +``` + +Evaluate WER, speaker similarity and streaming latency for your workload; +throughput gains do not establish quality equivalence. Do not apply partial +MPS SM quotas: BF16 GEMM outputs were observed incomplete in that configuration +on the validation environment. MPS uses an owned control socket or an +explicitly supplied operator socket; only the owned daemon is stopped at shutdown. + +The explicit C64 MRV2 and capped-C128 throughput profiles remain available. Both profiles below preserve 1-frame initial and 15-frame steady codec chunks and the model's sampling defaults. diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 5df8f01f0..070f35d64 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -2985,6 +2985,9 @@ class TestPlatformOverrides: ("moss_tts_local_mrv2.yaml", "moss_tts_local"), ("moss_tts_local_mrv2_high_concurrency.yaml", "moss_tts_local"), ("moss_tts_local_mrv2_low_latency.yaml", "moss_tts_local"), + ("moss_tts_local_mrv2_optimized.yaml", "moss_tts_local"), + ("moss_tts_local.yaml", "moss_tts_local"), + ("moss_tts_local_mrv2_high_concurrency_mps.yaml", "moss_tts_local"), ], ) @pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"]) @@ -3004,9 +3007,142 @@ class TestPlatformOverrides: # v2 only engages the native plane on stages declaring support. assert all(ps.supports_native_mrv2_data_plane for ps in pipeline.stages) + @pytest.mark.parametrize("memory_gib", [24, 80, 139, 140, 141, None]) + @pytest.mark.parametrize("mps_available", [False, True]) + def test_moss_local_cuda_default_selects_mrv2_capacity_and_mps(self, monkeypatch, memory_gib, mps_available): + from vllm.platforms import current_platform + + from vllm_omni.platforms import current_omni_platform + + monkeypatch.setattr(current_omni_platform, "device_name", "cuda") + + def memory(): + if memory_gib is None: + raise RuntimeError("NVML unavailable") + return memory_gib * 1024**3 + + monkeypatch.setattr(current_platform, "get_device_total_memory", memory) + monkeypatch.setattr( + "shutil.which", lambda command: "/usr/bin/nvidia-cuda-mps-control" if mps_available else None + ) + pipeline = resolve_pipeline_config("moss_tts_local") + high_capacity = mps_available and memory_gib is not None and memory_gib >= 140 + capacity = 128 if high_capacity else 64 + config = VllmOmniConfig.from_pipeline_config(pipeline) + talker, codec = config.stage_configs + assert talker.scheduler_config.max_num_seqs == codec.scheduler_config.max_num_seqs == capacity + assert talker.cache_config.kv_cache_memory_bytes == (32 * 1024**3 if high_capacity else None) + assert talker.cache_config.gpu_memory_utilization == 0.60 + assert codec.cache_config.gpu_memory_utilization == 0.30 + assert talker.cache_config.enable_prefix_caching is high_capacity + assert not codec.cache_config.enable_prefix_caching + for stage in config.stage_configs: + assert stage.model_config.use_v2_model_runner + assert stage.runtime_config.cuda_mps is mps_available + assert not stage.runtime_config.env + overrides = talker.model_config.hf_overrides or {} + assert overrides.get("mrv2_batch_prefill", False) is high_capacity + assert overrides.get("mrv2_direct_tokens", False) is high_capacity + assert overrides.get("local_compile_audio_sampler", False) is high_capacity + + # The serving engine still consumes the legacy representation. + stages, _ = StageConfigFactory._create_legacy_from_registry(pipeline, {}) + for stage in stages: + assert stage.yaml_engine_args["max_num_seqs"] == capacity + assert stage.yaml_engine_args["use_v2_model_runner"] + assert stage.yaml_runtime.get("cuda_mps", False) is mps_available + assert stages[0].yaml_engine_args.get("kv_cache_memory_bytes") == (32 * 1024**3 if high_capacity else None) + assert stages[0].yaml_engine_args["enable_prefix_caching"] is high_capacity + + @pytest.mark.parametrize("platform", ["cpu", "npu", "xpu", "rocm", "musa"]) + def test_moss_local_non_cuda_default_preserves_v1(self, monkeypatch, platform): + from vllm.platforms import current_platform + + from vllm_omni.platforms import current_omni_platform + + monkeypatch.setattr(current_omni_platform, "device_name", platform) + + def unexpected_probe(*args): + pytest.fail("Non-CUDA selection must not probe CUDA memory or MPS") + + monkeypatch.setattr(current_platform, "get_device_total_memory", unexpected_probe) + monkeypatch.setattr("shutil.which", unexpected_probe) + pipeline = resolve_pipeline_config("moss_tts_local") + config = VllmOmniConfig.from_pipeline_config(pipeline) + for stage in config.stage_configs: + assert not stage.model_config.use_v2_model_runner + assert not stage.runtime_config.cuda_mps + assert stage.scheduler_config.max_num_seqs == 64 + assert config.stage_configs[0].cache_config.kv_cache_memory_bytes is None + stages, _ = StageConfigFactory._create_legacy_from_registry(pipeline, {}) + assert all(not stage.yaml_engine_args["use_v2_model_runner"] for stage in stages) + assert all(not stage.yaml_runtime.get("cuda_mps", False) for stage in stages) + + @pytest.mark.parametrize( + "filename,runner,mps,capacity", + [ + ("moss_tts_local_v1.yaml", False, False, 64), + ("moss_tts_local_mrv2.yaml", True, False, 64), + ("moss_tts_local_mrv2_high_concurrency.yaml", True, False, 128), + ], + ) + def test_moss_local_explicit_profile_overrides_cuda_default(self, monkeypatch, filename, runner, mps, capacity): + from vllm.platforms import current_platform + + from vllm_omni.platforms import current_omni_platform + + monkeypatch.setattr(current_omni_platform, "device_name", "cuda") + monkeypatch.setattr(current_platform, "get_device_total_memory", lambda: 141 * 1024**3) + monkeypatch.setattr("shutil.which", lambda command: "/usr/bin/nvidia-cuda-mps-control") + pipeline = resolve_pipeline_config("moss_tts_local") + path = get_deploy_config_path(filename) + config = VllmOmniConfig.from_pipeline_config(pipeline, deploy_config_path=path) + for stage in config.stage_configs: + assert stage.model_config.use_v2_model_runner is runner + assert stage.runtime_config.cuda_mps is mps + assert stage.scheduler_config.max_num_seqs == capacity + stages, _ = StageConfigFactory._create_legacy_from_registry(pipeline, {}, deploy_config_path=path) + for stage in stages: + assert stage.yaml_engine_args["use_v2_model_runner"] is runner + assert stage.yaml_runtime.get("cuda_mps", False) is mps + assert stage.yaml_engine_args["max_num_seqs"] == capacity + + @pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"]) + def test_moss_local_system_profile_reaches_native_runner_and_cuda_only_mps(self, platform): + pipeline = resolve_pipeline_config("moss_tts_local") + path = Path(get_deploy_config_path("moss_tts_local_mrv2_optimized.yaml")) + deploy = _apply_platform_overrides(load_deploy_config(path), platform=platform) + stages = merge_pipeline_deploy(pipeline, deploy) + is_cuda = platform == "cuda" + assert deploy.cuda_mps is is_cuda + assert all(stage.yaml_runtime.get("cuda_mps", False) is is_cuda for stage in stages) + assert all(stage.yaml_engine_args["use_v2_model_runner"] is is_cuda for stage in stages) + if is_cuda: + args = stages[0].yaml_engine_args + assert args["enable_prefix_caching"] is True + assert args["hf_overrides"]["mrv2_gpu_slot_state"] is True + assert args["hf_overrides"]["mrv2_batch_prefill"] is True + assert args["hf_overrides"]["mrv2_direct_tokens"] is True + assert deploy.connectors["shm"]["extra"]["codec_first_chunk_fast_path"] == 1 + assert args["hf_overrides"]["local_compile_audio_sampler"] is True + extra = deploy.connectors["shm"]["extra"] + assert extra["codec_first_chunk_max_active_streams"] == 32 + assert extra["generation_min_batch_size"] == 16 + assert extra["generation_max_wait_ms"] == 6 + assert not stages[0].yaml_runtime.get("env") + else: + assert all(not stage.yaml_engine_args["enable_prefix_caching"] for stage in stages) + assert all(not stage.yaml_runtime.get("env") for stage in stages) + + def test_platform_mps_rejects_non_boolean(self): + deploy = load_deploy_config(get_deploy_config_path("moss_tts_local_mrv2_optimized.yaml")) + deploy.platforms["cuda"]["cuda_mps"] = "true" + with pytest.raises(ValueError, match="platform cuda_mps must be a boolean"): + _apply_platform_overrides(deploy, platform="cuda") + def test_moss_local_mrv2_preserves_base_profile_and_variant_scope(self): pipeline = resolve_pipeline_config("moss_tts_local") - base_path = Path(get_deploy_config_path(pipeline.default_deploy_config_name)) + base_path = Path(get_deploy_config_path("moss_tts_local_v1.yaml")) candidate_path = Path(get_deploy_config_path("moss_tts_local_mrv2.yaml")) base = load_deploy_config(base_path) candidate = load_deploy_config(candidate_path) @@ -3052,7 +3188,7 @@ class TestPlatformOverrides: assert compilation["cudagraph_mode"] == "FULL" assert compilation["cudagraph_capture_sizes"] == [1, 2, 4, 6, 8, 12, 16, 24, 32, 64, 128] assert compilation["inductor_compile_config"] == {"combo_kernels": False, "benchmark_combo_kernel": False} - base_path = Path(get_deploy_config_path("moss_tts_local.yaml")) + base_path = Path(get_deploy_config_path("moss_tts_local_v1.yaml")) for platform in ("npu", "xpu", "rocm", "musa"): actual = _apply_platform_overrides(load_deploy_config(path), platform=platform) expected = _apply_platform_overrides(load_deploy_config(base_path), platform=platform) diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 7027b5596..26b85a39c 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -1887,7 +1887,7 @@ def test_from_pipeline_config_rejects_reserved_diffusion_kv_mode(tmp_path): @pytest.mark.parametrize("source", ["default", "topology", "deploy", "stage-cli"]) @pytest.mark.parametrize("key_container", [list, tuple]) -def test_diffusion_stage_payload_keys_roundtrip(source, key_container): +def test_diffusion_stage_payload_keys_roundtrip(source, key_container, tmp_path): from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config @@ -1920,7 +1920,7 @@ def test_diffusion_stage_payload_keys_roundtrip(source, key_container): legacy_stage = merge_pipeline_deploy(pipeline, deploy)[0] legacy_args = {**legacy_stage.yaml_engine_args, **(override_keys if source == "stage-cli" else {})} restored_stage = ForkingPickler.loads(ForkingPickler.dumps(stage)) - engine_args = build_engine_args_dict_from_omni_stage_config(restored_stage, model="test-model") + engine_args = build_engine_args_dict_from_omni_stage_config(restored_stage, model=str(tmp_path)) diffusion_kwargs = omni_config_module.extract_diffusion_stage_config_kwargs( engine_args, stage_id=restored_stage.stage_id, include_engine_adapter_metadata=True ) diff --git a/tests/core/sched/test_generation_scheduler_init.py b/tests/core/sched/test_generation_scheduler_init.py index db0285e38..772d9683d 100644 --- a/tests/core/sched/test_generation_scheduler_init.py +++ b/tests/core/sched/test_generation_scheduler_init.py @@ -44,10 +44,13 @@ def _config(*, native=True, stateful=True, tp=1, pp=1, extras=None, capacity=128 ) -@pytest.mark.parametrize("profile", ["high_concurrency", "low_latency"]) +@pytest.mark.parametrize( + "profile,min_batch,max_wait_ms", + [("high_concurrency", 32, 12), ("low_latency", 32, 12), ("optimized", 16, 6)], +) @pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"]) def test_moss_profile_generation_constructor_after_platform_resolution( - construct_scheduler, monkeypatch, mocker, profile, platform + construct_scheduler, monkeypatch, mocker, profile, min_batch, max_wait_ms, platform ): # merge_pipeline_deploy owns platform resolution. Mock the detected device # instead of applying an override first and then resolving the worker's @@ -66,11 +69,12 @@ def test_moss_profile_generation_constructor_after_platform_resolution( warning = mocker.patch("vllm_omni.core.sched.omni_generation_scheduler.logger.warning") scheduler = construct_scheduler(config) # The raw extras persist through fallback; the live constructor must cope. - assert extras["generation_min_batch_size"] == 32 and extras["generation_max_wait_ms"] == 12 + assert extras["generation_min_batch_size"] == min_batch + assert extras["generation_max_wait_ms"] == max_wait_ms if platform == "cuda": assert scheduler._native_data_plane and scheduler.input_coordinator is not None - assert scheduler._generation_min_batch_size == 32 - assert scheduler._generation_max_wait_s == 0.012 + assert scheduler._generation_min_batch_size == min_batch + assert scheduler._generation_max_wait_s == max_wait_ms / 1000 assert scheduler._generation_max_regular_batch == (16 if profile == "low_latency" else 0) warning.assert_not_called() if profile == "low_latency": diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index dc49a7e21..f3ad3bf78 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -61,6 +61,7 @@ from vllm_omni.entrypoints.openai.tts_adapters.base import ( ) from vllm_omni.entrypoints.openai.tts_adapters.capabilities import load_supported_speakers from vllm_omni.entrypoints.openai.tts_adapters.ming_tts import MingTTSAdapter +from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import MossTTSAdapter from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSAdapter, Qwen3TTSCodecLimitError from vllm_omni.entrypoints.openai.tts_adapters.voxtral import VoxtralTTSAdapter from vllm_omni.entrypoints.serve.utils import errors as serve_errors @@ -84,6 +85,42 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] logger = logging.getLogger(__name__) +@pytest.mark.parametrize("sse", [False, True]) +@pytest.mark.parametrize("payload", [None, "empty", "valid"]) +@pytest.mark.asyncio +async def test_moss_empty_stream_is_an_error(mocker, sse, payload): + serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech) + serving._tts_model_type = "moss_tts" + serving.engine_client = SimpleNamespace( + model_config=SimpleNamespace(model="OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5") + ) + serving._adapter = MossTTSAdapter(SpeechServingContext(server=serving, engine_client=serving.engine_client)) + ready = mocker.patch.object(serving, "_mark_ref_audio_artifact_ready_for_request") + discard = mocker.patch.object(serving, "_discard_ref_audio_artifact_warmup") + + async def results(): + if payload is not None: + audio = torch.zeros(320) if payload == "valid" else torch.empty(0) + yield SimpleNamespace(multimodal_output={"model_outputs": audio, "sr": 24000}) + yield SimpleNamespace(multimodal_output={"model_outputs": torch.empty(0), "sr": 24000}) + + if sse: + events = [event async for event in serving._generate_audio_sse_events(results(), "moss-empty")] + assert any("speech.audio.done" in event for event in events) == (payload == "valid") + assert any("speech.audio.error" in event for event in events) == (payload != "valid") + else: + chunks = serving._generate_audio_chunks(results(), "moss-empty", response_format="wav") + if payload == "valid": + output = [chunk async for chunk in chunks] + assert len(output) == 2 and output[0].startswith(b"RIFF") + assert len(output[1]) == 640 + else: + with pytest.raises(TTSGenerationError, match="no audio"): + await anext(chunks) + assert ready.call_count == int(payload == "valid") + assert discard.call_count == int(payload != "valid") + + class TestAudioMixin: @pytest.fixture def audio_mixin(self): diff --git a/tests/entrypoints/openai_api/test_serving_speech_metrics.py b/tests/entrypoints/openai_api/test_serving_speech_metrics.py index a1fb74a46..a5e39872d 100644 --- a/tests/entrypoints/openai_api/test_serving_speech_metrics.py +++ b/tests/entrypoints/openai_api/test_serving_speech_metrics.py @@ -283,7 +283,7 @@ async def test_streaming_speech_does_not_count_empty_payload_as_first_packet(mon request_id="speech-test", request_arrival_ts=100.0, ) - assert len([chunk async for chunk in chunks]) == 2 + assert [chunk async for chunk in chunks] == [b"\0\0" * 320] assert metrics.ttfp_calls == [("1", "2", pytest.approx(0.25))] assert metrics.underrun_calls == [("1", "2", pytest.approx(0.0))] assert metrics.continuity_calls == [("1", "2", 100)] @@ -319,7 +319,11 @@ async def test_streaming_speech_does_not_finalize_continuity_on_validation_error def reject_generation(_tts_params, **_kwargs): raise RuntimeError("generation validation failed") - adapter = SimpleNamespace(validates_generation=True, validate_generation=reject_generation) + adapter = SimpleNamespace( + validates_generation=True, + validate_generation=reject_generation, + validate_stream_audio=lambda **_kwargs: None, + ) serving = _serving(metrics, adapter=adapter) monkeypatch.setattr("vllm_omni.entrypoints.openai.serving_speech.time.time", lambda: 100.25) diff --git a/tests/model_executor/models/moss_tts/test_first_chunk_fast_path.py b/tests/model_executor/models/moss_tts/test_first_chunk_fast_path.py index 85e53e80b..c8e1bd582 100644 --- a/tests/model_executor/models/moss_tts/test_first_chunk_fast_path.py +++ b/tests/model_executor/models/moss_tts/test_first_chunk_fast_path.py @@ -318,26 +318,45 @@ def test_failed_output_handoff_retains_decoded_samples(cuda): fast.close() +@pytest.mark.cpu +@pytest.mark.parametrize("enabled,complete", [(False, False), (True, False), (True, True), (True, None)]) +def test_gate_waits_only_for_enabled_inflight_decode(mocker, enabled, complete): + # A fixed GPU sleep cannot prove that an event is still pending when the + # host reaches gate(), particularly across CUDA and ROCm implementations. + fast = object.__new__(MossFirstChunkFastPath) + fast._device = torch.device("cuda") + fast._gate_main = enabled + event = None if complete is None else mocker.Mock() + if event is not None: + event.query.return_value = complete + fast._inflight = event + stream = mocker.Mock() + current_stream = mocker.patch("torch.cuda.current_stream", return_value=stream) + fast.gate() + if enabled and complete is False: + current_stream.assert_called_once_with(fast._device) + stream.wait_event.assert_called_once_with(event) + else: + current_stream.assert_not_called() + stream.wait_event.assert_not_called() + + @pytest.mark.cuda -def test_gate_orders_main_stream_after_inflight_fast_decode(cuda): +def test_gate_hands_inflight_gpu_output_to_main_stream(cuda): fast, _, _ = _fast(cuda) side = torch.cuda.Stream() + output = torch.zeros(16, device=cuda) + side.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(side): - torch.cuda._sleep(200_000_000) + output.fill_(7) pending = torch.cuda.Event() pending.record(side) fast._inflight = pending - fast.gate() # disabled: no wait - unordered = torch.cuda.Event() - unordered.record() fast._gate_main = True fast.gate() - ordered = torch.cuda.Event() - ordered.record() - unordered.synchronize() - assert not pending.query() and not ordered.query() - ordered.synchronize() - assert pending.query() + received = output.clone() + torch.testing.assert_close(received, torch.full_like(received, 7)) + fast.close() def _cpu_fast(monkeypatch, *, scheduler=None): @@ -545,3 +564,41 @@ def test_platform_failure_wakes_pending_handoff(monkeypatch): assert fast._closed with pytest.raises(RuntimeError, match="decode failed"): fast.order_after(0) + + +@pytest.mark.cpu +@pytest.mark.tts +@pytest.mark.parametrize("limit", [1, 2]) +def test_congested_admission_preserves_routes_and_slots_then_recovers(monkeypatch, limit): + import sys + from functools import partial + + helpers = sys.modules[__name__] + from vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path import MossFirstChunkFastPath + + monkeypatch.setattr(helpers, "MossFirstChunkFastPath", partial(MossFirstChunkFastPath, max_active_streams=limit)) + fast, session, sink, _ = helpers._cpu_fast(monkeypatch) + calls = [] + + class Sink: + def prepare(self, ids): + calls.append(ids) + return sink.prepare(ids) + + fast._sink = Sink() + thread = threading.Thread() + monkeypatch.setattr(thread, "is_alive", lambda: True) + fast._thread = thread + slots = {f"busy{i}": session.acquire() for i in range(limit)} + before_slots, before_free = slots.copy(), session.free.copy() + codes = torch.tensor([1, 2]) + assert not fast.submit("r", "r", codes, slots) + assert slots == before_slots and session.free == before_free + assert not calls and not fast._decoded and not fast._handoffs and fast._jobs.empty() + released = slots.pop("busy0") + session.release(released) + assert fast.submit("r", "r", codes, slots) + codes.fill_(99) + job = fast._jobs.get_nowait() + assert job.slot == slots["r"] and calls == [["r"]] + assert job.codes.tolist() == [[1], [2]] diff --git a/tests/model_executor/models/moss_tts/test_local_compile_config.py b/tests/model_executor/models/moss_tts/test_local_compile_config.py new file mode 100644 index 000000000..ab204d53b --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_local_compile_config.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Deployment sampler settings override the legacy environment switch.""" + +import pytest +from transformers import GPT2Config + +from vllm_omni.model_executor.models.moss_tts import modeling_moss_tts_local_depth as depth + +pytestmark = [pytest.mark.core_model, pytest.mark.tts, pytest.mark.cpu] + + +@pytest.mark.parametrize( + "configured,env,expected", [(None, "0", False), (None, "1", True), (True, "0", True), (False, "1", False)] +) +def test_sampler_compile_setting_and_idempotent_setup(monkeypatch, configured, env, expected): + monkeypatch.setenv("VLLM_OMNI_MOSS_LOCAL_COMPILE_AUDIO_SAMPLER", env) + monkeypatch.setattr(depth.current_omni_platform, "supports_torch_inductor", lambda: True) + calls = [] + + def compile_fn(fn, **kwargs): + calls.append((fn, kwargs)) + return fn + + monkeypatch.setattr(depth.torch, "compile", compile_fn) + model = depth.MossTTSLocalDepthTransformer( + GPT2Config(n_embd=80, n_head=1, n_inner=160), compile_audio_sampler=configured + ) + model.setup_compile() + model.setup_compile() + assert len(calls) == 1 + expected + assert (model._compiled_audio_sampler is depth._sample_token) is expected + sampler_calls = [kwargs for fn, kwargs in calls if fn is depth._sample_token] + assert len(sampler_calls) == int(expected) + if expected: + assert sampler_calls[0] == {"fullgraph": True, "dynamic": True, "options": {"fallback_random": True}} + + +def test_sampler_compile_setting_respects_platform_support(monkeypatch): + monkeypatch.setattr(depth.current_omni_platform, "supports_torch_inductor", lambda: False) + + def forbidden(*args, **kwargs): + raise AssertionError("compile called on unsupported platform") + + monkeypatch.setattr(depth.torch, "compile", forbidden) + model = depth.MossTTSLocalDepthTransformer(GPT2Config(n_embd=80, n_head=1, n_inner=160), compile_audio_sampler=True) + model.setup_compile() + assert model._compiled_audio_sampler is None diff --git a/tests/model_executor/models/moss_tts/test_local_model_state.py b/tests/model_executor/models/moss_tts/test_local_model_state.py index 4956d5c89..60d4269be 100644 --- a/tests/model_executor/models/moss_tts/test_local_model_state.py +++ b/tests/model_executor/models/moss_tts/test_local_model_state.py @@ -40,14 +40,18 @@ def _model(device): model.audio_assistant_slot_token_id, model.im_end_token_id = 2, 3 model.talker_mtp_output_key = ("audio_codes", "current") model.talker_mtp_graph_safe = False + model.talker_mtp_accepts_per_row_generators = True model.gpu_resident_buffer_keys = {("hidden_states", "last"), ("audio_codes", "current")} model.audio_lm_heads = model.audio_embeddings = model.local_text_lm_head = None model._audio_embed = lambda codes: codes[:, :1].expand(-1, 4).to(dtype) / 8 model.force_stop = False - def frame(hidden, *_args, generator=None, temperature=None, top_k=None, top_p=None, **_kwargs): + def frame(hidden, *_args, generator=None, generators=None, temperature=None, top_k=None, top_p=None, **_kwargs): assert (temperature, top_k, top_p) == (1.7, 25, 0.8) - noise = torch.randint(0, 8, (hidden.shape[0], 2), device=device, generator=generator) + if generators is not None: + noise = torch.cat([torch.randint(0, 8, (1, 2), device=device, generator=g) for g in generators]) + else: + noise = torch.randint(0, 8, (hidden.shape[0], 2), device=device, generator=generator) codes = (hidden[:, :2].mul(16).long() + noise) % 8 return torch.full((hidden.shape[0],), not model.force_stop, device=device), codes @@ -85,11 +89,11 @@ def _admit(state, slot, name, seed): state.intermediate_buffer.buffers[slot]["codes"] = {"ref": torch.tensor([[1, 2], [3, 4], [5, 6], [2, 3]])} -def _batch(device, slots, counts): +def _batch(device, slots, counts, index_dtype=torch.int32): starts = np.array([0, *np.cumsum(counts)], dtype=np.int32) return SimpleNamespace( idx_mapping_np=np.array(slots), - idx_mapping=torch.tensor(slots, device=device), + idx_mapping=torch.tensor(slots, device=device, dtype=index_dtype), num_reqs=len(slots), num_tokens=int(starts[-1]), num_scheduled_tokens=np.array(counts), @@ -122,7 +126,8 @@ def _step(state, batch, req_states, dispatcher=None): @pytest.mark.parametrize("seed", [None, 17]) @pytest.mark.parametrize("batch_prefill", [False, True]) -def test_slot_matches_canonical_mixed_prefill_reorder_stop_and_reuse(device, seed, batch_prefill): +@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) +def test_slot_matches_canonical_mixed_prefill_reorder_stop_and_reuse(device, seed, batch_prefill, index_dtype): reference, candidate = (_state(cls, device) for cls in (OmniModelState, MossLocalModelState)) candidate._batch_prefill = batch_prefill for state in (reference, candidate): @@ -154,7 +159,7 @@ def test_slot_matches_canonical_mixed_prefill_reorder_stop_and_reuse(device, see for state in (reference, candidate): state.model.force_stop = index == 3 torch.manual_seed(91 + index) - outputs.append(_step(state, _batch(device, slots, counts), req_states)) + outputs.append(_step(state, _batch(device, slots, counts, index_dtype), req_states)) a, b = outputs torch.testing.assert_close(a[0], b[0], rtol=0, atol=0) torch.testing.assert_close(a[2], b[2], rtol=0, atol=0) @@ -190,6 +195,23 @@ def test_slot_honors_external_stopping_control(device): assert output[1]["codes"]["audio"][0].tolist() == [[8, 8]] +@pytest.mark.cpu +@pytest.mark.parametrize("state_cls", [OmniModelState, MossLocalModelState]) +def test_local_seeded_mtp_keeps_per_row_generators_in_one_batch(state_cls): + state = _state(state_cls, torch.device("cpu")) + generators = [torch.Generator().manual_seed(seed) for seed in (17, 29)] + inputs = [torch.zeros(2, dtype=torch.long), *[torch.ones(2, 4) for _ in range(3)]] + frame = MagicMock(wraps=state.model.local_transformer.generate_frame) + state.model.local_transformer.generate_frame = frame + + state._call_mtp_with_sampling(*inputs, buffers=[{}, {}], req_ids=["first", "second"], generators=generators) + + frame.assert_called_once() + assert frame.call_args.args[0].shape[0] == 2 + assert frame.call_args.kwargs["generators"] == generators + assert frame.call_args.kwargs["generator"] is None + + @pytest.mark.cuda def test_slot_mtp_graph_padding_reorder_and_owned_output(): if not torch.cuda.is_available(): diff --git a/tests/model_executor/models/moss_tts/test_local_prefill_batch.py b/tests/model_executor/models/moss_tts/test_local_prefill_batch.py index 8b6856f4a..792253ddb 100644 --- a/tests/model_executor/models/moss_tts/test_local_prefill_batch.py +++ b/tests/model_executor/models/moss_tts/test_local_prefill_batch.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Reference batching preserves coordinates and asynchronous source ownership.""" -from types import SimpleNamespace +from dataclasses import dataclass import numpy as np import pytest @@ -15,6 +15,12 @@ from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import Mo pytestmark = pytest.mark.core_model +@dataclass +class _PrefillPositions: + prompt_len: np.ndarray + num_computed_tokens: np.ndarray + + @pytest.mark.cpu @pytest.mark.parametrize("flat", [False, True]) def test_prefill_reuses_text_embeddings_and_noncontiguous_reference_offsets(mocker, flat): @@ -28,7 +34,8 @@ def test_prefill_reuses_text_embeddings_and_noncontiguous_reference_offsets(mock "ref_offset": offset, } batch = _batch(torch.device("cpu"), [3, 0], [2, 1]) - req = SimpleNamespace(prompt_len=np.full(5, 8), num_computed_tokens=np.zeros(5)) + req = _PrefillPositions(prompt_len=np.full(5, 8), num_computed_tokens=np.zeros(5)) + req.num_computed_tokens[[3, 0]] = [1, 2] embeds = state._static_inputs_embeds[:3] expected = state.model.embed_input_ids(batch.input_ids) + state.model._audio_embed(codes[[1, 2, 2]]) embed = mocker.spy(state.model.model.embed_tokens, "forward") @@ -40,6 +47,65 @@ def test_prefill_reuses_text_embeddings_and_noncontiguous_reference_offsets(mock assert state.intermediate_buffer.buffers[0]["ref_offset"] == 3 +@pytest.mark.cpu +@pytest.mark.parametrize("batch_prefill", [False, True]) +@pytest.mark.parametrize("cached,count", [(2, 1), (1, 2)]) +def test_cached_prefix_prefill_uses_absolute_reference_position(batch_prefill, cached, count): + state = _state(MossLocalModelState, torch.device("cpu")) + state._batch_prefill = batch_prefill + codes = torch.tensor([[1, 2], [3, 4], [5, 6], [2, 3]]) + state.intermediate_buffer.buffers[0] = {"req_id": "prefix-hit", "codes": {"ref": codes}} + batch = _batch(torch.device("cpu"), [0], [count]) + req = _PrefillPositions(prompt_len=np.full(5, 4), num_computed_tokens=np.full(5, cached)) + embeds = state._static_inputs_embeds[:count] + expected = state.model.embed_input_ids(batch.input_ids) + state.model._audio_embed(codes[cached : cached + count]) + with torch.inference_mode(): + state.run_preprocess(batch, {"input_ids": batch.input_ids, "inputs_embeds": embeds}, req) + torch.testing.assert_close(embeds, expected, rtol=0, atol=0) + assert state.intermediate_buffer.buffers[0]["ref_offset"] == cached + count + + +@pytest.mark.cpu +def test_scalar_prefill_uses_cached_prompt_position(): + state = _state(MossLocalModelState, torch.device("cpu")) + codes = torch.tensor([[1, 2], [3, 4], [5, 6], [2, 3]]) + ids = torch.tensor([2]) + expected = state.model.embed_input_ids(ids) + state.model._audio_embed(codes[2:3]) + with torch.inference_mode(): + _, actual, updates = state.model.preprocess( + ids, None, codes={"ref": codes}, _omni_is_prefill=True, _omni_num_computed_tokens=2 + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert updates["ref_offset"] == 3 + + +@pytest.mark.cpu +@pytest.mark.parametrize("batch_prefill", [False, True]) +@pytest.mark.parametrize("omit_trailing_pad", [False, True]) +def test_prefill_chunk_crosses_reference_end(batch_prefill, omit_trailing_pad): + state = _state(MossLocalModelState, torch.device("cpu")) + state._batch_prefill = batch_prefill + # codes.ref uses prompt coordinates, including PAD rows for text tokens. + codes = torch.tensor([[8, 8], [1, 2], [3, 4], [5, 6], [8, 8], [8, 8]]) + if omit_trailing_pad: + codes = codes[:4] + state.model._audio_embed = lambda rows: rows[:, :1].masked_fill(rows[:, :1] == 8, 0).expand(-1, 4) / 8 + state.intermediate_buffer.buffers[0] = {"req_id": "reference-boundary", "codes": {"ref": codes}} + req = _PrefillPositions(prompt_len=np.full(5, 6), num_computed_tokens=np.zeros(5)) + outputs = [] + with torch.inference_mode(): + for start, count in [(0, 2), (2, 4)]: + req.num_computed_tokens[0] = start + batch = _batch(torch.device("cpu"), [0], [count]) + embeds = state._static_inputs_embeds[:count] + state.run_preprocess(batch, {"input_ids": batch.input_ids, "inputs_embeds": embeds}, req) + outputs.append(embeds.clone()) + expected = state.model.embed_input_ids(torch.full((6,), 2)) + expected[1:4] += torch.tensor([1, 3, 5]).reshape(-1, 1) / 8 + torch.testing.assert_close(torch.cat(outputs), expected, rtol=0, atol=0) + assert state.intermediate_buffer.buffers[0]["ref_offset"] == 6 + + @pytest.mark.cuda def test_pending_reference_uploads_survive_staging_reuse_and_overflow(): if not torch.cuda.is_available(): diff --git a/tests/worker_v2/test_omni_data_plane.py b/tests/worker_v2/test_omni_data_plane.py index 978346070..62898a456 100644 --- a/tests/worker_v2/test_omni_data_plane.py +++ b/tests/worker_v2/test_omni_data_plane.py @@ -11,6 +11,7 @@ from typing import Any import pytest import torch +from vllm import SamplingParams from vllm_omni.worker_v2.delivery import DeliveryCancelledError, DeliveryState, OmniDeliveryManager from vllm_omni.worker_v2.native_output_worker import NativeOutputWorker @@ -126,6 +127,60 @@ def test_full_payload_abort_discards_partial_and_late_outputs(plane): assert not plane._pending_full_payload_send +@pytest.mark.parametrize("async_chunk", [False, True]) +@pytest.mark.parametrize("stop", ["token", "length"]) +def test_non_resumable_output_stops_before_scheduler_terminal(plane, async_chunk, stop): + plane._async_chunk = async_chunk + plane._full_payload_replace_keys_cached = frozenset() + request = _new_request() + request.resumable = False + request.sampling_params = SamplingParams(stop_token_ids=[2150], max_tokens=2 if stop == "length" else 10) + plane.register_request(request) + first, last, stale = (torch.tensor([[value]]) for value in (1, 2, 99)) + plane.reserve_outputs(["internal"]) + _complete(plane, [{"codes.audio": first}], token=21) + plane.reserve_outputs(["internal"]) + _complete(plane, [{"codes.audio": last}], token=2150 if stop == "token" else 22) + plane.reserve_outputs(["internal"]) + assert plane.request_terminal({"internal"}) == 0 + _complete(plane, [{"codes.audio": stale}], token=23) + + payloads = [ + payload["codes"]["audio"] if async_chunk else payload["codes.audio"] + for batch in plane.record.batches + for _, payload in batch + if payload + ] + torch.testing.assert_close(torch.cat(payloads), torch.cat([first, last])) + terminal, _ = plane.record.batches[-1][0] + assert terminal.is_finished() + assert terminal.output_token_count == 2 + assert plane.record.cleaned == ["internal"] + assert not plane._native_outputs_in_flight and not plane._native_terminal_pending + + +@pytest.mark.parametrize("ignore_eos", [False, True]) +def test_native_output_honors_ignore_eos_and_min_tokens(plane, ignore_eos): + request = _new_request() + request.resumable = False + request.sampling_params = SamplingParams(ignore_eos=ignore_eos, min_tokens=2, max_tokens=10) + request.sampling_params.update_from_generation_config({}, eos_token_id=99) + plane.register_request(request) + for token in (99, 99, 21): + _complete(plane, [{"codes.audio": torch.tensor([[token]])}], token=token) + state = plane._native_requests["internal"] + assert state.output_token_ids == ([99, 99, 21] if ignore_eos else [99, 99]) + assert len(plane.record.batches) == (3 if ignore_eos else 2) + + +def test_resumable_output_keeps_scheduler_owned_segment_boundary(plane): + plane.register_request(_new_request()) + for token in (21, 2150, 22): + _complete(plane, [{"codes.audio": torch.tensor([[token]])}], token=token) + assert plane._native_requests["internal"].output_token_ids == [21, 2150, 22] + assert len(plane.record.batches) == 3 + + def test_full_payload_qwen_builder_receives_complete_codec(raw_plane, monkeypatch): from vllm_omni.model_executor.stage_input_processors.qwen3_tts import talker2code2wav_full_payload diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 1aa80ed92..80b7b9194 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -99,9 +99,9 @@ from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO_4_5_P from vllm_omni.model_executor.models.minimax_h3.pipeline import MINIMAX_H3_PIPELINE from vllm_omni.model_executor.models.minimax_music3.pipeline import MINIMAX_MUSIC3_PIPELINE from vllm_omni.model_executor.models.moss_tts.pipeline import ( - MOSS_TTS_LOCAL_PIPELINE, MOSS_TTS_PIPELINE, MOSS_TTS_REALTIME_PIPELINE, + resolve_moss_tts_local_pipeline, ) from vllm_omni.model_executor.models.moss_tts_nano.pipeline import MOSS_TTS_NANO_PIPELINE from vllm_omni.model_executor.models.nemotron_voicechat.pipeline import ( @@ -183,7 +183,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "minimax_h3_disaggregated": MINIMAX_H3_PIPELINE, "minimax_music3": MINIMAX_MUSIC3_PIPELINE, "moss_tts_delay": MOSS_TTS_PIPELINE, - "moss_tts_local": MOSS_TTS_LOCAL_PIPELINE, + "moss_tts_local": resolve_moss_tts_local_pipeline, "moss_tts_nano": MOSS_TTS_NANO_PIPELINE, "moss_tts_realtime": MOSS_TTS_REALTIME_PIPELINE, # Alias: the Nemotron-Labs-Audex-2B repo-root config.json reports diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index b243ea882..ea73c339a 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -914,6 +914,11 @@ def _apply_platform_overrides( device_name = current_omni_platform.device_name platform = device_name.lower() if device_name is not None else None platform_section = (deploy.platforms or {}).get(platform) if platform is not None else None + if platform_section is not None and "cuda_mps" in platform_section: + cuda_mps = platform_section["cuda_mps"] + if not isinstance(cuda_mps, bool): + raise ValueError("platform cuda_mps must be a boolean") + deploy.cuda_mps = cuda_mps if platform_section is not None and "model_runner" in platform_section: model_runner = platform_section["model_runner"] if model_runner not in ("v1", "v2"): diff --git a/vllm_omni/deploy/moss_tts_local.yaml b/vllm_omni/deploy/moss_tts_local.yaml index c351041f5..c2947f727 100644 --- a/vllm_omni/deploy/moss_tts_local.yaml +++ b/vllm_omni/deploy/moss_tts_local.yaml @@ -1,88 +1,9 @@ -# MOSS-TTS-Local-Transformer-v1.5 (MossTTSLocalModel, n_vq=12, 48 kHz stereo) -# HF: OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 -# Codec: OpenMOSS-Team/MOSS-Audio-Tokenizer-v2 (48 kHz, stereo, channel-interleaved RVQ) -# Uses a 1-layer GPT2-style local depth transformer (no delay warm-up), reset every -# frame — same per-frame timing (80 ms/frame, 12.5 fps) as MOSS-TTS-Realtime. -async_chunk: true +# CUDA MRV2/C64 default with full-quota private MPS. +# Automatic selection uses C128 on >=140 GiB CUDA GPUs, or this C64 profile +# on smaller/unknown-memory GPUs. Missing MPS selects moss_tts_local_mrv2.yaml. +# Non-CUDA platforms retain V1. Explicit deploy configs override selection. +base_config: moss_tts_local_mrv2.yaml -speech_cache: - resolve_max_bytes: 2147483648 # 2 GiB - resolve_max_entries: 2048 - -connectors: - shm: - name: SharedMemoryConnector - extra: - connector_get_sleep_s: 0.005 - connector_get_max_wait_first_chunk: 1000 - connector_get_max_wait: 300 - initial_codec_chunk_frames: 1 # 1 frame * 80 ms/frame = 80 ms first audio chunk - codec_chunk_frames: 15 # 15 frames * 80 ms/frame = 1.2 s steady chunks - # Optional startup ramp; its first entry overrides initial chunk size. - # Extra graph shapes/calls trade codec cost and latency for less underrun. - # codec_chunk_ramp: [1, 4, 15] - codec_left_context_frames: 0 - # Chunk ramp (opt-in): gradual chunk-size ladder replacing the - # initial_codec_chunk_frames -> codec_chunk_frames cliff. The last - # entry should equal codec_chunk_frames (steady state). - # codec_chunk_ramp: [2, 4, 8, 15] - -stages: - - stage_id: 0 - max_num_seqs: 64 - gpu_memory_utilization: 0.60 - enable_prefix_caching: false - async_scheduling: true - max_num_batched_tokens: 65536 - max_model_len: 4096 - devices: "0" - # profiler_config: - # profiler: torch - # torch_profiler_dir: /root/vllm-omni-workspace/vllm-omni/results/moss_local_stage0/batch_8 - output_connectors: - to_stage_1: shm - default_sampling_params: - temperature: 1.7 - top_p: 0.8 - top_k: 25 - max_tokens: 4096 - seed: 42 - repetition_penalty: 1.0 - - - stage_id: 1 - max_num_seqs: 64 - gpu_memory_utilization: 0.30 - enforce_eager: false - enable_prefix_caching: false - async_scheduling: true - compilation_config: - # Reused by the codec as execution-batch buckets; T is derived from - # initial_codec_chunk_frames and codec_chunk_frames. - cudagraph_capture_sizes: [1, 2, 4, 8, 16, 32, 64] - cudagraph_num_of_warmups: 1 - # profiler_config: - # profiler: torch - # torch_profiler_dir: /root/vllm-omni-workspace/vllm-omni/results/moss_local_stage1/batch_8 - max_num_batched_tokens: 65536 - max_model_len: 65536 - devices: "0" - input_connectors: - from_stage_0: shm - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 65536 - seed: 42 - repetition_penalty: 1.0 - -# Stage 1 NPU override: PIECEWISE cudagraph for NPUGraph capture; set enforce_eager: true to disable. platforms: - npu: - stages: - - stage_id: 1 - compilation_config: - cudagraph_mode: PIECEWISE - # Re-declared: _apply_platform_overrides() wholesale-replaces compilation_config. - cudagraph_capture_sizes: [1, 2, 4, 8, 16, 32, 64] - cudagraph_num_of_warmups: 1 + cuda: + cuda_mps: true diff --git a/vllm_omni/deploy/moss_tts_local_mrv2.yaml b/vllm_omni/deploy/moss_tts_local_mrv2.yaml index 4b25a41c7..60cabdfd4 100644 --- a/vllm_omni/deploy/moss_tts_local_mrv2.yaml +++ b/vllm_omni/deploy/moss_tts_local_mrv2.yaml @@ -1,10 +1,11 @@ -# Opt-in CUDA MRV2 for MOSS-TTS-Local-Transformer-v1.5. -# Keep the base profile's batching, chunk geometry and V1 default. -base_config: moss_tts_local.yaml -model_runner: v2 +# CUDA MRV2/C64 for MOSS-TTS-Local-Transformer-v1.5, without MPS. +# Also selected automatically when nvidia-cuda-mps-control is unavailable. +base_config: moss_tts_local_v1.yaml +model_runner: v1 platforms: cuda: + model_runner: v2 stages: - stage_id: 0 # Bound prefill work so new prompts do not monopolize streaming decode. diff --git a/vllm_omni/deploy/moss_tts_local_mrv2_high_concurrency_mps.yaml b/vllm_omni/deploy/moss_tts_local_mrv2_high_concurrency_mps.yaml new file mode 100644 index 000000000..6b11bab05 --- /dev/null +++ b/vllm_omni/deploy/moss_tts_local_mrv2_high_concurrency_mps.yaml @@ -0,0 +1,8 @@ +# Automatic CUDA MRV2/C128 default on GPUs with at least 140 GiB total memory +# when nvidia-cuda-mps-control is available. Keep the published C128 baseline +# and its original Local projection/sampler; add full-quota private MPS only. +base_config: moss_tts_local_mrv2_high_concurrency.yaml + +platforms: + cuda: + cuda_mps: true diff --git a/vllm_omni/deploy/moss_tts_local_mrv2_optimized.yaml b/vllm_omni/deploy/moss_tts_local_mrv2_optimized.yaml new file mode 100644 index 000000000..2588eeef4 --- /dev/null +++ b/vllm_omni/deploy/moss_tts_local_mrv2_optimized.yaml @@ -0,0 +1,43 @@ +# CUDA system profile for MOSS-TTS-Local-Transformer-v1.5. +# Selected automatically with at least 140 GiB GPU memory and MPS available. +# Requires H200-class memory and nvidia-cuda-mps-control on PATH. +# Adds prefix caching, Triton backbone attention and faster codec dispatch; +# keeps the original Local depth projection and native Torch sampler. +# The throughput baseline remains moss_tts_local_mrv2_high_concurrency.yaml. +base_config: moss_tts_local_mrv2_high_concurrency.yaml +model_runner: v1 + +connectors: + shm: + extra: + codec_first_chunk_fast_path: 1 + # Keep the fast first chunk for low load; coalesce crowded streams. + codec_first_chunk_max_active_streams: 32 + generation_min_batch_size: 16 + generation_max_wait_ms: 6 + +platforms: + cuda: + model_runner: v2 + cuda_mps: true + stages: + - stage_id: 0 + hf_overrides: + mrv2_gpu_slot_state: true + mrv2_batch_prefill: true + mrv2_direct_tokens: true + # Compile the native Torch sampler; explicit generators keep its fallback. + local_compile_audio_sampler: true + enable_prefix_caching: true + attention_config: + backend: TRITON_ATTN + - stage_id: 1 + compilation_config: + mode: 3 + backend: inductor + cudagraph_mode: FULL + cudagraph_capture_sizes: [1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32, 40, 48, 56, 64, 80, 96, 112, 128] + cudagraph_num_of_warmups: 1 + inductor_compile_config: + combo_kernels: false + benchmark_combo_kernel: false diff --git a/vllm_omni/deploy/moss_tts_local_v1.yaml b/vllm_omni/deploy/moss_tts_local_v1.yaml new file mode 100644 index 000000000..c351041f5 --- /dev/null +++ b/vllm_omni/deploy/moss_tts_local_v1.yaml @@ -0,0 +1,88 @@ +# MOSS-TTS-Local-Transformer-v1.5 (MossTTSLocalModel, n_vq=12, 48 kHz stereo) +# HF: OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 +# Codec: OpenMOSS-Team/MOSS-Audio-Tokenizer-v2 (48 kHz, stereo, channel-interleaved RVQ) +# Uses a 1-layer GPT2-style local depth transformer (no delay warm-up), reset every +# frame — same per-frame timing (80 ms/frame, 12.5 fps) as MOSS-TTS-Realtime. +async_chunk: true + +speech_cache: + resolve_max_bytes: 2147483648 # 2 GiB + resolve_max_entries: 2048 + +connectors: + shm: + name: SharedMemoryConnector + extra: + connector_get_sleep_s: 0.005 + connector_get_max_wait_first_chunk: 1000 + connector_get_max_wait: 300 + initial_codec_chunk_frames: 1 # 1 frame * 80 ms/frame = 80 ms first audio chunk + codec_chunk_frames: 15 # 15 frames * 80 ms/frame = 1.2 s steady chunks + # Optional startup ramp; its first entry overrides initial chunk size. + # Extra graph shapes/calls trade codec cost and latency for less underrun. + # codec_chunk_ramp: [1, 4, 15] + codec_left_context_frames: 0 + # Chunk ramp (opt-in): gradual chunk-size ladder replacing the + # initial_codec_chunk_frames -> codec_chunk_frames cliff. The last + # entry should equal codec_chunk_frames (steady state). + # codec_chunk_ramp: [2, 4, 8, 15] + +stages: + - stage_id: 0 + max_num_seqs: 64 + gpu_memory_utilization: 0.60 + enable_prefix_caching: false + async_scheduling: true + max_num_batched_tokens: 65536 + max_model_len: 4096 + devices: "0" + # profiler_config: + # profiler: torch + # torch_profiler_dir: /root/vllm-omni-workspace/vllm-omni/results/moss_local_stage0/batch_8 + output_connectors: + to_stage_1: shm + default_sampling_params: + temperature: 1.7 + top_p: 0.8 + top_k: 25 + max_tokens: 4096 + seed: 42 + repetition_penalty: 1.0 + + - stage_id: 1 + max_num_seqs: 64 + gpu_memory_utilization: 0.30 + enforce_eager: false + enable_prefix_caching: false + async_scheduling: true + compilation_config: + # Reused by the codec as execution-batch buckets; T is derived from + # initial_codec_chunk_frames and codec_chunk_frames. + cudagraph_capture_sizes: [1, 2, 4, 8, 16, 32, 64] + cudagraph_num_of_warmups: 1 + # profiler_config: + # profiler: torch + # torch_profiler_dir: /root/vllm-omni-workspace/vllm-omni/results/moss_local_stage1/batch_8 + max_num_batched_tokens: 65536 + max_model_len: 65536 + devices: "0" + input_connectors: + from_stage_0: shm + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + seed: 42 + repetition_penalty: 1.0 + +# Stage 1 NPU override: PIECEWISE cudagraph for NPUGraph capture; set enforce_eager: true to disable. +platforms: + npu: + stages: + - stage_id: 1 + compilation_config: + cudagraph_mode: PIECEWISE + # Re-declared: _apply_platform_overrides() wholesale-replaces compilation_config. + cudagraph_capture_sizes: [1, 2, 4, 8, 16, 32, 64] + cudagraph_num_of_warmups: 1 diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 3a9520574..42afdbeb0 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -1698,7 +1698,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): chunk_np = resampler.process(chunk_np) if chunk_np.size == 0: continue - if self._tts_model_type in _AUDEX_NO_AUDIO_GUARD_MODEL_TYPES and int(np.size(chunk_np)) == 0: + if int(np.size(chunk_np)) == 0: # Zero-size chunks must not emit a WAV header or count # as first audio; the post-loop guard below needs to # see an audio-less stream to fail the request. @@ -1762,6 +1762,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # Audex contract: zero codec tokens must abort the stream, not # complete it cleanly with zero audio bytes. raise ValueError("Audex produced no audio (the thinker emitted zero or invalid codec tokens)") + if adapter is not None: + adapter.validate_stream_audio(has_audio=first_audio_chunk_s is not None) # Check before committing the reference-audio artifact or logging # success. Streaming protocols may already have emitted partial # bytes, but they must terminate as an error rather than cleanly. diff --git a/vllm_omni/entrypoints/openai/tts_adapters/base.py b/vllm_omni/entrypoints/openai/tts_adapters/base.py index 89eb463d9..4d0d3098f 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/base.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/base.py @@ -331,6 +331,9 @@ class TTSModelAdapter(ABC): to unrelated TTS models. """ + def validate_stream_audio(self, *, has_audio: bool) -> None: + """Validate audio availability before recording streaming success.""" + def collect_response_metadata(self, audio_output: Mapping[str, Any], collect: dict) -> None: """Fold engine-side metadata from the mm payload into ``collect``. diff --git a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py index ed7865789..c91b5a8f3 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py @@ -21,6 +21,7 @@ from vllm_omni.entrypoints.openai.tts_adapters.base import ( ARTTSAdapter, OutputPolicy, PreparedRequest, + TTSGenerationError, apply_max_new_tokens, conditioning_cache_salt, ) @@ -519,6 +520,11 @@ class _MossTTSAdapterBase(ARTTSAdapter): params["ref_audio_2_cache_key"] = resolve_keys[1] return params + def validate_stream_audio(self, *, has_audio: bool) -> None: + # Local's sampled binary gate can stop before its first codec frame. + if self._moss_variant == "local" and not has_audio: + raise TTSGenerationError("MOSS-TTS Local produced no audio output.") + def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: """Validate any MOSS-TTS-family request (nano + 5 full variants). diff --git a/vllm_omni/model_executor/models/moss_tts/first_chunk_fast_path.py b/vllm_omni/model_executor/models/moss_tts/first_chunk_fast_path.py index 34361e547..cde07c1f6 100644 --- a/vllm_omni/model_executor/models/moss_tts/first_chunk_fast_path.py +++ b/vllm_omni/model_executor/models/moss_tts/first_chunk_fast_path.py @@ -64,6 +64,7 @@ class MossFirstChunkFastPath: sample_rate: torch.Tensor, device: torch.device, gate_main: bool = False, + max_active_streams: int = 0, handoff_timeout_s: float = 30.0, ) -> None: if not math.isfinite(handoff_timeout_s) or handoff_timeout_s <= 0: @@ -74,6 +75,7 @@ class MossFirstChunkFastPath: # before its next replay, so the small first-chunk graph runs mostly # uncontended instead of interleaving with a large regular replay. self._gate_main = bool(gate_main) + self._max_active_streams = max(0, int(max_active_streams)) self._inflight: torch.cuda.Event | None = None self._wrapper = wrapper self._n_vq = int(n_vq) @@ -154,6 +156,7 @@ class MossFirstChunkFastPath: or self._thread is None or not self._thread.is_alive() or request_key in req_slots + or (self._max_active_streams and len(req_slots) >= self._max_active_streams) ): return False delivery = self._sink.prepare([request_id]) diff --git a/vllm_omni/model_executor/models/moss_tts/local_model_state.py b/vllm_omni/model_executor/models/moss_tts/local_model_state.py index 5a1c8f153..83d1a7237 100644 --- a/vllm_omni/model_executor/models/moss_tts/local_model_state.py +++ b/vllm_omni/model_executor/models/moss_tts/local_model_state.py @@ -208,9 +208,12 @@ class MossLocalModelState(OmniModelState): def _select_rows(self, tensor, rows): if rows == list(range(len(rows))): - return tensor[: len(rows)] - indices = _metadata_to_device(np.asarray(rows, dtype=np.int64), tensor.device) - return tensor.index_select(0, indices) + selected = tensor[: len(rows)] + else: + indices = _metadata_to_device(np.asarray(rows, dtype=np.int64), tensor.device) + selected = tensor.index_select(0, indices) + # vLLM 0.31 uses int32 slot mappings; index_copy_ requires int64. + return selected.long() def run_preprocess(self, input_batch, model_inputs, req_states=None, mtp_batch_descriptor_dispatcher=None): input_ids = model_inputs.get("input_ids") @@ -255,7 +258,9 @@ class MossLocalModelState(OmniModelState): continue if self._batch_prefill: ref = (buf.get("codes", {}) or {}).get("ref") - ref_offset = int(buf.get("ref_offset", 0)) + # Prefix hits skip hooks for cached prompt tokens. Use the + # scheduler's absolute prompt position, also on first admission. + ref_offset = int(computed) if computed is not None else int(buf.get("ref_offset", 0)) # Device-side references keep the canonical path; do not add # an implicit D2H just to place them in a host staging buffer. if not isinstance(ref, torch.Tensor) or ref.device.type == "cpu": @@ -264,7 +269,7 @@ class MossLocalModelState(OmniModelState): ref = ref.view(-1, self.model.n_vq) if ref.dim() == 2: chunk = ref[ref_offset : ref_offset + count] - if chunk.numel() and chunk.shape[0] == count: + if chunk.numel(): references.append((start, chunk)) # Text embeddings have already been computed for all input # tokens. Match the canonical prefill state transition. @@ -419,7 +424,7 @@ class MossLocalModelState(OmniModelState): # mixed/chunked-prefill batches. index_copy snapshots graph outputs. last = input_batch.query_start_loc[1 : input_batch.num_reqs + 1] - 1 self._hidden_pool.index_copy_( - 0, input_batch.idx_mapping[: input_batch.num_reqs], hidden_states.index_select(0, last) + 0, input_batch.idx_mapping[: input_batch.num_reqs].long(), hidden_states.index_select(0, last) ) if self._local_eager_mtp: completing = self._completing_rows diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py index a526d8e33..b2ef81208 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py @@ -802,6 +802,7 @@ class MossTTSCodecDecoder(nn.Module): sample_rate=self._sr_tensor, device=next(self._codec.parameters()).device, gate_main=bool(self._connector_int("codec_first_chunk_gate", default=0)), + max_active_streams=self._connector_int("codec_first_chunk_max_active_streams", default=0), ) logger.info( "MOSS-TTS codec first-chunk fast path ready: B=%s T=%d", diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py index ae837eddb..523fe174d 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py @@ -197,7 +197,9 @@ class MossTTSLocalDepthTransformer(nn.Module): every cache position before reading it, including CUDA Graph replay. """ - def __init__(self, gpt2_config, hidden_size: int | None = None) -> None: + def __init__( + self, gpt2_config, hidden_size: int | None = None, *, compile_audio_sampler: bool | None = None + ) -> None: super().__init__() self.hidden_size = int(hidden_size if hidden_size is not None else gpt2_config.n_embd) n_head = int(gpt2_config.n_head) @@ -208,6 +210,7 @@ class MossTTSLocalDepthTransformer(nn.Module): self.ln_f = nn.LayerNorm(self.hidden_size, eps=eps) self._compiled_forward_prefix = None self._compiled_audio_sampler = None + self._compile_audio_sampler = compile_audio_sampler def _forward_prefix( self, @@ -231,7 +234,10 @@ class MossTTSLocalDepthTransformer(nn.Module): dynamic=True, options={"epilogue_fusion": False}, ) - if os.environ.get("VLLM_OMNI_MOSS_LOCAL_COMPILE_AUDIO_SAMPLER", "0") == "1": + compile_sampler = self._compile_audio_sampler + if compile_sampler is None: + compile_sampler = os.environ.get("VLLM_OMNI_MOSS_LOCAL_COMPILE_AUDIO_SAMPLER", "0") == "1" + if compile_sampler: # Keep the existing top-k/top-p algorithm and torch RNG. Explicit # per-request generators use the original helper below. The # binary continue/stop head is deliberately unchanged. diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py index 303724e2e..378f14824 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py @@ -1349,7 +1349,11 @@ class MossTTSLocalTalkerForGeneration(nn.Module): [nn.Linear(hidden_size, self.audio_vocab_size, bias=False) for _ in range(self.n_vq)] ) self.local_text_lm_head = nn.Linear(hidden_size, 2, bias=False) - self.local_transformer = MossTTSLocalDepthTransformer(self.config.gpt2_config, hidden_size=hidden_size) + self.local_transformer = MossTTSLocalDepthTransformer( + self.config.gpt2_config, + hidden_size=hidden_size, + compile_audio_sampler=getattr(self.config, "local_compile_audio_sampler", None), + ) self._batch_state: list[dict[str, Any]] | None = None # Stacked embedding cache built after load_weights() for vectorised @@ -1520,16 +1524,18 @@ class MossTTSLocalTalkerForGeneration(nn.Module): embeds = self.model.embed_tokens(input_ids) ref_codes = (info_dict.get("codes", {}) or {}).get("ref") - ref_offset = int(info_dict.get("ref_offset", 0)) + # Cached prompt tokens already consumed their reference rows even + # though this request's preprocessing hook never saw those tokens. + ref_offset = int(info_dict.get("_omni_num_computed_tokens", info_dict.get("ref_offset", 0))) if isinstance(ref_codes, torch.Tensor) and ref_codes.numel() > 0: if ref_codes.dim() == 1 and ref_codes.numel() % self.n_vq == 0: ref_codes = ref_codes.view(-1, self.n_vq) if isinstance(ref_codes, torch.Tensor) and ref_codes.dim() == 2: end_off = ref_offset + span_len chunk = ref_codes[ref_offset:end_off] - if chunk.numel() > 0 and chunk.shape[0] == span_len: + if chunk.numel() > 0: codes = chunk.to(device=device, dtype=torch.long) - embeds = embeds + self._audio_embed(codes) + embeds[: chunk.shape[0]] += self._audio_embed(codes) info_update: dict[str, Any] = { "audio_state": {"is_stopping": False}, @@ -1701,6 +1707,10 @@ class MossTTSLocalTalkerForGeneration(nn.Module): """Honor the same platform graph-safety override as the V1 runner.""" return self.talker_mtp_graph_safe + @property + def mtp_accepts_per_row_generators(self) -> bool: + return self.talker_mtp_accepts_per_row_generators + # Package runner-generated audio frames # ------------------------------------------------------------------ diff --git a/vllm_omni/model_executor/models/moss_tts/pipeline.py b/vllm_omni/model_executor/models/moss_tts/pipeline.py index 766c0d136..19701dd30 100644 --- a/vllm_omni/model_executor/models/moss_tts/pipeline.py +++ b/vllm_omni/model_executor/models/moss_tts/pipeline.py @@ -6,6 +6,11 @@ # Licensed under the Apache License, Version 2.0 (the "License"). """Pipeline topology for all MOSS-TTS variants (2-stage: talker → codec).""" +import shutil +from dataclasses import replace + +from transformers import PretrainedConfig +from vllm.logger import init_logger from vllm.sampling_params import RequestOutputKind from vllm_omni.config.stage_config import ( @@ -13,8 +18,10 @@ from vllm_omni.config.stage_config import ( StageExecutionType, StagePipelineConfig, ) +from vllm_omni.model_executor.models.moss_tts.configuration_moss_tts import MossTTSLocalConfig _PROC = "vllm_omni.model_executor.stage_input_processors.moss_tts" +logger = init_logger(__name__) # --------------------------------------------------------------------------- # Shared 2-stage pipeline (used by all 5 MOSS-TTS variants) @@ -135,4 +142,40 @@ MOSS_TTS_LOCAL_PIPELINE = PipelineConfig( # they have different talker architectures from the delay variant # (MossTTSRealtime / MossTTSLocalModel vs MossTTSDelayModel). -__all__ = ["MOSS_TTS_PIPELINE", "MOSS_TTS_REALTIME_PIPELINE", "MOSS_TTS_LOCAL_PIPELINE"] + +def resolve_moss_tts_local_pipeline(hf_config: PretrainedConfig | None = None) -> PipelineConfig | None: + """Select CUDA MRV2, with the C128 system profile when prerequisites are met. + + The platform memory query uses NVML, avoiding CUDA initialization before + workers spawn. Explicit deploy configs override this pipeline default. + """ + from vllm.platforms import current_platform + + from vllm_omni.platforms import current_omni_platform + + if hf_config is not None and not isinstance(hf_config, MossTTSLocalConfig): + return None + if current_omni_platform.device_name != "cuda": + return replace(MOSS_TTS_LOCAL_PIPELINE, default_deploy_config_name="moss_tts_local_v1.yaml") + if shutil.which("nvidia-cuda-mps-control") is None: + logger.info("MOSS Local defaults to CUDA MRV2/C64 without MPS: nvidia-cuda-mps-control is unavailable") + return replace(MOSS_TTS_LOCAL_PIPELINE, default_deploy_config_name="moss_tts_local_mrv2.yaml") + try: + memory_bytes = current_platform.get_device_total_memory() + except Exception as exc: + logger.warning("Cannot query CUDA memory for MOSS Local; using MRV2/C64 with MPS: %s", exc) + return MOSS_TTS_LOCAL_PIPELINE + if memory_bytes >= 140 * 1024**3: + return replace( + MOSS_TTS_LOCAL_PIPELINE, + default_deploy_config_name="moss_tts_local_mrv2_optimized.yaml", + ) + return MOSS_TTS_LOCAL_PIPELINE + + +__all__ = [ + "MOSS_TTS_PIPELINE", + "MOSS_TTS_REALTIME_PIPELINE", + "MOSS_TTS_LOCAL_PIPELINE", + "resolve_moss_tts_local_pipeline", +] diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/streaming_decoder.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/streaming_decoder.py index b1dffed00..a7d8a2473 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/streaming_decoder.py +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/streaming_decoder.py @@ -45,6 +45,7 @@ from vllm.triton_utils import tl, triton from vllm_omni.platforms import current_omni_platform _KERNEL = 7 +_CUDA_SIN = current_omni_platform.is_cuda() # --------------------------------------------------------------------------- kernels @@ -68,13 +69,15 @@ def _rvq_gather_kernel(codes_ptr, books_ptr, out_ptr, NQ: tl.constexpr, CB: tl.c @triton.jit -def _fast_sin(y): +def _fast_sin(y, CUDA_SIN: tl.constexpr = _CUDA_SIN): # Two-step reduction to [-pi, pi], then the hardware approximation # (absolute error ~4e-7 there): accurate sin's instruction count, not # memory, bounds every kernel that applies SnakeBeta. - k = tl.extra.cuda.libdevice.rint(y * 0.15915494309189535) + k = tl.extra.libdevice.rint(y * 0.15915494309189535) r = tl.fma(-k, -1.7484555314695172e-07, tl.fma(-k, 6.2831854820251465, y)) - return tl.inline_asm_elementwise("sin.approx.f32 $0, $1;", "=r,r", [r], dtype=tl.float32, is_pure=True, pack=1) + if CUDA_SIN: + return tl.inline_asm_elementwise("sin.approx.f32 $0, $1;", "=r,r", [r], dtype=tl.float32, is_pure=True, pack=1) + return tl.sin(r) @triton.jit diff --git a/vllm_omni/worker_v2/omni_data_plane.py b/vllm_omni/worker_v2/omni_data_plane.py index d1ace2a9a..d82e913dc 100644 --- a/vllm_omni/worker_v2/omni_data_plane.py +++ b/vllm_omni/worker_v2/omni_data_plane.py @@ -38,6 +38,30 @@ class _NativeRequestState: resumable: bool = False output_token_ids: list[int] = field(default_factory=list) finished: bool = False + output_stopped: bool = False + + def accept_tokens(self, token_ids: list[int]) -> None: + """Fence publications past the sampled stop, before scheduler ACK. + + Async scheduling may already have executed another step when the + scheduler consumes EOS. Keep the terminating step's payload, but do + not publish the later steps while lifecycle reservations drain. + Resumable requests keep their segment boundaries owned by the scheduler. + """ + params = self.sampling_params + stop_ids = set(getattr(params, "stop_token_ids", None) or ()) + if not getattr(params, "ignore_eos", False): + stop_ids.update(getattr(params, "_all_stop_token_ids", ()) or ()) + min_tokens = getattr(params, "min_tokens", 0) or 0 + max_tokens = getattr(params, "max_tokens", None) + for token_id in token_ids: + self.output_token_ids.append(int(token_id)) + if not self.resumable and ( + (len(self.output_token_ids) >= min_tokens and token_id in stop_ids) + or (max_tokens is not None and len(self.output_token_ids) >= max_tokens) + ): + self.output_stopped = True + break def snapshot(self, *, include_token_history: bool) -> SimpleNamespace: prompt = list(self.prompt_token_ids) if include_token_history else [] @@ -506,9 +530,15 @@ class OmniRunnerDataPlane(OmniConnectorModelRunnerMixin): state = self._native_requests.get(req_id) if state is None: continue - state.output_token_ids.extend(int(token_id) for token_id in sampled_by_req.get(req_id, [])) state.finished = req_id in terminal_req_ids - payload = payload_by_req.get(req_id) + # Native transport runs ahead of scheduler consumption. Its + # terminal notification fences resource cleanup, not the sampled + # token boundary: a later in-flight frame must not reach the codec. + if state.output_stopped: + payload = None + else: + state.accept_tokens(sampled_by_req.get(req_id, [])) + payload = payload_by_req.get(req_id) if not self._async_chunk: # Full-payload consumers execute once, after the producer has # finished. Reuse V1's concat/replace contract for flat outputs