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

recipes/NVIDIA/NemotronLabs-VoiceChat.md

@@ -31,7 +31,8 @@ NeMo modules (`nemo_vendored/`), so no `nemo_toolkit` install is needed.
   [`examples/offline_inference/nemotron_voicechat/end2end.py`](../../examples/offline_inference/nemotron_voicechat/end2end.py)
 - Model modules (thinker / talker / code2wav / vendored NeMo):
   [`vllm_omni/model_executor/models/nemotron_voicechat/`](../../vllm_omni/model_executor/models/nemotron_voicechat/)
-- Staged pipeline config:
+- Staged pipeline config (fast default; "PARITY:" comments restore the
+  NeMo-bit-exact reference):
   [`vllm_omni/deploy/nemotron_labs_voicechat.yaml`](../../vllm_omni/deploy/nemotron_labs_voicechat.yaml)
 - Native duplex config and E2E probe:
   [`vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml`](../../vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml),
@@ -46,14 +47,60 @@ NeMo modules (`nemo_vendored/`), so no `nemo_toolkit` install is needed.
 
 | stage | arch | dtype | role |
 |---|---|---|---|
-| 0 thinker | `NemotronVoiceChatThinkerForConditionalGeneration` (LLM_AR) | fp32 | WAV + system prompt -> frame-locked text-token timeline (+ function channel) |
+| 0 thinker | `NemotronVoiceChatThinkerForConditionalGeneration` (LLM_AR) | bf16 (PARITY: fp32) | WAV + system prompt -> frame-locked text-token timeline (+ function channel) |
 | 1 talker | `NemotronVoiceChatTalker` (LLM_AR) | fp32 | text timeline -> 31-quantizer RVQ code stacks (one per 80 ms frame) |
 | 2 code2wav | `NemotronVoiceChatCode2Wav` (LLM_GENERATION) | fp32 | RVQ-VAE decode -> 22.05 kHz PCM |
 
-Stages 0/1 default to fp32 for exact parity with the NeMo reference
-implementation (greedy decoding matches it token for token on the acceptance
-fixture). The deploy yaml documents a ~2x-faster bf16 thinker option whose
-output stayed within one word of the reference in testing.
+The deploy yamls default to a fast execution profile; every fast setting
+carries a "PARITY:" comment whose edits restore the eager fp32 reference
+execution that matches the NeMo implementation token for token (see
+Performance below).
+
+## Performance
+
+Each deploy yaml defaults to a fast execution profile, mirroring how NVIDIA's
+own serving stack runs this model (bf16 LLM + CUDA-graph decode, EAR-TTS on a
+vLLM engine, incremental codec cache):
+
+- `nemotron_labs_voicechat.yaml` (offline sync): bf16 + CUDA-graph thinker and
+  the talker's Gemma3 backbone as a native vLLM model
+  (`hf_overrides.use_native_talker`) — PagedAttention KV, vLLM decode CUDA
+  graphs, and the MoG sampling step captured into a single graph. RTF ~0.23 on
+  the 196-frame acceptance fixture (1x H100, warm).
+- `nemotron_labs_voicechat_streaming.yaml` (offline streaming): the vendored
+  per-frame EAR-TTS step captured into CUDA graphs
+  (`hf_overrides.use_talker_cuda_graphs`), 5-frame codec chunks decoded
+  incrementally with a per-request causal-conv cache
+  (`use_incremental_codec_cache`, O(T) instead of O(T^2) codec work). First
+  audio chunk ~1.5 s after submission on the same fixture.
+- `nemotron_labs_voicechat_duplex.yaml` (realtime): native talker + bf16
+  thinker; audio is delivered as one 80 ms packet per frame at the frame clock
+  (median inter-chunk gap ~81 ms on the bundled fixture) and the session
+  length is unbounded (paged KV, no StaticCache bucket cap).
+
+Notes on the fast settings:
+
+- Thinker: `dtype: bfloat16` + `enforce_eager: false`. The 9B NemotronH decode
+  is memory- and launch-bound, so both are enabled together; weights fit in
+  ~21 GB, so the pipeline also runs on sub-80 GB parts. Text stays coherent
+  but greedy decoding can legitimately diverge from the fp32 reference
+  (verified turn-for-turn coherent and ASR-clean on the bundled fixtures).
+- Talker (captured-step, streaming yaml): the whole EAR-TTS frame step
+  (Gemma3 backbone via HF StaticCache + CFG batch doubling + MoG iterative
+  unmasking) replays as one CUDA graph per cache-size bucket. MoG noise draws
+  from the CUDA-graph Philox stream: audio is equivalent but not
+  bit-identical to the eager path. Capture failures and oversize sessions
+  fall back to the eager step automatically.
+- Talker (native, sync/duplex yamls): engine positions map 1:1 onto timeline
+  steps; fused input embeddings are numerically verified against the
+  reference forward, and the reference's EOS-silence code feedback
+  (`inference_force_speech_silence_on_eos`) is applied. Supports
+  multi-session batching (`max_num_seqs > 1`).
+
+Bit parity with the NeMo reference is preserved behind the `PARITY:` comments
+in each yaml: applying those edits runs the eager fp32 reference execution
+(md5-anchored WAV on the acceptance fixture; the drain-all talker scheduling
+added for performance was verified WAV-byte identical against it).
 
 ## Hardware Support
 
@@ -110,7 +157,7 @@ vllm-omni serve /path/to/NVIDIA-NemotronLabs-VoiceChat-11B \
   --deploy-config vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml
 ```
 
-The single-GPU profile runs three eager fp32 engine processes on one device.
+The single-GPU profile runs all three engine processes on one device.
 
 In another shell, stream only the user channel of one of NVIDIA's bundled
 stereo conversation fixtures:
@@ -142,14 +189,53 @@ python tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py \
   --expected-post-tool-text "random number"
 ```
 
+Probe caveats observed on the bundled `turn_taking.wav`: stream the FULL
+fixture (a `--max-frames` truncation cuts the conversation mid-turn, the model
+then never closes its last response, and the probe times out waiting for a
+`response.done`), and note that the pass/fail condition "a response completes
+after the final commit" is timing-sensitive — a faster pipeline can legitimately
+complete every turn before the commit lands (use
+`--allow-incomplete-response` when measuring latency rather than protocol
+completion).
+
+#### Duplex performance profile
+
+`nemotron_labs_voicechat_duplex.yaml` defaults to the fast profile: the
+talker's Gemma3 backbone on PagedAttention + vLLM CUDA graphs (MoG sampling
+in a captured step graph, unbounded session length) and a bf16 + CUDA-graph
+thinker inherited from the streaming base. On the bundled `turn_taking.wav`
+fixture (1x H100) the client receives one 80 ms packet per frame at the frame
+clock (median inter-chunk gap ~81 ms / p95 ~110 ms) with sub-second first-packet
+latency per turn; playback clients should still keep a small (~200-300 ms)
+prebuffer to absorb scheduling jitter.
+
+The bf16 thinker's greedy text stream can legitimately diverge from fp32
+mid-conversation. On the bundled fixture it was verified turn-for-turn:
+coherent transcripts, every turn reaching EOS/`response.done`, ASR-clean
+output. Apply the PARITY edits in the yaml to bit-match the fp32 reference
+thinker instead.
+
+Classifier-free guidance on the native talker
+(`hf_overrides.native_talker_guidance: true` + `gpu_memory_utilization: 0.18`
+on stage 1): the conditional stream stays on vLLM paged KV while the
+unconditional stream mirrors it on the vendored HF backbone (a captured
+StaticCache step), and the pair is blended with NeMo's `generate_step` math
+(hidden-space blends in `lm_head` and inside `MoGHead.infer`). The per-frame
+cost stays within the 80 ms duplex budget, and a paired 4-fixture ASR study
+matched the guided reference. Off by default since guidance at this
+checkpoint's scale (0.2) is ASR-indistinguishable; enable it for strict parity
+with NVIDIA's guided deployment or for listening tests. Under CFG the
+unconditional stream's StaticCache is bucketed (largest 4096 positions ≈ 5.4
+minutes); one session at a time uses the captured step, extra concurrent
+sessions fall back to eager stepping.
+
 #### Notes
 
-- Memory usage: the shipped yaml runs all three stages on one GPU
-  (`gpu_memory_utilization` 0.62 / 0.12 / 0.06); peak usage is dominated by the
-  fp32 thinker. The fp32 default has a hard floor of roughly 43 GB of thinker
-  weights alone (9B backbone + 587M `embed_tokens` + 587M `function_head` +
-  0.6B Conformer), so 48 GB cards cannot run it — use the bf16 thinker option
-  documented in the deploy yaml on anything smaller than an 80 GB part.
+- Memory usage: the shipped yamls run all three stages on one GPU
+  (`gpu_memory_utilization` 0.62 / 0.12 / 0.06). The default bf16 thinker fits
+  in ~21 GB of weights; the fp32 PARITY thinker needs roughly 43 GB of weights
+  alone (9B backbone + 587M `embed_tokens` + 587M `function_head` + 0.6B
+  Conformer), so the PARITY execution requires an 80 GB part.
 - Input sizing: the timeline is frame-locked, so the reply budget IS the input
   duration. The acoustic channel trails the text channel; if the WAV does not
   carry enough trailing silence for the reply to finish, the spoken answer is
@@ -161,15 +247,17 @@ python tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py \
   set `min_tokens` on the thinker: the tokenizer's EOS token is also the
   frame-locked PAD/silence token, so masking it forces the model to babble
   instead of pausing.
-- The talker's `max_tokens` is 16383 (its stage prompt is one placeholder
-  token, and the stage context is 16384).
+- The talker's `max_tokens` is 16347: its stage prompt is the 37-token
+  speaker-prompt prefill (`talker_init_len: 37`), and prompt + max_tokens
+  must fit the 16384 stage context.
 - The 80 ms frame period is a model/protocol contract, not a throughput
-  guarantee. In an eager fp32 H200 profile, Stage 2 emitted 80 ms packets at
-  roughly 0.87--0.88 s intervals (RTF about 10.8--10.9); placing the three
-  stages on separate GPUs did not remove the serial 8-step EAR-TTS bottleneck.
-  The current native duplex path is functionally streaming but not wall-clock
-  realtime.
-- Known limitations: batch=1 only. The native duplex deployment allows one
+  guarantee. The default duplex profile keeps up with the frame clock on 1x
+  H100; the eager fp32 PARITY execution is functionally streaming but not
+  wall-clock realtime (the serial 8-step EAR-TTS unmasking dominates, and
+  placing the three stages on separate GPUs does not remove it).
+- Known limitations: the offline pipeline is single-request per call
+  (the native talker supports engine-level multi-session batching via
+  max_num_seqs > 1). The native duplex deployment allows one
   active session and does not support barge-in. Tool execution remains
   client-owned: the server emits function-call events, accepts a validated
   `function_call_output`, and resumes the live model with the returned result;

tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py

@@ -216,7 +216,10 @@ async def run(args: argparse.Namespace) -> dict[str, object]:
                     if args.expect_function_call
                     else (
                         client.events.count("response.audio.delta") >= args.minimum_audio_chunks
-                        and (client.events.count("response.done") > completed_responses_at_commit)
+                        and (
+                            args.allow_incomplete_response
+                            or client.events.count("response.done") > completed_responses_at_commit
+                        )
                     )
                 )
             ),
@@ -227,7 +230,7 @@ async def run(args: argparse.Namespace) -> dict[str, object]:
         if client.events.errors():
             raise AssertionError(f"Realtime session emitted errors: {client.events.errors()}")
         done_events = _events(client, "response.done")
-        if not args.expect_function_call:
+        if not args.expect_function_call and not args.allow_incomplete_response:
             response = done_events[-1].get("response") if done_events else None
             status = response.get("status") if isinstance(response, dict) else None
             if status != "completed":
@@ -344,6 +347,17 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
     parser.add_argument("--instructions", default=DEFAULT_INSTRUCTIONS)
     parser.add_argument("--max-frames", type=int)
     parser.add_argument("--minimum-audio-chunks", type=int, default=1)
+    parser.add_argument(
+        "--allow-incomplete-response",
+        action="store_true",
+        help=(
+            "Accept a session whose responses all completed before the final "
+            "commit. A realtime-paced pipeline delivers each turn's audio and "
+            "response.done as the turn happens, so the strict "
+            "done-after-commit gate only holds when delivery lags the frame "
+            "clock; use this flag when measuring latency on fast configs."
+        ),
+    )
     parser.add_argument("--minimum-audio-rms", type=float, default=1e-4)
     parser.add_argument("--expect-function-call", action="store_true")
     parser.add_argument("--expected-function-name", default="generate_random_number")

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

vllm_omni/config/stage_config.py

@@ -506,6 +506,11 @@ class DuplexSessionRuntimeConfig:
     max_pending_turns_per_session: int = 4
     max_sessions: int = 1
     completed_append_cache_size: int = 256
