[Model] Support NVIDIA-NemotronLabs-VoiceChat-11B offline speech-to-speech (#5842)

Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

tests/model_executor/models/test_nemotron_voicechat_registration.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""L1 registration + config-shim tests for NemotronVoiceChat (CPU, no weights)."""
+
+import subprocess
+from pathlib import Path
+
+import pytest
+
+from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
+from vllm_omni.config.stage_config import StageExecutionType
+from vllm_omni.model_executor.models.nemotron_voicechat.configuration_nemotron_voicechat import (
+    NemotronVoiceChatConfig,
+)
+from vllm_omni.model_executor.models.registry import _OMNI_MODELS
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_ARCHS = (
+    "NemotronVoiceChatThinkerForConditionalGeneration",
+    "NemotronVoiceChatTalkerForConditionalGeneration",
+    "NemotronVoiceChatCode2Wav",
+)
+
+
+def _minimal_nemo_dict() -> dict:
+    return {
+        "data": {"frame_length": 0.08, "source_sample_rate": 16000, "target_sample_rate": 22050},
+        "model": {
+            "inference_speaker_name": "Aria",
+            "stt": {
+                "model": {
+                    "perception": {"encoder": {}, "preprocessor": {}, "modality_adapter": {}},
+                    "pretrained_llm": "nvidia/NVIDIA-Nemotron-Nano-9B-v2",
+                    "use_function_head": True,
+                }
+            },
+            "speech_generation": {
+                "model": {
+                    "tts_config": {"backbone_config": {"hidden_size": 1152, "sliding_window": 7500}},
+                    "codec_config": {"num_quantizers": 31, "codebook_size": 1024},
+                }
+            },
+        },
+    }
+
+
+def test_pipeline_registered_with_three_stages() -> None:
+    pipeline = OMNI_PIPELINES["nemotron_voicechat"]
+    assert OMNI_PIPELINES["nemotron_labs_voicechat"] is pipeline
+    assert len(pipeline.stages) == 3
+    thinker, talker, code2wav = pipeline.stages
+    assert thinker.execution_type == StageExecutionType.LLM_AR
+    assert thinker.final_output and thinker.final_output_type == "text"
+    assert thinker.owns_tokenizer
+    assert talker.execution_type == StageExecutionType.LLM_AR
+    assert talker.model_arch == "NemotronVoiceChatTalkerForConditionalGeneration"
+    assert talker.input_sources == (0,)
+    assert code2wav.execution_type == StageExecutionType.LLM_GENERATION
+    assert code2wav.final_output and code2wav.final_output_type == "audio"
+    assert code2wav.input_sources == (1,)
+    assert pipeline.validate() == []
+
+
+def test_architectures_registered() -> None:
+    for arch in _ARCHS:
+        assert arch in _OMNI_MODELS, f"{arch} missing from _OMNI_MODELS"
+
+
+def test_full_payload_whitelists_cover_code2wav_only() -> None:
+    from vllm_omni.core.sched.omni_scheduling_coordinator import _FULL_PAYLOAD_INPUT_STAGES
+
+    assert ("NemotronVoiceChatCode2Wav", "code2wav") in _FULL_PAYLOAD_INPUT_STAGES
+    # The thinker->talker hop is token-path only: the talker must NOT wait on a
+    # connector payload or stage 1 hangs.
+    assert ("NemotronVoiceChatTalkerForConditionalGeneration", "talker") not in _FULL_PAYLOAD_INPUT_STAGES
+
+
+def test_config_parses_nemo_layout_without_hub_access(monkeypatch) -> None:
+    # Provide the thinker text config explicitly so the test never touches HF.
+    cfg = NemotronVoiceChatConfig(
+        thinker_text_config={"model_type": "llama", "hidden_size": 4480, "num_hidden_layers": 56},
+        **_minimal_nemo_dict(),
+    )
+    assert cfg.get_text_config().hidden_size == 4480
+    assert cfg.talker_config.hidden_size == 1152
+    assert cfg.talker_config.sliding_window == 7500
+    assert cfg.code2wav_config.num_quantizers == 31
+    assert cfg.inference_speaker_name == "Aria"
+
+
+@pytest.mark.parametrize(
+    "mutation",
+    ["drop_stt", "drop_speech_generation", "drop_perception", "drop_codec"],
+)
+def test_config_rejects_non_voicechat_checkpoints(mutation: str) -> None:
+    raw = _minimal_nemo_dict()
+    if mutation == "drop_stt":
+        del raw["model"]["stt"]
+    elif mutation == "drop_speech_generation":
+        del raw["model"]["speech_generation"]
+    elif mutation == "drop_perception":
+        del raw["model"]["stt"]["model"]["perception"]
+    else:
+        del raw["model"]["speech_generation"]["model"]["codec_config"]
+    with pytest.raises(ValueError, match="NemotronVoiceChat|missing"):
+        NemotronVoiceChatConfig(
+            thinker_text_config={"model_type": "llama", "hidden_size": 4480},
+            **raw,
+        )
+
+
+def test_no_nemo_imports_in_nemotron_voicechat_tree() -> None:
+    """Dependency guard: the vendored tree must not import the nemo package.
+
+    Scoped to the nemotron_voicechat model tree (other models may have their
+    own optional nemo integrations).
+    """
+    repo_root = Path(__file__).resolve().parents[3]
+    targets = [
+        str(repo_root / "vllm_omni/model_executor/models/nemotron_voicechat"),
+        str(repo_root / "vllm_omni/model_executor/stage_input_processors/nemotron_voicechat.py"),
+    ]
+    result = subprocess.run(
+        ["grep", "-rnE", r"^\s*(import nemo\b|from nemo(\.|\s))", *targets, "--include=*.py"],
+        capture_output=True,
+        text=True,
+    )
+    assert result.returncode != 0, f"Found nemo imports in the nemotron_voicechat tree:\n{result.stdout}"

tests/model_executor/models/test_nemotron_voicechat_shapes.py

@@ -0,0 +1,227 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Golden shape test for the NemotronVoiceChat frame-locked timeline contract.
+
+Pins the exact prefill/off-by-one arithmetic against values verified with the
+NeMo reference run on sample_general.wav (fp32, greedy): perception produced
+196 acoustic frames for the 249,734-sample 16 kHz fixture (the NeMo dump's
+250,760 was the loader's padded batch width; perception consumed the true
+length); the system prompt tokenized to 56 ids ([bos] + 54 + [eos]); NeMo's
+timeline T = 252 = 56 + 196; and our fp32 pipeline matched the reference text
+channel 196/196 tokens.
+"""
+
+import pytest
+
+from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import (
+    compute_acoustic_frame_count,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Perception config subset matching the checkpoint (mel 10 ms hop, n_fft 512,
+# dw-striding 8x CAUSAL subsampling, conv kernel 3). causal_downsampling is
+# what the real checkpoint sets; omitting it here is exactly what let the
+# non-causal padding formula pass review round 0 (helper vs itself).
+_STT_CFG = {
+    "perception": {
+        "preprocessor": {"sample_rate": 16000, "window_stride": 0.01, "window_size": 0.025, "n_fft": 512},
+        "encoder": {
+            "subsampling": "dw_striding",
+            "subsampling_factor": 8,
+            "subsampling_conv_kernel_size": 3,
+            "causal_downsampling": True,
+        },
+    }
+}
+
+# (num_16k_samples, expected 12.5 Hz frames), all verified against the REAL
+# vendored AudioPerceptionModule built from the checkpoint config. The first
+# row is the sample_general.wav acceptance fixture (196 frames -> the NeMo
+# T=252 timeline). 71271 is the reviewer's 4.45 s repro that the non-causal
+# formula undersized by one; 250760 is the NeMo loader's padded width, which
+# genuinely yields 197 (not the fixture's 196).
+_GOLDEN = [
+    (249734, 196),
+    (250760, 197),
+    (71271, 57),
+    (16000, 14),
+    (16000 * 2, 26),
+    (16000 * 4, 51),
+    (16000 * 8, 101),
+]
+
+
+@pytest.mark.parametrize(("num_samples", "expected_frames"), _GOLDEN)
+def test_acoustic_frame_count_golden(num_samples: int, expected_frames: int) -> None:
+    assert compute_acoustic_frame_count(_STT_CFG, num_samples) == expected_frames
+
+
+def test_frame_count_matches_vendored_subsampling_geometry() -> None:
+    """Pin the helper against the real ConvSubsampling, not against itself.
+
+    Builds the vendored subsampling module with the checkpoint geometry (both
+    causal and non-causal) and checks the helper's padding arithmetic against
+    the module's actual ``_left_padding + _right_padding`` and its
+    ``calc_length`` output over a spread of mel lengths.
+    """
+    import torch
+
+    from vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored.asr.subsampling import (
+        ConvSubsampling,
+        calc_length,
+    )
+
+    for causal in (True, False):
+        sub = ConvSubsampling(
+            subsampling="dw_striding",
+            subsampling_factor=8,
+            feat_in=128,
+            feat_out=8,
+            conv_channels=8,
+            is_causal=causal,
+        )
+        expected_paddings = sub._left_padding + sub._right_padding
+        cfg = {
+            "perception": {
+                "preprocessor": _STT_CFG["perception"]["preprocessor"],
+                "encoder": {
+                    "subsampling": "dw_striding",
+                    "subsampling_factor": 8,
+                    "subsampling_conv_kernel_size": 3,
+                    "causal_downsampling": causal,
+                },
+            }
+        }
+        for num_samples in (16000, 71271, 249734, 250760):
+            # Recompute the helper's mel length, then ask the module's own
+            # calc_length what the subsampled length should be.
+            mel_len = (num_samples + 512 - 512) // 160 + 1
+            module_frames = int(
+                calc_length(
+                    torch.tensor([mel_len]),
+                    all_paddings=expected_paddings,
+                    kernel_size=3,
+                    stride=2,
+                    ceil_mode=False,
+                    repeat_num=3,
+                )[0]
+            )
+            assert compute_acoustic_frame_count(cfg, num_samples) == module_frames, (
+                f"helper diverged from vendored geometry (causal={causal}, samples={num_samples})"
+            )
+
+
+def test_timeline_max_model_len_validation() -> None:
+    from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import (
+        validate_timeline_fits_max_model_len,
+    )
+
+    # The acceptance fixture (56 + 1 + 196 = 253) fits the shipped 8192 limit.
+    validate_timeline_fits_max_model_len(56, 196, 8192)
+    # Exactly at the limit: the vLLM request occupies prompt + 1 (frame-0
+    # placeholder) + frames tokens, so 57 + 8135 == 8192 is allowed...
+    validate_timeline_fits_max_model_len(56, 8135, 8192)
+    # ...and one more frame must be rejected, naming every offending length.
+    with pytest.raises(ValueError) as excinfo:
+        validate_timeline_fits_max_model_len(56, 8136, 8192)
+    message = str(excinfo.value)
+    for fragment in (
+        "vllm_prefill_len=57",
+        "acoustic_frame_count=8136",
+        "vllm_total_len=8193",
+        "max_model_len=8192",
+        "logical_prompt_token_len=56",
+        "logical_total_len=8192",
+    ):
+        assert fragment in message
+
+
+def _bare_thinker():
+    import torch
+    from torch import nn
+
+    from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import (
+        NemotronVoiceChatThinkerForConditionalGeneration,
+    )
+
+    thinker = NemotronVoiceChatThinkerForConditionalGeneration.__new__(NemotronVoiceChatThinkerForConditionalGeneration)
+    nn.Module.__init__(thinker)
+    thinker._use_function_head = True
+    thinker._dtype = torch.float32
+    thinker.function_head = nn.Linear(8, 16, bias=False)
+    return thinker
+
+
+def test_postprocess_appends_function_token_to_timeline(monkeypatch: pytest.MonkeyPatch) -> None:
+    import torch
+
+    thinker = _bare_thinker()
+    hidden = torch.randn(3, 8)
+    expected = int(thinker.function_head(hidden[-1:]).argmax(dim=-1))
+
+    # Default: only the load-bearing prev-token; the cumulative debug timeline
+    # (a per-step host sync + growing concat) stays off.
+    update = thinker.postprocess(hidden)
+    assert int(update["nvc_prev_function_token"]) == expected
+    assert "nvc_function_tokens" not in update
+
+    monkeypatch.setenv("NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE", "1")
+    # No existing timeline: starts one.
+    update = thinker.postprocess(hidden)
+    assert update["nvc_function_tokens"].tolist() == [expected]
+    # Existing timeline: appended, not replaced.
+    existing = torch.tensor([7, 8, 9], dtype=torch.long)
+    update = thinker.postprocess(hidden, nvc_function_tokens=existing)
+    assert update["nvc_function_tokens"].tolist() == [7, 8, 9, expected]
+
+
+def test_postprocess_skips_intermediate_prefill_chunks(monkeypatch: pytest.MonkeyPatch) -> None:
+    import torch
+
+    monkeypatch.setenv("NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE", "1")
+    thinker = _bare_thinker()
+    hidden = torch.randn(4, 8)
+    # Intermediate chunk: 0 computed + 4 rows < prompt_len 10 -> no token.
+    assert thinker.postprocess(hidden, _omni_is_prefill=True, _omni_num_computed_tokens=0, _omni_prompt_len=10) == {}
+    # Final chunk reaches the last prefill position (acoustic frame 0).
+    update = thinker.postprocess(hidden, _omni_is_prefill=True, _omni_num_computed_tokens=6, _omni_prompt_len=10)
+    assert "nvc_prev_function_token" in update
+    assert update["nvc_function_tokens"].numel() == 1
+
+
+def test_prefill_contract_arithmetic() -> None:
+    # vllm_prefill_len = logical_prompt_token_len + 1 (the +1 is acoustic
+    # frame 0); decode steps == acoustic_frame_count; NeMo timeline
+    # T == prompt + frames. Values from the verified sample_general.wav run.
+    logical_prompt_len = 56
+    frames = compute_acoustic_frame_count(_STT_CFG, 249734)
+    vllm_prefill_len = logical_prompt_len + 1
+    assert vllm_prefill_len == 57
+    assert logical_prompt_len + frames == 252  # NeMo reference T
+    # The talker consumes the full timeline from t=1 and the producer trims
+    # the prompt region: rows P-1 onward of the (T-1)-row code stack.
+    talker_rows = (logical_prompt_len + frames) - 1
+    trimmed = talker_rows - (logical_prompt_len - 1)
+    assert trimmed == frames
+
+
+def test_torchaudio_mel_filterbank_matches_librosa() -> None:
+    """The perception frontend builds its mel filterbank with torchaudio.
+
+    torchaudio's melscale_fbanks(norm="slaney", mel_scale="slaney") implements
+    the same slaney-convention math as librosa.filters.mel (the NeMo reference
+    dependency) up to float32 rounding. The 1e-6 tolerance is ~4x the observed
+    2.6e-7 max diff; e2e this stays token- and WAV-identical to the reference
+    on the parity fixture (per-feature normalization absorbs it).
+    """
+    librosa = pytest.importorskip("librosa")
+    torchaudio = pytest.importorskip("torchaudio")
+
+    # The checkpoint's perception preprocessor config (16 kHz, n_fft=512, 128 mels).
+    ta = torchaudio.functional.melscale_fbanks(
+        n_freqs=257, f_min=0.0, f_max=8000.0, n_mels=128, sample_rate=16000, norm="slaney", mel_scale="slaney"
+    ).T.numpy()
+    reference = librosa.filters.mel(sr=16000, n_fft=512, n_mels=128, fmin=0.0, fmax=8000.0, norm="slaney")
+    assert ta.shape == reference.shape
+    assert abs(ta - reference).max() < 1e-6

tests/model_executor/stage_input_processors/test_nemotron_voicechat.py

@@ -0,0 +1,422 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import (
+    _FULL_PAYLOAD_REPLACE_KEYS,
+    talker2code2wav_full_payload,
+    talker2code2wav_token_only,
+    thinker2talker_token_only,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_PAD = 12
+
+
+def _thinker_output(prompt_len: int, generated: list[int]):
+    return SimpleNamespace(
+        finished=True,
+        prompt_token_ids=list(range(100, 100 + prompt_len)) + [_PAD],
+        additional_information=None,
+        outputs=[SimpleNamespace(token_ids=generated, cumulative_token_ids=None)],
+    )
+
+
+def test_thinker2talker_builds_pad_prompt_timeline() -> None:
+    generated = [7, 8, 9, 10]
+    inputs = thinker2talker_token_only([_thinker_output(prompt_len=5, generated=generated)])
+    assert len(inputs) == 1
+    prompt = inputs[0]
+    # 1-token placeholder prompt; the timeline rides additional_information.
+    assert prompt["prompt_token_ids"] == [0]
+    info = prompt["additional_information"]
+    assert info["nvc_logical_prompt_len"] == 5
+    assert info["nvc_frame_count"] == 4
+    assert info["nvc_text_timeline"] == [_PAD] * 5 + generated
+
+
+def test_thinker2talker_skips_unfinished() -> None:
+    out = _thinker_output(3, [1])
+    out.finished = False
+    assert thinker2talker_token_only([out]) == []
+
+
+def test_full_payload_trims_prompt_region_rows() -> None:
+    # Rows correspond to timeline steps t=1..T-1; with prompt_len P the NeMo
+    # trim keeps rows for t >= P, i.e. row index P-1 onward.
+    prompt_len = 4
+    frames, q = 10, 31
+    codes = torch.arange(frames * q, dtype=torch.long).reshape(frames, q) % 1024
+    request = SimpleNamespace(
+        additional_information={
+            "codes": {"audio": codes},
+            "nvc_logical_prompt_len": prompt_len,
+        },
+        additional_information_cpu=None,
+    )
+    payload = talker2code2wav_full_payload(request=request, is_finished=True)
+    assert payload.codes.audio.shape == (frames - (prompt_len - 1), q)
+    assert torch.equal(payload.codes.audio, codes[prompt_len - 1 :])
+    assert bool(payload.meta.finished)
+
+
+def test_full_payload_without_codes_is_empty_finished() -> None:
+    # The sync producer must never emit a meta-only "pending" struct (it would
+    # be shipped as the real payload); with no codes it ships an empty finished
+    # payload regardless of the request's is_finished flag.
+    payload = talker2code2wav_full_payload(request=SimpleNamespace(), is_finished=False)
+    assert payload is not None
+    assert bool(payload.meta.finished)
+    assert payload.codes.audio.numel() == 0
+
+
+def test_full_payload_replace_semantics_declared() -> None:
+    # Full-sequence tensors must REPLACE in the accumulator, not concatenate.
+    # The AR runner flattens nested payloads before accumulation, so the
+    # flattened "codes.audio" key is the one that matters.
+    assert "codes.audio" in _FULL_PAYLOAD_REPLACE_KEYS
+    assert "codes" in _FULL_PAYLOAD_REPLACE_KEYS
+
+
+def test_full_payload_missing_prompt_len_fails_explicitly() -> None:
+    codes = torch.zeros((5, 31), dtype=torch.long)
+    request = SimpleNamespace(
+        additional_information={"codes": {"audio": codes}},
+        additional_information_cpu=None,
+    )
+    with pytest.raises(ValueError, match="nvc_logical_prompt_len"):
+        talker2code2wav_full_payload(request=request, is_finished=True)
+
+
+def test_code2wav_distinguishes_empty_payload_from_missing_info() -> None:
+    from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_code2wav import (
+        NemotronVoiceChatCode2Wav,
+    )
+
+    extract = NemotronVoiceChatCode2Wav._codes_from_runtime_info
+    # Explicit empty payload (e.g. zero frames after the prompt trim): the
+    # empty tensor comes back as-is -> empty audio, never a placeholder-id
+    # decode.
+    empty = extract(None, {"codes": {"audio": torch.empty(0, dtype=torch.long)}})
+    assert isinstance(empty, torch.Tensor) and empty.numel() == 0
+    # Missing codes info entirely -> None (forward raises: placeholder
+    # input_ids are never decoded).
+    assert extract(None, {"meta": {}}) is None
+    assert extract(None, None) is None
+    # Normal payload passes through (nested and flattened key forms).
+    codes = torch.zeros((3, 31), dtype=torch.long)
+    assert extract(None, {"codes": {"audio": codes}}) is codes
+    assert extract(None, {"codes.audio": codes}) is codes
+
+
+def test_validate_code_stack_negatives() -> None:
+    from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_code2wav import (
+        validate_code_stack,
+    )
+
+    good = torch.zeros((4, 31), dtype=torch.long)
+    assert validate_code_stack(good, 31, 1024).shape == (4, 31)
+    with pytest.raises(ValueError, match="not divisible"):
+        validate_code_stack(torch.zeros(100, dtype=torch.long), 31, 1024)
+    with pytest.raises(ValueError, match="shaped"):
+        validate_code_stack(torch.zeros((4, 30), dtype=torch.long), 31, 1024)
+    bad = good.clone()
+    bad[1, 5] = 1024
+    with pytest.raises(ValueError, match="out-of-range"):
+        validate_code_stack(bad, 31, 1024)
+
+
+def test_tokenizer_resolution_failure_is_actionable(monkeypatch) -> None:
+    from transformers import AutoConfig
+
+    from vllm_omni.model_executor.models.nemotron_voicechat.configuration_nemotron_voicechat import (
+        NemotronVoiceChatConfig,
+    )
+
+    def _boom(*args, **kwargs):
+        raise OSError("offline")
+
+    monkeypatch.setattr(AutoConfig, "from_pretrained", _boom)
+    with pytest.raises(RuntimeError, match="NEMOTRON_VOICECHAT_LLM_PATH"):
+        NemotronVoiceChatConfig(
+            data={"frame_length": 0.08, "source_sample_rate": 16000, "target_sample_rate": 22050},
+            model={
+                "inference_speaker_name": "Aria",
+                "stt": {"model": {"perception": {"encoder": {}}, "pretrained_llm": "nvidia/nope"}},
+                "speech_generation": {"model": {"tts_config": {}, "codec_config": {}}},
+            },
+        )
+
+
+def test_sanitize_tts_model_cfg_disables_pretrained_loads(monkeypatch) -> None:
+    from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_talker import (
+        sanitize_tts_model_cfg,
+    )
+
+    original = {
+        "pretrained_lm_name": "nvidia/NVIDIA-Nemotron-Nano-9B-v2",
+        "pretrained_codec_model": "/some/external/codec.ckpt",
+        "tts_config": {
+            "pretrained_text_name": "some/hub-llm",
+            "backbone_type": "gemma3_text",
+            "cas_config": {"pretrained_tokenizer_name": "some/hub-tokenizer", "keep": 1},
+        },
+    }
+    monkeypatch.delenv("NEMOTRON_VOICECHAT_LLM_PATH", raising=False)
+    cfg = sanitize_tts_model_cfg(original)
+    # External weight-load triggers are gone; construction is config-only.
+    assert "pretrained_text_name" not in cfg["tts_config"]
+    assert "pretrained_codec_model" not in cfg
+    assert "pretrained_tokenizer_name" not in cfg["tts_config"]["cas_config"]
+    assert cfg["tts_config"]["cas_config"]["keep"] == 1
+    # Tokenizer reference survives.
+    assert cfg["pretrained_lm_name"] == "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
+    # The shared stage config is never mutated (deep copy).
+    assert original["tts_config"]["pretrained_text_name"] == "some/hub-llm"
+
+    # Air-gapped override reroutes the tokenizer reference.
+    monkeypatch.setenv("NEMOTRON_VOICECHAT_LLM_PATH", "/local/nemotron")
+    cfg = sanitize_tts_model_cfg(original)
+    assert cfg["pretrained_lm_name"] == "/local/nemotron"
+
+
+def test_code2wav_token_only_sizes_placeholder() -> None:
+    talker_output = SimpleNamespace(
+        finished=True,
+        outputs=[SimpleNamespace(token_ids=[0] * 9, cumulative_token_ids=None)],
+    )
+    inputs = talker2code2wav_token_only([talker_output])
+    assert len(inputs) == 1
+    assert len(inputs[0]["prompt_token_ids"]) == 9
+
+
+# =========================
+# Async-chunk (streaming) producers
+# =========================
+
+
+def _fake_transfer_manager(codec_chunk_frames: int | None = None):
+    extra = {}
+    if codec_chunk_frames is not None:
+        extra["codec_chunk_frames"] = codec_chunk_frames
+    return SimpleNamespace(
+        request_payload={},
+        connector=SimpleNamespace(config={"extra": extra}),
+        put_req_chunk={},
+    )
+
+
+def _streaming_request(prompt_len: int, generated: list[int], info: dict | None = None):
+    return SimpleNamespace(
+        external_req_id="req-0",
+        prompt_token_ids=list(range(100, 100 + prompt_len)) + [_PAD],
+        output_token_ids=generated,
+        additional_information=info,
+    )
+
+
+def test_thinker2talker_async_chunk_cumulative_timeline() -> None:
+    from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import (
+        thinker2talker_async_chunk,
+    )
+
+    tm = _fake_transfer_manager()
+    req = _streaming_request(prompt_len=3, generated=[7])
+    payload = thinker2talker_async_chunk(tm, {}, req, is_finished=False)
+    assert payload is not None
+    assert payload.ids.all == [_PAD] * 3 + [7]
+    assert payload.meta.num_processed_tokens == 3
+    assert not bool(payload.meta.finished)
+
+    # No new tokens since the last chunk -> skip.
+    assert thinker2talker_async_chunk(tm, {}, req, is_finished=False) is None
+
+    # More tokens -> longer cumulative timeline; final chunk sets finished.
+    req.output_token_ids = [7, 8, 9]
+    payload = thinker2talker_async_chunk(tm, {}, req, is_finished=True)
+    assert payload.ids.all == [_PAD] * 3 + [7, 8, 9]
+    assert bool(payload.meta.finished)
+
+
+def test_talker2code2wav_async_chunk_cadence_and_left_context() -> None:
+    from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import (
+        talker2code2wav_async_chunk,
+    )
+
+    tm = _fake_transfer_manager(codec_chunk_frames=2)
+    prompt_len = 3
+    info = {"nvc_logical_prompt_len": prompt_len}
+    req = _streaming_request(prompt_len=prompt_len, generated=[], info=info)
+
+    def cumulative(rows: int) -> dict:
+        # Talker rows are timeline steps t=1..; prompt region occupies the
+        # first prompt_len-1 rows before acoustic frames start.
+        return {"codes": {"audio": torch.arange((prompt_len - 1 + rows) * 31).reshape(-1, 31)}}
+
+    # Only prompt-region rows so far -> nothing to ship.
+    assert talker2code2wav_async_chunk(tm, cumulative(0), req, is_finished=False) is None
+    # One new frame < chunk size 2 -> hold.
+    assert talker2code2wav_async_chunk(tm, cumulative(1), req, is_finished=False) is None
+    # Two frames -> first chunk, no left context.
+    payload = talker2code2wav_async_chunk(tm, cumulative(2), req, is_finished=False)
+    assert payload is not None
+    assert payload.meta.left_context_size == 0
+    assert payload.codes.audio.shape == (2, 31)
+    # Third frame + finished -> cumulative history with left context 2.
+    payload = talker2code2wav_async_chunk(tm, cumulative(3), req, is_finished=True)
+    assert payload.meta.left_context_size == 2
+    assert payload.codes.audio.shape == (3, 31)
+    assert bool(payload.meta.finished)
+
+
+def test_talker2code2wav_async_chunk_requires_prompt_len() -> None:
+    from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import (
+        talker2code2wav_async_chunk,
+    )
+
+    tm = _fake_transfer_manager()
+    req = _streaming_request(prompt_len=3, generated=[], info=None)
+    with pytest.raises(ValueError, match="logical prompt length"):
+        talker2code2wav_async_chunk(tm, {"codes": {"audio": torch.zeros(5, 31, dtype=torch.long)}}, req)
+
+
+# =========================
+# Talker drain-per-wake (async scheduling review P1/P2)
+# =========================
+
+
+class _FakeTTSModelOutput:
+    past_key_values = None
+
+
+class _FakeTTS:
+    """Stub DuplexEARTTS: infer_codes_one_step returns a constant code row."""
+
+    def __init__(self, q: int = 31):
+        self.q = q
+        self.steps = 0
+
+    def set_init_inputs(self, speaker_name: str) -> None:
+        del speaker_name
+
+    def get_init_inputs(self, B: int) -> dict:  # noqa: N803 — mirrors the vendored DuplexEARTTS signature
+        return {
+            "code": torch.zeros((B, 2, self.q), dtype=torch.long),
+            "subword_ids": torch.zeros((B, 3), dtype=torch.long),
+        }
+
+    def _get_generation_config(self, guidance_enabled: bool):
+        del guidance_enabled
+        return {}
+
+    def tts_model(self, **kwargs):
+        del kwargs
+        return _FakeTTSModelOutput()
+
+    def infer_codes_one_step(self, **kwargs):
+        del kwargs
+        self.steps += 1
+        return torch.full((1, 1, self.q), self.steps, dtype=torch.long), None
+
+
+def _bare_talker():
+    from torch import nn
+
+    from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_talker import (
+        NemotronVoiceChatTalkerForConditionalGeneration,
+    )
+
+    talker = NemotronVoiceChatTalkerForConditionalGeneration.__new__(NemotronVoiceChatTalkerForConditionalGeneration)
+    nn.Module.__init__(talker)
+    talker.tts = _FakeTTS()
+    talker._sessions = {}
+    talker._hidden = 8
+    talker._dtype = torch.float32
+    talker._speaker_name = "Aria"
+    talker.config = SimpleNamespace(tts_cfg={"inference_guidance_enabled": False})
+    return talker
+
+
+def _async_info(timeline: list[int], prompt_len: int, finished: bool, request_id: str = "req-0") -> dict:
+    return {
+        "request_id": request_id,
+        "additional_information": {
+            "ids": {"all": list(timeline)},
+            "meta": {
+                "num_processed_tokens": prompt_len,
+                "finished": torch.tensor(finished),
+                "request_id": request_id,
+            },
+        },
+    }
+
+
+def test_talker_drains_all_received_positions_per_wake() -> None:
+    """P2: chunk 0's PAD prompt region must not cost one upstream chunk per step,
+    and P1: a coalesced chunk (2+ new positions, 1 wake) must be fully drained."""
+    talker = _bare_talker()
+    pad, prompt_len = 12, 4
+    ids = torch.zeros(1, dtype=torch.long)
+
+    # Prefill wake: chunk 0 carries only the PAD prompt region.
+    info = _async_info([pad] * prompt_len, prompt_len, finished=False)
+    talker.preprocess(ids, None, _omni_is_prefill=True, **info)
+    session = talker._sessions["req-0"]
+    assert session["step"] == 1 and not session["sync_mode"]
+
+    # First decode wake: drains ALL prompt-region steps (t=1..P-1) at once.
+    _, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info)
+    assert session["step"] == prompt_len  # t advanced 1 -> P in one wake
+    assert update["codes"]["audio"].shape == (prompt_len - 1, 31)
+    assert float(embeds[0, 0]) == 0.0  # not finished
+
+    # Coalesced chunk: TWO new frame tokens arrive in ONE wake (delayed save
+    # thread merged two thinker steps). Both must be consumed by this wake.
+    info2 = _async_info([pad] * prompt_len + [7, 8], prompt_len, finished=False)
+    _, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info2)
+    assert session["step"] == prompt_len + 2
+    assert update["codes"]["audio"].shape == (prompt_len + 1, 31)
+    assert float(embeds[0, 0]) == 0.0
+
+    # Zero-progress wake (no new positions, not finished): plain CONTINUE, no
+    # fabricated frames, no codes update.
+    _, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info2)
+    assert session["step"] == prompt_len + 2
+    assert update == {}
+    assert float(embeds[0, 0]) == 0.0
+
+    # Final chunk: one more token + finished marker -> drain + stop flag.
+    info3 = _async_info([pad] * prompt_len + [7, 8, 9], prompt_len, finished=True)
+    _, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info3)
+    assert session["step"] == prompt_len + 3
+    assert update["codes"]["audio"].shape == (prompt_len + 2, 31)
+    assert float(embeds[0, 0]) == 1.0  # stop
+
+
+def test_talker_sync_mode_steps_once_per_wake() -> None:
+    """The verified sync parity path keeps its one-TTS-step-per-engine-step cadence."""
+    talker = _bare_talker()
+    pad, prompt_len = 12, 3
+    ids = torch.zeros(1, dtype=torch.long)
+    info = {
+        "request_id": "req-1",
+        "additional_information": {
+            "nvc_text_timeline": [pad] * prompt_len + [7, 8],
+            "nvc_logical_prompt_len": prompt_len,
+        },
+    }
+    talker.preprocess(ids, None, _omni_is_prefill=True, **info)
+    session = talker._sessions["req-1"]
+    assert session["sync_mode"] and session["upstream_finished"]
+    for expected_step in (2, 3, 4):
+        _, embeds, _ = talker.preprocess(ids, None, _omni_is_prefill=False, **info)
+        assert session["step"] == expected_step
+    # Last position -> stop flag on the final step.
+    _, embeds, _ = talker.preprocess(ids, None, _omni_is_prefill=False, **info)
+    assert session["step"] == 5
+    assert float(embeds[0, 0]) == 1.0