[Perf][NemotronVoiceChat] Fast defaults for all three deploy profiles: 2.8x offline, realtime full-duplex audio (#6354)

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

tests/model_executor/models/test_nemotron_voicechat_perception_window.py

@@ -0,0 +1,135 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Rolling-window mel slicing must match full-history slicing bit-for-bit.
+
+The duplex thinker keeps a bounded trailing audio window instead of the full
+session history (the mel featurizer is local: hop-aligned STFT columns, no
+cross-stream normalization).  These tests replay a stream both ways through
+the vendored preprocessor and require identical chunk tensors.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import (
+    _DUPLEX_MEL_MARGIN_COLS,
+    perception_chunk_geometry,
+    slice_perception_streaming_mel,
+)
+
+FRAME_SAMPLES = 1280  # 80 ms @ 16 kHz
+HOP = 160  # 10 ms window stride
+
+pytestmark = pytest.mark.core_model
+
+
+def _make_preprocessor():
+    from vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored.asr.audio_preprocessing import (
+        AudioToMelSpectrogramPreprocessor,
+    )
+
+    # Mirrors model.stt.model.perception.preprocessor in the checkpoint with
+    # the thinker's streaming overrides (dither=0, pad_to=0).
+    return (
+        AudioToMelSpectrogramPreprocessor(
+            sample_rate=16000,
+            window_size=0.025,
+            window_stride=0.01,
+            n_fft=512,
+            features=128,
+            dither=0.0,
+            pad_to=0,
+            normalize="NA",
+            log=True,
+            frame_splicing=1,
+        )
+        .to(torch.float32)
+        .eval()
+    )
+
+
+def _mel(preprocessor, audio: torch.Tensor) -> torch.Tensor:
+    with torch.inference_mode():
+        processed, _ = preprocessor(
+            input_signal=audio.unsqueeze(0),
+            length=torch.tensor([audio.numel()], dtype=torch.long),
+        )
+    return processed
+
+
+@pytest.mark.parametrize(
+    "streaming_cfg",
+    [
+        SimpleNamespace(
+            chunk_size=[9, 8],
+            shift_size=[9, 8],
+            pre_encode_cache_size=[0, 9],
+            drop_extra_pre_encoded=3,
+        ),
+        SimpleNamespace(
+            chunk_size=[17, 16],
+            shift_size=[9, 8],
+            pre_encode_cache_size=[0, 17],
+            drop_extra_pre_encoded=5,
+        ),
+    ],
+)
+def test_windowed_mel_chunks_match_full_history(streaming_cfg):
+    torch.manual_seed(0)
+    preprocessor = _make_preprocessor()
+    num_frames = 80  # 6.4 s: far past the point where trimming kicks in
+    stream = torch.randn(num_frames * FRAME_SAMPLES, dtype=torch.float32) * 0.1
+
+    window = None
+    window_col0 = 0
+    trimmed = False
+    for frame_idx in range(num_frames):
+        frame = stream[frame_idx * FRAME_SAMPLES : (frame_idx + 1) * FRAME_SAMPLES]
+        full_audio = stream[: (frame_idx + 1) * FRAME_SAMPLES]
+        window = frame if window is None else torch.cat([window, frame])
+
+        if frame_idx > 0:
+            total_samples = window_col0 * HOP + int(window.numel())
+            stream_total_cols = total_samples // HOP + 1
+            assert stream_total_cols == full_audio.numel() // HOP + 1
+            cache_start, _, _ = perception_chunk_geometry(frame_idx, stream_total_cols, streaming_cfg)
+            desired_col0 = max(0, cache_start - _DUPLEX_MEL_MARGIN_COLS)
+            if desired_col0 > window_col0:
+                window = window[(desired_col0 - window_col0) * HOP :]
+                window_col0 = desired_col0
+                trimmed = True
+
+        expected, expected_drop = slice_perception_streaming_mel(
+            _mel(preprocessor, full_audio), frame_idx, streaming_cfg
+        )
+        actual, actual_drop = slice_perception_streaming_mel(
+            _mel(preprocessor, window), frame_idx, streaming_cfg, window_col0=window_col0
+        )
+        assert actual_drop == expected_drop
+        assert actual.shape == expected.shape, f"frame {frame_idx}"
+        assert torch.equal(actual, expected), f"frame {frame_idx} mel chunk diverged (window_col0={window_col0})"
+
+    assert trimmed, "test never exercised the trimmed-window path"
+    # The rolling window must stay bounded instead of growing with the stream.
+    assert window.numel() < 4 * 16000, f"window grew to {window.numel()} samples"
+
+
+def test_stream_start_requires_untrimmed_window():
+    streaming_cfg = SimpleNamespace(
+        chunk_size=[9, 8], shift_size=[9, 8], pre_encode_cache_size=[0, 9], drop_extra_pre_encoded=3
+    )
+    mel = torch.zeros(1, 128, 9)
+    with pytest.raises(ValueError, match="untrimmed"):
+        slice_perception_streaming_mel(mel, 0, streaming_cfg, window_col0=4)
+
+
+def test_window_missing_needed_columns_rejected():
+    streaming_cfg = SimpleNamespace(
+        chunk_size=[9, 8], shift_size=[9, 8], pre_encode_cache_size=[0, 9], drop_extra_pre_encoded=3
+    )
+    mel = torch.zeros(1, 128, 24)
+    with pytest.raises(ValueError, match="needs column"):
+        slice_perception_streaming_mel(mel, 40, streaming_cfg, window_col0=340)

tests/model_executor/models/test_nemotron_voicechat_shapes.py

@@ -249,7 +249,7 @@ def test_code2wav_reuses_and_clears_per_request_streaming_cache() -> None:
     model = NemotronVoiceChatCode2Wav.__new__(NemotronVoiceChatCode2Wav)
     nn.Module.__init__(model)
     model._sample_rate, model._num_quantizers, model._codebook_size, model._wav_per_frame = 22050, 31, 1024, 1764
-    model._duplex_codec_caches = {}
+    model._stream_caches = {}
     model.audio_codec = Codec()
 
     def info(codes):
@@ -258,6 +258,6 @@ def test_code2wav_reuses_and_clears_per_request_streaming_cache() -> None:
     model.forward(runtime_additional_information=info(torch.zeros((1, 31), dtype=torch.long)))
     model.forward(runtime_additional_information=info(torch.zeros((2, 31), dtype=torch.long)))
     assert model.audio_codec.caches[0] is model.audio_codec.caches[1]
-    assert "req" in model._duplex_codec_caches
+    assert "req" in model._stream_caches
     model.on_requests_finished({"req"})
-    assert "req" not in model._duplex_codec_caches
+    assert "req" not in model._stream_caches

tests/model_executor/models/test_nemotron_voicechat_talker_replay.py

@@ -0,0 +1,134 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Pin the native talker's KV-recompute replay arithmetic.
+
+Replay turns what used to be a loud desync failure into a silent
+reconstruction, so these tests pin the exact index math: ``codes_rows[i]``
+holds the codes sampled at step ``i + 1``, replaying step ``t`` feeds the
+codes from step ``t - 1`` (``initial_code`` for ``t == 1``), the live step
+(if any) is the last row of the span, and a prefill-boundary-crossing span
+prepends the prefill embeds. ``build_decode_embeds`` is faked to encode
+``(timeline step, prev-code sum, uncond flag)`` into each returned row so the
+composition is directly assertable.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.nemotron_voicechat import talker_native
+from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_talker import (
+    NemotronVoiceChatTalkerForConditionalGeneration,
+)
+
+pytestmark = pytest.mark.core_model
+
+INIT_LEN = 4
+HID = 4
+Q = 2
+
+
+def _fake_build_decode_embeds(model, *, prev_codes, subword_embed, uncond):
+    t = float(subword_embed.reshape(-1)[0])
+    row = torch.tensor([[[t, float(prev_codes.float().sum()), 1.0 if uncond else 0.0, 0.0]]])
+    return row
+
+
+def _codes_for_step(step: int) -> torch.Tensor:
+    """codes_rows entry for ``step`` (appended after the step ran)."""
+    return torch.full((1, Q), 100.0 + step)
+
+
+def _make_talker(step: int, rows_through_step: int):
+    talker = object.__new__(NemotronVoiceChatTalkerForConditionalGeneration)
+    talker._dtype = torch.float32
+    talker._use_native_backbone = True
+    talker.tts = SimpleNamespace(tts_model=object())
+    session = {
+        "use_native": True,
+        "step": step,
+        "init_len": INIT_LEN,
+        "prefill_embeds": torch.arange(INIT_LEN * HID, dtype=torch.float32).reshape(INIT_LEN, HID),
+        # cas_table[t] == [t] so the fake embed builder can recover the step.
+        "cas_table": torch.arange(64, dtype=torch.float32).reshape(-1, 1),
+        "timeline": torch.zeros(64, dtype=torch.long),
+        "codes_rows": [_codes_for_step(s) for s in range(1, rows_through_step + 1)],
+        "code": torch.full((1, 1, Q), 999.0),
+        "initial_code": torch.full((1, 1, Q), 7.0),
+        "silence_codes": None,
+        "text_eos_id": 1,
+        "uncond_stream": None,
+        "uncond_hidden": None,
+        "codec_streaming": True,
+        "upstream_finished": False,
+    }
+    talker._sessions = {"req": session}
+    return talker, session
+
+
+def _preprocess(talker, *, offset: int, span: int):
+    info = {"_omni_num_computed_tokens": offset, "request_id": "req"}
+    input_ids = torch.zeros(span, dtype=torch.long)
+    _, embeds, _ = talker._native_preprocess(input_ids, info, torch.device("cpu"), span, "req")
+    return embeds
+
+
+@pytest.fixture(autouse=True)
+def fake_embed_builder(monkeypatch):
+    monkeypatch.setattr(talker_native, "build_decode_embeds", _fake_build_decode_embeds)
+
+
+def test_pure_replay_reads_previous_steps_codes():
+    talker, session = _make_talker(step=5, rows_through_step=4)
+    # Engine rewound to timeline steps 1..3 (offset init_len => t_first = 1).
+    embeds = _preprocess(talker, offset=INIT_LEN, span=3)
+    assert embeds.shape == (3, HID)
+    assert embeds[:, 0].tolist() == [1.0, 2.0, 3.0]  # timeline steps
+    # step 1 feeds initial_code; steps 2..3 feed codes sampled at steps 1..2.
+    assert embeds[:, 1].tolist() == [7.0 * Q, (100.0 + 1) * Q, (100.0 + 2) * Q]
+    assert embeds[:, 2].tolist() == [0.0, 0.0, 0.0]  # uncond never used
+    assert session.get("pending_step") is None  # nothing to resample
+    assert session["replaying"] is True
+
+
+def test_replay_plus_live_step_appends_live_row_last():
+    talker, session = _make_talker(step=5, rows_through_step=4)
+    # Steps 3..4 replayed, step 5 is live (t_last == step).
+    embeds = _preprocess(talker, offset=INIT_LEN + 2, span=3)
+    assert embeds.shape == (3, HID)
+    assert embeds[:, 0].tolist() == [3.0, 4.0, 5.0]
+    # Replays feed stored history; the live step feeds the latest code.
+    assert embeds[:, 1].tolist() == [(100.0 + 2) * Q, (100.0 + 3) * Q, 999.0 * Q]
+    assert session["pending_step"] == 5
+    assert "replaying" not in session
+
+
+def test_replay_span_crossing_prefill_boundary_prepends_prefill():
+    talker, session = _make_talker(step=4, rows_through_step=3)
+    # Positions 2..3 are speaker-prompt prefill, then timeline steps 1..2.
+    embeds = _preprocess(talker, offset=2, span=4)
+    assert embeds.shape == (4, HID)
+    assert torch.equal(embeds[:2], session["prefill_embeds"][2:4])
+    assert embeds[2:, 0].tolist() == [1.0, 2.0]
+    assert embeds[2:, 1].tolist() == [7.0 * Q, (100.0 + 1) * Q]
+    assert session.get("pending_step") is None
+
+
+def test_live_first_step_crossing_prefill_boundary_still_rejected():
+    talker, _ = _make_talker(step=1, rows_through_step=0)
+    with pytest.raises(ValueError, match="crossing"):
+        _preprocess(talker, offset=2, span=3)
+
+
+def test_short_code_history_is_rejected():
+    talker, _ = _make_talker(step=5, rows_through_step=1)
+    with pytest.raises(RuntimeError, match="cannot replay"):
+        _preprocess(talker, offset=INIT_LEN, span=3)
+
+
+def test_position_past_session_step_is_rejected():
+    talker, _ = _make_talker(step=3, rows_through_step=2)
+    with pytest.raises(RuntimeError, match="outpaced"):
+        _preprocess(talker, offset=INIT_LEN + 3, span=1)

tests/model_executor/stage_input_processors/test_nemotron_voicechat.py

@@ -131,6 +131,88 @@ def test_validate_code_stack_negatives() -> None:
         validate_code_stack(bad, 31, 1024)
 
 
+def test_scalar_bool_normalizes_tensor_metadata() -> None:
+    from vllm_omni.model_executor.models.nemotron_voicechat.runtime_info import scalar_bool
+
+    assert scalar_bool(True)
+    assert scalar_bool(torch.tensor(True))
+    assert scalar_bool(torch.tensor([True]))
+    assert not scalar_bool(torch.empty(0, dtype=torch.bool))
+    with pytest.raises(ValueError, match="scalar boolean metadata"):
+        scalar_bool(torch.tensor([True, False]))
+
+
+def _bare_code2wav(use_codec_cache: bool):
+    from torch import nn
+
+    from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_code2wav import (
+        NemotronVoiceChatCode2Wav,
+    )
+
+    class _FakeCodec(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.dummy = nn.Parameter(torch.zeros(1))
+            self.calls: list[tuple[int, bool]] = []
+
+        def decode(self, codes, lens, cache=None, flush=False):
+            frames = int(codes.shape[1])
+            self.calls.append((frames, cache is not None))
+            wav = torch.zeros(1, 1, frames * 4, dtype=torch.float32)
+            return wav, lens * 4
+
+    model = NemotronVoiceChatCode2Wav.__new__(NemotronVoiceChatCode2Wav)
+    nn.Module.__init__(model)
+    model._num_quantizers = 31
+    model._codebook_size = 1024
+    model._sample_rate = 22050
+    model._wav_per_frame = 4
+    model._use_codec_cache = use_codec_cache
+    model.requires_request_ids = use_codec_cache
+    model._stream_caches = {}
+    model.audio_codec = _FakeCodec()
+    return model
+
+
+def test_code2wav_incremental_cache_decodes_only_new_frames() -> None:
+    model = _bare_code2wav(use_codec_cache=True)
+    q = 31
+
+    def chunk_info(frames: int, left: int) -> dict:
+        return {
+            "codes": {"audio": torch.zeros((frames, q), dtype=torch.long)},
+            "meta": {"left_context_size": left},
+        }
+
+    # Chunk 0: 13 cumulative frames, none emitted yet -> decode 13 with a cache.
+    out = model.forward(runtime_additional_information=[chunk_info(13, 0)], request_ids=["r0"])
+    assert model.audio_codec.calls == [(13, True)]
+    assert out.multimodal_outputs["model_outputs"][0].numel() == 13 * 4
+    assert "r0" in model._stream_caches
+    # Chunk 1: 20 cumulative, 13 emitted -> decode only the 7 NEW frames, same cache.
+    cache = model._stream_caches["r0"]
+    out = model.forward(runtime_additional_information=[chunk_info(20, 13)], request_ids=["r0"])
+    assert model.audio_codec.calls[-1] == (7, True)
+    assert model._stream_caches["r0"] is cache
+    assert out.multimodal_outputs["model_outputs"][0].numel() == 7 * 4
+    # Request end drops the per-request cache.
+    model.on_requests_finished(["r0"])
+    assert "r0" not in model._stream_caches
+
+
+def test_code2wav_prefix_redecode_unchanged_without_cache_flag() -> None:
+    model = _bare_code2wav(use_codec_cache=False)
+    q = 31
+    info = {
+        "codes": {"audio": torch.zeros((20, q), dtype=torch.long)},
+        "meta": {"left_context_size": 13},
+    }
+    out = model.forward(runtime_additional_information=[info], request_ids=["r0"])
+    # Full cumulative decode (no cache), then slice off the emitted prefix.
+    assert model.audio_codec.calls == [(20, False)]
+    assert out.multimodal_outputs["model_outputs"][0].numel() == (20 - 13) * 4
+
+
 def test_tokenizer_resolution_failure_is_actionable(monkeypatch) -> None:
     from transformers import AutoConfig
 
@@ -375,6 +457,8 @@ def _bare_talker():
     talker._hidden = 8
     talker._dtype = torch.float32
     talker._speaker_name = "Aria"
+    talker._step_graph = None
+    talker._use_native_backbone = False
     talker.config = SimpleNamespace(tts_cfg={"inference_guidance_enabled": False})
     return talker
 
@@ -401,6 +485,112 @@ def _async_info(
     }
 
 
+def test_talker_output_keeps_streaming_codec_mode_on_wire() -> None:
+    talker = _bare_talker()
+    output = talker.make_omni_output(
+        torch.zeros((1, talker._hidden)),
+        model_intermediate_buffer=[
+            {
+                "codes": {"audio": torch.zeros((1, 31), dtype=torch.long)},
+                "meta": {
+                    "nvc_logical_prompt_len": 3,
+                    "codec_streaming": torch.tensor(True),
+                    "request_id": "req-duplex",
+                },
+            }
+        ],
+    )
+
+    meta = output.multimodal_outputs["meta"]
+    assert meta["nvc_logical_prompt_len"].item() == 3
+    assert meta["codec_streaming"].item() is True
+
+
+def test_code2wav_accepts_tensor_streaming_flag() -> None:
+    from torch import nn
+
+    from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_code2wav import (
+        NemotronVoiceChatCode2Wav,
+    )
+
+    class _FakeCodec(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.anchor = nn.Parameter(torch.zeros(1))
+            self.cache = None
+
+        def decode(self, codes, lens, cache=None):
+            del codes, lens
+            self.cache = cache
+            return torch.zeros((1, 4)), torch.tensor([4])
+
+    model = NemotronVoiceChatCode2Wav.__new__(NemotronVoiceChatCode2Wav)
+    nn.Module.__init__(model)
+    model._sample_rate = 22050
+    model._num_quantizers = 31
+    model._codebook_size = 1024
+    model._wav_per_frame = 1764
+    model._use_codec_cache = False
+    model._stream_caches = {}
+    model.audio_codec = _FakeCodec()
+
+    output = model.forward(
+        runtime_additional_information=[
+            {
+                "codes": {"audio": torch.zeros((1, 31), dtype=torch.long)},
+                "meta": {"codec_streaming": torch.tensor([True]), "request_id": "req-codec"},
+            }
+        ]
+    )
+
+    assert "req-codec" in model._stream_caches
+    assert model.audio_codec.cache is model._stream_caches["req-codec"]
+    assert output.multimodal_outputs["model_outputs"][0].shape == (4,)
+
+
+def test_talker_duplex_drains_each_wake_and_keeps_cumulative_code_history() -> 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)
+    info["additional_information"]["meta"]["codec_streaming"] = torch.tensor([True])
+    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]) == 1.0  # current scheduler segment is drained
+
+    # 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]) == 1.0
+
+    # Zero-progress wake: never fabricate frames. The stop flag parks this
+    # resumable scheduler request until another connector chunk arrives.
+    _, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info2)
+    assert session["step"] == prompt_len + 2
+    assert update == {}
+    assert float(embeds[0, 0]) == 1.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_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."""
@@ -456,8 +646,10 @@ def test_talker_native_codec_stream_ends_each_received_segment() -> None:
     assert float(embeds[0, 0]) == 1.0
 
 
-def test_talker_sync_mode_steps_once_per_wake() -> None:
-    """The verified sync parity path keeps its one-TTS-step-per-engine-step cadence."""
+def test_talker_sync_mode_drains_whole_timeline_in_one_wake() -> None:
+    """Sync mode ships the whole timeline up front: the first decode wake runs
+    every TTS step back-to-back (no per-frame engine round trip) and emits the
+    stop flag with the full cumulative code stack."""
     talker = _bare_talker()
     pad, prompt_len = 12, 3
     ids = torch.zeros(1, dtype=torch.long)
@@ -471,10 +663,7 @@ def test_talker_sync_mode_steps_once_per_wake() -> None:
     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
+    _, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info)
+    assert session["step"] == 5  # t advanced 1 -> T in one wake
+    assert update["codes"]["audio"].shape == (4, 31)
+    assert float(embeds[0, 0]) == 1.0  # stop flag on the same wake