+    # Startup warmup: run this many silent 80 ms-style frames through a
+    # throwaway realtime session before real clients are admitted, so
+    # one-time costs (kernel JIT, first prefill/decode paths, codec caches)
+    # never land on the first user. 0 disables the warmup.
+    warmup_frames: int = 0
 
     def __post_init__(self) -> None:
         positive = {

vllm_omni/data_entry_keys.py

@@ -204,6 +204,12 @@ class MetaStruct(_StructBase):
     decode_flag: bool | None = None
     codec_streaming: bool | None = None
     ref_code_len: int | None = None
+    # Expected FINAL length of a growing async-chunk sequence, when the
+    # producer knows it up front (e.g. a frame-locked AR stage whose
+    # max_tokens is exact). Lets consumers size fixed-capacity state (such as
+    # a CUDA-graph StaticCache bucket) optimally instead of assuming the
+    # worst case.
+    expected_total_tokens: int | None = None
     ref_context_size: int | None = None
     ref_context_request_id: str | None = None
     ref_context_included: bool | None = None

vllm_omni/deploy/nemotron_labs_voicechat.yaml

@@ -2,6 +2,23 @@
 # 3 stages on one GPU: thinker (NemotronH AR) -> talker (EAR-TTS AR) -> code2wav
 # (RVQ-VAE decode). thinker->talker is token-path only; talker->code2wav ships
 # the [frames, 31] code stack via the shared-memory connector.
+#
+# The default is the FAST profile (2.8x on the acceptance fixture, 10.3 s ->
+# 3.65 s / RTF 0.23 on 1x H100): a bf16 + CUDA-graph thinker and the talker's
+# Gemma3 backbone running as a native vLLM model (PagedAttention KV + vLLM
+# decode CUDA graphs, MoG sampling in one captured graph) — the way NVIDIA's
+# own deployment serves EAR-TTS. Every fast setting is marked with a
+# "PARITY:" comment; apply those edits to bit-match the eager fp32 NeMo
+# reference (196/196 token parity, md5-anchored WAV on the acceptance
+# fixture).
+#
+# Notes on the native talker:
+#   * classifier-free guidance is OFF by default (a paired ASR study over 4
+#     fixtures at this checkpoint's guidance_scale=0.2 found no systematic
+#     difference). Opt in via hf_overrides.native_talker_guidance below.
+#   * the talker's engine prompt covers the speaker-prompt prefill
+#     (hf_overrides.talker_init_len; 37 for the shipped "Aria" speaker). The
+#     talker validates the value and reports the right one on mismatch.
 pipeline: nemotron_voicechat
 async_chunk: false
 
@@ -15,19 +32,19 @@ stages:
   - stage_id: 0  # thinker: speech -> frame-locked text (final text output)
     max_num_seqs: 1
     gpu_memory_utilization: 0.62
-    enforce_eager: true
+    # PARITY: set enforce_eager: true and dtype: float32. bf16 + CUDA graphs
+    # is ~4x faster per frame and halves the weights (~43 GB -> ~21 GB); its
+    # greedy text can legitimately diverge from the fp32 reference (measured
+    # within one word on the acceptance fixture, ASR-clean).
+    enforce_eager: false
+    dtype: bfloat16
+    # NemotronH mamba SSM state in fp32 (matches NVIDIA's own vLLM deployment
+    # of this model; bf16 SSM state accumulates drift over long sequences).
+    mamba_ssm_cache_dtype: float32
     enable_prefix_caching: false
     async_scheduling: false
     max_num_batched_tokens: 8192
     max_model_len: 8192
-    # fp32 matches the NeMo reference exactly (196/196 token parity on the
-    # acceptance fixture). bf16 is ~2x faster and measured within one word of
-    # the reference on the same fixture — flip to bfloat16 if you prefer speed
-    # over exact parity.
-    dtype: float32
-    # NemotronH mamba SSM state in fp32 (matches NVIDIA's own vLLM deployment
-    # of this model; bf16 SSM state accumulates drift over long sequences).
-    mamba_ssm_cache_dtype: float32
     devices: "0"
     default_sampling_params:
       # Greedy, frame-locked: the example sets max_tokens to the
@@ -39,7 +56,11 @@ stages:
   - stage_id: 1  # talker: text timeline -> 31-quantizer code stacks
     max_num_seqs: 1
     gpu_memory_utilization: 0.12
-    enforce_eager: true
+    # PARITY: set enforce_eager: true and delete the hf_overrides block. The
+    # vendored eager fp32 EAR-TTS step is the NeMo-bit-exact path; the native
+    # backbone replays the same math on PagedAttention with fresh
+    # CUDA-graph-safe RNG draws (audio equivalent, not bit-identical).
+    enforce_eager: false
     enable_prefix_caching: false
     async_scheduling: false
     skip_tokenizer_init: true
@@ -49,15 +70,26 @@ stages:
     # fp32 buffers into the forward, and 1B params in fp32 are cheap.
     dtype: float32
     devices: "0"
+    hf_overrides:
+      use_native_talker: true
+      talker_init_len: 37
+      # Classifier-free guidance (checkpoint guidance_scale=0.2): mirrors the
+      # stream on the vendored HF backbone (captured StaticCache step) and
+      # blends per NeMo's generate_step math. 196-frame fixture: 3.65 s ->
+      # 5.3 s (RTF 0.23 -> 0.34). Requires gpu_memory_utilization ~0.18 on
+      # this stage (the vendored backbone copy stays loaded). Off by default:
+      # measured ASR-equivalent at this checkpoint's scale.
+      # native_talker_guidance: true
     output_connectors:
       to_stage_2: connector_of_shared_memory
     default_sampling_params:
       # Placeholder-token loop; the stage emits stop token 1 after the last
-      # timeline position. 16383 (not 16384): the request also carries one
-      # placeholder prompt token, and prompt + max_tokens must fit
+      # timeline position. 16347 = max_model_len (16384) minus the native
+      # talker's 37-token speaker-prompt prefill: the request carries the
+      # placeholder prompt tokens, and prompt + max_tokens must fit
       # max_model_len.
       temperature: 0.0
-      max_tokens: 16383
+      max_tokens: 16347
       stop_token_ids: [1]
       detokenize: false
       seed: 0

vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml

@@ -3,6 +3,35 @@
 # One scheduler segment consumes one 1280-sample (80 ms) microphone frame.
 # Stage 0 remains one resumable request, so its NemotronH KV/Mamba state is
 # preserved between frame appends. Stage 1 and the codec overlap per frame.
+#
+# The default is the FAST profile (TTFP 712 ms -> 408 ms and realtime audio
+# delivery — one 80 ms packet per frame, median inter-chunk gap ~81 ms,
+# RTF ~1.01 on the bundled turn_taking fixture, 1x H100): the thinker runs
+# bf16 + CUDA graphs (inherited from the streaming base) and the talker's
+# Gemma3 backbone runs as a native vLLM model on PagedAttention, so the
+# duplex session length is unbounded (no StaticCache bucket cap) and the
+# per-frame talker cost drops below the vendored step. Every fast setting is
+# marked with a "PARITY:" comment here and in the streaming base; apply those
+# edits for the eager fp32 reference execution.
+#
+# Native-talker semantics:
+#   * stage 1 sets max_num_batched_tokens: 64 so the speaker-prompt prefill
+#     (hf_overrides.talker_init_len positions, embeddings precomputed — no
+#     step-to-step feedback) runs as ONE chunked prefill at session start.
+#     The per-frame code feedback IS strictly sequential, so the
+#     thinker->talker producer ships at most one new timeline token per
+#     chunk (drip; see thinker2talker_async_chunk) — at most one timeline
+#     position is ever schedulable per engine step.
+#   * the connector extra `native_talker: true` makes the thinker->talker
+#     producer append one engine position PER NEW timeline token (engine
+#     positions map 1:1 onto timeline steps).
+#   * matching NeMo's realtime pipeline, the talker never sees the system
+#     prompt (it conditions the thinker only); the duplex timeline is one
+#     anchor PAD + the per-frame text tokens.
+#   * classifier-free guidance is OFF by default (measured ASR-equivalent at
+#     this checkpoint's guidance_scale=0.2); opt in via
+#     hf_overrides.native_talker_guidance below (duplex pacing stays
+#     realtime with it on).
 base_config: nemotron_labs_voicechat_streaming.yaml
 pipeline: nemotron_voicechat
 session_mode: duplex
@@ -18,6 +47,10 @@ duplex_session:
   max_pending_turns_per_session: 4
   max_sessions: 1
   completed_append_cache_size: 256
+  # Startup warmup: run two silent frames through a throwaway session before
+  # real clients are admitted, so one-time costs (kernel JIT, first
+  # prefill/decode paths, codec caches) never land on the first user.
+  warmup_frames: 2
 
 async_chunk: true
 
@@ -27,6 +60,10 @@ connectors:
     extra:
       shm_threshold_bytes: 65536
       codec_chunk_frames: 1
+      # PARITY: delete this line together with the stage-1 hf_overrides
+      # below (the vendored talker drains chunks model-side and keeps the
+      # 1-position-per-chunk default).
+      native_talker: true
 
 stages:
   - stage_id: 0
@@ -41,6 +78,22 @@ stages:
   - stage_id: 1
     max_num_seqs: 1
     async_scheduling: false
+    # PARITY: set enforce_eager: true, delete max_num_batched_tokens and the
+    # hf_overrides block (and the connector's native_talker flag above) for
+    # the vendored eager talker of the reference duplex deployment.
+    enforce_eager: false
+    max_num_batched_tokens: 64
+    hf_overrides:
+      use_native_talker: true
+      talker_init_len: 37
+      # The prewarmed placeholder prompt must cover the speaker-prompt
+      # prefill positions (read by the orchestrator's async-chunk prewarm).
+      async_chunk_prewarm_prompt_len: 37
+      # Classifier-free guidance: mirrors the stream on the vendored HF
+      # backbone (captured StaticCache step) and blends per NeMo's
+      # generate_step math (~+8.5 ms/frame, still realtime). Requires
+      # gpu_memory_utilization ~0.18 on this stage.
+      # native_talker_guidance: true
 
   - stage_id: 2
     max_num_seqs: 1

vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml

@@ -9,9 +9,12 @@
 # seam-free. First audio arrives after ~codec_chunk_frames thinker steps
 # instead of after the whole pipeline completes.
 #
-# The non-streaming yaml remains the parity-acceptance path (bit-exact WAV);
-# streamed audio is validated by ASR/WER instead (prefix decodes are not
-# bit-identical to the one-shot decode of the final stack).
+# The default is the FAST profile (TTFP 1.47 s / RTF 0.33 on the 196-frame
+# fixture, 1x H100): bf16 + CUDA-graph thinker, CUDA-graph EAR-TTS talker
+# step, and incremental codec decode. Every fast setting is marked with a
+# "PARITY:" comment; apply those edits for the eager fp32 reference
+# execution. Streamed audio is validated by ASR/WER, not bit parity (prefix
+# decodes are not bit-identical to the one-shot decode of the final stack).
 pipeline: nemotron_voicechat
 async_chunk: true
 
@@ -21,19 +24,25 @@ connectors:
     extra:
       shm_threshold_bytes: 65536
       # Streaming cadence for talker -> code2wav: emit a chunk once this many
-      # NEW frames (80 ms each) have accumulated; ~1 s of audio per chunk.
-      codec_chunk_frames: 13
+      # NEW frames (80 ms each) have accumulated (5 frames = 0.4 s of audio).
+      # Small chunks are cheap because the incremental codec cache decodes
+      # only each chunk's new frames; on the 196-frame fixture 5-frame chunks
+      # cut TTFP by ~0.4 s vs 13-frame chunks at unchanged total time.
+      # PARITY: 13 (~1 s per chunk) was the reference cadence.
+      codec_chunk_frames: 5
 
 stages:
   - stage_id: 0  # thinker: speech -> frame-locked text (final text output)
     max_num_seqs: 1
     gpu_memory_utilization: 0.62
-    enforce_eager: true
+    # PARITY: set enforce_eager: true and dtype: float32 for the fp32
+    # reference thinker (greedy text can legitimately diverge under bf16).
+    enforce_eager: false
     enable_prefix_caching: false
     async_scheduling: false
     max_num_batched_tokens: 8192
     max_model_len: 8192
-    dtype: float32
+    dtype: bfloat16
     mamba_ssm_cache_dtype: float32
     devices: "0"
     output_connectors:
@@ -54,6 +63,12 @@ stages:
     max_model_len: 16384
     dtype: float32
     devices: "0"
+    hf_overrides:
+      # Capture the vendored per-frame EAR-TTS step (backbone + MoG, CFG
+      # batch layout) into CUDA graphs (~3x per-frame). PARITY: delete this
+      # block for the eager reference step (bit-exact vs NeMo).
+      use_talker_cuda_graphs: true
+      talker_graph_max_cache_len: 4096
     input_connectors:
       from_stage_0: connector_of_shared_memory
     output_connectors:
@@ -75,6 +90,12 @@ stages:
     max_num_batched_tokens: 65536
     max_model_len: 65536
     devices: "0"
+    hf_overrides:
+      # Decode only each chunk's NEW frames with a per-request causal-conv
+      # cache (NVIDIA's S2S_USE_CODEC_CACHE path) instead of re-decoding the
+      # cumulative prefix per chunk (O(T^2) in the reply length).
+      # PARITY: delete this block for the prefix re-decode reference.
+      use_incremental_codec_cache: true
     input_connectors:
       from_stage_1: connector_of_shared_memory
     default_sampling_params:

vllm_omni/engine/orchestrator.py

@@ -2936,6 +2936,19 @@ class Orchestrator:
                     next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids))
                 except Exception:
                     next_prompt_len = max(1, len(prompt_token_ids))
+                # A stage may size its own prewarmed placeholder prompt via
+                # hf_overrides.async_chunk_prewarm_prompt_len (e.g. a talker
+                # whose engine positions must cover a speaker-prompt prefill;
+                # the model validates the value and reports the right one on
+                # mismatch).
+                prewarm_len = getattr(
+                    getattr(next_pool.stage_vllm_config, "model_config", None),
+                    "hf_config",
+                    None,
+                )
+                prewarm_len = getattr(prewarm_len, "async_chunk_prewarm_prompt_len", None)
+                if prewarm_len is not None and int(prewarm_len) > 1:
+                    next_prompt_len = int(prewarm_len)
 
                 original_prompt = req_state.prompt
                 if isinstance(original_prompt, dict):

vllm_omni/engine/stage_engine_core_client.py

@@ -150,6 +150,7 @@ class StageEngineCoreClientBase(StageClientBase):
         self.engine_outputs: Any = None
         self.client_addresses = dict(client_addresses or {})
         self._omni_kv_config = getattr(getattr(vllm_config, "model_config", None), "omni_kv_config", None)
+        self._stage_hf_config = getattr(getattr(vllm_config, "model_config", None), "hf_config", None)
         self._kv_sender_host = self._resolve_contact_host()
         self._kv_sender_info: dict[str, Any] | None = None
         self._kv_sender_initialized = False
@@ -396,17 +397,29 @@ class StageEngineCoreClientBase(StageClientBase):
         """
         if self.custom_process_input_func is not None:
             signature = inspect.signature(self.custom_process_input_func)
-            if len(signature.parameters) >= 4:
+            extra_kwargs: dict[str, Any] = {}
+            if "next_stage_hf_config" in signature.parameters:
+                # Let a processor size the next stage's prompt from that
+                # stage's model config (e.g. a talker whose engine positions
+                # must cover a speaker-prompt prefill).
+                extra_kwargs["next_stage_hf_config"] = self._stage_hf_config
+            # Match the context parameter by name, including the
+            # underscore-prefixed spelling some processors use (e.g.
+            # MiniCPM-o's ``llm2tts(..., _streaming_context)``), so bridge
+            # state keeps flowing to them.
+            if "streaming_context" in signature.parameters or "_streaming_context" in signature.parameters:
                 return self.custom_process_input_func(
                     source_outputs,
                     prompt,
                     self.requires_multimodal_data,
                     streaming_context,
+                    **extra_kwargs,
                 )
             return self.custom_process_input_func(
                 source_outputs,
                 prompt,
                 self.requires_multimodal_data,
+                **extra_kwargs,
             )
 
         if not self.engine_input_source:

vllm_omni/entrypoints/openai/api_server.py

@@ -13,6 +13,7 @@ import os
 import random
 import tempfile
 import time
+import uuid
 from argparse import Namespace
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
@@ -445,6 +446,116 @@ class _DiffusionServingModels:
 # Server entry points
 
 
+async def _warmup_duplex_realtime(app, args, warmup_frames: int) -> None:
+    """Run silent frames through a throwaway realtime session at startup.
+
+    Self-connects over the server's own websocket endpoint so the warmup
+    exercises the exact production path (serving adapter, duplex control
+    plane, every pipeline stage). One-time costs — kernel JIT compilation,
+    first prefill/decode paths, codec caches — are paid here instead of by
+    the first real client; ``/v1/realtime`` holds other connections until
+    the warmup finishes (see ``duplex_warmup_done``).
+    """
+    import base64
+
+    warmup_done = getattr(app.state, "duplex_warmup_done", None)
+    try:
+        try:
+            import websockets
+        except ImportError:
+            logger.warning("Duplex warmup skipped: the 'websockets' package is not installed.")
+            return
+        served = getattr(args, "served_model_name", None)
+        if isinstance(served, list | tuple) and served:
+            model_name = served[0]
+        elif isinstance(served, str) and served:
+            model_name = served
+        else:
+            model_name = args.model
+        adapter = getattr(getattr(app.state, "openai_serving_duplex", None), "_serving_runtime_adapter", None)
+        frame_samples = int(getattr(adapter, "silence_continuation_samples", 16000))
+        silence = base64.b64encode(bytes(frame_samples * 4)).decode("ascii")
+        from vllm_omni.experimental.fullduplex.client import build_realtime_url
+
+        url = (
+            build_realtime_url(f"ws://127.0.0.1:{args.port}/v1/realtime", model_name, autostart=False)
+            + "&vllm_omni_warmup=1"
+        )
+        logger.info("Duplex warmup starting: %d silent frames of %d samples.", warmup_frames, frame_samples)
+        start = time.time()
+        # The socket is bound before uvicorn finishes starting, so a plain
+        # connect can be accepted and then reset. Wait for /health first.
+        async with httpx.AsyncClient() as http:
+            for _ in range(120):
+                try:
+                    if (await http.get(f"http://127.0.0.1:{args.port}/health", timeout=2)).status_code == 200:
+                        break
+                except httpx.HTTPError:
+                    pass
+                await asyncio.sleep(1)
+            else:
+                logger.warning("Duplex warmup: /health never became ready; skipping warmup.")
+                return
+        async with websockets.connect(url, max_size=64 * 1024 * 1024, open_timeout=10) as ws:
+            await ws.send(
+                json.dumps(
+                    {
+                        "type": "session.update",
+                        "session": {
+                            "session_id": f"warmup-{uuid.uuid4().hex[:8]}",
+                            "model": model_name,
+                            "modalities": ["audio", "text"],
+                            "input_audio_format": "pcm_f32le",
+                            "output_audio_format": "pcm16",
+                            "idle_timeout_s": 60,
+                            "turn_detection": None,
+                            "extra_body": {"auto_response": True},
+                        },
+                    }
+                )
+            )
+            saw_audio = False
+            sent = 0
+
+            async def _recv_until(predicate, timeout_s: float) -> bool:
+                deadline = time.monotonic() + timeout_s
+                while time.monotonic() < deadline:
+                    try:
+                        raw = await asyncio.wait_for(ws.recv(), timeout=max(0.05, deadline - time.monotonic()))
+                    except (TimeoutError, asyncio.TimeoutError):
+                        return False
+                    event = json.loads(raw)
+                    if predicate(event):
+                        return True
+                return False
+
+            if not await _recv_until(lambda e: e.get("type") == "session.created", 30):
+                logger.warning("Duplex warmup: no session.created within 30 s; aborting warmup.")
+                return
+            while sent < warmup_frames:
+                await ws.send(json.dumps({"type": "input_audio_buffer.append", "audio": silence}))
+                sent += 1
+                await asyncio.sleep(0.08)
+            # Wait for the pipeline's first audio output so every stage ran.
+            saw_audio = await _recv_until(lambda e: e.get("type") == "response.audio.delta", 30)
+            # Close the session EXPLICITLY: a bare websocket close parks the
+            # session in its disconnect grace window, where it would keep
+            # occupying the max_sessions=1 slot and block the first client.
+            await ws.send(json.dumps({"type": "session.close"}))
+            await _recv_until(lambda e: e.get("type") == "session.closed", 10)
+        logger.info(
+            "Duplex warmup finished in %.1f s (%d silent frames, first audio %s).",
+            time.time() - start,
+            sent,
+            "received" if saw_audio else "NOT received",
+        )
+    except Exception:
+        logger.exception("Duplex warmup failed; continuing to serve (first session pays cold-start costs).")
+    finally:
+        if warmup_done is not None:
+            warmup_done.set()
+
+
 async def omni_run_server(args, **uvicorn_kwargs) -> None:
     """Run a single-worker API server.
 
@@ -585,6 +696,21 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
                     scope["state"]["request_timestamp"] = time.time()
                 await self._inner(scope, receive, send)
 
