[Model] Add single-stage Qwen3-TTS with in-Talker streaming codec decode

docs/configuration/stage_configs.md

@@ -319,7 +319,9 @@ ROCm and MUSA. Select one of these deployment profiles:
 | `qwen3_tts_mrv2.yaml` | V2 | B1 | Explicit MRV2 profile (same runner selection as the default) |
 | `qwen3_tts_high_concurrency_mrv2.yaml` | V2 | B1, B2 | Opt-in throughput tuning |
 | `qwen3_tts_high_concurrency_mrv2_b4.yaml` | V2 | B1, B2, B3, B4 | Experimental throughput / buffered playback |
+| `qwen3_tts_high_concurrency_mrv2_single_gpu.yaml` | V2 | B1–B8 | Experimental two-stage deployment on one GPU, with MPS |
 | `qwen3_tts_high_concurrency.yaml` | V1 | Existing defaults | V1 high-concurrency control |
+| `qwen3_tts_fused_single_gpu.yaml` | V2 | N/A (in-Talker decoder) | Opt-in single-stage CUDA pipeline on one GPU |
 
 ```bash
 # MRV2 is the default; pass a copy with `model_runner: v1` to force V1.
@@ -350,12 +352,73 @@ supply `mtp_sampling_params` and `get_mtp_seed(sampling_params)` for model-local
 (with `mtp_sample_steps` and `mtp_sample_vocab_size`) as needed. Qwen3-TTS retains
 its existing `talker_mtp` entry point for V1.
 
+### Selecting one or two Qwen3-TTS stages
+
+Choose the deployment profile explicitly with `--deploy-config`. The default
+`qwen3_tts` pipeline retains two stages: the Talker publishes codec chunks to
+Code2Wav, which produces audio. Its stage devices can be placed on the same
+GPU or on separate GPUs. `qwen3_tts_fused` is a separate, opt-in pipeline for
+one CUDA GPU: the Talker owns the stateful codec decoder and emits PCM directly,
+without a Code2Wav engine or stage connector.
+
+```bash
+# Two stages on one GPU; the two-stage pipeline remains independently selectable.
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \
+  --deploy-config vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml
+
+# One stage on one GPU; no inter-stage transport or Code2Wav scheduling.
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \
+  --deploy-config vllm_omni/deploy/qwen3_tts_fused_single_gpu.yaml
+```
+
+The single-stage profile requires CUDA, V2, asynchronous chunks, TP=PP=1,
+an in-process worker and disabled prefix caching. Select the full profile;
+adding `talker_stream_decode: true` to a two-stage Talker is rejected because
+its latent output belongs to the Code2Wav input contract. Single-stage options
+live in `stages[].additional_config`; two-stage transport and first-frame
+options live in `connectors.*.extra`.
+
+The single-stage profile defaults `talker_stream_first_audio` to `false` for
+throughput. To opt into a separate first-frame delivery queue, save this
+overlay beside the bundled deployment files and pass its path to
+`--deploy-config`:
+
+```yaml
+base_config: qwen3_tts_fused_single_gpu.yaml
+stages:
+  - stage_id: 0
+    additional_config:
+      talker_stream_decode: true
+      talker_stream_first_audio: true
+      ref_code_context_frames: 72
+      code_predictor_kv_cache: true
+      code_predictor_fused_sampling: true
+      code_predictor_fused: true
+```
+
+This queue uses PCM from the single-stage decoder. The two-stage
+`talker_first_audio` option below uses its own first-frame decoder and does
+not require single-stage execution. The two options cannot be enabled
+together. Single-stage prefix-cache reset with
+`reset_running_requests=True` returns `False` while requests are active;
+wait for completion or abort them before retrying.
+
+Reference-cloning context is selected independently by each profile. The
+single-stage profile explicitly uses 72 reference frames, matching the
+`qwen3_tts.yaml` context boundary. A shorter
+`additional_config.ref_code_context_frames` (for example 25) can change
+timbre continuity even when WER is similar;
+the throughput measurements for CustomVoice do not qualify Base voice
+cloning quality. Reference layout, grouping and priming are model hooks;
+the shared worker does not select Qwen's context length or codebook layout.
+
 ### Qwen3-TTS first-frame delivery and rollback
 
 The standard `qwen3_tts.yaml`, `qwen3_tts_high_concurrency.yaml`,
 `qwen3_tts_mrv2.yaml` and `qwen3_tts_high_concurrency_mrv2.yaml` profiles
 enable the `talker_first_audio` connector option, as does the experimental
-single-GPU profile. This changes the default CUDA MRV2 streaming path:
+two-stage `qwen3_tts_high_concurrency_mrv2_single_gpu.yaml` profile.
+This changes the default CUDA MRV2 streaming path:
 residual prediction runs eagerly after the Talker sample, the Talker loads
 an additional first-frame decoder with its weights and CUDA graphs, and the
 orchestrator orders its audio before subsequent Code2Wav chunks.

tests/config/test_config_merge.py

@@ -35,3 +35,161 @@ def test_recursive_deploy_inheritance(tmp_path):
     merged = resolve_deploy_yaml(leaf)
     assert merged["speech_cache"] == {"nested": {"a": 1, "b": 0}}
     assert merged["async_chunk"] is False
+
+
+def test_single_stage_streaming_capability_uses_central_async_resolver():
+    from vllm_omni.config.stage_config import DeployConfig, _resolve_pipeline_async_chunk_enabled
+    from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_FUSED_PIPELINE
+
+    assert _resolve_pipeline_async_chunk_enabled(QWEN3_TTS_FUSED_PIPELINE, DeployConfig())
+    with pytest.raises(ValueError, match="requires async_chunk=True"):
+        _resolve_pipeline_async_chunk_enabled(QWEN3_TTS_FUSED_PIPELINE, DeployConfig(async_chunk=False))
+
+
+@pytest.mark.parametrize(
+    "pipeline_key,profile,stage_count",
+    [
+        ("qwen3_tts", "qwen3_tts_high_concurrency_mrv2_single_gpu.yaml", 2),
+        ("qwen3_tts_fused", "qwen3_tts_fused_single_gpu.yaml", 1),
+    ],
+)
+def test_qwen3_tts_deployment_selection_preserves_output_and_transport_contract(pipeline_key, profile, stage_count):
+    from tests.helpers.stage_config import get_deploy_config_path
+    from vllm_omni.config.pipeline_registry import resolve_pipeline_config
+    from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
+
+    pipeline = resolve_pipeline_config(pipeline_key)
+    assert pipeline is not None
+    deploy = load_deploy_config(get_deploy_config_path(profile))
+    stages = merge_pipeline_deploy(pipeline, deploy)
+    assert len(stages) == stage_count
+    talker = stages[0]
+    if stage_count == 1:
+        assert talker.final_output and talker.final_output_type == "audio"
+        assert talker.yaml_engine_args["engine_output_type"] == "audio"
+        assert "custom_process_next_stage_input_func" not in talker.yaml_engine_args
+        assert not talker.yaml_extras.get("output_connectors")
+        assert talker.yaml_engine_args["supports_running_prefix_cache_reset"] is False
+        assert talker.yaml_engine_args["additional_config"]["ref_code_context_frames"] == 72
+    else:
+        assert not talker.final_output
+        assert talker.yaml_engine_args["engine_output_type"] == "latent"
+        assert talker.yaml_engine_args["custom_process_next_stage_input_func"].endswith("talker2code2wav_async_chunk")
+        assert talker.yaml_extras["output_connectors"]
+        assert stages[1].input_sources == [0] and stages[1].final_output_type == "audio"
+        assert talker.yaml_engine_args.get("supports_running_prefix_cache_reset", True)
+
+
+def test_single_stage_running_reset_capability_cannot_be_overridden():
+    from vllm_omni.config.omni_config import VllmOmniConfig
+    from vllm_omni.config.stage_config import DeployConfig, StageDeployConfig, merge_pipeline_deploy
+    from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_FUSED_PIPELINE
+
+    deploy = DeployConfig(
+        stages=[StageDeployConfig(stage_id=0, engine_extras={"supports_running_prefix_cache_reset": True})]
+    )
+    legacy = merge_pipeline_deploy(QWEN3_TTS_FUSED_PIPELINE, deploy)[0]
+    legacy.runtime_overrides["supports_running_prefix_cache_reset"] = True
+    assert legacy.to_omegaconf().engine_args.supports_running_prefix_cache_reset is False
+    with pytest.raises(ValueError, match="no structured config owner: supports_running_prefix_cache_reset"):
+        VllmOmniConfig.from_pipeline_config(QWEN3_TTS_FUSED_PIPELINE, user_deploy_config=deploy)
+    structured = VllmOmniConfig.from_pipeline_config(QWEN3_TTS_FUSED_PIPELINE)
+    assert structured.stage_by_id(0).model_config.supports_running_prefix_cache_reset is False
+
+
+@pytest.fixture
+def single_stage_tts_config(mocker):
+    from vllm.config import VllmConfig
+
+    config = mocker.Mock(spec=VllmConfig)
+    config.additional_config = {"talker_stream_decode": True}
+    config.model_config = mocker.Mock(
+        stage_connector_config={},
+        engine_output_type="audio",
+        final_output=True,
+        supports_running_prefix_cache_reset=False,
+        use_v2_model_runner=True,
+        async_chunk=True,
+    )
+    config.parallel_config = mocker.Mock(
+        tensor_parallel_size=1, pipeline_parallel_size=1, distributed_executor_backend="uni"
+    )
+    config.cache_config = mocker.Mock(enable_prefix_caching=False)
+    config.device_config = mocker.Mock(device="cuda")
+    mocker.patch(
+        "vllm_omni.model_executor.models.qwen3_tts.stream_decode.current_omni_platform.is_cuda", return_value=True
+    )
+    return config
+
+
+def test_single_stage_tts_rejects_unsupported_decoder_configuration(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.cache_config.enable_prefix_caching = True
+    with pytest.raises(ValueError, match="enable_prefix_caching=False"):
+        talker_stream_decode_enabled(config)
+    config.cache_config.enable_prefix_caching = False
+    assert talker_stream_decode_enabled(config)
+
+
+@pytest.mark.parametrize("option_location", ["additional_config", "connector"])
+@pytest.mark.parametrize("output_type", ["latent", "audio"])
+def test_two_stage_talker_rejects_single_stage_pcm_option(single_stage_tts_config, option_location, output_type):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.model_config.engine_output_type = output_type
+    config.model_config.final_output = False
+    if option_location == "connector":
+        config.model_config.stage_connector_config = {"extra": config.additional_config}
+        config.additional_config = {}
+    with pytest.raises(ValueError, match="final audio-output Talker"):
+        talker_stream_decode_enabled(config)
+
+
+def test_two_stage_first_audio_remains_independent(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.first_audio import talker_first_audio_enabled
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.model_config.engine_output_type = "latent"
+    config.additional_config = {}
+    config.model_config.stage_connector_config = {"extra": {"talker_first_audio": True}}
+    assert not talker_stream_decode_enabled(config)
+    assert talker_first_audio_enabled(config)
+    config.model_config.stage_connector_config["extra"]["talker_stream_first_audio"] = True
+    with pytest.raises(ValueError, match="requires the single-stage"):
+        talker_stream_decode_enabled(config)
+
+
+def test_single_stage_pcm_requires_cuda_device_and_own_first_audio_option(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.device_config.device = "cpu"
+    with pytest.raises(ValueError, match="requires CUDA"):
+        talker_stream_decode_enabled(config)
+    config.device_config.device = "cuda"
+    config.additional_config["talker_first_audio"] = True
+    with pytest.raises(ValueError, match="mutually exclusive"):
+        talker_stream_decode_enabled(config)
+    config.additional_config.pop("talker_first_audio")
+    config.additional_config["talker_stream_first_audio"] = True
+    assert talker_stream_decode_enabled(config)
+
+
+def test_single_stage_audio_talker_requires_pcm_option(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    single_stage_tts_config.additional_config = {}
+    with pytest.raises(ValueError, match="requires talker_stream_decode=True"):
+        talker_stream_decode_enabled(single_stage_tts_config)
+
+
+def test_single_stage_pcm_requires_reset_capability_declaration(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    single_stage_tts_config.model_config.supports_running_prefix_cache_reset = True
+    with pytest.raises(ValueError, match="must disable running prefix-cache reset"):
+        talker_stream_decode_enabled(single_stage_tts_config)

tests/config/test_omni_config.py

@@ -1046,6 +1046,7 @@ def test_sub_config_fields_match_structured_scopes():
         "interleave_mm_strings",
         "media_io_kwargs",
         "final_output",
+        "supports_running_prefix_cache_reset",
         "active_stream_window",
         "session_mode",
         "duplex_max_sessions",

tests/core/sched/test_omni_ar_scheduler_stale_drain.py

@@ -106,15 +106,15 @@ def _make_drain_sched(session: Request) -> MagicMock:
     return sched
 
 
-def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: int) -> bool:
+def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: int, spec=None) -> bool:
     """Feed one model-runner frame through update_from_output.
 
     Returns True when the frame's tokens were delivered (the append ran),
     False when the frame was swallowed by the stale drain.
     """
     scheduler_output = MagicMock(spec=SchedulerOutput)
     scheduler_output.num_scheduled_tokens = {session.request_id: num_scheduled}
-    scheduler_output.scheduled_spec_decode_tokens = {}
+    scheduler_output.scheduled_spec_decode_tokens = {session.request_id: spec} if spec else {}
     scheduler_output.num_invalid_spec_tokens = 0
 
     model_runner_output = MagicMock(spec=ModelRunnerOutput)
@@ -134,6 +134,27 @@ def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token:
     return sched._update_request_with_output.called
 
 
+def test_connector_prompt_replacement_drops_old_frame_and_delivers_new_frame() -> None:
+    from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+
+    session = _make_session()
+    session.num_computed_tokens = 6
+    session.num_in_flight_tokens = 1
+    session.num_output_placeholders = 1
+    sched = _make_drain_sched(session)
+    sched.chunk_transfer_adapter.replaced_streaming_prompt_ids = {session.request_id}
+    sched.chunk_transfer_adapter.requests_with_ready_chunks = {session.request_id}
+    sched.chunk_transfer_adapter.requests_num_chunks_sent = {session.external_req_id: 2}
+
+    OmniSchedulerMixin._reset_ready_async_chunk_replacements(sched)
+    assert session.num_stale_output_tokens == 1 and session.drop_stale_output
+    assert session.num_output_placeholders == 0
+    sched._release_replaced_streaming_prompt_cache.assert_called_once_with(session)
+    assert _run_step(sched, session, num_scheduled=1, token=42) is False
+    assert session.num_stale_output_tokens == 0
+    assert _run_step(sched, session, num_scheduled=1, token=43) is True
+
+
 def test_exact_drain_delivers_new_segment_frame() -> None:
     """One in-flight decode frame at replacement: its late output drains the
     counter exactly and the new segment's first frame is delivered."""
@@ -225,3 +246,80 @@ def test_in_flight_prefill_chunk_drains_exactly_without_underflow() -> None:
     assert _run_step(sched, session, num_scheduled=5, token=42) is False  # late prefill chunk dropped
     assert session.num_stale_output_tokens == 0
     assert _run_step(sched, session, num_scheduled=1, token=43) is True  # new segment survives
+
+
+@pytest.mark.parametrize("drop", [False, True])
+def test_preemption_stale_output_follows_upstream_delivery_policy(drop):
+    request = _make_session()
+    request.status = RequestStatus.PREEMPTED
+    request.num_stale_output_tokens = 1
+    request.num_in_flight_tokens = 1
+    request.drop_stale_output = drop
+    sched = _make_drain_sched(request)
+    assert _run_step(sched, request, num_scheduled=1, token=42) is (not drop)
+    assert request.num_stale_output_tokens == 0
+    if not drop:
+        sched._update_request_with_output.assert_called_once_with(request, [42], is_stale=True)
+
+
+def test_preemption_stale_spec_rejection_does_not_roll_back_resumed_counters():
+    request = _make_session()
+    request.num_stale_output_tokens = 2
+    request.num_in_flight_tokens = 2
+    request.num_computed_tokens = 3
+    request.num_output_placeholders = 2
+    request.drop_stale_output = False
+    sched = _make_drain_sched(request)
+    assert _run_step(sched, request, num_scheduled=2, token=42, spec=[41])
+    assert request.num_computed_tokens == 3
+    assert request.num_output_placeholders == 2
+
+
+@pytest.mark.parametrize(
+    "arch,supports_reset,running,reset_running,blocked",
+    [
+        ("Qwen3TTSTalkerForConditionalGeneration", False, True, True, True),
+        ("Qwen3TTSTalkerForConditionalGeneration", False, False, True, False),
+        ("Qwen3TTSTalkerForConditionalGeneration", False, True, False, False),
+        ("Qwen3TTSTalkerForConditionalGeneration", True, True, True, False),
+        ("AnotherStatefulAudioModel", False, True, True, True),
+        ("AnotherStatefulAudioModel", True, True, True, False),
+    ],
+)
+def test_reset_running_streaming_codec_is_rejected_before_preemption(
+    mocker, arch, supports_reset, running, reset_running, blocked
+):
+    from vllm.config import VllmConfig
+    from vllm.v1.core.sched.scheduler import Scheduler
+
+    from vllm_omni.config.model import OmniModelConfig
+
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler.vllm_config = mocker.Mock(
+        spec=VllmConfig,
+        model_config=mocker.Mock(
+            spec=OmniModelConfig,
+            model_arch=arch,
+            supports_running_prefix_cache_reset=supports_reset,
+        ),
+    )
+    request = _make_session()
+    request.num_in_flight_tokens = 1
+    request.num_output_placeholders = 1
+    scheduler.running = [request] if running else []
+
+    def reset_cache(_self, reset_running_requests=False, reset_connector=False):
+        assert reset_running_requests is reset_running and reset_connector is True
+        return True
+
+    reset = mocker.patch.object(Scheduler, "reset_prefix_cache", autospec=True, side_effect=reset_cache)
+
+    assert scheduler.reset_prefix_cache(reset_running, reset_connector=True) is (not blocked)
+    if blocked:
+        reset.assert_not_called()
+        assert scheduler.running == [request]
+        assert request.status == RequestStatus.RUNNING
+        assert request.num_in_flight_tokens == request.num_output_placeholders == 1
+        assert not request.drop_stale_output
+    else:
+        reset.assert_called_once()