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