+        # Startup duplex warmup (duplex_session.warmup_frames in the deploy
+        # yaml): real /v1/realtime connections wait on this event so the
+        # first client never pays cold-start costs.
+        duplex_warmup_frames = 0
+        if getattr(app.state, "openai_serving_duplex", None) is not None:
+            duplex_cfg = getattr(engine_client, "duplex_session_config", None)
+            duplex_warmup_frames = int(getattr(duplex_cfg, "warmup_frames", 0) or 0)
+        app.state.duplex_warmup_done = asyncio.Event() if duplex_warmup_frames > 0 else None
+        warmup_task: asyncio.Task | None = None
+        if duplex_warmup_frames > 0:
+            # Scheduled before serve_http (which may not return until
+            # shutdown); the coroutine retries its self-connect until the
+            # server socket is accepting.
+            warmup_task = asyncio.create_task(_warmup_duplex_realtime(app, args, duplex_warmup_frames))
+
         shutdown_task = await serve_http(
             _TimestampMiddleware(app),
             sock=sock,
@@ -609,6 +735,8 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
         try:
             await shutdown_task
         finally:
+            if warmup_task is not None:
+                warmup_task.cancel()
             state = getattr(app, "state", None)
             serving_video = getattr(state, "openai_serving_video", None) if state is not None else None
             if serving_video is not None:
@@ -1709,6 +1837,14 @@ async def streaming_video_output(websocket: WebSocket):
 @router.websocket("/v1/realtime")
 async def realtime_websocket(websocket: WebSocket):
     """WebSocket endpoint for OpenAI-style realtime interactions."""
+    # Hold real clients until the startup duplex warmup finishes (the warmup
+    # connection marks itself with vllm_omni_warmup=1 and passes through).
+    warmup_done = getattr(websocket.app.state, "duplex_warmup_done", None)
+    if warmup_done is not None and not warmup_done.is_set() and websocket.query_params.get("vllm_omni_warmup") != "1":
+        try:
+            await asyncio.wait_for(warmup_done.wait(), timeout=120)
+        except (TimeoutError, asyncio.TimeoutError):
+            logger.warning("Duplex warmup still running after 120 s; admitting the client anyway.")
     duplex_handler = getattr(websocket.app.state, "openai_serving_duplex", None)
     duplex_query = websocket.query_params.get("duplex")
     use_duplex_realtime = (

vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/duplex_ear_tts.py

@@ -705,6 +705,10 @@ class DuplexEARTTS(nn.Module):
         generation_config=None,
         ignore_eos_flag_stop=True,
         request_id=None,
+        current_subword_embed=None,
+        cache_position=None,
+        attention_mask=None,
+        position_ids=None,
     ):
         """
         Runs a single autoregressive prediction step to infer audio codec codes.
@@ -756,6 +760,20 @@ class DuplexEARTTS(nn.Module):
         # request_id is only used by vLLM backend — skip for native PyTorch inference
         if request_id is not None:
             inputs["request_id"] = request_id
+        # Fast-path extras (see RVQEARTTSModel.forward): a precomputed subword
+        # embedding bypasses the CAS encoder (host-sync free); cache_position
+        # enables StaticCache decoding; an explicit (4D) attention_mask and
+        # position_ids keep the mask and RoPE fully tensor-driven, which CUDA
+        # graph capture requires (the HF-derived paths read python-int cache
+        # lengths that would be baked into the captured graph).
+        if current_subword_embed is not None:
+            inputs["subword_embeds"] = current_subword_embed
+        if cache_position is not None:
+            inputs["cache_position"] = cache_position
+        if attention_mask is not None:
+            inputs["attention_mask"] = attention_mask
+        if position_ids is not None:
+            inputs["position_ids"] = position_ids
 
         outputs = self.tts_model(**inputs)
 

vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_model.py

@@ -1141,14 +1141,24 @@ class RVQEARTTSModel(nn.Module):
         subword_ids: Tensor | None,
         subword_mask: Tensor | None,
         uncond_dec_flag: Tensor,
+        subword_embeds: Tensor | None = None,
     ) -> Tensor:
-        """Computes the final conditioning tensor by combining all sources."""
+        """Computes the final conditioning tensor by combining all sources.
+
+        ``subword_embeds`` short-circuits the CharAwareSubwordEncoder with a
+        precomputed embedding (same shape as its output). Besides skipping the
+        char-level backbone, this path is free of host syncs
+        (``subword_mask.any()`` / the encoder's ``.cpu().tolist()``), which
+        makes it CUDA-graph capturable.
+        """
         cond = torch.zeros((1, 1, self.hidden_size), device=uncond_dec_flag.device)
 
         if self.embed_context is not None and context_hidden_state is not None:
             cond = cond + self.embed_context(context_hidden_state)
 
-        if self.embed_subword is not None and subword_ids is not None:
+        if subword_embeds is not None:
+            cond = cond + subword_embeds
+        elif self.embed_subword is not None and subword_ids is not None:
             # Infer subword mask from context if not provided
             if subword_mask is None and context_hidden_state is not None:
                 subword_mask = torch.any(context_hidden_state != 0, dim=-1)
@@ -1244,6 +1254,8 @@ class RVQEARTTSModel(nn.Module):
         ignore_eos_flag_stop: bool = False,
         asr_speech_tokens_emb: Tensor | None = None,
         audio_prompt_latent: Tensor | None = None,
+        subword_embeds: Tensor | None = None,
+        cache_position: Tensor | None = None,
     ) -> RVQEARTTSOutput:
         """
         Performs a forward pass handling training, generation, or single-step inference.
@@ -1352,13 +1364,17 @@ class RVQEARTTSModel(nn.Module):
                 subword_ids = torch.cat([subword_ids] * 2, 0)
                 if subword_mask is not None:
                     subword_mask = torch.cat([subword_mask] * 2, 0)
+            if subword_embeds is not None:
+                subword_embeds = torch.cat([subword_embeds] * 2, 0)
             if asr_speech_tokens_emb is not None:
                 asr_speech_tokens_emb = torch.cat([asr_speech_tokens_emb] * 2, 0)
 
             uncond_dec_flag = torch.cat([uncond_dec_flag, torch.ones_like(uncond_dec_flag)], 0)
 
         # Prepare conditioning
-        cond = self._prepare_conditioning(context_hidden_state, subword_ids, subword_mask, uncond_dec_flag)
+        cond = self._prepare_conditioning(
+            context_hidden_state, subword_ids, subword_mask, uncond_dec_flag, subword_embeds=subword_embeds
+        )
 
         if asr_speech_tokens_emb is not None:
             cond = cond + asr_speech_tokens_emb
@@ -1368,13 +1384,20 @@ class RVQEARTTSModel(nn.Module):
         else:
             inputs_embeds = code_embeds + cond
 
-        # Main backbone pass
+        # Main backbone pass. cache_position is threaded only when provided so
+        # backbones without the parameter keep working (it is required for
+        # StaticCache-based decoding, where the write index cannot be derived
+        # from the cache contents).
+        backbone_kwargs: dict[str, Any] = {}
+        if cache_position is not None:
+            backbone_kwargs["cache_position"] = cache_position
         backbone_outputs = self.backbone(
             inputs_embeds=inputs_embeds,
             attention_mask=attention_mask,
             position_ids=position_ids,
             past_key_values=past_key_values,
             use_cache=use_cache,
+            **backbone_kwargs,
         )
         hidden_states = backbone_outputs.last_hidden_state
 
@@ -1564,7 +1587,9 @@ class RVQEARTTSModel(nn.Module):
             else:
                 eos_flag = lm_logits.argmax(-1) == 1
 
-            if torch.all(eos_flag) and not ignore_eos_flag_stop:
+            # Check the cheap Python flag first: evaluating torch.all() forces a
+            # host sync per frame even when the result is ignored.
+            if not ignore_eos_flag_stop and torch.all(eos_flag):
                 return None, lm_logits, eos_flag
         else:
             lm_logits = None
@@ -1573,17 +1598,18 @@ class RVQEARTTSModel(nn.Module):
         # Initialize the full code tensor
         code = torch.zeros((b, t, d), dtype=torch.long, device=device) + self.config.codebook_size
 
-        # 3. Set up the iterative denoising schedule for the continuous part
-        rates = torch.linspace(0.0, 1.0, num_iter + 1, device=device)[:-1].unsqueeze(-1)
-        masking_rates = get_masking_rate(rates, exponent=exponent)
-        num_maskings = torch.ceil(masking_rates * self.config.num_quantizers).long()
-
-        ks = num_maskings - F.pad(num_maskings[1:], [0, 0, 0, 1])
+        # 3. Set up the iterative denoising schedule for the continuous part.
+        # The schedule is a pure function of (num_iter, exponent,
+        # num_quantizers) — compute it once as Python ints and cache it, so the
+        # per-frame loop below runs without any GPU->CPU sync (`.item()` /
+        # `torch.all(k == 0)` on a device tensor would stall the dispatch
+        # pipeline several times per 80 ms frame).
+        ks = self._unmask_schedule(num_iter, float(exponent))
 
         # 4. Iteratively unmask the continuous part of the code
         cnt = 0
         for i, k in enumerate(ks):
-            if torch.all(k == 0):
+            if k == 0:
                 continue
 
             # Prepare input for the MoG head
@@ -1593,7 +1619,7 @@ class RVQEARTTSModel(nn.Module):
 
             mog_input_embeds = self.embed_code(self.depthsum_embedding(code).to(self.embed_code.weight.dtype))
             if self.config.random_target_masking:
-                mog_input_embeds += self.embed_target_mask(cnt + k - 1)
+                mog_input_embeds += self.embed_target_mask(torch.tensor(cnt + k - 1, device=device))
             if guidance_scale_i > 0.0:
                 mog_input_embeds = torch.cat(
                     [mog_input_embeds + hidden_states, mog_input_embeds + uncond_hidden_states], 0
@@ -1607,10 +1633,26 @@ class RVQEARTTSModel(nn.Module):
                 top_p_or_k=top_p_or_k_i,
             )
             z = mog_mu + torch.exp(mog_logs) * torch.randn_like(mog_mu) * noise_scale_i
-            code = depthsum_encoding_step(self.rvq_embs, z, code, cnt, k[0].item())
-            cnt += k[0].item()
+            code = depthsum_encoding_step(self.rvq_embs, z, code, cnt, k)
+            cnt += k
         return code, lm_logits, eos_flag
 
+    def _unmask_schedule(self, num_iter: int, exponent: float) -> list[int]:
+        """Per-iteration unmask counts as Python ints (cached; sync-free loop)."""
+        cache = getattr(self, "_unmask_schedule_cache", None)
+        if cache is None:
+            cache = {}
+            self._unmask_schedule_cache = cache
+        key = (int(num_iter), float(exponent), int(self.config.num_quantizers))
+        ks = cache.get(key)
+        if ks is None:
+            rates = torch.linspace(0.0, 1.0, num_iter + 1)[:-1].unsqueeze(-1)
+            masking_rates = get_masking_rate(rates, exponent=exponent)
+            num_maskings = torch.ceil(masking_rates * self.config.num_quantizers).long()
+            ks = (num_maskings - F.pad(num_maskings[1:], [0, 0, 0, 1])).reshape(-1).tolist()
+            cache[key] = ks
+        return ks
+
     def load_state_dict(self, state_dict, strict: bool = True):
         try:
             super().load_state_dict(state_dict, strict=strict)

vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_code2wav.py

@@ -82,9 +82,22 @@ class NemotronVoiceChatCode2Wav(nn.Module):
         self.enable_update_additional_information = True
         self.requires_raw_input_tokens = True
 
+        # Incremental streaming decode (hf_overrides.use_incremental_codec_cache):
+        # keep a CausalConv1dCache per request so each async chunk decodes ONLY
+        # its new frames — O(T) codec work — instead of re-decoding the whole
+        # cumulative prefix per chunk (O(T^2)). This is the codec-cache path
+        # NVIDIA's own serving stack runs (S2S_USE_CODEC_CACHE); the default
+        # stays the prefix re-decode, whose concatenation is bit-equal to the
+        # one-shot decode of the final stack.
+        self._use_codec_cache = bool(getattr(self.config, "use_incremental_codec_cache", False))
+        self.requires_request_ids = self._use_codec_cache
+        # Shared per-request causal-conv caches; ALSO used by duplex serving,
+        # whose per-segment payloads mark themselves with meta.codec_streaming
+        # + meta.request_id (independent of the hf_override).
+        self._stream_caches: dict[str, Any] = {}
+
         # Constructed in load_weights (owns the tts_model.audio_codec.* subtree).
         self.audio_codec: nn.Module | None = None
-        self._duplex_codec_caches: dict[str, object] = {}
 
     # ------------------------------------------------------------------
     # Runner-facing placeholder hooks.
@@ -130,6 +143,7 @@ class NemotronVoiceChatCode2Wav(nn.Module):
         intermediate_tensors: Any = None,
         inputs_embeds: torch.Tensor | None = None,
         runtime_additional_information: list[dict[str, Any]] | None = None,
+        request_ids: list[str] | None = None,
         **kwargs: Any,
     ) -> OmniOutput:
         sr_tensor = torch.tensor(self._sample_rate, dtype=torch.int32)
@@ -189,25 +203,40 @@ class NemotronVoiceChatCode2Wav(nn.Module):
                 left_context_frames = max(int(runtime_info["meta.left_context_size"]), 0)
             if left_context_frames >= frames:
                 continue  # terminal empty chunk: all frames already emitted
-            lens = torch.tensor([frames], device=device, dtype=torch.long)
+            is_stream_chunk = (isinstance(meta, Mapping) and "left_context_size" in meta) or (
+                isinstance(runtime_info, Mapping) and "meta.left_context_size" in runtime_info
+            )
+            request_id = str(request_ids[i]) if request_ids is not None and i < len(request_ids) else None
             codec_cache = None
+            decode_codes = codes
             if codec_streaming:
+                # Duplex serving: the per-segment payload carries ONLY the new
+                # frames; continuity lives in the per-request causal-conv cache
+                # keyed by the talker-forwarded meta.request_id.
                 if not isinstance(cache_request_id, str) or not cache_request_id:
                     raise ValueError("Nemotron VoiceChat incremental codec input lacks meta.request_id")
                 from vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored.ear_tts_vae_codec import (
                     CausalConv1dCache,
                 )
 
-                codec_cache = self._duplex_codec_caches.setdefault(cache_request_id, CausalConv1dCache())
+                codec_cache = self._stream_caches.setdefault(cache_request_id, CausalConv1dCache())
+            elif self._use_codec_cache and is_stream_chunk and request_id is not None:
+                # Async-chunk streaming: cumulative payload; decode ONLY the
+                # frames past the already-emitted prefix, with the causal-conv
+                # left context carried in the per-request cache.
+                from vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored.ear_tts_vae_codec import (
+                    CausalConv1dCache,
+                )
+
+                codec_cache = self._stream_caches.setdefault(request_id, CausalConv1dCache())
+                decode_codes = codes[left_context_frames:]
+            lens = torch.tensor([decode_codes.shape[0]], device=device, dtype=torch.long)
             # NeMo decodes the codec strictly in fp32.
             with torch.autocast(device_type=device.type, enabled=False):
-                wav, wav_len = self.audio_codec.decode(
-                    codes.unsqueeze(0),
-                    lens,
-                    cache=codec_cache,
-                )
+                wav, wav_len = self.audio_codec.decode(decode_codes.unsqueeze(0), lens, cache=codec_cache)
             wav = wav.reshape(-1)[: int(wav_len.reshape(-1)[0])]
-            if left_context_frames > 0 and not codec_streaming:
+            if left_context_frames > 0 and codec_cache is None:
+                # Prefix re-decode path: slice off the samples already emitted.
                 wav = wav[left_context_frames * self._wav_per_frame :]
             audios[i] = wav.detach().to(dtype=torch.float32).cpu()
             logger.debug(
@@ -222,15 +251,15 @@ class NemotronVoiceChatCode2Wav(nn.Module):
             multimodal_outputs={"model_outputs": audios, "sr": srs},
         )
 
-    def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None:
-        for request_id in finished_req_ids:
-            self._duplex_codec_caches.pop(str(request_id), None)
-
     def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput | tuple, **kwargs: Any) -> OmniOutput:
         if isinstance(model_outputs, OmniOutput):
             return model_outputs
         raise TypeError(f"NemotronVoiceChatCode2Wav expected OmniOutput, got {type(model_outputs)}")
 
+    def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None:
+        for request_id in finished_req_ids:
+            self._stream_caches.pop(str(request_id), None)
+
     # ------------------------------------------------------------------
     # Weights.
     # ------------------------------------------------------------------

vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_talker.py

@@ -113,13 +113,71 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module):
         self.has_postprocess = False
         self.requires_full_prefix_cached_hidden_states = False
 
-        # Stateful per-request execution is only validated at batch size 1
-        # (the shipped offline scope); fail fast on silent multi-request use.
+        # The vendored (eager / captured-step) paths share single-session
+        # buffers (StaticCache, capture-time addresses) and stay batch=1; the
+        # NATIVE path keeps all per-request state in _sessions (paged KV,
+        # per-request MoG on each request's own hidden slice) and supports
+        # multi-session batching (validated at max_num_seqs<=4). The native
+        # flag is read below, so peek at it here for the guard.
         max_num_seqs = int(getattr(vllm_config.scheduler_config, "max_num_seqs", 1))
-        if max_num_seqs != 1:
+        if max_num_seqs != 1 and not bool(getattr(self.config, "use_native_talker", False)):
             raise NotImplementedError(
-                f"NemotronVoiceChat talker supports max_num_seqs=1 only (got {max_num_seqs}); "
-                "the vendored EAR-TTS session state is not batch-safe."
+                f"NemotronVoiceChat talker supports max_num_seqs=1 only (got {max_num_seqs}) on the "
+                "vendored paths; the EAR-TTS session buffers are not batch-safe. Use "
+                "hf_overrides.use_native_talker for multi-session batching."
+            )
+
+        # CUDA-graph fast path (hf_overrides.use_talker_cuda_graphs). Off by
+        # default: the eager path is the NeMo bit-parity acceptance path, while
+        # the captured step replays the same math ~3x faster with a
+        # graph-managed RNG (audio equivalent but not bit-identical).
+        self._use_cuda_graphs = bool(getattr(self.config, "use_talker_cuda_graphs", False))
+        self._graph_max_cache_len = int(getattr(self.config, "talker_graph_max_cache_len", 4096))
+        self._step_graph: Any | None = None
+
+        # Native-vLLM path (hf_overrides.use_native_talker, experimental): the
+        # Gemma3 backbone runs as a real vLLM model on PagedAttention (KV and
+        # per-step attention cost scale with the ACTUAL session length; vLLM's
+        # own decode CUDA graphs apply), mirroring NVIDIA's deployment. The
+        # per-frame MoG sampling runs on the backbone's hidden state in
+        # make_omni_output. Classifier-free guidance is NOT applied yet on
+        # this path (single conditional stream; the paired-request CFG design
+        # is a documented follow-up), so audio quality trails the vendored
+        # CFG paths — use for development only.
+        self._use_native_backbone = bool(getattr(self.config, "use_native_talker", False))
+        # Classifier-free guidance on the native path
+        # (hf_overrides.native_talker_guidance): the conditional stream stays
+        # on vLLM paged KV; the unconditional stream mirrors it on the vendored
+        # HF backbone (see talker_native.UncondStream) and the pair is blended
+        # per NeMo's generate_step math. Off by default (measured ASR-equivalent
+        # at this checkpoint's guidance_scale=0.2, but kept available for
+        # listening-test parity with the vendored paths).
+        self._native_guidance = self._use_native_backbone and bool(
+            getattr(self.config, "native_talker_guidance", False)
+        )
+        # Captured MoG sampling step (built lazily on the first decode frame;
+        # falls back to eager sampling if capture fails).
+        self._mog_graph: Any | None = None
+        self._uncond_graph: Any | None = None
+        self.backbone: nn.Module | None = None
+        if self._use_native_backbone:
+            if self._use_cuda_graphs:
+                raise ValueError(
+                    "NemotronVoiceChat talker: use_native_talker and use_talker_cuda_graphs are mutually exclusive."
+                )
+            from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
+
+            from vllm_omni.model_executor.models.nemotron_voicechat.talker_native import (
+                synthesize_backbone_config,
+            )
+
+            tts_config = dict(getattr(self.config, "tts_cfg", {}) or {}).get("tts_config") or {}
+            backbone_cfg = synthesize_backbone_config(dict(tts_config))
+            self.backbone = init_vllm_registered_model(
+                vllm_config=vllm_config,
+                hf_config=backbone_cfg,
+                prefix=maybe_prefix(prefix, "backbone"),
+                architectures=["Gemma3ForCausalLM"],
             )
 
         # Vendored DuplexEARTTS; constructed in load_weights.
