[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

@@ -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.

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.
 

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)

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
     )

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":

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):

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)
 

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]]

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

tests/model_executor/models/moss_tts/test_local_model_state.py

@@ -40,13 +40,17 @@ 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)
+        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():

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():

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
 

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

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"):

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

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.

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

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

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

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.

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``.
 

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).
 

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])

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)]
+            selected = tensor[: len(rows)]
+        else:
             indices = _metadata_to_device(np.asarray(rows, dtype=np.int64), tensor.device)
-        return tensor.index_select(0, indices)
+            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

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",

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.

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
     # ------------------------------------------------------------------
 

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",
+]

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))
+    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

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,8 +530,14 @@ 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
+            # 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