@@ -141,6 +199,16 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module):
         inputs_embeds: torch.Tensor | None = None,
         **_: Any,
     ) -> torch.Tensor:
+        if self._use_native_backbone:
+            # Real Gemma3 forward on PagedAttention; preprocess injected the
+            # fused embeddings, make_omni_output turns the hidden state into
+            # codes (MoG) and the CONTINUE/STOP flag.
+            return self.backbone(
+                input_ids=None,
+                positions=positions,
+                intermediate_tensors=intermediate_tensors,
+                inputs_embeds=inputs_embeds,
+            )
         # The per-frame compute happens in preprocess; the "hidden state" only
         # transports the stop flag (column 0) to compute_logits.
         if inputs_embeds is not None:
@@ -161,6 +229,8 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module):
     def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput:
         if isinstance(model_outputs, OmniOutput):
             return model_outputs
+        if self._use_native_backbone:
+            return self._native_make_omni_output(model_outputs, **kwargs)
         hidden = model_outputs
         info_dicts = kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information") or []
         codes_list: list[torch.Tensor] = []
@@ -275,6 +345,36 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module):
             )
         return timeline
 
+    @staticmethod
+    def _expected_total_tokens(info: dict[str, Any]) -> int | None:
+        """Producer-announced FINAL timeline length (meta.expected_total_tokens)."""
+        meta = info.get("meta")
+        value = meta.get("expected_total_tokens") if isinstance(meta, dict) else None
+        if value is None:
+            value = info.get("meta.expected_total_tokens")
+        return int(value) if value is not None else None
+
+    def _graph_session_usable(self, timeline: torch.Tensor, expected_len: int | None) -> bool:
+        """Whether the CUDA-graph fast path can serve this session.
+
+        ``expected_len`` is the known FINAL timeline length (sync mode, or a
+        frame-locked async producer's meta.expected_total_tokens); the graph
+        picks the smallest fitting cache bucket from it, since per-step
+        attention cost scales with the bucket size. ``None`` = unbounded
+        (duplex) -> largest bucket.
+        """
+        if self._step_graph is None:
+            return False
+        needed = max(int(timeline.numel()), expected_len or 0)
+        if not self._step_graph.fits(needed):
+            logger.warning(
+                "NemotronVoiceChat talker: timeline of %d positions exceeds the CUDA-graph "
+                "capacity; using the eager per-frame step for this request.",
+                needed,
+            )
+            return False
+        return self._step_graph.ensure_captured(expected_len)
+
     def _init_session(self, request_id: str, info: dict[str, Any], device: torch.device) -> dict[str, Any]:
         assert self.tts is not None
         timeline = self._timeline(info).to(device)
@@ -283,15 +383,6 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module):
                 f"NemotronVoiceChat talker timeline for request {request_id} has "
                 f"{timeline.numel()} positions; need at least 2 (prompt + one frame)."
             )
-        # Warmup: pre-baked speaker latent + init prompt forward (NeMo
-        # offline_inference lines around get_init_inputs/first forward).
-        self.tts.set_init_inputs(speaker_name=self._speaker_name)
-        init_inputs = self.tts.get_init_inputs(B=1)
-        guidance_enabled = self._guidance_enabled()
-        generation_config = self.tts._get_generation_config(guidance_enabled)
-        init_inputs.update({"use_cache": True, "past_key_values": None, "guidance_enabled": guidance_enabled})
-        with torch.inference_mode():
-            outputs = self.tts.tts_model(**init_inputs)
         prompt_len = self._logical_prompt_len(info)
         if prompt_len is None:
             raise ValueError(
@@ -304,17 +395,17 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module):
         codec_streaming = scalar_bool(meta.get("codec_streaming")) if isinstance(meta, dict) else False
         codec_streaming = codec_streaming or scalar_bool(info.get("meta.codec_streaming"))
         codec_request_id = meta.get("request_id") if isinstance(meta, dict) else None
+        # Warmup: pre-baked speaker latent (NeMo offline_inference lines around
+        # get_init_inputs); the init prompt forward runs in the eager branch
+        # below, or inside the CUDA-graph session's own StaticCache prefill.
+        self.tts.set_init_inputs(speaker_name=self._speaker_name)
+        guidance_enabled = self._guidance_enabled()
         session = {
             "timeline": timeline,
             "prompt_len": prompt_len,
             "step": 1,  # NeMo's loop starts at t=1
-            "past_key_values": outputs.past_key_values,
-            "code": init_inputs["code"][:, -1:],
-            "first_context_subword_id": init_inputs["subword_ids"][:, -1].unsqueeze(-1),
             "current_subword_mask": torch.ones(1, 1, device=device, dtype=torch.bool),
-            "tts_initialized": False,
             "guidance_enabled": guidance_enabled,
-            "generation_config": generation_config,
             "codes_rows": [],
             "codec_streaming": codec_streaming,
             "codec_request_id": codec_request_id,
@@ -324,14 +415,434 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module):
             "sync_mode": info.get("nvc_text_timeline") is not None,
             "upstream_finished": self._upstream_finished(info) or info.get("nvc_text_timeline") is not None,
         }
+        if session["sync_mode"]:
+            expected_len: int | None = int(timeline.numel())
+        else:
+            expected_len = self._expected_total_tokens(info)
+        if self._graph_session_usable(timeline, expected_len):
+            self._step_graph.start_session(timeline, expected_len=expected_len)
+            session["use_graph"] = True
+        else:
+            init_inputs = self.tts.get_init_inputs(B=1)
+            generation_config = self.tts._get_generation_config(guidance_enabled)
+            init_inputs.update({"use_cache": True, "past_key_values": None, "guidance_enabled": guidance_enabled})
+            with torch.inference_mode():
+                outputs = self.tts.tts_model(**init_inputs)
+            session.update(
+                {
+                    "use_graph": False,
+                    "past_key_values": outputs.past_key_values,
+                    "code": init_inputs["code"][:, -1:],
+                    "first_context_subword_id": init_inputs["subword_ids"][:, -1].unsqueeze(-1),
+                    "tts_initialized": False,
+                    "generation_config": generation_config,
+                }
+            )
         self._sessions[request_id] = session
         return session
 
+    # ------------------------------------------------------------------
+    # Native-vLLM path (experimental; see __init__).
+    # ------------------------------------------------------------------
+    def _encode_cas_rows(self, ids: torch.Tensor) -> torch.Tensor:
+        """Batched CharAwareSubwordEncoder rows for timeline tokens ([n, H])."""
+        model = self.tts.tts_model
+        mask = torch.ones(1, ids.numel(), dtype=torch.bool, device=ids.device)
+        with torch.inference_mode():
+            return model.embed_subword(ids.unsqueeze(0), mask)[0]
+
+    def _init_native_session(self, request_id: str, info: dict[str, Any], device: torch.device) -> dict[str, Any]:
+        from vllm_omni.model_executor.models.nemotron_voicechat.talker_native import (
+            build_prefill_embeds,
+        )
+
+        timeline = self._timeline(info).to(device)
+        prompt_len = self._logical_prompt_len(info)
+        if prompt_len is None:
+            raise ValueError(
+                "NemotronVoiceChat talker request is missing its logical prompt length "
+                "('nvc_logical_prompt_len' or async meta.num_processed_tokens); the code2wav "
+                "producer cannot trim the prompt-region codes without it."
+            )
+        meta = info.get("meta")
+        codec_streaming = scalar_bool(meta.get("codec_streaming")) if isinstance(meta, dict) else False
+        codec_streaming = codec_streaming or scalar_bool(info.get("meta.codec_streaming"))
+        codec_request_id = meta.get("request_id") if isinstance(meta, dict) else None
+        sync_mode = info.get("nvc_text_timeline") is not None
+        if not sync_mode and not codec_streaming:
+            raise NotImplementedError(
+                "NemotronVoiceChat native talker supports the sync pipeline and duplex "
+                "(resumable) serving; the offline async-chunk pipeline would free-run past "
+                "the received timeline (non-resumable requests cannot park per segment) — "
+                "use the captured-step fast_streaming yaml for offline streaming."
+            )
+        self.tts.set_init_inputs(speaker_name=self._speaker_name)
+        init_inputs = self.tts.get_init_inputs(B=1)
+        init_len = int(init_inputs["code"].shape[1])
+        model = self.tts.tts_model
+        with torch.inference_mode():
+            prefill = build_prefill_embeds(
+                model,
+                code=init_inputs["code"],
+                subword_ids=init_inputs["subword_ids"],
+                subword_mask=init_inputs["subword_mask"],
+                audio_mask=init_inputs["audio_mask"],
+                audio_prompt_latent=init_inputs.get("audio_prompt_latent"),
+                uncond=False,
+            )[0]  # [init_len, H]
+        uncond_stream = None
+        if self._native_guidance:
+            from vllm_omni.model_executor.models.nemotron_voicechat.talker_native import (
+                UncondStepGraph,
+                UncondStream,
+            )
+
+            with torch.inference_mode():
+                uncond_prefill = build_prefill_embeds(
+                    model,
+                    code=init_inputs["code"],
+                    subword_ids=init_inputs["subword_ids"],
+                    subword_mask=init_inputs["subword_mask"],
+                    audio_mask=init_inputs["audio_mask"],
+                    audio_prompt_latent=init_inputs.get("audio_prompt_latent"),
+                    uncond=True,
+                ).to(self._dtype)
+            # Captured StaticCache step when the graph is free (one session at
+            # a time; sync sessions pick the smallest fitting bucket, duplex
+            # takes the largest). Concurrent sessions step eagerly.
+            if self._uncond_graph is None:
+                self._uncond_graph = UncondStepGraph(
+                    model,
+                    hidden_size=int(uncond_prefill.shape[-1]),
+                    dtype=self._dtype,
+                    device=timeline.device,
+                    max_cache_len=self._graph_max_cache_len,
+                )
+            total_positions = (init_len + int(timeline.numel()) + 8) if sync_mode else None
+            if self._uncond_graph.start_session(request_id, uncond_prefill, total_positions):
+                uncond_stream = self._uncond_graph
+            else:
+                uncond_stream = UncondStream(model)
+                uncond_stream.prefill(uncond_prefill)
+        session = {
+            "use_native": True,
+            "timeline": timeline,
+            "cas_table": self._encode_cas_rows(timeline),
+            "prompt_len": int(prompt_len),
+            "init_len": init_len,
+            "prefill_embeds": prefill,
+            "step": 1,  # NeMo's loop starts at t=1
+            "code": init_inputs["code"][:, -1:],
+            # Kept for KV-recompute replay: step t=1's feedback code, before
+            # session["code"] starts tracking the latest sampled frame.
+            "initial_code": init_inputs["code"][:, -1:],
+            # CFG: native_talker_guidance mirrors the stream on the vendored
+            # HF backbone (UncondStream); otherwise a single conditional
+            # stream (measured ASR-equivalent at guidance_scale=0.2; see yaml).
+            "generation_config": self.tts._get_generation_config(self._native_guidance),
+            "uncond_stream": uncond_stream,
+            "uncond_hidden": None,
+            # NeMo's inference_force_speech_silence_on_eos (default on): EOS
+            # text frames feed the silence codec frame back instead of the
+            # sampled codes.
+            "silence_codes": (
+                self.tts.codec_silence_tokens.reshape(1, 1, -1).to(device=timeline.device, dtype=torch.long)
+                if bool(self.tts.cfg.get("inference_force_speech_silence_on_eos", True))
+                else None
+            ),
+            "text_eos_id": int(self.tts.text_eos_id),
+            "codes_rows": [],
+            "pending_step": None,
+            "codec_streaming": codec_streaming,
+            "codec_request_id": codec_request_id,
+            "sync_mode": info.get("nvc_text_timeline") is not None,
+            "upstream_finished": self._upstream_finished(info) or info.get("nvc_text_timeline") is not None,
+        }
+        self._sessions[request_id] = session
+        return session
+
+    def _refresh_native_timeline(self, session: dict[str, Any], info: dict[str, Any], device: torch.device) -> None:
+        """Adopt a longer timeline from the latest chunk; CAS-encode new rows."""
+        latest = self._try_timeline(info)
+        old_len = int(session["timeline"].numel())
+        if latest is not None and latest.numel() > old_len:
+            latest = latest.to(device)
+            new_rows = self._encode_cas_rows(latest[old_len:].to(torch.long))
+            session["timeline"] = latest
+            session["cas_table"] = torch.cat([session["cas_table"], new_rows.to(session["cas_table"].dtype)], dim=0)
+        if self._upstream_finished(info):
+            session["upstream_finished"] = True
+
+    def _native_preprocess(
+        self,
+        input_ids: torch.Tensor,
+        info: dict[str, Any],
+        device: torch.device,
+        span: int,
+        request_id: str,
+    ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
+        from vllm_omni.model_executor.models.nemotron_voicechat.talker_native import (
+            build_decode_embeds,
+        )
+
+        # NOTE: resumable (duplex) segments report _omni_is_prefill=True on
+        # EVERY wake (each appended placeholder extends the prompt), so the
+        # branch is decided by the engine POSITION, never by the prefill flag:
+        # positions [0, init_len) carry the speaker-prompt prefill embeddings,
+        # every later position consumes exactly one timeline step.
+        session = self._sessions.get(request_id)
+        if session is None:
+            session = self._init_native_session(request_id, info, device)
+        else:
+            self._refresh_native_timeline(session, info, device)
+        offset = max(0, int(info.get("_omni_num_computed_tokens", 0) or 0))
+        init_len = int(session["init_len"])
+        step = int(session["step"])
+        total = int(session["timeline"].numel())
+        prefill_head: torch.Tensor | None = None
+        if offset < init_len:
+            # Speaker-prompt prefill region (possibly chunked across steps).
+            if offset + span <= init_len:
+                prefill: torch.Tensor = session["prefill_embeds"]
+                return input_ids, prefill[offset : offset + span].to(self._dtype), {}
+            # A span crossing the boundary is only legal while a KV-cache
+            # recomputation replays already-generated positions in bulk; a
+            # LIVE first timeline step sharing a chunk with the prefill means
+            # a timeline chunk outpaced the speaker prefill.
+            if step <= 1:
+                raise ValueError(
+                    f"NemotronVoiceChat native talker got a span [{offset}, {offset + span}) crossing "
+                    f"the speaker-prompt boundary ({init_len}); the stage prompt must be exactly "
+                    f"{init_len} placeholder tokens (set hf_overrides.talker_init_len / "
+                    f"async_chunk_prewarm_prompt_len to {init_len} in the deploy yaml). A timeline "
+                    "chunk extended this request before its speaker prefill finished — timeline "
+                    "positions are strictly sequential and cannot share a step with the prefill."
+                )
+            prefill_head = session["prefill_embeds"][offset:init_len]
+            span -= init_len - offset
+            offset = init_len
+
+        t_first = offset - init_len + 1
+        t_last = t_first + span - 1
+        if t_last > step:
+            raise RuntimeError(
+                f"NemotronVoiceChat native talker was scheduled timeline steps {t_first}..{t_last} but "
+                f"the session is at step {step}; the per-frame code feedback is strictly sequential. "
+                "The duplex producer drips one new timeline token per chunk — a position past the "
+                "session step means chunks outpaced engine steps (raise the frame budget or "
+                "investigate the engine stall)."
+            )
+        if t_last >= total:
+            raise RuntimeError(f"NemotronVoiceChat native talker stepped past the timeline: step {t_last} of {total}.")
+        replay_last = min(t_last, step - 1)
+        if replay_last >= t_first and len(session["codes_rows"]) < replay_last - 1:
+            raise RuntimeError(
+                f"NemotronVoiceChat native talker cannot replay steps {t_first}..{replay_last}: only "
+                f"{len(session['codes_rows'])} stored code rows."
+            )
+        replay_rows: list[torch.Tensor] = []
+        if replay_last >= t_first:
+            # KV-cache recomputation replays already-generated engine
+            # positions.  Rebuild their decode embeds deterministically from
+            # the stored code history: no resampling (pending_step is only set
+            # for a live step) and the manual unconditional stream is not
+            # advanced (it was never rewound).
+            if not session.get("replaying"):
+                session["replaying"] = True
+                logger.warning(
+                    "NemotronVoiceChat native talker replaying timeline steps %d..%d after KV recompute.",
+                    t_first,
+                    replay_last,
+                )
+            replay_rows = [self._native_replay_embeds(session, rt) for rt in range(t_first, replay_last + 1)]
+        if t_last < step:
+            # Pure replay chunk: no live position this step.
+            embeds = torch.cat(replay_rows, dim=0)
+            if prefill_head is not None:
+                embeds = torch.cat([prefill_head.to(embeds.dtype), embeds], dim=0)
+            return input_ids, embeds.to(self._dtype), {}
+        session.pop("replaying", None)
+        t = step
+        model = self.tts.tts_model
+        with torch.inference_mode():
+            prev_codes = session["code"]
+            if session["silence_codes"] is not None:
+                # NeMo's inference_force_speech_silence_on_eos: on an EOS text
+                # frame the code feedback becomes the silence frame, so the
+                # turn ends in clean silence instead of a trailing artifact.
+                # Tensor-level select — no host sync on the timeline token.
+                is_eos = session["timeline"][t : t + 1].reshape(1, 1, 1) == session["text_eos_id"]
+                prev_codes = torch.where(is_eos, session["silence_codes"].expand_as(prev_codes), prev_codes)
+            embeds = build_decode_embeds(
+                model,
+                prev_codes=prev_codes,
+                subword_embed=session["cas_table"][t : t + 1].unsqueeze(0),
+                uncond=False,
+            )[0]  # [1, H]
+            if session["uncond_stream"] is not None:
+                # Advance the unconditional stream in lockstep: same prev codes,
+                # null conditioning. Its hidden row pairs with this step's
+                # conditional row in make_omni_output.
+                uncond_embed = build_decode_embeds(
+                    model,
+                    prev_codes=prev_codes,
+                    subword_embed=session["cas_table"][t : t + 1].unsqueeze(0),
+                    uncond=True,
+                ).to(self._dtype)
+                session["uncond_hidden"] = session["uncond_stream"].step(uncond_embed)
+        session["pending_step"] = t
+        embeds = embeds.to(self._dtype)
+        if replay_rows:
+            embeds = torch.cat([*replay_rows, embeds], dim=0)
+        if prefill_head is not None:
+            embeds = torch.cat([prefill_head.to(embeds.dtype), embeds], dim=0)
+        return input_ids, embeds, {}
+
+    def _native_replay_embeds(self, session: dict[str, Any], t: int) -> torch.Tensor:
+        """Decode embeds for an already-generated step ``t`` (KV recompute).
+
+        Reconstructs the exact feedback the original step saw: the codes
+        sampled at step ``t-1`` (or the speaker-prompt code for ``t == 1``)
+        with the same silence-on-EOS substitution.
+        """
+        from vllm_omni.model_executor.models.nemotron_voicechat.talker_native import (
+            build_decode_embeds,
+        )
+
+        model = self.tts.tts_model
+        with torch.inference_mode():
+            template: torch.Tensor = session["code"]
+            if t == 1:
+                prev_codes = session["initial_code"]
+            else:
+                prev_codes = session["codes_rows"][t - 2].reshape(template.shape).to(dtype=template.dtype)
+            if session["silence_codes"] is not None:
+                is_eos = session["timeline"][t : t + 1].reshape(1, 1, 1) == session["text_eos_id"]
+                prev_codes = torch.where(is_eos, session["silence_codes"].expand_as(prev_codes), prev_codes)
+            embeds = build_decode_embeds(
+                model,
+                prev_codes=prev_codes,
+                subword_embed=session["cas_table"][t : t + 1].unsqueeze(0),
+                uncond=False,
+            )[0]
+        return embeds.to(self._dtype)
+
+    def _native_make_omni_output(self, hidden: torch.Tensor, **kwargs: Any) -> OmniOutput:
+        infos = kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information") or []
+        spans = kwargs.get("request_token_spans")
+        flag = torch.zeros(hidden.shape[0], 2, dtype=torch.float32, device=hidden.device)
+        # Batch-aligned per-request lists: the generic mm splitter
+        # (to_payload_element) hands list element ``idx`` to request ``idx``
+        # but clones a shared tensor to EVERY request, so a single shared
+        # payload would leak the last request's codes into the others'
+        # streams.  Requests with nothing to ship this step keep placeholder
+        # entries.
+        empty_codes = torch.empty(0, dtype=torch.long)
+        audio_list: list[torch.Tensor] = [empty_codes for _ in infos]
+        prompt_len_list: list[torch.Tensor] = [torch.zeros(1, dtype=torch.long) for _ in infos]
+        streaming_list: list[torch.Tensor] = [torch.zeros(1, dtype=torch.bool) for _ in infos]
+        emitted = False
+
+        def _emit(index: int, session_state: dict[str, Any]) -> None:
+            nonlocal emitted
+            audio_list[index] = torch.cat(session_state["codes_rows"], dim=0)
+            # 1-D tensors on purpose: the wire-payload builder indexes
+            # element.shape[0] (see the eager make_omni_output).
+            prompt_len_list[index] = torch.tensor([session_state["prompt_len"]], dtype=torch.long)
+            streaming_list[index] = torch.tensor([bool(session_state["codec_streaming"])], dtype=torch.bool)
+            emitted = True
+
+        for i, info in enumerate(infos):
+            if not isinstance(info, dict):
+                continue
+            request_id = info.get("request_id")
+            session = self._sessions.get(str(request_id)) if request_id is not None else None
+            if not session or not session.get("use_native"):
+                continue
+            pending = session.get("pending_step")
+            if pending is None:
+                # Prefill forward: no new codes this step, but keep the
+                # request's cumulative history in its own slot so the
+                # splitter never falls back to a shared tensor.
+                if session.get("codes_rows"):
+                    _emit(i, session)
+                continue
+            session["pending_step"] = None
+            if spans is None:
+                raise RuntimeError("NemotronVoiceChat native talker needs request_token_spans from the runner.")
+            start, end = spans[i]
+            model = self.tts.tts_model
+            if session["uncond_stream"] is not None:
+                # CFG: pair the conditional row with this frame's uncond row
+                # ([cond; uncond] order — generate_step chunks in that order)
+                # and let generate_step apply NeMo's hidden-space blends
+                # (lm_head EOS blend + the post-MLP blend inside MoGHead.infer).
+                uncond_hidden = session["uncond_hidden"]
+                if uncond_hidden is None:
+                    raise RuntimeError(
+                        "NemotronVoiceChat native talker CFG desynchronized: no unconditional "
+                        "hidden row for this frame (preprocess/make_omni_output mismatch)."
+                    )
+                session["uncond_hidden"] = None
+                pair = torch.cat(
+                    [hidden[end - 1 : end].reshape(1, 1, -1), uncond_hidden.reshape(1, 1, -1)],
+                    dim=0,
+                ).to(self._dtype)
+                if self._mog_graph is None:
+                    from vllm_omni.model_executor.models.nemotron_voicechat.talker_native import (
+                        MoGStepGraph,
+                    )
+
+                    self._mog_graph = MoGStepGraph(model, session["generation_config"], batch=2)
+                codes = self._mog_graph.run(pair)
+            else:
+                if self._mog_graph is None:
+                    from vllm_omni.model_executor.models.nemotron_voicechat.talker_native import (
+                        MoGStepGraph,
+                    )
+
+                    self._mog_graph = MoGStepGraph(model, session["generation_config"], batch=1)
+                codes = self._mog_graph.run(hidden[end - 1 : end].reshape(1, 1, -1))
+            session["code"] = codes
+            session["codes_rows"].append(codes.reshape(1, -1).to(torch.long))
+            session["step"] = int(pending) + 1
+            # Duplex (codec_streaming) sessions stop per segment when the
+            # received timeline is exhausted and resume with KV intact when
+            # the next thinker frame arrives; sync sessions stop at the end.
+            finished = int(session["step"]) >= int(session["timeline"].numel()) and self._segment_stops_at_exhaustion(
+                session
+            )
+            if finished:
+                flag[start:end, 0] = 1.0
+            _emit(i, session)
+            # Belt and braces: also mirror into the live per-request buffer so
+            # the flush-time full-payload producer can read it from
+            # additional_information as well as from the wire payload.
+            info["codes"] = {"audio": audio_list[i]}
+            info["meta"] = {
+                "nvc_tts_step": int(session["step"]),
+                "nvc_logical_prompt_len": int(session["prompt_len"]),
+                "codec_streaming": bool(session["codec_streaming"]),
+                "request_id": session["codec_request_id"],
+            }
+        outputs: dict[str, Any] = {}
+        if emitted:
+            outputs = {
+                "codes": {"audio": audio_list},
+                "meta": {
+                    "nvc_logical_prompt_len": prompt_len_list,
+                    "codec_streaming": streaming_list,
+                },
+            }
+        return OmniOutput(text_hidden_states=flag, multimodal_outputs=outputs)
+
     def _refresh_session_timeline(self, session: dict[str, Any], info: dict[str, Any], device: torch.device) -> None:
         """Adopt a longer timeline from the latest async-chunk payload."""
         latest = self._try_timeline(info)
         if latest is not None and latest.numel() > int(session["timeline"].numel()):
             session["timeline"] = latest.to(device)
+            if session.get("use_graph"):
+                self._step_graph.extend_timeline(session["timeline"])
         if self._upstream_finished(info):
             session["upstream_finished"] = True
 
@@ -374,15 +885,23 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module):
         session["code"] = code
         session["past_key_values"] = past_key_values
         session["step"] = t + 1
-        # Native codec streaming uses one resumable scheduler segment per
-        # received timeline suffix. The pre-existing offline async-chunk path,
-        # however, must stay alive until the thinker sends its terminal marker;
-        # stopping at an intermediate cumulative snapshot starves Stage 2.
-        finished = (t + 1) >= int(session["timeline"].numel()) and (
-            bool(session["codec_streaming"]) or bool(session["upstream_finished"])
-        )
+        finished = (t + 1) >= int(session["timeline"].numel()) and self._segment_stops_at_exhaustion(session)
         return code.reshape(1, -1).to(torch.long), finished
 
+    @staticmethod
+    def _segment_stops_at_exhaustion(session: dict[str, Any]) -> bool:
+        """Whether exhausting the received timeline should emit the stop token.
+
+        Duplex serving submits RESUMABLE talker requests and marks them with
+        ``codec_streaming``: exhausting the timeline ends the current scheduler
+        segment, and the parked request resumes (TTS state intact) when the
+        next thinker frame arrives. Offline requests are NOT resumable — a
+        stop token would end them permanently — so mid-stream exhaustion emits
+        CONTINUE instead, and the stop waits for the upstream's final chunk
+        (``meta.finished``; sync mode ships the whole timeline up front, so it
+        counts as finished from the start)."""
+        return bool(session["codec_streaming"]) or bool(session["upstream_finished"])
+
     def preprocess(
         self,
         input_ids: torch.Tensor,
@@ -394,6 +913,9 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module):
         span = int(input_ids.shape[0])
         request_id = require_request_id(info, "talker")
 
+        if self._use_native_backbone:
+            return self._native_preprocess(input_ids, info, device, span, request_id)
+
         embeds = torch.zeros((span, self._hidden), device=device, dtype=self._dtype)
         if request_id not in self._sessions:
             # The request id, not vLLM's per-segment prefill flag, is the
@@ -404,47 +926,56 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module):
 
         session = self._sessions[request_id]
         self._refresh_session_timeline(session, info, device)
-        # Sync mode: one NeMo TTS step per engine step (the verified parity
-        # path, unchanged). Async-chunk mode: DRAIN every received timeline
-        # position on each wake. The engine tokens are placeholders
-        # (CONTINUE/STOP) and the codes ride the cumulative payload, so
-        # multiple TTS steps per engine step are safe — and necessary:
+        # DRAIN every received timeline position on each wake, in both modes.
+        # The engine tokens are placeholders (CONTINUE/STOP) and the codes ride
+        # the cumulative payload, so multiple TTS steps per engine step are
+        # safe — and necessary:
+        #   * sync mode ships the WHOLE timeline up front, so the first decode
+        #     wake runs every TTS step back-to-back instead of paying a full
+        #     engine round trip (scheduler + sampler + IPC) per 80 ms frame;
+        #     the codes are computed in the same order as the stepwise loop,
+        #     so the output is bit-identical,
         #   * the save_async background thread can coalesce thinker steps into
         #     one chunk (fewer wakes than new positions), and
         #   * the prompt-region PAD steps (logical_prompt_len - 1 of them) all
         #     become available with chunk 0 and must not cost one upstream
         #     chunk each.
+        # Rows stay on the GPU (tiny [1, Q] longs): a per-frame .cpu() would
+        # stall the dispatch pipeline once per frame; the stage producers move
+        # the final payload to CPU when shipping.
         finished = False
         steps_run = 0
-        new_code_rows: list[torch.Tensor] = []
-        codec_streaming = bool(session["codec_streaming"])
-        while int(session["step"]) < int(session["timeline"].numel()):
-            codes_row, finished = self._step_session(session)
-            if codec_streaming:
-                # One D2H copy per wake is enough; avoid synchronizing after every drained TTS step.
-                new_code_rows.append(codes_row)
-            else:
-                session["codes_rows"].append(codes_row.cpu())
-            steps_run += 1
-            if session["sync_mode"]:
-                break
+        total = int(session["timeline"].numel())
+        if session.get("use_graph"):
+            steps_run = total - int(session["step"])
+            if steps_run > 0:
+                self._step_graph.run_steps(steps_run)
+                session["step"] = total
+                finished = self._segment_stops_at_exhaustion(session)
+        else:
+            while int(session["step"]) < total:
+                codes_row, finished = self._step_session(session)
+                session["codes_rows"].append(codes_row)
+                steps_run += 1
         if steps_run == 0:
             # Zero-progress wake (duplicate/terminal-marker chunk): no TTS step
-            # was run, so never fabricate a frame. Stop the current scheduler
-            # segment only for native codec streaming, or once the ordinary
-            # offline async producer has sent its terminal marker.
-            finished = int(session["step"]) >= int(session["timeline"].numel()) and (
-                bool(session["codec_streaming"]) or bool(session["upstream_finished"])
+            # was run, so never fabricate a frame. Stop only when this session
+            # may end at exhaustion (see _segment_stops_at_exhaustion);
+            # otherwise emit CONTINUE and wait for more upstream input.
+            finished = int(session["step"]) >= int(session["timeline"].numel()) and self._segment_stops_at_exhaustion(
+                session
             )
             if finished:
                 embeds[:, 0] = 1.0
             return input_ids, embeds, {}
-        if codec_streaming:
-            # The code2wav producer owns suffix slicing and therefore requires
-            # cumulative history on every wake. Keep the history on CPU and do
-            # one D2H transfer for all rows produced by this wake.
-            session["codes_rows"].append(torch.cat(new_code_rows, dim=0).cpu())
-        cumulative = torch.cat(session["codes_rows"], dim=0)
+        # The code2wav producers own the suffix slicing and therefore receive
+        # the cumulative history on every wake. Rows stay on the GPU (tiny
+        # [1, Q] longs; a per-step .cpu() would stall dispatch once per frame);
+        # the producers move the shipped payload to CPU.
+        if session.get("use_graph"):
+            cumulative = self._step_graph.codes_rows(1, int(session["step"]))
+        else:
+            cumulative = torch.cat(session["codes_rows"], dim=0)
         if finished:
             embeds[:, 0] = 1.0  # stop flag -> compute_logits emits the stop token
         info_update: dict[str, Any] = {
@@ -466,6 +997,8 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module):
     def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None:
         for request_id in finished_req_ids:
             self._sessions.pop(str(request_id), None)
+            if self._uncond_graph is not None:
+                self._uncond_graph.release(str(request_id))
 
     # ------------------------------------------------------------------
     # Weights.
@@ -497,22 +1030,49 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module):
         tts = DuplexEARTTS({"data": tts_data, "model": tts_model_cfg})
 
         prefix = "tts_model."
+        backbone_prefix = "tts_model.tts_model.backbone."
         state: dict[str, torch.Tensor] = {}
+        backbone_weights: list[tuple[str, torch.Tensor]] = []
         for name, tensor in weights:
             if not name.startswith(prefix):
                 continue
+            if self._use_native_backbone and name.startswith(backbone_prefix):
+                # The backbone runs as a native vLLM Gemma3 model instead.
+                backbone_weights.append((name, tensor))
+                if not self._native_guidance:
+                    continue
+                # CFG keeps a second copy in the vendored module: the
+                # unconditional stream decodes on the HF backbone.
             target = name.removeprefix(prefix)
             if target.startswith("audio_codec."):
                 state[target] = tensor.to(dtype=torch.float32)
             else:
                 state[target] = tensor.to(dtype=self._dtype)
+        if self._use_native_backbone:
+            from vllm_omni.model_executor.models.nemotron_voicechat.talker_native import (
+                iter_backbone_weights,
+            )
+
+            loaded = self.backbone.load_weights(
+                iter_backbone_weights(backbone_weights, hidden_size=self._hidden, dtype=self._dtype)
+            )
+            logger.info("NemotronVoiceChat native talker backbone loaded %d vLLM tensors", len(loaded))
         if not state:
             raise ValueError(
                 "No 'tts_model.*' tensors found in the checkpoint; the NemotronVoiceChat "
                 "unified model.safetensors is required for the talker stage."
             )
+        if self._use_native_backbone and not self._native_guidance:
+            # The HF backbone inside the vendored model is unused on the native
+            # path (only the embedding/CAS/MoG modules are consulted); drop it
+            # before loading so its ~2.4 GiB of fp32 parameters are freed and
+            # its (rerouted) weights are not reported missing. With CFG the
+            # backbone stays: the unconditional stream decodes on it.
+            tts.tts_model.backbone = nn.Identity()
         missing, unexpected = tts.load_state_dict(state, strict=False)
         real_missing = [m for m in missing if not m.startswith("audio_codec.")]
+        if self._use_native_backbone and not self._native_guidance:
+            real_missing = [m for m in real_missing if not m.startswith("tts_model.backbone.")]
         if real_missing:
             raise ValueError(f"NemotronVoiceChat talker is missing TTS weights: {sorted(real_missing)[:10]} ...")
         if unexpected:
@@ -533,4 +1093,40 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module):
             self._speaker_name,
             self._guidance_enabled(),
         )
+        if self._use_cuda_graphs:
+            from vllm_omni.model_executor.models.nemotron_voicechat.talker_graph import (
+                TalkerStepGraph,
+            )
+
+            step_graph = TalkerStepGraph(
+                self.tts,
+                dtype=self._dtype,
+                device=torch.device(device) if not isinstance(device, torch.device) else device,
+                max_cache_len=self._graph_max_cache_len,
+                guidance_enabled=self._guidance_enabled(),
+            )
+            supported, reason = step_graph.config_supported()
+            if supported:
+                self._step_graph = step_graph
+            else:
+                logger.warning(
+                    "NemotronVoiceChat talker CUDA graphs requested but unsupported (%s); "
+                    "using the eager per-frame step.",
+                    reason,
+                )
+        if self._use_native_backbone:
+            # Capture the per-frame MoG sampling graph at load time; a lazy
+            # first-frame capture would otherwise cost the first session a
+            # few hundred ms of first-packet latency.
+            from vllm_omni.model_executor.models.nemotron_voicechat.talker_native import (
+                MoGStepGraph,
+            )
+
+            mog_batch = 2 if self._native_guidance else 1
+            self._mog_graph = MoGStepGraph(
+                self.tts.tts_model,
+                self.tts._get_generation_config(self._native_guidance),
+                batch=mog_batch,
+            )
+            self._mog_graph.ensure_captured(torch.zeros(mog_batch, 1, self._hidden, dtype=self._dtype, device=device))
         return {name for name, _ in self.named_parameters()}

vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_thinker.py

@@ -61,43 +61,74 @@ from vllm_omni.model_executor.models.output_templates import OmniOutput
 
 logger = init_logger(__name__)
 
+# Mel columns kept ahead of the earliest column a perception chunk consumes,
+# so reflect-pad / preemphasis boundary effects at the trimmed rolling-window
+# head (which only reach the first ~2 columns) never touch a consumed column.
+_DUPLEX_MEL_MARGIN_COLS = 8
+
 _LM_ARCHITECTURE = "NemotronHForCausalLM"
 
 
-def slice_perception_streaming_mel(
-    processed_signal: torch.Tensor,
+def _streaming_cfg_at(value: Any, index: int) -> int:
+    return int(value[index] if isinstance(value, list | tuple) else value)
+
+
+def perception_chunk_geometry(
     frame_idx: int,
+    stream_total_cols: int,
     streaming_cfg: Any,
-) -> tuple[torch.Tensor, int]:
-    """Slice one 80 ms mel chunk using NeMo's cache-manager geometry."""
-
-    def _at(value: Any, index: int) -> int:
-        return int(value[index] if isinstance(value, list | tuple) else value)
-
-    chunk_size_first = _at(streaming_cfg.chunk_size, 0)
-    chunk_size = _at(streaming_cfg.chunk_size, 1)
-    shift_size_first = _at(streaming_cfg.shift_size, 0)
-    shift_size = _at(streaming_cfg.shift_size, 1)
-    pre_encode_first = _at(streaming_cfg.pre_encode_cache_size, 0)
-    pre_encode_size = _at(streaming_cfg.pre_encode_cache_size, 1)
-
-    if frame_idx == 0:
-        mel_chunk = processed_signal[:, :, :chunk_size_first]
-        if pre_encode_first > 0:
-            mel_chunk = torch.nn.functional.pad(mel_chunk, (pre_encode_first, 0))
-        return mel_chunk, 0
+) -> tuple[int, int, int]:
+    """Stream-global (cache_start, chunk_start, chunk_end) mel columns for ``frame_idx > 0``."""
+    chunk_size = _streaming_cfg_at(streaming_cfg.chunk_size, 1)
+    shift_size_first = _streaming_cfg_at(streaming_cfg.shift_size, 0)
+    shift_size = _streaming_cfg_at(streaming_cfg.shift_size, 1)
+    pre_encode_size = _streaming_cfg_at(streaming_cfg.pre_encode_cache_size, 1)
 
     chunk_start = shift_size_first + (frame_idx - 1) * shift_size
     chunk_end = chunk_start + chunk_size
     # Match NeMo's tail-boundary rule.  With one 80 ms frame per call this
     # normally leaves the expected trailing STFT columns unused.
     offset = chunk_size - shift_size_first
-    if chunk_end > processed_signal.shape[-1] - offset:
-        chunk_end = processed_signal.shape[-1] - offset
+    if chunk_end > stream_total_cols - offset:
+        chunk_end = stream_total_cols - offset
         chunk_start = chunk_end - chunk_size
-    main_chunk = processed_signal[:, :, chunk_start:chunk_end]
     cache_start = max(0, chunk_start - pre_encode_size)
-    cache_mel = processed_signal[:, :, cache_start:chunk_start]
+    return cache_start, chunk_start, chunk_end
+
+
+def slice_perception_streaming_mel(
+    processed_signal: torch.Tensor,
+    frame_idx: int,
+    streaming_cfg: Any,
+    *,
+    window_col0: int = 0,
+) -> tuple[torch.Tensor, int]:
+    """Slice one 80 ms mel chunk using NeMo's cache-manager geometry.
+
+    ``processed_signal`` may be the mel of a trailing window of the stream
+    instead of the full history; ``window_col0`` is the stream-global column
+    index of its first column (the window must end at the stream tail).
+    """
+    chunk_size_first = _streaming_cfg_at(streaming_cfg.chunk_size, 0)
+    pre_encode_first = _streaming_cfg_at(streaming_cfg.pre_encode_cache_size, 0)
+    pre_encode_size = _streaming_cfg_at(streaming_cfg.pre_encode_cache_size, 1)
+
+    if frame_idx == 0:
+        if window_col0 != 0:
+            raise ValueError("perception stream start requires an untrimmed mel window")
+        mel_chunk = processed_signal[:, :, :chunk_size_first]
+        if pre_encode_first > 0:
+            mel_chunk = torch.nn.functional.pad(mel_chunk, (pre_encode_first, 0))
+        return mel_chunk, 0
+
+    stream_total_cols = window_col0 + int(processed_signal.shape[-1])
+    cache_start, chunk_start, chunk_end = perception_chunk_geometry(frame_idx, stream_total_cols, streaming_cfg)
+    if cache_start < window_col0:
+        raise ValueError(
+            f"perception mel window starts at column {window_col0} but frame {frame_idx} needs column {cache_start}"
+        )
+    main_chunk = processed_signal[:, :, chunk_start - window_col0 : chunk_end - window_col0]
+    cache_mel = processed_signal[:, :, cache_start - window_col0 : chunk_start - window_col0]
     if cache_mel.shape[-1] < pre_encode_size:
         cache_mel = torch.nn.functional.pad(
             cache_mel,
@@ -300,12 +331,14 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState,
 
         # Stateful per-request execution is only validated at batch size 1
         # (the shipped offline scope); fail fast on silent multi-request use.
+        # Multi-session batching: all per-request state (audio frames, fused
+        # prefill embeds, function-token feedback) lives in _sessions keyed by
+        # request id; the runner drives preprocess per request and hands
+        # postprocess each request's own hidden-state slice, and the NemotronH
+        # backbone batches natively on paged KV. Validated at max_num_seqs<=4.
         max_num_seqs = int(getattr(vllm_config.scheduler_config, "max_num_seqs", 1))
-        if max_num_seqs != 1:
-            raise NotImplementedError(
-                f"NemotronVoiceChat thinker supports max_num_seqs=1 only (got {max_num_seqs}); "
-                "multi-request batching of the per-request fusion state is not implemented."
-            )
+        if max_num_seqs > 1:
+            logger.info("NemotronVoiceChat thinker: multi-session batching enabled (max_num_seqs=%d)", max_num_seqs)
 
         # request_id -> per-request state (audio frames, prefill embeds, counters).
         self._sessions: dict[str, dict[str, Any]] = {}
@@ -367,11 +400,16 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState,
         if isinstance(model_outputs, OmniOutput):
             return model_outputs
         multimodal_outputs: dict[str, torch.Tensor] = {}
-        if self._use_function_head and model_outputs.numel():
+        info_dicts = kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information")
+        single_request = isinstance(info_dicts, list) and len(info_dicts) == 1
+        if self._use_function_head and model_outputs.numel() and single_request:
+            # Only a single-request batch can attribute the batch's last row
+            # to a request.  With several concurrent requests this key is
+            # omitted and postprocess computes each request's token from its
+            # own hidden-state slice instead (and applies any forced token).
             with torch.inference_mode():
                 function_token = self.function_head(model_outputs[-1:, :].to(self._dtype)).argmax(dim=-1)
-            info_dicts = kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information")
-            info = info_dicts[0] if isinstance(info_dicts, list) and len(info_dicts) == 1 else None
+            info = info_dicts[0]
             request_id = info.get("request_id") if isinstance(info, dict) else None
             session = self._sessions.get(request_id) if isinstance(request_id, str) else None
             forced_token = session.get("forced_function_token") if isinstance(session, dict) else None
@@ -474,25 +512,51 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState,
                 )
             audio = torch.cat([audio, frame_audio])
         else:
+            # Stream start.  A session normally starts at input sequence 1;
+            # seeing a later sequence means model-side session state was
+            # dropped mid-stream (request abort racing an already-queued step,
+            # or a stage-request re-open).  Restart perception from this frame
+            # instead of failing the engine step.
             if source_input_seq != 1:
-                raise ValueError(
-                    f"Nemotron VoiceChat cache-aware perception must start at input sequence 1, got {source_input_seq}"
+                logger.warning(
+                    "Nemotron VoiceChat perception session reset mid-stream; restarting at input sequence %d.",
+                    source_input_seq,
                 )
+            session["duplex_seq_base"] = source_input_seq - 1
+            session["duplex_window_col0"] = 0
             audio = frame_audio
 
+        frame_idx = source_input_seq - 1 - int(session.get("duplex_seq_base", 0))
+
         encoder = self.perception.encoder
         streaming_cfg = encoder.streaming_cfg
+        window_col0 = int(session.get("duplex_window_col0", 0))
+        if frame_idx > 0:
+            # Bound the rolling audio window: each mel column is local (STFT
+            # centered on hop-aligned samples, no cross-stream normalization),
+            # so only the columns the chunk slice consumes need to exist.
+            # Keep a margin ahead of cache_start so reflect-pad/preemphasis
+            # boundary effects at the trimmed window head never reach a
+            # consumed column.
+            hop = int(getattr(self._streaming_preprocessor.featurizer, "hop_length", 160))
+            total_samples = window_col0 * hop + int(audio.numel())
+            stream_total_cols = total_samples // 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:
+                audio = audio[(desired_col0 - window_col0) * hop :]
+                window_col0 = desired_col0
         with torch.inference_mode():
             processed_signal, _ = self._streaming_preprocessor(
                 input_signal=audio.unsqueeze(0),
                 length=torch.tensor([audio.numel()], device=device, dtype=torch.long),
             )
 
-            frame_idx = source_input_seq - 1
             mel_chunk, drop_extra_pre_encoded = slice_perception_streaming_mel(
                 processed_signal,
                 frame_idx,
                 streaming_cfg,
+                window_col0=window_col0,
             )
             if frame_idx == 0:
                 caches = encoder.get_initial_cache_state(
@@ -538,6 +602,7 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState,
             )
         stable = stable[0].to(self._dtype)
         session["duplex_audio"] = audio
+        session["duplex_window_col0"] = window_col0
         session["perception_cache_last_channel"] = cache_channel
         session["perception_cache_last_time"] = cache_time
         session["perception_cache_last_channel_len"] = cache_channel_len
@@ -599,12 +664,14 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState,
 
         self._sync_forced_function_response(session, runtime_config)
 
-        if source_input_seq <= 1:
-            offset = max(0, int(info.get("_omni_num_computed_tokens", 0) or 0))
-            prefill = session["prefill_embeds"]
-            span = int(input_ids.shape[0])
-            if offset + span > int(prefill.shape[0]):
-                raise ValueError("Nemotron VoiceChat duplex first append exceeds its fused prompt")
+        offset = max(0, int(info.get("_omni_num_computed_tokens", 0) or 0))
+        span = int(input_ids.shape[0])
+        prefill = session["prefill_embeds"]
+        if offset + span <= int(prefill.shape[0]):
+            # Engine positions inside the fused prompt prefill (possibly
+            # chunked).  The position decides, not the duplex input sequence:
+            # after a mid-stream session reset the re-opened request prefills
+            # its prompt again at a sequence far beyond 1.
             return (
                 input_ids,
                 prefill[offset : offset + span].to(self._dtype),
@@ -887,14 +954,18 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState,
         if isinstance(request_id, str):
             session = self._sessions.get(request_id)
             if session is not None:
-                # StreamingUpdate replaces the runner payload at each append;
-                # retain the function channel in model-owned session state.
-                session["func_token"] = int(token.item())
                 forced = session.pop("forced_function_token", None)
                 if isinstance(forced, int):
                     queue = session.get("forced_function_tokens")
                     if not isinstance(queue, list) or not queue or queue.pop(0) != forced:
                         raise RuntimeError("Nemotron VoiceChat forced function-token queue lost alignment")
+                    # Apply the override here too: with several concurrent
+                    # requests make_omni_output omits the shared token, so the
+                    # token above is this request's own computed one.
+                    token = torch.full_like(token, forced)
+                # StreamingUpdate replaces the runner payload at each append;
+                # retain the function channel in model-owned session state.
+                session["func_token"] = int(token.item())
         update: dict[str, Any] = {"nvc_prev_function_token": token}
         if os.environ.get("NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE", "0") == "1":
             existing = kwargs.get("nvc_function_tokens")

vllm_omni/model_executor/models/nemotron_voicechat/talker_graph.py

@@ -0,0 +1,329 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CUDA-graph fast path for the NemotronVoiceChat talker per-frame step.
+
+The eager talker step is launch-bound: one 80 ms frame costs ~3200 kernel
+launches (28-layer Gemma3 backbone + 5 MoG unmasking passes + a 31-level
+depthsum loop) for ~9 ms of actual GPU work, so the eager wall time is
+~25-34 ms/frame. This module captures the WHOLE step into one CUDA graph so
+the drain loop degenerates to ``graph.replay()`` per frame (~12 ms/frame,
+GPU-bound). This mirrors NVIDIA's own deployment of this model, which serves
+EAR-TTS through vLLM decode CUDA graphs with CFG baked into the capture.
+
+What makes the step capturable:
+
+* ``StaticCache`` (preallocated KV, write index via ``cache_position``)
+  instead of a growing HF ``DynamicCache``;
+* the CharAwareSubwordEncoder — dynamic char-length shapes plus per-call
+  ``.cpu().tolist()`` / ``subword_mask.any()`` host syncs — is hoisted OUT of
+  the step: the frame-locked text timeline is known ahead of the audio, so its
+  subword embeddings are precomputed in one batched call per session and the
+  in-graph step just gathers a row (``subword_embeds`` fast path in the
+  vendored model);
+* the code feedback, output row write, step counter, and cache position all
+  advance INSIDE the graph, so a replay is self-contained.
+
+The captured state machine (all device tensors, no host round trips):
+
+    cur_id  = timeline_ids[step_ctr]        # for EOS-silence gating
+    cur_emb = cas_table[step_ctr]           # precomputed CAS row
+    code    = one EAR-TTS step (backbone + MoG, CFG batch=2, StaticCache)
+    code_buf, codes_out[step_ctr] = code    # feedback + output
+    step_ctr += 1; cache_pos += 1
+
+MoG sampling noise draws from the CUDA graph-safe Philox generator, which
+advances per replay — audio is NOT bit-identical to the eager path (nor is the
+eager path bit-stable across process RNG states); the frame-locked text channel
+is unaffected. Sessions longer than ``max_cache_len`` fall back to the eager
+path (the talker decides per session).
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class TalkerStepGraph:
+    """Owns the StaticCache, the per-session buffers, and the captured graph.
+
+    Sessions are strictly serial (the talker enforces ``max_num_seqs=1``); the
+    same buffers and capture are reused across requests via ``start_session``.
+    """
+
+    def __init__(
+        self,
+        tts: Any,
+        *,
+        dtype: torch.dtype,
+        device: torch.device,
+        max_cache_len: int,
+        guidance_enabled: bool,
+    ) -> None:
+        self.tts = tts
+        self.model = tts.tts_model
+        self.device = device
+        self.dtype = dtype
+        self.max_cache_len = int(max_cache_len)
+        self.guidance_enabled = bool(guidance_enabled)
+        self.generation_config = tts._get_generation_config(self.guidance_enabled)
+
+        # Attention in StaticCache mode runs over the FULL padded cache length,
+        # so per-step cost scales with the bucket size (~11 ms/step at 512
+        # positions vs ~20 ms/step at 4096 in fp32 on H100). Sessions with a
+        # known timeline (sync mode) pick the smallest fitting bucket; growing
+        # (async) sessions use the largest, since codes already shipped cannot
+        # be regenerated under a different cache length mid-request.
+        self._bucket_lens = [b for b in (512, 1024, 2048) if b < self.max_cache_len] + [self.max_cache_len]
+        self._buckets: dict[int, dict[str, Any]] = {}
+        self._active_bucket: int | None = None
+        self._broken = False
+        self._init_len: int | None = None
+        self._timeline_len = 0
+
+        hidden = int(self.model.hidden_size)
+        quantizers = int(self.model.config.num_quantizers)
+        with torch.inference_mode():
+            self._timeline_ids = torch.zeros(self.max_cache_len, dtype=torch.long, device=device)
+            self._cas_table = torch.zeros(self.max_cache_len, hidden, dtype=dtype, device=device)
+            self._codes_out = torch.zeros(self.max_cache_len, quantizers, dtype=torch.long, device=device)
+            self._code_buf = torch.zeros(1, 1, quantizers, dtype=torch.long, device=device)
+            self._step_ctr = torch.ones(1, dtype=torch.long, device=device)
+            self._cache_pos = torch.zeros(1, dtype=torch.long, device=device)
+            self._ones_mask = torch.ones(1, 1, dtype=torch.bool, device=device)
+            # Prepared 4D attention mask over the cache slots (True = attend).
+            # Passing a ready 4D mask makes transformers' mask pipeline return
+            # it as-is; the HF-derived mask (and anything else that reads the
+            # cache's python-int lengths) would otherwise bake the capture-time
+            # position into the graph. Per-bucket graphs slice a leading view.
+            self._attn_mask = torch.zeros(1, 1, 1, self.max_cache_len, dtype=torch.bool, device=device)
+        self._cache: Any = None  # the ACTIVE bucket's StaticCache
+
+    # ------------------------------------------------------------------
+    # Support / capacity checks.
+    # ------------------------------------------------------------------
+    def config_supported(self) -> tuple[bool, str]:
+        """Static-config gate for the graph path (independent of the request)."""
+        if self.device.type != "cuda":
+            return False, f"device {self.device} is not CUDA"
+        if not self.guidance_enabled:
+            return False, "classifier-free guidance disabled (capture assumes the CFG batch layout)"
+        if self.tts.cfg.tts_config.context_hidden_size is not None:
+            return False, "context-LM conditioning enabled (prev-subword context is not captured)"
+        if self.model.embed_subword is None:
+            return False, "no CharAwareSubwordEncoder to precompute"
+        sliding_window = getattr(self.model.backbone.config, "sliding_window", None)
+        if sliding_window is not None and self.max_cache_len >= int(sliding_window):
+            # The single prepared mask stands in for both the full and the
+            # sliding mask, which is only equivalent while no position can
+            # fall outside the window.
+            return False, (
+                f"max_cache_len {self.max_cache_len} >= backbone sliding_window {sliding_window}; "
+                "the prepared causal mask would diverge from the sliding mask"
+            )
+        return True, ""
+
+    def fits(self, timeline_len: int) -> bool:
+        init_len = self._init_len if self._init_len is not None else 64
+        return timeline_len <= self.max_cache_len - init_len - 1
+
+    def _bucket_len_for(self, timeline_len: int | None) -> int:
+        """Smallest bucket that fits ``init_len + timeline_len + 1`` positions.
+
+        ``None`` (growing async timeline) selects the largest bucket: codes
+        already shipped cannot be regenerated under a different cache length.
+        """
+        if timeline_len is None:
+            return self._bucket_lens[-1]
+        init_len = self._init_len if self._init_len is not None else 64
+        needed = init_len + timeline_len + 1
+        for bucket in self._bucket_lens:
+            if needed <= bucket:
+                return bucket
+        return self._bucket_lens[-1]
+
+    # ------------------------------------------------------------------
+    # Capture.
+    # ------------------------------------------------------------------
+    def _get_bucket(self, bucket_len: int) -> dict[str, Any]:
+        bucket = self._buckets.get(bucket_len)
+        if bucket is None:
+            from transformers import StaticCache
+
+            bucket = {
+                "cache": StaticCache(
+                    config=self.model.backbone.config,
+                    max_batch_size=2 if self.guidance_enabled else 1,
+                    max_cache_len=bucket_len,
+                    device=self.device,
+                    dtype=self.dtype,
+                ),
+                "graph": None,
+            }
+            self._buckets[bucket_len] = bucket
+        return bucket
+
+    def _init_forward(self) -> None:
+        """Eager speaker-prompt prefill into the StaticCache (per session)."""
+        init_inputs = self.tts.get_init_inputs(B=1)
+        init_len = int(init_inputs["code"].shape[1])
+        if self._init_len is None:
+            self._init_len = init_len
+        elif init_len != self._init_len:
+            raise RuntimeError(
+                f"NemotronVoiceChat talker graph: init prompt length changed ({self._init_len} -> {init_len}); "
+                "the captured graph assumes a fixed speaker prompt."
+            )
+        init_inputs.update(
+            {
+                "use_cache": True,
+                "past_key_values": self._cache,
+                "guidance_enabled": self.guidance_enabled,
+                "cache_position": torch.arange(init_len, device=self.device),
+            }
+        )
+        self.model(**init_inputs)
+        self._cache_pos.fill_(init_len)
+        self._code_buf.copy_(init_inputs["code"][:, -1:])
+        # Seed the prepared mask: the speaker-prompt slots are attendable, the
+        # per-step slot bits are set inside the captured step.
+        self._attn_mask.zero_()
+        self._attn_mask[..., :init_len] = True
+
+    def _one_step(self) -> None:
+        """The captured body: one EAR-TTS frame step + in-graph state advance."""
+        cur_id = self._timeline_ids.index_select(0, self._step_ctr).reshape(1, 1)
+        cur_emb = self._cas_table.index_select(0, self._step_ctr).reshape(1, 1, -1)
+        # Open this step's own cache slot before attending (a token attends to
+        # itself), then run the step with fully tensor-driven mask/positions.
+        self._attn_mask.index_fill_(3, self._cache_pos, True)
+        code, _ = self.tts.infer_codes_one_step(
+            current_subword_id=cur_id,
+            # prev_subword_id only feeds the context-LM channel, which
+            # config_supported() guarantees is disabled.
+            prev_subword_id=cur_id,
+            current_subword_mask=self._ones_mask,
+            prev_audio_tokens=self._code_buf,
+            past_key_values=self._cache,
+            guidance_enabled=self.guidance_enabled,
+            generation_config=self.generation_config,
+            ignore_eos_flag_stop=True,
+            current_subword_embed=cur_emb,
+            cache_position=self._cache_pos,
+            attention_mask=self._attn_mask[..., : self._active_bucket],
+            position_ids=self._cache_pos.reshape(1, 1),
+        )
+        self._code_buf.copy_(code)
+        self._codes_out.index_copy_(0, self._step_ctr, code.reshape(1, -1))
+        self._step_ctr.add_(1)
+        self._cache_pos.add_(1)
+
+    def ensure_captured(self, timeline_len: int | None = None) -> bool:
+        """Warm up and capture the bucket for this timeline (throwaway session)."""
+        if self._broken:
+            return False
+        bucket_len = self._bucket_len_for(timeline_len)
+        bucket = self._get_bucket(bucket_len)
+        if bucket["graph"] is not None:
+            return True
+        try:
+            with torch.inference_mode():
+                # Throwaway session: buffer VALUES are irrelevant for capture,
+                # only shapes/addresses matter. PAD-filled timeline rows and a
+                # zero CAS table are fine.
+                self._cache = bucket["cache"]
+                self._active_bucket = bucket_len
+                self._cache.reset()
+                self._step_ctr.fill_(1)
+                self._cache_pos.zero_()
+                self._init_forward()
+                side = torch.cuda.Stream()
+                side.wait_stream(torch.cuda.current_stream())
+                with torch.cuda.stream(side):
+                    for _ in range(3):
+                        self._one_step()
+                torch.cuda.current_stream().wait_stream(side)
+                graph = torch.cuda.CUDAGraph()
+                with torch.cuda.graph(graph):
+                    self._one_step()
+                bucket["graph"] = graph
+            logger.info(
+                "NemotronVoiceChat talker CUDA graph captured (bucket=%d, init_len=%d)",
+                bucket_len,
+                self._init_len,
+            )
+            return True
+        except Exception:
+            self._broken = True
+            logger.exception(
+                "NemotronVoiceChat talker CUDA graph capture failed; falling back to the eager per-frame step."
+            )
+            return False
+
+    # ------------------------------------------------------------------
+    # Session lifecycle.
+    # ------------------------------------------------------------------
+    def start_session(self, timeline: torch.Tensor, *, expected_len: int | None) -> None:
+        """Reset all captured state and prefill for a new request.
+
+        ``expected_len`` is the FINAL timeline length when known up front
+        (sync mode ships the whole timeline; frame-locked async producers ship
+        meta.expected_total_tokens) and selects the smallest fitting bucket;
+        ``None`` (unbounded growth, e.g. duplex) selects the largest bucket.
+        """
+        bucket_len = self._bucket_len_for(expected_len)
+        bucket = self._buckets[bucket_len]
+        assert bucket["graph"] is not None, "ensure_captured() must succeed before start_session()"
+        with torch.inference_mode():
+            self._cache = bucket["cache"]
+            self._active_bucket = bucket_len
+            self._cache.reset()
+            self._step_ctr.fill_(1)
+            self._cache_pos.zero_()
+            self._codes_out.zero_()
+            self._timeline_len = 0
+            self._init_forward()
+            self.extend_timeline(timeline)
+
+    def extend_timeline(self, timeline: torch.Tensor) -> None:
+        """Adopt a (possibly longer) timeline; CAS-encode only the new rows."""
+        new_len = int(timeline.numel())
+        if new_len <= self._timeline_len:
+            return
+        if not self.fits(new_len):
+            raise RuntimeError(
+                f"NemotronVoiceChat talker graph: timeline of {new_len} positions exceeds the "
+                f"captured capacity ({self.max_cache_len - (self._init_len or 0) - 1} frames). Raise "
+                "hf_overrides.talker_graph_max_cache_len or disable hf_overrides.use_talker_cuda_graphs."
+            )
+        active = self._active_bucket
+        if active is not None and (self._init_len or 0) + new_len + 1 > active:
+            raise RuntimeError(
+                f"NemotronVoiceChat talker graph: timeline of {new_len} positions exceeds the active "
+                f"cache bucket ({active} positions). Raise hf_overrides.talker_graph_max_cache_len or "
+                "disable hf_overrides.use_talker_cuda_graphs."
+            )
+        start = self._timeline_len
+        with torch.inference_mode():
+            new_ids = timeline[start:new_len].to(device=self.device, dtype=torch.long).reshape(-1)
+            self._timeline_ids[start:new_len] = new_ids
+            mask = torch.ones(1, new_ids.numel(), dtype=torch.bool, device=self.device)
+            # One batched CAS call per chunk of new tokens (the encoder's host
+            # syncs and dynamic char shapes stay OUTSIDE the captured step).
+            embeds = self.model.embed_subword(new_ids.unsqueeze(0), mask)[0]
+            self._cas_table[start:new_len] = embeds.to(self.dtype)
+        self._timeline_len = new_len
+
+    def run_steps(self, num_steps: int) -> None:
+        graph = self._buckets[self._active_bucket]["graph"]
+        for _ in range(num_steps):
+            graph.replay()
+
+    def codes_rows(self, first_step: int, last_step: int) -> torch.Tensor:
+        """Codes for NeMo steps ``first_step..last_step-1`` (row t = step t)."""
+        with torch.inference_mode():
+            return self._codes_out[first_step:last_step].clone()

vllm_omni/model_executor/models/nemotron_voicechat/talker_native.py

@@ -0,0 +1,429 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Native-vLLM building blocks for the NemotronVoiceChat talker.
+
+The eager/CUDA-graph talker executes the vendored HF Gemma3 backbone inside the
+per-request ``preprocess`` hook, next to (not inside) the vLLM engine. The
+native path (``hf_overrides.use_native_talker``) instead runs the backbone as a
+real vLLM model — PagedAttention KV cache (capacity and per-step attention cost
+scale with the ACTUAL session length, no StaticCache buckets), vLLM's own decode
+CUDA graphs, and, later, multi-session batching. This mirrors NVIDIA's own
+deployment, whose EAR-TTS is served by a vLLM engine.
+
+This module hosts the pieces that are independent of the talker class:
+
+* ``synthesize_backbone_config``: the checkpoint stores the backbone as a NeMo
+  ``tts_config.backbone_config`` dict; vLLM needs a real ``Gemma3TextConfig``.
+  The embedding table is never used (the fused embeddings are built from
+  codes + text conditioning), so a tiny dummy vocab keeps it negligible.
+* ``iter_backbone_weights``: remaps the checkpoint's
+  ``tts_model.tts_model.backbone.*`` HF-Gemma3 tensors onto the names vLLM's
+  ``Gemma3ForCausalLM`` loader expects, and fabricates the dummy
+  ``model.embed_tokens.weight``.
+* ``build_prefill_embeds`` / ``build_decode_embeds``: the per-token fused input
+  embeddings, numerically verified against the vendored
+  ``RVQEARTTSModel.forward`` embedding section for both CFG roles (the
+  speaker-prompt prefill must use the checkpoint's pre-baked
+  ``audio_prompt_latent``, not the frozen projection).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+
+_DUMMY_VOCAB_SIZE = 8
+
+
+def synthesize_backbone_config(tts_config: dict[str, Any]) -> Any:
+    """Build a ``Gemma3TextConfig`` from the checkpoint's NeMo backbone dict."""
+    from transformers import AutoConfig
+
+    backbone_type = tts_config.get("backbone_type")
+    if backbone_type != "gemma3_text":
+        raise NotImplementedError(
+            f"NemotronVoiceChat native talker supports a gemma3_text backbone only (got {backbone_type!r})."
+        )
+    backbone_cfg = dict(tts_config.get("backbone_config") or {})
+    backbone_cfg.pop("use_cache", None)
+    config = AutoConfig.for_model(
+        backbone_type,
+        **backbone_cfg,
+        vocab_size=_DUMMY_VOCAB_SIZE,
+        tie_word_embeddings=True,
+    )
+    return config
+
+
+def iter_backbone_weights(
+    weights: Iterable[tuple[str, torch.Tensor]],
+    *,
+    hidden_size: int,
+    dtype: torch.dtype,
+) -> Iterable[tuple[str, torch.Tensor]]:
+    """Yield vLLM-Gemma3 weights from the unified checkpoint stream.
+
+    Checkpoint names: ``tts_model.tts_model.backbone.<hf gemma3 names>``.
+    vLLM ``Gemma3ForCausalLM`` expects ``model.<hf gemma3 names>`` plus an
+    embedding table (tied lm_head); the embedding is never consulted (fused
+    embeddings are injected via ``inputs_embeds``), so a dummy is fabricated.
+    """
+    prefix = "tts_model.tts_model.backbone."
+    for name, tensor in weights:
+        if name.startswith(prefix):
+            yield "model." + name.removeprefix(prefix), tensor.to(dtype=dtype)
+    yield (
+        "model.embed_tokens.weight",
+        torch.zeros(_DUMMY_VOCAB_SIZE, hidden_size, dtype=dtype),
+    )
+
+
+def build_prefill_embeds(
+    model: Any,
+    *,
+    code: torch.Tensor,
+    subword_ids: torch.Tensor,
+    subword_mask: torch.Tensor,
+    audio_mask: torch.Tensor,
+    audio_prompt_latent: torch.Tensor | None,
+    uncond: bool,
+) -> torch.Tensor:
+    """Speaker-prompt prefill embeddings for one CFG role.
+
+    Mirrors the ``audio_mask`` inference branch of the vendored
+    ``RVQEARTTSModel.forward`` (verified to ~4e-6 max abs diff against the
+    embeddings the vendored forward feeds its backbone). ``model`` is the
+    vendored ``RVQEARTTSModel``. Returns ``[1, T, H]``.
+    """
+    shifted_code = F.pad(code[:, :-1], [0, 0, 1, 0])
+    code_embed = model.depthsum_embedding(shifted_code)
+    bos_mask = (audio_mask & (~F.pad(audio_mask[:, :-1], [1, 0]))).unsqueeze(-1)
+    pre_bos_mask = bos_mask.cumsum(dim=1) == 0
+    code_embed = model.embed_code(code_embed)
+    if model.config.get("use_audio_prompt_frozen_projection", False):
+        if audio_prompt_latent is None:
+            weight = model.audio_prompt_projection_W.to(code_embed.device, code_embed.dtype)
+            audio_prompt_latent = torch.nn.functional.linear(code_embed, weight.T)
+        code_embed = torch.where(pre_bos_mask, audio_prompt_latent, code_embed)
+    code_embeds = code_embed + bos_mask * model.bos_emb
+    flag = torch.full((code.size(0), 1, 1), bool(uncond), dtype=torch.bool, device=code.device)
+    cond = model._prepare_conditioning(None, subword_ids, subword_mask, flag)
+    if model.config.use_gated_fusion_for_text_audio:
+        return model.gated_fusion_audio_text(code_embeds, cond)
+    return code_embeds + cond
+
+
+class MoGStepGraph:
+    """One captured CUDA graph for the per-frame MoG sampling step.
+
+    On the native path the backbone runs under vLLM's own decode CUDA graphs,
+    but the MoG iterative unmasking (``RVQEARTTSModel.generate_step``: 5
+    effective MoG-head passes + a 31-level depthsum quantization loop) would
+    otherwise run eagerly in ``make_omni_output`` — ~1000 kernel launches for
+    ~2 ms of GPU work per 80 ms frame. All shapes are static and the unmask
+    schedule is Python ints, so the whole step captures into one graph whose
+    replay is a single launch. MoG noise draws from the CUDA-graph-safe Philox
+    stream (fresh randomness per replay, not bit-identical to eager).
+
+    ``batch`` is the leading hidden dim (1 for the single conditional stream;
+    2 once paired-CFG lands). Capture failures fall back to eager permanently.
+    """
+
+    def __init__(self, model: Any, generation_config: dict[str, Any], *, batch: int = 1) -> None:
+        self.model = model
+        self.generation_config = dict(generation_config)
+        self.batch = int(batch)
+        self._graph: torch.cuda.CUDAGraph | None = None
+        self._broken = False
+        self._hidden_buf: torch.Tensor | None = None
+        self._code_buf: torch.Tensor | None = None
+
+    def _step(self) -> None:
+        codes, _, _ = self.model.generate_step(self._hidden_buf, ignore_eos_flag_stop=True, **self.generation_config)
+        self._code_buf.copy_(codes)
+
+    def ensure_captured(self, hidden: torch.Tensor) -> bool:
+        if self._graph is not None:
+            return True
+        if self._broken:
+            return False
+        try:
+            with torch.inference_mode():
+                self._hidden_buf = hidden.clone()
+                # Under CFG (guidance_scale set) generate_step chunks the
+                # [cond; uncond] pair and returns codes for the effective
+                # batch only.
+                out_batch = self.batch
+                if self.generation_config.get("guidance_scale") is not None:
+                    out_batch = max(self.batch // 2, 1)
+                self._code_buf = torch.zeros(
+                    out_batch, 1, int(self.model.config.num_quantizers), dtype=torch.long, device=hidden.device
+                )
+                side = torch.cuda.Stream()
+                side.wait_stream(torch.cuda.current_stream())
+                with torch.cuda.stream(side):
+                    for _ in range(3):
+                        self._step()
+                torch.cuda.current_stream().wait_stream(side)
+                graph = torch.cuda.CUDAGraph()
+                with torch.cuda.graph(graph):
+                    self._step()
+                self._graph = graph
+            return True
+        except Exception:
+            self._broken = True
+            self._graph = None
+            import logging
+
+            logging.getLogger(__name__).exception(
+                "NemotronVoiceChat native talker: MoG step graph capture failed; sampling eagerly."
+            )
+            return False
+
+    def run(self, hidden: torch.Tensor) -> torch.Tensor:
+        """MoG-sample one frame from ``hidden`` ([batch, 1, H]) -> codes [batch, 1, Q]."""
+        if self.ensure_captured(hidden):
+            with torch.inference_mode():
+                self._hidden_buf.copy_(hidden)
+                self._graph.replay()
+                return self._code_buf.clone()
+        with torch.inference_mode():
+            codes, _, _ = self.model.generate_step(hidden, ignore_eos_flag_stop=True, **self.generation_config)
+            return codes
+
+
+def build_decode_embeds(
+    model: Any,
+    *,
+    prev_codes: torch.Tensor,
+    subword_embed: torch.Tensor,
+    uncond: bool,
+) -> torch.Tensor:
+    """One decode-step fused embedding for one CFG role (``[1, 1, H]``).
+
+    ``subword_embed`` is the precomputed CharAwareSubwordEncoder row for the
+    current timeline token (host-sync free; see the CAS table precompute).
+    """
+    code_embeds = model.embed_code(model.depthsum_embedding(prev_codes))
+    flag = torch.full((prev_codes.size(0), 1, 1), bool(uncond), dtype=torch.bool, device=prev_codes.device)
+    cond = model._prepare_conditioning(None, None, None, flag, subword_embeds=subword_embed)
+    if model.config.use_gated_fusion_for_text_audio:
+        return model.gated_fusion_audio_text(code_embeds, cond)
+    return code_embeds + cond
+
+
+class UncondStream:
+    """The unconditional CFG stream, run on the vendored HF Gemma3 backbone.
+
+    NeMo applies classifier-free guidance by duplicating every backbone input
+    into a [cond; uncond] batch (``RVQEARTTSModel.forward``) whose second half
+    replaces the whole conditioning sum with ``null_emb``, then blending the
+    two hidden rows inside ``generate_step`` (lm_head EOS blend and the
+    post-MLP blend inside ``MoGHead.infer``).
+
+    On the native path the conditional stream's KV lives inside vLLM's paged
+    cache, where a second same-length sequence cannot ride along without
+    scheduler-level pair locking. This class instead mirrors the stream
+    position-for-position on the vendored HF backbone with its own
+    ``DynamicCache`` — one extra [1, 1, H] decode step per frame. Both streams
+    are fed the SAME sampled code row each frame (prev-codes conditioning is
+    shared in NeMo's layout; only the text conditioning differs).
+    """
+
+    def __init__(self, model: Any) -> None:
+        self.model = model
+        self.cache: Any = None
+        self.pos = 0
+
+    @torch.inference_mode()
+    def prefill(self, embeds: torch.Tensor) -> None:
+        """Consume the speaker-prompt prefill embeddings (``[1, T, H]``)."""
+        from transformers import DynamicCache
+
+        self.cache = DynamicCache()
+        length = int(embeds.shape[1])
+        position_ids = torch.arange(length, device=embeds.device).unsqueeze(0)
+        outputs = self.model.backbone(
+            inputs_embeds=embeds,
+            attention_mask=None,
+            position_ids=position_ids,
+            past_key_values=self.cache,
+            use_cache=True,
+        )
+        self.cache = outputs.past_key_values
+        self.pos = length
+
+    @torch.inference_mode()
+    def step(self, embed: torch.Tensor) -> torch.Tensor:
+        """One decode step (``[1, 1, H]`` in) -> uncond hidden (``[1, 1, H]``)."""
+        if self.cache is None:
+            raise RuntimeError("UncondStream.step called before prefill")
+        position_ids = torch.tensor([[self.pos]], device=embed.device)
+        outputs = self.model.backbone(
+            inputs_embeds=embed,
+            attention_mask=None,
+            position_ids=position_ids,
+            past_key_values=self.cache,
+            use_cache=True,
+        )
+        self.cache = outputs.past_key_values
+        self.pos += 1
+        return outputs.last_hidden_state
+
+
+class UncondStepGraph:
+    """Captured per-frame decode step for the unconditional CFG stream.
+
+    Backbone-only StaticCache stepping with the same capture hygiene as
+    ``talker_graph.TalkerStepGraph``: a prepared 4D boolean attention mask
+    (triggers the HF early-exit so no python-int cache length is baked into
+    the graph), explicit ``position_ids``/``cache_position`` tensors, and the
+    in-graph slot-open ``index_fill_``. One session owns the graph at a time
+    (StaticCache addresses are capture-bound); concurrent sessions fall back
+    to the eager ``UncondStream``.
+    """
+
+    _BUCKETS = (512, 1024, 2048, 4096)
+
+    def __init__(
+        self,
+        model: Any,
+        *,
+        hidden_size: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        max_cache_len: int = 4096,
+    ) -> None:
+        self.model = model
+        self.dtype = dtype
+        self.device = device
+        self.max_cache_len = int(max_cache_len)
+        self._buckets: dict[int, dict[str, Any]] = {}
+        self._broken = False
+        self.owner: str | None = None
+        self._embed_buf = torch.zeros(1, 1, hidden_size, dtype=dtype, device=device)
+        self._hidden_buf = torch.zeros(1, 1, hidden_size, dtype=dtype, device=device)
+        self._cache_pos = torch.zeros(1, dtype=torch.long, device=device)
+        self._attn_mask = torch.zeros(1, 1, 1, self.max_cache_len, dtype=torch.bool, device=device)
+        self._cache: Any = None
+        self._active_bucket = 0
+
+    def bucket_for(self, total_positions: int | None) -> int | None:
+        """Smallest bucket holding ``total_positions``; None if it cannot fit."""
+        for bucket in self._BUCKETS:
+            if bucket > self.max_cache_len:
+                break
+            if total_positions is not None and total_positions <= bucket:
+                return bucket
+        largest = max(b for b in self._BUCKETS if b <= self.max_cache_len)
+        return largest if total_positions is None or total_positions <= largest else None
+
+    def _get_bucket(self, bucket_len: int) -> dict[str, Any]:
+        state = self._buckets.get(bucket_len)
+        if state is None:
+            from transformers import StaticCache
+
+            state = {
+                "cache": StaticCache(
+                    config=self.model.backbone.config,
+                    max_batch_size=1,
+                    max_cache_len=bucket_len,
+                    device=self.device,
+                    dtype=self.dtype,
+                ),
+                "graph": None,
+            }
+            self._buckets[bucket_len] = state
+        return state
+
+    def _one_step(self) -> None:
+        self._attn_mask.index_fill_(3, self._cache_pos, True)
+        outputs = self.model.backbone(
+            inputs_embeds=self._embed_buf,
+            attention_mask=self._attn_mask[..., : self._active_bucket],
+            position_ids=self._cache_pos.reshape(1, 1),
+            past_key_values=self._cache,
+            use_cache=True,
+            cache_position=self._cache_pos,
+        )
+        self._hidden_buf.copy_(outputs.last_hidden_state)
+        self._cache_pos.add_(1)
+
+    def start_session(self, request_id: str, prefill_embeds: torch.Tensor, total_positions: int | None) -> bool:
+        """Claim the graph for one session and prefill its cache. False -> use eager."""
+        if self._broken or self.owner is not None:
+            return False
+        bucket_len = self.bucket_for(total_positions)
+        if bucket_len is None:
+            return False
+        try:
+            with torch.inference_mode():
+                state = self._get_bucket(bucket_len)
+                self._cache = state["cache"]
+                self._active_bucket = bucket_len
+                self._cache.reset()
+                init_len = int(prefill_embeds.shape[1])
+                self.model.backbone(
+                    inputs_embeds=prefill_embeds.to(self.dtype),
+                    position_ids=torch.arange(init_len, device=self.device).unsqueeze(0),
+                    past_key_values=self._cache,
+                    use_cache=True,
+                    cache_position=torch.arange(init_len, device=self.device),
+                )
+                self._attn_mask.zero_()
+                self._attn_mask[..., :init_len] = True
+                self._cache_pos.fill_(init_len)
+                if state["graph"] is None:
+                    side = torch.cuda.Stream()
+                    side.wait_stream(torch.cuda.current_stream())
+                    saved_pos = int(self._cache_pos.item())
+                    with torch.cuda.stream(side):
+                        for _ in range(3):
+                            self._one_step()
+                    torch.cuda.current_stream().wait_stream(side)
+                    graph = torch.cuda.CUDAGraph()
+                    with torch.cuda.graph(graph):
+                        self._one_step()
+                    state["graph"] = graph
+                    # Rewind the warmup/capture side effects: re-prefill so the
+                    # cache holds only the speaker prompt again.
+                    self._cache.reset()
+                    self.model.backbone(
+                        inputs_embeds=prefill_embeds.to(self.dtype),
+                        position_ids=torch.arange(init_len, device=self.device).unsqueeze(0),
+                        past_key_values=self._cache,
+                        use_cache=True,
+                        cache_position=torch.arange(init_len, device=self.device),
+                    )
+                    self._attn_mask.zero_()
+                    self._attn_mask[..., :init_len] = True
+                    self._cache_pos.fill_(saved_pos)
+            self.owner = request_id
+            return True
+        except Exception:
+            self._broken = True
+            import logging
+
+            logging.getLogger(__name__).exception(
+                "NemotronVoiceChat native talker: uncond step graph capture failed; stepping eagerly."
+            )
+            return False
+
+    def step(self, embed: torch.Tensor) -> torch.Tensor:
+        """One captured decode step: embed [1,1,H] -> uncond hidden [1,1,H]."""
+        if int(self._cache_pos.item()) >= self._active_bucket:
+            raise RuntimeError(
+                f"NemotronVoiceChat uncond CFG stream exceeded its StaticCache bucket "
+                f"({self._active_bucket} positions)."
+            )
+        with torch.inference_mode():
+            self._embed_buf.copy_(embed.to(self.dtype))
+            self._buckets[self._active_bucket]["graph"].replay()
+            return self._hidden_buf.clone()
+
+    def release(self, request_id: str) -> None:
+        if self.owner == request_id:
+            self.owner = None

vllm_omni/model_executor/stage_input_processors/nemotron_voicechat.py

@@ -79,6 +79,7 @@ def thinker2talker_token_only(
     source_outputs: list,
     prompt: Any = None,
     _requires_multimodal_data: bool = False,
+    next_stage_hf_config: Any = None,
 ) -> list:
     """Build the talker's engine inputs from the thinker's finished outputs.
 
@@ -91,6 +92,16 @@ def thinker2talker_token_only(
     from vllm_omni.inputs.data import OmniTokensPrompt
 
     del prompt, _requires_multimodal_data
+    # Default: one placeholder token (timeline position 0; the NeMo TTS loop
+    # starts at t=1) — the vendored talker paths ignore the vLLM positions.
+    # The NATIVE talker (hf_overrides.use_native_talker) instead maps vLLM
+    # positions onto real PagedAttention KV: its prompt must be exactly the
+    # speaker-prompt length (hf_overrides.talker_init_len; 37 for the shipped
+    # "Aria" speaker) so the prefill covers the speaker-init embeddings — the
+    # talker validates the value and reports the right one on mismatch.
+    talker_prompt_len = 1
+    if next_stage_hf_config is not None and bool(getattr(next_stage_hf_config, "use_native_talker", False)):
+        talker_prompt_len = max(int(getattr(next_stage_hf_config, "talker_init_len", 37)), 1)
     inputs: list = []
     for thinker_output in source_outputs:
         if not getattr(thinker_output, "finished", False):
@@ -107,11 +118,10 @@ def thinker2talker_token_only(
         timeline = [pad_id] * logical_prompt_len + [int(t) for t in generated]
         inputs.append(
             OmniTokensPrompt(
-                # The vLLM-side prompt is one placeholder (timeline position 0;
-                # the NeMo TTS loop starts at t=1). Each decode step consumes
-                # one timeline position from additional_information, so the
-                # talker needs len(timeline) - 1 decode steps.
-                prompt_token_ids=[0],
+                # Each decode step consumes one timeline position from
+                # additional_information, so the talker needs
+                # len(timeline) - 1 decode steps.
+                prompt_token_ids=[0] * talker_prompt_len,
                 additional_information={
                     "nvc_text_timeline": timeline,
                     "nvc_logical_prompt_len": logical_prompt_len,
@@ -154,11 +164,14 @@ def thinker2talker_async_chunk(
     generated_now = [int(token_id) for token_id in (getattr(request, "output_token_ids", None) or [])]
     prompt_ids = list(getattr(request, "prompt_token_ids", None) or [])
     if is_duplex:
-        runtime_config = duplex_info.get("runtime_config")
-        runtime_prompt = runtime_config.get("nvc_prompt_token_ids") if isinstance(runtime_config, dict) else None
-        if not isinstance(runtime_prompt, list | tuple) or not runtime_prompt:
-            raise ValueError("Nemotron VoiceChat duplex stage transfer lacks nvc_prompt_token_ids")
-        logical_prompt_len = len(runtime_prompt)
+        # NeMo's realtime pipeline never feeds the system prompt to the TTS
+        # model — the prompt conditions the thinker's KV only, and the talker
+        # sees just the speaker prefill plus the per-frame text tokens. A
+        # single leading PAD anchors timeline position 0 (the NeMo per-frame
+        # loop reads timeline[t-1] as the previous subword at t=1). The
+        # OFFLINE pipelines keep the full PAD prompt region for bit parity
+        # with NeMo's offline recipe.
+        logical_prompt_len = 1
         previous = list(state.get("nvc_duplex_text_tokens", ())) if isinstance(state, dict) else []
         if has_step_snapshot:
             generated = previous + generated_delta
@@ -178,24 +191,69 @@ def thinker2talker_async_chunk(
         pad_id = int(reported_pad)
     timeline = [pad_id] * logical_prompt_len + generated
 
+    # The NATIVE talker (connector extra ``native_talker: true``, matching the
+    # stage's hf_overrides.use_native_talker) maps engine positions 1:1 onto
+    # timeline steps and consumes them strictly sequentially (per-frame code
+    # feedback), so a chunk must never make more than ONE new timeline step
+    # schedulable at once when the stage batches prompt tokens
+    # (max_num_batched_tokens > 1 lets the speaker prefill run as one chunked
+    # prefill). The vendored talker drains chunks model-side instead.
+    connector = getattr(transfer_manager, "connector", None)
+    raw_cfg = getattr(connector, "config", {}) or {}
+    cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
+    native_talker = bool(cfg.get("native_talker", False))
+
     # Skip no-progress wakeups (same length as the last emitted chunk).
     last_len = int(state.get("nvc_timeline_len", -1)) if isinstance(state, dict) else -1
-    if len(timeline) <= last_len and not is_finished:
+    if is_duplex and native_talker:
+        # Drip: ship at most one NEW timeline token per producer invocation.
+        # The producer fires on every thinker step (one per 80 ms frame, and
+        # the duplex thinker never stops emitting), so a save-thread burst
+        # that coalesced K thinker steps drains within K subsequent wakes —
+        # audio trails text by at most one frame per coalesced token.
+        ship_len = min(len(timeline), max(last_len, logical_prompt_len) + 1)
+        if ship_len <= last_len and not is_finished:
+            return None
+        timeline = timeline[:ship_len]
+    elif len(timeline) <= last_len and not is_finished:
         return None
     transfer_manager.request_payload[request_id] = {
         "nvc_timeline_len": len(timeline),
         "nvc_duplex_text_tokens": generated if is_duplex else (),
     }
 
+    # Offline timelines are frame-locked: the thinker emits exactly max_tokens
+    # tokens (ignore_eos), so the FINAL timeline length is known up front.
+    # Shipping it lets the talker's CUDA-graph fast path pick the smallest
+    # fitting StaticCache bucket (per-step attention cost scales with the
+    # bucket size) instead of assuming an unbounded session. Duplex sessions
+    # really are unbounded — no expectation is shipped for them.
+    expected_total = None
+    if not is_duplex:
+        max_tokens = getattr(getattr(request, "sampling_params", None), "max_tokens", None)
+        if max_tokens is not None:
+            expected_total = logical_prompt_len + int(max_tokens)
+
+    new_positions = 1
+    if native_talker:
+        # A timeline of length L needs L-1 engine positions (the NeMo loop
+        # starts at t=1), so grow by new_len - previous_len against a baseline
+        # of logical_prompt_len for the first chunk. Under the duplex drip
+        # above this is exactly 1 per chunk.
+        new_positions = max(len(timeline) - max(last_len, logical_prompt_len), 1)
+
     return OmniPayloadStruct(
         ids=IdsStruct(all=timeline, prompt=[0]),
         meta=MetaStruct(
             finished=torch.tensor(bool(is_finished and not is_duplex), dtype=torch.bool),
             num_processed_tokens=logical_prompt_len,
-            # The talker's vLLM prompt stays a single placeholder token.
-            next_stage_prompt_len=1,
+            # Vendored talker: one placeholder position per chunk (it drains
+            # all new timeline tokens model-side). Native talker: one position
+            # PER NEW TOKEN (see native_talker above).
+            next_stage_prompt_len=new_positions,
             codec_streaming=is_duplex,
             request_id=request_id if is_duplex else None,
+            expected_total_tokens=expected_total,
         ),
     )