From 432b2c6ca3a0bc29c9cf65b6de11196d13c2ddf7 Mon Sep 17 00:00:00 2001
From: Sy03 <1370724210@qq.com>
Date: Wed, 26 Aug 2026 13:19:20 -0700
Subject: [PATCH] [Model] Add native full-duplex Nemotron VoiceChat serving
(#6089)
Signed-off-by: Sysy <1370724210@qq.com>
Signed-off-by: sy03 <137072421@qq.com>
Co-authored-by: sy03 <137072421@qq.com>
---
.buildkite/amd/test-amd-ready.yml | 2 +-
recipes/NVIDIA/NemotronLabs-VoiceChat.md | 77 +++-
.../sched/test_omni_ar_scheduler_streaming.py | 31 ++
.../test_chunk_transfer_adapter.py | 101 +++++
.../nemotron_voicechat/test_data_plane.py | 94 ++++
.../nemotron_voicechat/test_input.py | 73 ++++
.../test_runtime_contract.py | 115 +++++
.../nemotron_voicechat_realtime_duplex.py | 364 +++++++++++++++
.../test_cosyvoice3_tts_expansion.py | 5 +-
.../test_nemotron_voicechat_duplex.py | 56 +++
.../openai_api/test_duplex_handler.py | 71 ++-
.../models/test_nemotron_voicechat_shapes.py | 36 ++
.../test_nemotron_voicechat.py | 74 +++-
tests/worker/test_gpu_ar_model_runner.py | 1 +
vllm_omni/core/sched/omni_ar_scheduler.py | 7 +-
vllm_omni/core/sched/omni_scheduler_mixin.py | 33 +-
.../nemotron_labs_voicechat_duplex.yaml | 47 ++
.../chunk_transfer_adapter.py | 208 +++++++--
.../fullduplex/nemotron_voicechat/__init__.py | 4 +
.../nemotron_voicechat/data_plane.py | 403 +++++++++++++++++
.../fullduplex/nemotron_voicechat/input.py | 225 ++++++++++
.../fullduplex/nemotron_voicechat/runtime.py | 167 +++++++
.../nemotron_voicechat/serving_adapter.py | 363 +++++++++++++++
.../fullduplex/openai/realtime_input.py | 43 ++
.../fullduplex/openai/realtime_output.py | 61 +++
.../fullduplex/openai/runtime_adapter.py | 1 +
.../fullduplex/openai/runtime_bridge.py | 29 +-
.../experimental/fullduplex/openai/serving.py | 4 +
.../fullduplex/openai/session_runner.py | 79 +++-
.../fullduplex/openai/websocket.py | 3 +
.../nemo_vendored/asr/conformer_encoder.py | 69 +++
.../nemo_vendored/compat.py | 45 ++
.../nemo_vendored/ear_tts_model.py | 10 +-
.../nemotron_voicechat_code2wav.py | 41 +-
.../nemotron_voicechat_talker.py | 83 +++-
.../nemotron_voicechat_thinker.py | 413 +++++++++++++++++-
.../models/nemotron_voicechat/pipeline.py | 7 +
.../models/nemotron_voicechat/runtime_info.py | 21 +
.../nemotron_voicechat.py | 224 ++++++++--
vllm_omni/worker/gpu_ar_model_runner.py | 24 +-
40 files changed, 3560 insertions(+), 154 deletions(-)
create mode 100644 tests/e2e/features/fullduplex/nemotron_voicechat/test_data_plane.py
create mode 100644 tests/e2e/features/fullduplex/nemotron_voicechat/test_input.py
create mode 100644 tests/e2e/features/fullduplex/nemotron_voicechat/test_runtime_contract.py
create mode 100644 tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py
create mode 100644 tests/e2e/online_serving/test_nemotron_voicechat_duplex.py
create mode 100644 vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml
create mode 100644 vllm_omni/experimental/fullduplex/nemotron_voicechat/__init__.py
create mode 100644 vllm_omni/experimental/fullduplex/nemotron_voicechat/data_plane.py
create mode 100644 vllm_omni/experimental/fullduplex/nemotron_voicechat/input.py
create mode 100644 vllm_omni/experimental/fullduplex/nemotron_voicechat/runtime.py
create mode 100644 vllm_omni/experimental/fullduplex/nemotron_voicechat/serving_adapter.py
diff --git a/.buildkite/amd/test-amd-ready.yml b/.buildkite/amd/test-amd-ready.yml
index 68ceed718..53b39e031 100644
--- a/.buildkite/amd/test-amd-ready.yml
+++ b/.buildkite/amd/test-amd-ready.yml
@@ -192,7 +192,7 @@ steps:
grade: Blocking
commands:
- |
- timeout 20m bash -c '
+ timeout 45m bash -c '
pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "slow" --run-level "core_model"
'
diff --git a/recipes/NVIDIA/NemotronLabs-VoiceChat.md b/recipes/NVIDIA/NemotronLabs-VoiceChat.md
index 8faa268d2..dbe7d5431 100644
--- a/recipes/NVIDIA/NemotronLabs-VoiceChat.md
+++ b/recipes/NVIDIA/NemotronLabs-VoiceChat.md
@@ -1,7 +1,7 @@
# NemotronLabs VoiceChat 11B
-> Offline speech-to-speech: a 16 kHz user utterance in, the agent's spoken reply
-> (text + 22.05 kHz WAV) out, on a 3-stage vLLM-Omni pipeline.
+> Speech-to-speech on a 3-stage vLLM-Omni pipeline: offline single-turn
+> inference and experimental model-native, frame-locked Realtime serving.
## Summary
@@ -11,8 +11,8 @@
timeline: a Conformer + NemotronH hybrid-Mamba thinker emits one text token
per acoustic frame, a Gemma3-1B EAR-TTS talker turns the text timeline into
31-quantizer RVQ code stacks, and an RVQ-VAE codec decodes them to audio.
-- Mode: Offline single-turn inference (batch=1). Online/duplex serving,
- function calling, and batch>1 are documented follow-ups.
+- Mode: Offline single-turn inference (batch=1), plus experimental native
+ duplex Realtime serving at one 80 ms microphone frame per scheduler wake.
- Maintainer: [`@yuekaizhang`](https://github.com/yuekaizhang)
## When to use this recipe
@@ -33,6 +33,11 @@ NeMo modules (`nemo_vendored/`), so no `nemo_toolkit` install is needed.
[`vllm_omni/model_executor/models/nemotron_voicechat/`](../../vllm_omni/model_executor/models/nemotron_voicechat/)
- Staged pipeline config:
[`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),
+ [`tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py`](../../tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py)
+- Nightly model-level gate:
+ [`tests/e2e/online_serving/test_nemotron_voicechat_duplex.py`](../../tests/e2e/online_serving/test_nemotron_voicechat_duplex.py)
- Upstream model card:
[`nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`](https://huggingface.co/nvidia/NVIDIA-NemotronLabs-VoiceChat-11B)
- Reference implementation: NVIDIA-NeMo/Speech, branch `nemotron-labs-voicechat`
@@ -60,7 +65,7 @@ output stayed within one word of the reference in testing.
- OS: Linux
- Python: 3.12
-- vLLM version: 0.26.0
+- vLLM version: 0.27.0
- vLLM-Omni version or commit: this PR / current `main`
#### Command
@@ -88,6 +93,55 @@ The reply text should read as a coherent spoken-style answer to the question in
the input WAV, and the WAV should transcribe to (approximately) the same text
with any ASR model.
+#### Native duplex serving
+
+The checkpoint does not contain the underlying Nemotron text tokenizer. Set
+`NEMOTRON_VOICECHAT_LLM_PATH` to a local snapshot of
+`nvidia/NVIDIA-Nemotron-Nano-9B-v2` before starting an offline deployment.
+
+```bash
+export HF_HUB_OFFLINE=1
+export TRANSFORMERS_OFFLINE=1
+export NEMOTRON_VOICECHAT_LLM_PATH=/path/to/NVIDIA-Nemotron-Nano-9B-v2
+
+vllm-omni serve /path/to/NVIDIA-NemotronLabs-VoiceChat-11B \
+ --omni \
+ --served-model-name nemotron-voicechat \
+ --deploy-config vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml
+```
+
+The single-GPU profile runs three eager fp32 engine processes on one device.
+
+In another shell, stream only the user channel of one of NVIDIA's bundled
+stereo conversation fixtures:
+
+```bash
+python tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py \
+ --model nemotron-voicechat \
+ --input-wav /path/to/NVIDIA-NemotronLabs-VoiceChat-11B/turn_taking.wav \
+ --input-channel 0 \
+ --output-dir results/nemotron_voicechat_duplex \
+ --timeout-s 600
+```
+
+The probe requires a completed response, non-silent 22.05 kHz output, and the
+advertised native 80 ms append capabilities. Add `--expect-function-call` and
+use `tool_call.wav` to validate the function-call channel. To validate the
+Realtime tool round trip, also provide the expected arguments and return a
+tool result:
+
+```bash
+python tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py \
+ --model nemotron-voicechat \
+ --input-wav /path/to/NVIDIA-NemotronLabs-VoiceChat-11B/tool_call.wav \
+ --output-dir results/nemotron_voicechat_tool_call \
+ --expect-function-call \
+ --expected-function-name generate_random_number \
+ --expected-function-arguments '{"min":1,"max":50}' \
+ --function-output 20 \
+ --expected-post-tool-text "random number"
+```
+
#### Notes
- Memory usage: the shipped yaml runs all three stages on one GPU
@@ -109,5 +163,14 @@ with any ASR model.
instead of pausing.
- The talker's `max_tokens` is 16383 (its stage prompt is one placeholder
token, and the stage context is 16384).
-- Known limitations: batch=1 only; offline single turn (no online/duplex
- serving yet); the function-call channel is decoded but not yet acted on.
+- 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
+ 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;
+ it does not execute arbitrary tools itself.
diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py
index fb4188ced..bf35bfdbc 100644
--- a/tests/core/sched/test_omni_ar_scheduler_streaming.py
+++ b/tests/core/sched/test_omni_ar_scheduler_streaming.py
@@ -198,6 +198,7 @@ def test_resumable_segment_boundary_keeps_pre_transition_send_watermark() -> Non
inter_stage_output,
session,
True,
+ new_token_ids=[42],
confirmed_num_computed_tokens=26,
)
@@ -427,3 +428,33 @@ def test_ready_async_chunk_prompt_replacement_releases_stale_kv_once() -> None:
assert sched.chunk_transfer_adapter.replaced_streaming_prompt_ids == set()
assert sched.chunk_transfer_adapter.requests_with_ready_chunks == {session.request_id}
assert sched.chunk_transfer_adapter.requests_num_chunks_sent == {}
+
+
+def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None:
+ """A WAITING_FOR_CHUNK stop must not delete its newly parked session."""
+ session = _make_request()
+ session.resumable = True
+ session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+
+ def queue(*requests):
+ result = MagicMock()
+ result.requests = set(requests)
+ result.add_request.side_effect = result.requests.add
+ result.remove_requests.side_effect = result.requests.difference_update
+ return result
+
+ sched = OmniARScheduler.__new__(OmniARScheduler)
+ sched.running = []
+ sched.waiting = queue()
+ sched.skipped_waiting = queue(session)
+ sched.num_waiting_for_streaming_input = 1
+ sched.chunk_transfer_adapter = SimpleNamespace(segment_finished_requests={session.request_id})
+
+ sched._resume_downstream_chunk_receiver(session)
+ sched._remove_stopped_requests_from_queues(set(), {session})
+
+ assert session.status == RequestStatus.WAITING
+ assert session in sched.waiting.requests
+ assert session not in sched.skipped_waiting.requests
+ assert sched.num_waiting_for_streaming_input == 0
+ assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests
diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py
index 2dbe967e7..a5537bdae 100644
--- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py
+++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py
@@ -1209,6 +1209,67 @@ def test_cleanup_preserves_pending_save(build_adapter):
assert len(adapter._pending_save_reqs) == 1
+def test_abort_invalidates_queued_sender_task_before_external_id_reuse(build_adapter):
+ adapter, connector = build_adapter(stage_id=1)
+ request = _req("req-old", RequestStatus.WAITING, external_req_id="ext-reused")
+ adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
+ codes=CodesStruct(audio=torch.tensor([1], dtype=torch.long))
+ )
+
+ adapter.save_async(multimodal_output=None, request=request)
+ stale_task = adapter._pending_save_reqs.popleft()
+ adapter.cleanup(request.request_id, request.external_req_id)
+ adapter._send_single_request(stale_task)
+
+ connector.put.assert_not_called()
+ assert request.external_req_id not in adapter.put_req_chunk
+
+ replacement = _req("req-new", RequestStatus.WAITING, external_req_id=request.external_req_id)
+ adapter.save_async(multimodal_output=None, request=replacement)
+ current_task = adapter._pending_save_reqs.popleft()
+ adapter._send_single_request(current_task)
+
+ connector.put.assert_called_once()
+ assert connector.put.call_args.kwargs["put_key"] == "ext-reused_1_0"
+
+
+def test_finish_requests_does_not_wait_for_inflight_send(build_adapter):
+ adapter, connector = build_adapter(stage_id=1)
+ first = _req("req-first", RequestStatus.WAITING, external_req_id="ext-first")
+ adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct()
+ put_started = threading.Event()
+ release_put = threading.Event()
+
+ def blocking_put(**kwargs):
+ put_started.set()
+ release_put.wait(timeout=2)
+ return True, 1, {}
+
+ connector.put.side_effect = blocking_put
+ adapter.save_async(multimodal_output=None, request=first)
+ task = adapter._pending_save_reqs.popleft()
+ sender = threading.Thread(target=adapter._send_single_request, args=(task,))
+ sender.start()
+ assert put_started.wait(timeout=1)
+ cleaner = threading.Thread(
+ target=adapter.finish_requests,
+ args=([first.request_id], RequestStatus.FINISHED_ABORTED, {first.request_id: first}),
+ )
+ cleaner.start()
+ try:
+ cleaner.join(timeout=0.5)
+ assert not cleaner.is_alive()
+ assert task["sender_token"].cancelled
+ assert first.external_req_id in adapter._sender_tokens
+ finally:
+ release_put.set()
+ sender.join(timeout=1)
+ cleaner.join(timeout=1)
+ assert not sender.is_alive()
+ assert first.external_req_id not in adapter._sender_tokens
+ assert first.external_req_id not in adapter.put_req_chunk
+
+
def test_cleanup_only_affects_target_request(build_adapter):
"""Cleanup for one request must not affect another request's state."""
adapter, _ = build_adapter(stage_id=1)
@@ -2216,3 +2277,43 @@ def test_expiry_is_per_request(build_adapter):
assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == {"stalled"}
assert "healthy" in adapter._waiting_since
+
+
+def test_abort_clears_native_codec_state_before_external_id_reuse(build_adapter):
+ from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import (
+ talker2code2wav_async_chunk,
+ )
+
+ adapter, _ = build_adapter(stage_id=1, connector_extra={"codec_chunk_frames": 1})
+ external_req_id = "ext-native-abort"
+ first = _req("req-native-abort", RequestStatus.WAITING, external_req_id=external_req_id)
+ first.resumable = True
+ first.model_intermediate_buffer = None
+ first.additional_information = {"meta": {"codec_streaming": True}, "nvc_logical_prompt_len": 1}
+ first_frame = torch.full((1, 31), 101, dtype=torch.long)
+
+ talker2code2wav_async_chunk(
+ adapter,
+ {"codes": {"audio": first_frame}, "meta": {"codec_streaming": True}},
+ first,
+ is_finished=False,
+ )
+
+ adapter.requests_num_chunks_sent[external_req_id] = 1
+ adapter.finish_requests([first.request_id], RequestStatus.FINISHED_ABORTED, {first.request_id: first})
+
+ assert external_req_id not in adapter.request_payload
+ assert external_req_id not in adapter.requests_num_chunks_sent
+
+ replacement = _req("req-native-abort", RequestStatus.WAITING, external_req_id=external_req_id)
+ replacement.resumable = True
+ replacement.model_intermediate_buffer = None
+ replacement.additional_information = first.additional_information
+ replacement_payload = talker2code2wav_async_chunk(
+ adapter,
+ {"codes": {"audio": torch.full((1, 31), 202, dtype=torch.long)}, "meta": {"codec_streaming": True}},
+ replacement,
+ is_finished=False,
+ )
+
+ assert torch.equal(replacement_payload.codes.audio, torch.full((1, 31), 202, dtype=torch.long))
diff --git a/tests/e2e/features/fullduplex/nemotron_voicechat/test_data_plane.py b/tests/e2e/features/fullduplex/nemotron_voicechat/test_data_plane.py
new file mode 100644
index 000000000..5790b11f1
--- /dev/null
+++ b/tests/e2e/features/fullduplex/nemotron_voicechat/test_data_plane.py
@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+
+from vllm_omni.experimental.fullduplex.engine.duplex_runtime import duplex_resource_request_id
+from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
+from vllm_omni.experimental.fullduplex.nemotron_voicechat.data_plane import (
+ NemotronVoiceChatDataPlaneContext,
+ NemotronVoiceChatDataPlaneSession,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+_RUNTIME = {
+ "nvc_text_bos_id": 0,
+ "nvc_text_eos_id": 1,
+ "nvc_text_pad_id": 12,
+ "nvc_function_sotc_id": 20,
+ "nvc_function_eotc_id": 21,
+ "nvc_tokenizer_ref": "test-tokenizer",
+}
+
+
+def _projector(encode_audio=lambda *_: None) -> NemotronVoiceChatDataPlaneSession:
+ projector = NemotronVoiceChatDataPlaneSession(encode_audio)
+ projector.configure_runtime(_RUNTIME, tokenizer=SimpleNamespace(decode=lambda *_args, **_kwargs: "text"))
+ return projector
+
+
+def _stage0_output(request_id: str = "req-0", *, text_token: int = 12, function_token: int | None = None) -> object:
+ metadata: dict[str, object] = {"nvc_text_token_ids": [text_token]}
+ if function_token is not None:
+ metadata["nvc_function_token"] = [function_token]
+ completion = SimpleNamespace(multimodal_output=metadata)
+ return SimpleNamespace(
+ stage_id=0,
+ request_output=SimpleNamespace(request_id=request_id, outputs=[completion]),
+ )
+
+
+def test_function_channel_projects_completed_call_without_ending_speech() -> None:
+ projector = _projector()
+ projector._decode = lambda token_ids: (
+ '[{"name":"weather","arguments":{"city":"Shanghai"}}]' if token_ids == [99] else ""
+ )
+ context = NemotronVoiceChatDataPlaneContext(epoch=0)
+
+ events = []
+ for function_token in (20, 99, 21):
+ events.extend(projector.project_output(_stage0_output(function_token=function_token), context=context))
+
+ listen = [event for event in events if event.get("is_listen") is True]
+ function = [event for event in events if event.get("function_call") is True]
+ assert len(listen) == 3
+ assert len(function) == 1
+ assert function[0]["name"] == "weather"
+ assert function[0]["arguments"] == '{"city":"Shanghai"}'
+
+ projector._decode = lambda _token_ids: "not-json"
+ bad_events = []
+ for function_token in (20, 98, 21):
+ bad_events.extend(projector.project_output(_stage0_output(function_token=function_token)))
+ assert any(event.get("error_code") == "nemotron_function_call_parse_error" for event in bad_events)
+
+
+def test_new_epoch_request_does_not_inherit_partial_function_or_frame_state() -> None:
+ projector = _projector(lambda *_: "audio")
+ old_request = duplex_resource_request_id(DuplexFence("sid", epoch=1), "stage0")
+ new_request = duplex_resource_request_id(DuplexFence("sid", epoch=2), "stage0")
+
+ projector.begin_request(old_request)
+ list(projector.project_output(_stage0_output(old_request, text_token=42, function_token=20)))
+ projector.begin_request(new_request)
+
+ # EOTC in the new epoch must not close the old epoch's partial function.
+ new_events = list(projector.project_output(_stage0_output(new_request, text_token=1, function_token=21)))
+ assert not [event for event in new_events if event.get("function_call") is True]
+ assert not [event for event in new_events if event.get("end_of_turn") is True]
+
+ audio = SimpleNamespace(
+ stage_id=2,
+ request_output=SimpleNamespace(
+ request_id=new_request,
+ outputs=[
+ SimpleNamespace(multimodal_output={"model_outputs": [np.ones(1764, dtype=np.float32)], "sr": [22050]})
+ ],
+ ),
+ )
+ audio_events = list(projector.project_output(audio))
+ assert sum(event.get("end_of_turn") is True for event in audio_events) == 1
diff --git a/tests/e2e/features/fullduplex/nemotron_voicechat/test_input.py b/tests/e2e/features/fullduplex/nemotron_voicechat/test_input.py
new file mode 100644
index 000000000..84c8b123a
--- /dev/null
+++ b/tests/e2e/features/fullduplex/nemotron_voicechat/test_input.py
@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+import base64
+
+import numpy as np
+import pytest
+
+from vllm_omni.experimental.fullduplex.nemotron_voicechat.input import (
+ NemotronVoiceChatPcmAppendBuffer,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _payload(samples: np.ndarray, *, sample_rate_hz: int = 16000, fmt: str = "pcm_f32le") -> dict[str, object]:
+ return {
+ "type": "audio",
+ "audio": base64.b64encode(np.asarray(samples, dtype=np.float32).tobytes()).decode("ascii"),
+ "format": fmt,
+ "sample_rate_hz": sample_rate_hz,
+ }
+
+
+def _append(buffer, samples, operation_id):
+ return buffer.prepare_append(
+ _payload(np.asarray(samples, dtype=np.float32)),
+ operation_id=operation_id,
+ chunk_period_ms=80,
+ allow_emit=True,
+ )
+
+
+def test_irregular_browser_packets_emit_exact_ordered_80ms_frame() -> None:
+ buffer = NemotronVoiceChatPcmAppendBuffer()
+ assert _append(buffer, np.arange(1000), "packet-1") is None
+
+ reservation = _append(buffer, np.arange(1000, 1280), "packet-2")
+
+ assert reservation is not None
+ assert reservation.operation_id == "packet-2"
+ decoded = np.frombuffer(base64.b64decode(reservation.payload["audio"]), dtype=np.float32)
+ np.testing.assert_array_equal(decoded, np.arange(1280, dtype=np.float32))
+ assert buffer.pending_byte_count == 0
+
+
+def test_append_rejects_invalid_frame_contract() -> None:
+ cases = (
+ (1281, 16000, "pcm_f32le", 0.0, "at most 1280"),
+ (1, 8000, "pcm_f32le", 0.0, "sample_rate_hz"),
+ (1, 16000, "audio/pcm", 0.0, "format"),
+ (1, 16000, "pcm_f32le", np.nan, "finite"),
+ )
+ for samples, sample_rate_hz, fmt, value, message in cases:
+ with pytest.raises(ValueError, match=message):
+ NemotronVoiceChatPcmAppendBuffer().prepare_append(
+ _payload(np.full(samples, value, dtype=np.float32), sample_rate_hz=sample_rate_hz, fmt=fmt),
+ operation_id="invalid",
+ chunk_period_ms=80,
+ allow_emit=True,
+ )
+
+
+def test_terminal_commit_rejects_more_than_one_frame_tail() -> None:
+ buffer = NemotronVoiceChatPcmAppendBuffer()
+ for index, samples in enumerate((1000, 281)):
+ buffer.prepare_append(
+ _payload(np.zeros(samples, dtype=np.float32)),
+ operation_id=f"tail-{index}",
+ chunk_period_ms=80,
+ allow_emit=False,
+ )
+ with pytest.raises(ValueError, match="more than one 80 ms frame"):
+ buffer.prepare_commit(operation_id="commit", chunk_period_ms=80)
diff --git a/tests/e2e/features/fullduplex/nemotron_voicechat/test_runtime_contract.py b/tests/e2e/features/fullduplex/nemotron_voicechat/test_runtime_contract.py
new file mode 100644
index 000000000..0b32322c0
--- /dev/null
+++ b/tests/e2e/features/fullduplex/nemotron_voicechat/test_runtime_contract.py
@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import base64
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.experimental.fullduplex.engine.contracts import DuplexInputMode
+from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
+from vllm_omni.experimental.fullduplex.nemotron_voicechat.runtime import (
+ NemotronVoiceChatDuplexRuntimeExtension,
+)
+from vllm_omni.experimental.fullduplex.nemotron_voicechat.serving_adapter import (
+ NemotronVoiceChatServingRuntimeAdapter,
+ _render_tool_response,
+)
+from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import (
+ NemotronVoiceChatThinkerForConditionalGeneration,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _frame() -> dict[str, object]:
+ raw = np.zeros(1280, dtype=np.float32).tobytes()
+ return {
+ "type": "audio",
+ "audio": base64.b64encode(raw).decode("ascii"),
+ "format": "pcm_f32le",
+ "sample_rate_hz": 16000,
+ }
+
+
+def _plan(extension, *, input_seq: int):
+ return extension.plan_append(
+ request_id="req",
+ fence=DuplexFence("sid", incarnation=2, epoch=3),
+ session_config={},
+ runtime_config={
+ "nvc_prompt_token_ids": [0, 42, 1],
+ "nvc_text_pad_id": 12,
+ "nvc_max_model_len": 8192,
+ },
+ seq=input_seq,
+ turn_seq=input_seq,
+ mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
+ payload=_frame(),
+ final=False,
+ sampling_params=SamplingParams(),
+ )
+
+
+def test_first_append_prefills_prompt_then_each_append_consumes_one_frame() -> None:
+ extension = NemotronVoiceChatDuplexRuntimeExtension()
+
+ params = extension.configure_sampling_params(
+ runtime_config={}, defaults=(SamplingParams(temperature=0.7, max_tokens=4),)
+ )[0]
+ assert params.temperature == 0.0 and params.max_tokens == 1 and params.top_p == 1.0 and params.top_k == 0
+
+ first = _plan(extension, input_seq=1)
+ later = _plan(extension, input_seq=2)
+
+ assert first.prompt["prompt_token_ids"] == [0, 42, 1, 12]
+ assert later.prompt["prompt_token_ids"] == [12]
+ assert first.prompt["model_intermediate_buffer"]["duplex"]["source_input_seq"] == 1
+ assert later.prompt["model_intermediate_buffer"]["duplex"]["source_input_seq"] == 2
+
+
+def test_append_rejects_stage0_context_overflow() -> None:
+ with pytest.raises(ValueError, match="max_model_len"):
+ _plan(NemotronVoiceChatDuplexRuntimeExtension(), input_seq=8190)
+
+
+def test_function_output_becomes_versioned_nvidia_channel_tokens() -> None:
+ encoded: list[str] = []
+ adapter = NemotronVoiceChatServingRuntimeAdapter(lambda *_: None)
+ adapter._tokenizer = SimpleNamespace(
+ encode=lambda text, **_kwargs: encoded.append(text) or [31, 32, 33],
+ )
+
+ first = adapter.runtime_config_for_function_output(
+ {"type": "function_call_output", "call_id": "call-1", "output": '{"result":20}'},
+ {},
+ )
+ second = adapter.runtime_config_for_function_output(
+ {"type": "function_call_output", "call_id": "call-2", "output": "plain text"},
+ first,
+ )
+
+ assert _render_tool_response('{"result":20}') == '[{"result":20}]'
+ assert encoded == [
+ '[{"result":20}]',
+ '["plain text"]',
+ ]
+ assert first["nvc_function_response_generation"] == 1
+ assert second["nvc_function_response_generation"] == 2
+ assert second["nvc_function_response_token_ids"] == [31, 32, 33]
+ assert second["nvc_function_response_call_id"] == "call-2"
+ assert second["nvc_function_response_batches"] == [
+ {"generation": 1, "call_id": "call-1", "token_ids": [31, 32, 33]},
+ {"generation": 2, "call_id": "call-2", "token_ids": [31, 32, 33]},
+ ]
+
+ session = {"function_response_generation": 0, "forced_function_tokens": []}
+ NemotronVoiceChatThinkerForConditionalGeneration._sync_forced_function_response(session, second)
+
+ assert session["function_response_generation"] == 2
+ assert session["forced_function_token"] == 31
+ assert session["forced_function_tokens"] == [31, 32, 33, 31, 32, 33]
diff --git a/tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py b/tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py
new file mode 100644
index 000000000..b623446ad
--- /dev/null
+++ b/tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py
@@ -0,0 +1,364 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import argparse
+import asyncio
+import base64
+import json
+import math
+import uuid
+import wave
+from collections.abc import Sequence
+from contextlib import asynccontextmanager, suppress
+from pathlib import Path
+from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
+
+import numpy as np
+from scipy.signal import resample_poly
+
+from vllm_omni.experimental.fullduplex.client import RealtimeDuplexClient, wait_for, write_pcm16_wav
+
+INPUT_SAMPLE_RATE_HZ = 16_000
+OUTPUT_SAMPLE_RATE_HZ = 22_050
+FRAME_SAMPLES = 1_280
+FRAME_PERIOD_S = FRAME_SAMPLES / INPUT_SAMPLE_RATE_HZ
+DEFAULT_FUNCTION_TOOLS = [
+ {
+ "type": "function",
+ "name": "generate_random_number",
+ "description": "Generate a random integer between min and max (inclusive).",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "min": {"type": "integer", "description": "Minimum value (inclusive)"},
+ "max": {"type": "integer", "description": "Maximum value (inclusive)"},
+ },
+ "required": ["min", "max"],
+ },
+ }
+]
+DEFAULT_INSTRUCTIONS = "You are NVIDIA Voice Chat. Answer briefly. Start by greeting the user."
+DEFAULT_FUNCTION_INSTRUCTIONS = (
+ "You are NVIDIA Voice Chat. If the user's request matches an available tool, "
+ "you MUST call that tool instead of answering from your own knowledge. "
+ "Use only argument values spoken by the user and never invent missing values."
+)
+
+
+def _url(base_url: str, model: str, session_id: str) -> str:
+ parts = urlsplit(base_url)
+ query = dict(parse_qsl(parts.query, keep_blank_values=True))
+ query.update(duplex="1", model=model, autostart="0", session_id=session_id)
+ return urlunsplit((parts.scheme, parts.netloc, parts.path, urlencode(query), parts.fragment))
+
+
+def _read_wav(path: Path, *, input_channel: int = 0) -> np.ndarray:
+ with wave.open(str(path), "rb") as wav_file:
+ channels = wav_file.getnchannels()
+ width = wav_file.getsampwidth()
+ source_rate = wav_file.getframerate()
+ raw = wav_file.readframes(wav_file.getnframes())
+ dtypes = {1: np.uint8, 2: np.dtype(" 1:
+ pcm = pcm.reshape(-1, channels)[:, input_channel]
+ if source_rate != INPUT_SAMPLE_RATE_HZ:
+ divisor = math.gcd(source_rate, INPUT_SAMPLE_RATE_HZ)
+ pcm = resample_poly(pcm, up=INPUT_SAMPLE_RATE_HZ // divisor, down=source_rate // divisor)
+ return np.ascontiguousarray(pcm, dtype=" int:
+ count = math.ceil(pcm.size / FRAME_SAMPLES)
+ if max_frames is not None:
+ count = min(count, max_frames)
+ for seq in range(count):
+ frame = pcm[seq * FRAME_SAMPLES : (seq + 1) * FRAME_SAMPLES]
+ frame = np.pad(frame, (0, FRAME_SAMPLES - frame.size)).astype(" list[dict[str, object]]:
+ return [event for event in client.events.events if event.get("type") == event_type]
+
+
+async def _return_function_output_when_ready(
+ client: RealtimeDuplexClient,
+ *,
+ output: str,
+ timeout_s: float,
+) -> tuple[str, int]:
+ await wait_for(
+ lambda: bool(client.events.errors()) or client.events.count("response.function_call_arguments.done") > 0,
+ timeout_s=timeout_s,
+ label="function call to execute",
+ )
+ if client.events.errors():
+ raise AssertionError(f"function call failed before tool execution: {client.events.errors()}")
+ function_done = _events(client, "response.function_call_arguments.done")[-1]
+ call_id = function_done.get("call_id")
+ if not isinstance(call_id, str) or not call_id:
+ raise AssertionError(f"completed function call has no call_id: {function_done}")
+ event_count_before_output = len(client.events.events)
+ await client.send(
+ {
+ "type": "conversation.item.create",
+ "item": {
+ "type": "function_call_output",
+ "call_id": call_id,
+ "output": output,
+ },
+ }
+ )
+ return call_id, event_count_before_output
+
+
+def _write_events(path: Path, client: RealtimeDuplexClient) -> None:
+ path.write_text(
+ "".join(json.dumps(event, ensure_ascii=False) + "\n" for event in client.events.events),
+ encoding="utf-8",
+ )
+
+
+@asynccontextmanager
+async def _managed_client(client: RealtimeDuplexClient, *, timeout_s: float):
+ async with client:
+ try:
+ yield
+ finally:
+ if client.events.count("session.created") and not client.events.count("session.closed"):
+ with suppress(Exception):
+ await client.close_session(timeout_s=min(timeout_s, 30.0))
+
+
+async def run(args: argparse.Namespace) -> dict[str, object]:
+ output_dir = Path(args.output_dir)
+ output_dir.mkdir(parents=True, exist_ok=True)
+ if args.expect_function_call and args.instructions == DEFAULT_INSTRUCTIONS:
+ instructions = DEFAULT_FUNCTION_INSTRUCTIONS
+ else:
+ instructions = args.instructions
+ tools = DEFAULT_FUNCTION_TOOLS if args.expect_function_call else None
+ session_id = f"nemotron-voicechat-{uuid.uuid4().hex}"
+ client = RealtimeDuplexClient(_url(args.url, args.model, session_id))
+ async with _managed_client(client, timeout_s=args.timeout_s):
+ session_payload: dict[str, object] = {
+ "session_id": session_id,
+ "model": args.model,
+ "modalities": ["audio", "text"],
+ "input_audio_format": "pcm_f32le",
+ "output_audio_format": "pcm16",
+ "instructions": instructions,
+ "idle_timeout_s": args.timeout_s,
+ "turn_detection": None,
+ "extra_body": {"auto_response": True},
+ }
+ if tools is not None:
+ session_payload["tools"] = tools
+ await client.send({"type": "session.update", "session": session_payload})
+ await wait_for(
+ lambda: client.events.count("session.created") > 0 or bool(client.events.errors()),
+ timeout_s=args.timeout_s,
+ label="session.created",
+ )
+ if client.events.errors():
+ raise AssertionError(f"session setup failed: {client.events.errors()}")
+ created = _events(client, "session.created")[-1]
+ session = created.get("session")
+ capabilities = session.get("capabilities") if isinstance(session, dict) else None
+ expected = dict(
+ implementation_level="model_native_duplex",
+ chunk_period_ms=80,
+ supports_core_resumable_request=True,
+ supports_core_kv_lease=False,
+ supports_multi_session=False,
+ )
+ if not isinstance(capabilities, dict) or any(capabilities.get(key) != value for key, value in expected.items()):
+ raise AssertionError(f"unexpected capabilities: {capabilities}")
+
+ function_output_task = (
+ asyncio.create_task(
+ _return_function_output_when_ready(
+ client,
+ output=args.function_output,
+ timeout_s=args.timeout_s,
+ )
+ )
+ if args.function_output is not None
+ else None
+ )
+ pcm = _read_wav(Path(args.input_wav), input_channel=args.input_channel)
+ frame_count = await _stream(client, pcm, max_frames=args.max_frames, realtime=not args.no_realtime)
+ completed_responses_at_commit = client.events.count("response.done")
+ await client.send({"type": "input_audio_buffer.commit", "final": True})
+ await wait_for(
+ lambda: (
+ bool(client.events.errors())
+ or (
+ client.events.count("response.function_call_arguments.done") > 0
+ 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)
+ )
+ )
+ ),
+ timeout_s=args.timeout_s,
+ label="model output",
+ )
+ await asyncio.sleep(args.drain_s)
+ 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:
+ response = done_events[-1].get("response") if done_events else None
+ status = response.get("status") if isinstance(response, dict) else None
+ if status != "completed":
+ raise AssertionError(f"response did not complete successfully: {done_events[-1:]}")
+
+ function_events = [
+ event for event in client.events.events if str(event.get("type", "")).startswith("response.function_call")
+ ]
+ function_items = [
+ event
+ for event in _events(client, "response.output_item.done")
+ if isinstance(event.get("item"), dict) and event["item"].get("type") == "function_call"
+ ]
+ if args.expect_function_call and not any(
+ event.get("type") == "response.function_call_arguments.done" for event in function_events
+ ):
+ raise AssertionError(f"no completed function call: {function_events}")
+ if args.expect_function_call:
+ matching_items = [
+ event for event in function_items if event["item"].get("name") == args.expected_function_name
+ ]
+ if not matching_items:
+ raise AssertionError(f"expected {args.expected_function_name!r}, got {function_items}")
+ function_item = matching_items[-1]["item"]
+ try:
+ function_arguments = json.loads(str(function_item.get("arguments", "")))
+ except json.JSONDecodeError as exc:
+ raise AssertionError(f"function arguments are not JSON: {function_item}") from exc
+ if args.expected_function_arguments is not None:
+ expected_arguments = json.loads(args.expected_function_arguments)
+ if function_arguments != expected_arguments:
+ raise AssertionError(
+ f"function arguments differ: expected={expected_arguments!r}, actual={function_arguments!r}"
+ )
+
+ if args.function_output is not None:
+ call_id = function_item.get("call_id")
+ if not isinstance(call_id, str) or not call_id:
+ raise AssertionError(f"function item has no call_id: {function_item}")
+ assert function_output_task is not None
+ returned_call_id, event_count_before_output = await function_output_task
+ assert returned_call_id == call_id
+
+ def tool_result_completed() -> bool:
+ later = client.events.events[event_count_before_output:]
+ transcript = "".join(
+ str(event.get("delta", ""))
+ for event in later
+ if event.get("type") == "response.audio_transcript.delta"
+ ).lower()
+ return (
+ any(event.get("type") == "response.audio.delta" for event in later)
+ and any(event.get("type") == "response.done" for event in later)
+ and (args.expected_post_tool_text is None or args.expected_post_tool_text.lower() in transcript)
+ )
+
+ await wait_for(
+ lambda: bool(client.events.errors()) or tool_result_completed(),
+ timeout_s=args.timeout_s,
+ label="completed response after function output",
+ )
+ if client.events.errors():
+ raise AssertionError(f"function output failed: {client.events.errors()}")
+ await asyncio.sleep(args.drain_s)
+
+ audio = client.events.audio_bytes()
+ audio_events = _events(client, "response.audio.delta")
+ rates = {event.get("sample_rate_hz") for event in audio_events}
+ if not args.expect_function_call and audio and rates != {OUTPUT_SAMPLE_RATE_HZ}:
+ raise AssertionError(f"unexpected output sample rates: {rates}")
+ if not args.expect_function_call and args.minimum_audio_chunks and not audio:
+ raise AssertionError("model produced no audio")
+ expected_bytes = 2 * OUTPUT_SAMPLE_RATE_HZ * expected["chunk_period_ms"] // 1000
+ packet_sizes = [len(base64.b64decode(str(event.get("delta", "")), validate=True)) for event in audio_events]
+ if not args.expect_function_call and any(size != expected_bytes for size in packet_sizes):
+ raise AssertionError(f"audio deltas are not fixed 80 ms PCM16 packets: {packet_sizes}")
+ audio_pcm = np.frombuffer(audio, dtype=" argparse.Namespace:
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--url", default="ws://127.0.0.1:8125/v1/realtime")
+ parser.add_argument("--model", required=True)
+ parser.add_argument("--input-wav", required=True)
+ parser.add_argument("--input-channel", type=int, default=0)
+ parser.add_argument("--output-dir", default="/tmp/nemotron-voicechat-duplex")
+ 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("--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")
+ parser.add_argument("--expected-function-arguments")
+ parser.add_argument("--function-output")
+ parser.add_argument("--expected-post-tool-text")
+ parser.add_argument("--no-realtime", action="store_true")
+ parser.add_argument("--drain-s", type=float, default=2.0)
+ parser.add_argument("--timeout-s", type=float, default=600.0)
+ return parser.parse_args(argv)
+
+
+def main() -> None:
+ print(json.dumps(asyncio.run(run(parse_args())), ensure_ascii=False, indent=2))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py b/tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py
index cae854be3..544afc7bd 100644
--- a/tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py
+++ b/tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py
@@ -52,9 +52,8 @@ def get_prompt(prompt_type="zh"):
def _set_rocm_request_timeout(request_config: dict) -> None:
if current_omni_platform.is_rocm():
- # A cold AITER MHA kernel build took 301 seconds on MI300X.
- # Allow 600 seconds for the first request; later requests reuse the cache.
- request_config["timeout"] = 600.0
+ # Cold AITER MHA builds on loaded MI300 agents can exceed 20 minutes.
+ request_config["timeout"] = 1800.0
tts_server_params = [
diff --git a/tests/e2e/online_serving/test_nemotron_voicechat_duplex.py b/tests/e2e/online_serving/test_nemotron_voicechat_duplex.py
new file mode 100644
index 000000000..e82552750
--- /dev/null
+++ b/tests/e2e/online_serving/test_nemotron_voicechat_duplex.py
@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import asyncio
+import os
+from pathlib import Path
+
+import pytest
+from huggingface_hub import snapshot_download
+
+from tests.e2e.online_serving.nemotron_voicechat_realtime_duplex import parse_args, run
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams, get_model_prefix
+from tests.helpers.stage_config import get_deploy_config_path
+
+MODEL = os.environ.get("VLLM_TEST_NEMOTRON_VOICECHAT_MODEL", "nvidia/NVIDIA-NemotronLabs-VoiceChat-11B")
+DEPLOY_CONFIG = get_deploy_config_path("nemotron_labs_voicechat_duplex.yaml")
+TOKENIZER = os.environ.get("VLLM_TEST_NEMOTRON_VOICECHAT_LLM_PATH")
+pytestmark = [pytest.mark.full_model, pytest.mark.omni]
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize(
+ "omni_server",
+ [
+ pytest.param(
+ OmniServerParams(
+ model=MODEL,
+ stage_config_path=DEPLOY_CONFIG,
+ env_dict={"NEMOTRON_VOICECHAT_LLM_PATH": TOKENIZER} if TOKENIZER else None,
+ ),
+ id="native-duplex",
+ )
+ ],
+ indirect=True,
+)
+def test_native_duplex_turn_taking_streams_model_audio(omni_server, tmp_path: Path) -> None:
+ model_prefix = get_model_prefix()
+ root = Path(model_prefix) / MODEL if model_prefix else Path(MODEL)
+ if not root.is_dir():
+ root = Path(snapshot_download(MODEL, local_files_only=True))
+ args = parse_args(
+ (
+ f"--url ws://{omni_server.host}:{omni_server.port}/v1/realtime --model {omni_server.model} "
+ f"--input-wav {root / 'turn_taking.wav'} --input-channel 0 --max-frames 190 "
+ f"--minimum-audio-chunks 48 --minimum-audio-rms 0.001 --no-realtime --timeout-s 300 "
+ f"--output-dir {tmp_path / 'native_duplex'}"
+ ).split()
+ )
+ result = asyncio.run(run(args))
+ assert result["ok"] is True
+ assert result["input_frames"] == 190
+ assert result["event_counts"]["response.speak"] > 0
+ assert result["audio_bytes"] >= result["input_frames"] * 1764 * 2 // 4
diff --git a/tests/entrypoints/openai_api/test_duplex_handler.py b/tests/entrypoints/openai_api/test_duplex_handler.py
index 617a8a101..6a2b9d224 100644
--- a/tests/entrypoints/openai_api/test_duplex_handler.py
+++ b/tests/entrypoints/openai_api/test_duplex_handler.py
@@ -551,6 +551,54 @@ async def test_native_realtime_protocol_conversation_item_create_commits_user_te
}
+@pytest.mark.asyncio
+async def test_function_output_is_recorded_only_after_runtime_delivery():
+ ws = TimedWebSocket()
+ protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type]
+ protocol.bind_sender(ws.send_json)
+ protocol.encode_outbound_event(
+ {
+ "type": "conversation.item.created",
+ "item": {
+ "id": "item-call",
+ "type": "function_call",
+ "call_id": "call-1",
+ "name": "weather",
+ "arguments": "{}",
+ "status": "completed",
+ },
+ }
+ )
+ event = {
+ "type": "conversation.item.create",
+ "item": {
+ "id": "item-output",
+ "type": "function_call_output",
+ "call_id": "call-1",
+ "output": '{"temperature":20}',
+ },
+ }
+
+ first = await protocol._to_duplex_event(event)
+ retry = await protocol._to_duplex_event(event)
+
+ assert first is not None
+ assert retry is not None
+ assert "item-output" not in protocol._conversation_items
+
+ delivered_item = first["payload"]["item"]
+ protocol.encode_outbound_event(
+ {
+ "type": "conversation.item.created",
+ "item": delivered_item,
+ }
+ )
+
+ assert protocol._conversation_items["item-output"] == delivered_item
+ assert await protocol._to_duplex_event(event) is None
+ assert ws.sent[-1]["error"]["code"] == "invalid_function_call_output"
+
+
@pytest.mark.asyncio
async def test_native_realtime_protocol_audio_commit_requires_non_empty_buffer():
ws = TimedWebSocket()
@@ -5481,7 +5529,7 @@ async def test_failed_native_append_prevents_queued_append_from_running():
@pytest.mark.asyncio
-async def test_failed_final_append_keeps_committed_audio_for_response_retry():
+async def test_failed_final_append_discards_retained_audio_without_retry():
failed_result = {
"operation": "append",
"session_id": "sid-final-append-retry",
@@ -5490,15 +5538,7 @@ async def test_failed_final_append_keeps_committed_audio_for_response_retry():
"error_count": 1,
"stage_results": [],
}
- successful_result = {
- "operation": "append",
- "session_id": "sid-final-append-retry",
- "ok": True,
- "unsupported_count": 0,
- "error_count": 0,
- "stage_results": [],
- }
- engine = FakeEngineClient(append_results=[failed_result, successful_result])
+ engine = FakeEngineClient(append_result=failed_result)
handler = OmniDuplexSessionHandler(
chat_service=FakeChatService(engine),
config_timeout_s=0.1,
@@ -5523,15 +5563,18 @@ async def test_failed_final_append_keeps_committed_audio_for_response_retry():
handler_task = asyncio.create_task(handler.handle_session(ws))
for _ in range(100):
- if len(engine.appended) >= 2:
+ state = handler._serving_runtime_adapter.session_states.get("sid-final-append-retry")
+ if len(engine.appended) >= 1 and state is not None and state.committed_audio_payload is None:
break
await asyncio.sleep(0.01)
+ state = handler._serving_runtime_adapter.session_states["sid-final-append-retry"]
+ assert state.committed_audio_payload is None
+ assert state.committed_audio_reserved_bytes == 0
+ assert handler._registry.get("sid-final-append-retry").pending_input_bytes == 0
ws.put({"type": "session.close"})
await asyncio.wait_for(handler_task, timeout=2)
- assert len(engine.appended) == 2
- assert engine.appended[0][2]["audio"] == engine.appended[1][2]["audio"]
- assert engine.append_operation_ids[0] == engine.append_operation_ids[1]
+ assert len(engine.appended) == 1
@pytest.mark.asyncio
diff --git a/tests/model_executor/models/test_nemotron_voicechat_shapes.py b/tests/model_executor/models/test_nemotron_voicechat_shapes.py
index 8bd279bb8..9b3c01559 100644
--- a/tests/model_executor/models/test_nemotron_voicechat_shapes.py
+++ b/tests/model_executor/models/test_nemotron_voicechat_shapes.py
@@ -225,3 +225,39 @@ def test_torchaudio_mel_filterbank_matches_librosa() -> None:
reference = librosa.filters.mel(sr=16000, n_fft=512, n_mels=128, fmin=0.0, fmax=8000.0, norm="slaney")
assert ta.shape == reference.shape
assert abs(ta - reference).max() < 1e-6
+
+
+def test_code2wav_reuses_and_clears_per_request_streaming_cache() -> None:
+ import torch
+ from torch import nn
+
+ from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_code2wav import (
+ NemotronVoiceChatCode2Wav,
+ )
+
+ class Codec(nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.anchor = nn.Parameter(torch.zeros(()))
+ self.caches = []
+
+ def decode(self, codes, lengths, *, cache):
+ self.caches.append(cache)
+ wav = torch.zeros((1, int(codes.shape[1]) * 1764), device=codes.device)
+ return wav, lengths * 1764
+
+ 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.audio_codec = Codec()
+
+ def info(codes):
+ return [{"meta": {"codec_streaming": True, "request_id": "req"}, "codes": {"audio": codes}}]
+
+ 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
+ model.on_requests_finished({"req"})
+ assert "req" not in model._duplex_codec_caches
diff --git a/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py b/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py
index dcdff5bbb..aba69ad14 100644
--- a/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py
+++ b/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
@@ -285,9 +285,46 @@ def test_talker2code2wav_async_chunk_requires_prompt_len() -> None:
talker2code2wav_async_chunk(tm, {"codes": {"audio": torch.zeros(5, 31, dtype=torch.long)}}, req)
-# =========================
-# Talker drain-per-wake (async scheduling review P1/P2)
-# =========================
+def test_talker2code2wav_duplex_keeps_cumulative_history_across_wakes() -> None:
+ from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import (
+ talker2code2wav_async_chunk,
+ )
+
+ tm = _fake_transfer_manager(codec_chunk_frames=2)
+ prompt_len = 3
+ req = _streaming_request(
+ prompt_len=prompt_len,
+ generated=[],
+ info={"meta": {"codec_streaming": torch.tensor([True])}, "nvc_logical_prompt_len": prompt_len},
+ )
+ prompt_rows = torch.zeros((prompt_len - 1, 31), dtype=torch.long)
+ first = torch.full((1, 31), 101, dtype=torch.long)
+ second = torch.full((1, 31), 202, dtype=torch.long)
+
+ for rows, finished in ((prompt_rows, False), (torch.cat([prompt_rows, first]), False)):
+ payload = talker2code2wav_async_chunk(
+ tm,
+ {"codes": {"audio": rows}, "meta": {"codec_streaming": torch.tensor([True])}},
+ req,
+ is_finished=finished,
+ )
+ assert payload is not None and not bool(payload.meta.finished)
+
+ payload = talker2code2wav_async_chunk(
+ tm,
+ {
+ "codes": {"audio": torch.cat([prompt_rows, first, second])},
+ "meta": {"codec_streaming": torch.tensor([True])},
+ },
+ req,
+ is_finished=True,
+ )
+
+ assert payload is not None
+ assert not bool(payload.meta.finished)
+ assert not bool(payload.meta.is_segment_finished)
+ assert torch.equal(payload.codes.audio, torch.cat([first, second]))
+ assert tm.request_payload["req-0"]["nvc_frames_seen"] == 2
class _FakeTTSModelOutput:
@@ -342,7 +379,14 @@ def _bare_talker():
return talker
-def _async_info(timeline: list[int], prompt_len: int, finished: bool, request_id: str = "req-0") -> dict:
+def _async_info(
+ timeline: list[int],
+ prompt_len: int,
+ finished: bool,
+ request_id: str = "req-0",
+ *,
+ codec_streaming: bool = False,
+) -> dict:
return {
"request_id": request_id,
"additional_information": {
@@ -351,6 +395,7 @@ def _async_info(timeline: list[int], prompt_len: int, finished: bool, request_id
"num_processed_tokens": prompt_len,
"finished": torch.tensor(finished),
"request_id": request_id,
+ "codec_streaming": codec_streaming,
},
},
}
@@ -373,7 +418,7 @@ def test_talker_drains_all_received_positions_per_wake() -> None:
_, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info)
assert session["step"] == prompt_len # t advanced 1 -> P in one wake
assert update["codes"]["audio"].shape == (prompt_len - 1, 31)
- assert float(embeds[0, 0]) == 0.0 # not finished
+ assert float(embeds[0, 0]) == 0.0 # offline stream waits for upstream terminal
# Coalesced chunk: TWO new frame tokens arrive in ONE wake (delayed save
# thread merged two thinker steps). Both must be consumed by this wake.
@@ -383,14 +428,13 @@ def test_talker_drains_all_received_positions_per_wake() -> None:
assert update["codes"]["audio"].shape == (prompt_len + 1, 31)
assert float(embeds[0, 0]) == 0.0
- # Zero-progress wake (no new positions, not finished): plain CONTINUE, no
- # fabricated frames, no codes update.
+ # Zero-progress wake: never fabricate frames or end the offline stream
+ # before the thinker sends its terminal marker.
_, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info2)
assert session["step"] == prompt_len + 2
assert update == {}
assert float(embeds[0, 0]) == 0.0
- # Final chunk: one more token + finished marker -> drain + stop flag.
info3 = _async_info([pad] * prompt_len + [7, 8, 9], prompt_len, finished=True)
_, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info3)
assert session["step"] == prompt_len + 3
@@ -398,6 +442,18 @@ def test_talker_drains_all_received_positions_per_wake() -> None:
assert float(embeds[0, 0]) == 1.0 # stop
+def test_talker_native_codec_stream_ends_each_received_segment() -> None:
+ talker = _bare_talker()
+ ids = torch.zeros(1, dtype=torch.long)
+ info = _async_info([12, 12, 7], 2, finished=False, codec_streaming=True)
+ talker.preprocess(ids, None, _omni_is_prefill=True, **info)
+
+ _, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info)
+
+ assert update["codes"]["audio"].shape == (2, 31)
+ 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."""
talker = _bare_talker()
diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py
index 99f609ea5..539329d24 100644
--- a/tests/worker/test_gpu_ar_model_runner.py
+++ b/tests/worker/test_gpu_ar_model_runner.py
@@ -1832,6 +1832,7 @@ class TestPreferModelSamplerNoneFallback:
"hunyuan_image3",
"minicpmo_4_5",
"minimax_music3",
+ "nemotron_voicechat",
}
assert declarers == expected, (
"The set of models declaring `prefer_model_sampler` changed:\n"
diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py
index 64bd94916..abecfdeef 100644
--- a/vllm_omni/core/sched/omni_ar_scheduler.py
+++ b/vllm_omni/core/sched/omni_ar_scheduler.py
@@ -561,8 +561,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
if self.chunk_transfer_adapter:
if self.vllm_config.model_config.stage_id != 0:
# Downstream async-chunk stages receive real payloads from the
- # connector. This update only resumes polling for the next segment.
- self.chunk_transfer_adapter.segment_finished_requests.discard(request.request_id)
+ # connector, not an external StreamingUpdate. Move the
+ # resumable request back to the ordinary waiting queue so
+ # connector polling starts again immediately.
+ self._resume_downstream_chunk_receiver(request)
outstanding_async_tokens = request.num_output_placeholders
# Always record the discard signal (0 when nothing is in
# flight). Upstream a0c092ee72 removed the
@@ -641,6 +643,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
inter_stage_output,
request,
is_segment_finished,
+ new_token_ids=new_token_ids,
confirmed_num_computed_tokens=confirmed_num_computed_tokens,
)
diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py
index aef8332fe..4b7c89539 100644
--- a/vllm_omni/core/sched/omni_scheduler_mixin.py
+++ b/vllm_omni/core/sched/omni_scheduler_mixin.py
@@ -564,8 +564,37 @@ class OmniSchedulerMixin:
if stopped_running_reqs:
self.running = remove_all(self.running, stopped_running_reqs)
if stopped_preempted_reqs:
- self.waiting.remove_requests(stopped_preempted_reqs)
- self.skipped_waiting.remove_requests(stopped_preempted_reqs)
+ # ``_handle_stopped_request`` re-enqueues a resumable segment as
+ # WAITING_FOR_STREAMING_REQ before this cleanup runs. A
+ # downstream async-chunk request can have entered the update with
+ # the older WAITING_FOR_CHUNK status, which puts it in
+ # ``stopped_preempted_reqs``; deleting it here would immediately
+ # undo the requeue and the receiver would never poll chunk 1.
+ removable = {
+ request
+ for request in stopped_preempted_reqs
+ if not (
+ request.resumable
+ and request.status
+ in (
+ RequestStatus.WAITING,
+ RequestStatus.WAITING_FOR_STREAMING_REQ,
+ )
+ )
+ }
+ self.waiting.remove_requests(removable)
+ self.skipped_waiting.remove_requests(removable)
+
+ def _resume_downstream_chunk_receiver(self, request: Request) -> None:
+ """Resume connector polling without waiting for an external update."""
+ adapter = self.chunk_transfer_adapter
+ adapter.segment_finished_requests.discard(request.request_id)
+ if request.status != RequestStatus.WAITING_FOR_STREAMING_REQ:
+ return
+ self.num_waiting_for_streaming_input -= 1
+ request.status = RequestStatus.WAITING
+ self.skipped_waiting.remove_requests((request,))
+ self._enqueue_waiting_request(request)
def _aggregate_kv_connector_stats(
self,
diff --git a/vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml b/vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml
new file mode 100644
index 000000000..1a69b24e6
--- /dev/null
+++ b/vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml
@@ -0,0 +1,47 @@
+# Nemotron VoiceChat native Full-Duplex deployment.
+#
+# 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.
+base_config: nemotron_labs_voicechat_streaming.yaml
+pipeline: nemotron_voicechat
+session_mode: duplex
+active_stream_window: 1
+
+duplex_session:
+ idle_ttl_s: 300
+ disconnect_grace_s: 30
+ reaper_interval_s: 5
+ resume_replay_ttl_s: 60
+ resume_replay_max_bytes_per_session: 8388608
+ max_pending_input_bytes_per_session: 16777216
+ max_pending_turns_per_session: 4
+ max_sessions: 1
+ completed_append_cache_size: 256
+
+async_chunk: true
+
+connectors:
+ connector_of_shared_memory:
+ name: SharedMemoryConnector
+ extra:
+ shm_threshold_bytes: 65536
+ codec_chunk_frames: 1
+
+stages:
+ - stage_id: 0
+ max_num_seqs: 1
+ async_scheduling: false
+ default_sampling_params:
+ temperature: 0.0
+ max_tokens: 1
+ ignore_eos: true
+ seed: 0
+
+ - stage_id: 1
+ max_num_seqs: 1
+ async_scheduling: false
+
+ - stage_id: 2
+ max_num_seqs: 1
+ async_scheduling: false
diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py
index f79192faa..34cc4851a 100644
--- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py
+++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py
@@ -3,9 +3,11 @@
import copy
import importlib
+import inspect
+import threading
import time
from collections import defaultdict, deque
-from collections.abc import Callable, Mapping
+from collections.abc import Callable, Iterable, Mapping
from typing import Any
import torch
@@ -24,6 +26,14 @@ from .base import OmniTransferAdapterBase
logger = get_connector_logger(__name__)
+class _SenderGeneration:
+ """Fence one external request generation without blocking cleanup."""
+
+ def __init__(self) -> None:
+ self.in_flight = False
+ self.cancelled = False
+
+
class OmniChunkTransferAdapter(OmniTransferAdapterBase):
"""Chunk-level transfer adapter for Omni connector pipelines.
@@ -44,6 +54,12 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
def __init__(self, vllm_config: Any):
model_config = vllm_config.model_config
+ # The base constructor starts the save thread, so sender-generation
+ # state must exist before it can observe an enqueued task. The global
+ # lock only protects this map and short state transitions. Slow payload
+ # construction and connector writes never hold it.
+ self._sender_state_lock = threading.Lock()
+ self._sender_tokens: dict[str, _SenderGeneration] = {}
self.scheduler_max_num_seqs = vllm_config.scheduler_config.max_num_seqs
active_stream_window = int(getattr(model_config, "active_stream_window", 0) or 0)
model_max_num_seqs = int(getattr(model_config, "max_num_seqs", self.scheduler_max_num_seqs) or 0)
@@ -70,6 +86,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
module_path, func_name = custom_process_next_stage_input_func.rsplit(".", 1)
module = importlib.import_module(module_path)
self.custom_process_next_stage_input_func = getattr(module, func_name)
+ self._processor_accepts_step_tokens: dict[Callable[..., Any], bool] = {}
# mapping for request id and chunk id
self.put_req_chunk: dict[str, int] = defaultdict(int)
self.get_req_chunk: dict[str, int] = defaultdict(int)
@@ -198,11 +215,28 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
with self._recv_cond:
self._recv_cond.notify()
+ def _accepts_new_token_ids(self, processor: Callable[..., Any]) -> bool:
+ cached = self._processor_accepts_step_tokens.get(processor)
+ if cached is not None:
+ return cached
+ try:
+ parameters = inspect.signature(processor).parameters.values()
+ except (TypeError, ValueError):
+ accepts = False
+ else:
+ accepts = any(
+ parameter.name == "new_token_ids" or parameter.kind is inspect.Parameter.VAR_KEYWORD
+ for parameter in parameters
+ )
+ self._processor_accepts_step_tokens[processor] = accepts
+ return accepts
+
def save_async(
self,
multimodal_output: dict[str, Any] | None = None,
request: Request | None = None,
is_segment_finished: bool = False,
+ new_token_ids: Iterable[int] | None = None,
confirmed_num_computed_tokens: int | None = None,
):
"""Build and enqueue one chunk for asynchronous sending.
@@ -221,38 +255,56 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
multimodal_output: Per-request multimodal output dictionary
request: Request object
is_segment_finished: whether the segment of request is finished
+ new_token_ids: immutable snapshot of tokens generated by this
+ scheduler step, before resumable cleanup mutates the request
confirmed_num_computed_tokens: committed token count captured
before a streaming transition can mutate ``request``
"""
is_finished = request.is_finished() and not request.resumable
-
+ external_req_id = request.external_req_id
if confirmed_num_computed_tokens is None:
confirmed_num_computed_tokens = self._confirmed_num_computed_tokens(request)
-
- # If the request is preempted, skip the already saved chunks.
- if confirmed_num_computed_tokens < self.requests_num_chunks_sent.get(request.external_req_id, 0):
- logger.warning(
- f"Enqueue save_async for request {request.external_req_id}, "
- f"request.num_computed_tokens={request.num_computed_tokens}, "
- f"request.num_output_placeholders={getattr(request, 'num_output_placeholders', 0)}, "
- f"previous_chunks_sent={self.requests_num_chunks_sent.get(request.external_req_id, 0)}"
- )
- return
-
- self.requests_num_chunks_sent[request.external_req_id] = confirmed_num_computed_tokens
+ processor_request = (
+ self._snapshot_processor_request(request) if self.custom_process_next_stage_input_func else request
+ )
task = {
"multimodal_output": multimodal_output,
- "request": (
- self._snapshot_processor_request(request) if self.custom_process_next_stage_input_func else request
- ),
+ "request": processor_request,
"is_finished": is_finished,
"is_segment_finished": is_segment_finished,
+ "new_token_ids": tuple(int(token_id) for token_id in (new_token_ids or ())),
}
- self._pending_save_reqs.append(task)
- if is_segment_finished:
- # The queued FIFO item now owns the old segment. Start the next
- # segment's deduplication watermark before the worker sends it.
- self.requests_num_chunks_sent.pop(request.external_req_id, None)
+
+ reject_reason = None
+ with self._sender_state_lock:
+ sender_token = self._sender_tokens.get(external_req_id)
+ if sender_token is not None and sender_token.cancelled:
+ reject_reason = "previous sender generation is still draining"
+ # If the request is preempted, skip the already saved chunks.
+ elif confirmed_num_computed_tokens < self.requests_num_chunks_sent.get(external_req_id, 0):
+ logger.warning(
+ f"Enqueue save_async for request {external_req_id}, "
+ f"request.num_computed_tokens={request.num_computed_tokens}, "
+ f"request.num_output_placeholders={getattr(request, 'num_output_placeholders', 0)}, "
+ f"previous_chunks_sent={self.requests_num_chunks_sent.get(external_req_id, 0)}"
+ )
+ return
+
+ else:
+ self.requests_num_chunks_sent[external_req_id] = confirmed_num_computed_tokens
+ if sender_token is None:
+ sender_token = _SenderGeneration()
+ self._sender_tokens[external_req_id] = sender_token
+ task["sender_token"] = sender_token
+ self._pending_save_reqs.append(task)
+ if is_segment_finished:
+ # The queued FIFO item now owns the old segment. Start the next
+ # segment's deduplication watermark before the worker sends it.
+ self.requests_num_chunks_sent.pop(external_req_id, None)
+ if reject_reason is not None:
+ logger.error("Cannot enqueue %s: %s", external_req_id, reject_reason)
+ self.record_send_failure(request.request_id, reject_reason)
+ return
with self._save_cond:
self._save_cond.notify()
@@ -378,6 +430,44 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
return False
def _send_single_request(self, task: dict):
+ request = task["request"]
+ external_req_id = request.external_req_id
+ sender_token = task.get("sender_token")
+ if sender_token is None:
+ self._send_single_request_for_generation(task)
+ return
+ with self._sender_state_lock:
+ is_current = self._sender_tokens.get(external_req_id) is sender_token
+ if is_current and not sender_token.cancelled:
+ sender_token.in_flight = True
+ else:
+ is_current = False
+ if not is_current:
+ logger.debug("Discarding stale queued chunk for aborted request %s", external_req_id)
+ return
+ try:
+ self._send_single_request_for_generation(task, sender_token)
+ finally:
+ with self._sender_state_lock:
+ if self._sender_tokens.get(external_req_id) is sender_token:
+ sender_token.in_flight = False
+ if sender_token.cancelled:
+ self._sender_tokens.pop(external_req_id, None)
+ self._clear_sender_state_locked(external_req_id)
+
+ def _sender_generation_is_active(
+ self,
+ external_req_id: str,
+ sender_token: _SenderGeneration,
+ ) -> bool:
+ with self._sender_state_lock:
+ return self._sender_tokens.get(external_req_id) is sender_token and not sender_token.cancelled
+
+ def _send_single_request_for_generation(
+ self,
+ task: dict,
+ sender_token: _SenderGeneration | None = None,
+ ):
raw_mm = task["multimodal_output"]
multimodal_output = unflatten_payload(raw_mm) if isinstance(raw_mm, Mapping) else raw_mm
request = task["request"]
@@ -392,10 +482,11 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
payload_data: OmniPayloadStruct | None = None
if self.custom_process_next_stage_input_func:
try:
- payload_data = self.custom_process_next_stage_input_func(
- transfer_manager=self,
- multimodal_output=multimodal_output,
- request=request,
+ processor = self.custom_process_next_stage_input_func
+ processor_kwargs = {
+ "transfer_manager": self,
+ "multimodal_output": multimodal_output,
+ "request": request,
# Existing processors use is_finished as a flush signal.
# Terminal stops no longer count as segment boundaries
# (is_segment_finished is False when the request finishes,
@@ -403,8 +494,11 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
# accumulated tail on the terminal chunk — otherwise the
# downstream stage receives the finished marker without
# the final payload (#5413).
- is_finished=is_segment_finished or is_finished,
- )
+ "is_finished": is_segment_finished or is_finished,
+ }
+ if self._accepts_new_token_ids(processor):
+ processor_kwargs["new_token_ids"] = task.get("new_token_ids", ())
+ payload_data = processor(**processor_kwargs)
except Exception as e:
logger.error(f"Failed to use custom_process_input_func for payload extraction: {e}")
@@ -417,10 +511,23 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
payload_data = OmniPayloadStruct()
if payload_data.meta is None:
payload_data.meta = MetaStruct()
- payload_data.meta.finished = torch.tensor(is_finished, dtype=torch.bool)
+ # A model-native codec stream may span many resumable scheduler
+ # requests. Its processor is the only layer that knows whether a
+ # request finish is a transport wake or the lifetime of the codec
+ # state itself. Preserve an explicit streaming lifetime decision;
+ # all ordinary producers keep the scheduler-derived default.
+ processor_controls_codec_lifetime = payload_data.meta.finished is not None and self._is_truthy_scalar(
+ payload_data.meta.codec_streaming
+ )
+ if not processor_controls_codec_lifetime:
+ payload_data.meta.finished = torch.tensor(is_finished, dtype=torch.bool)
if payload_data.meta.is_segment_finished is None:
payload_data.meta.is_segment_finished = torch.tensor(is_segment_finished, dtype=torch.bool)
+ if sender_token is not None and not self._sender_generation_is_active(external_req_id, sender_token):
+ logger.debug("Skipping cancelled chunk for request %s before connector put", external_req_id)
+ return
+
success, size, metadata = self.connector.put(
from_stage=str(stage_id),
to_stage=str(next_stage_id),
@@ -428,6 +535,13 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
data=payload_data,
)
+ if sender_token is not None and not self._sender_generation_is_active(external_req_id, sender_token):
+ # cleanup_sender() may cancel this generation while put() is
+ # blocked. Do not let that stale completion recreate state for a
+ # request whose cleanup is already in progress.
+ logger.debug("Ignoring completed put for cancelled request %s", external_req_id)
+ return
+
if success:
self.put_req_chunk[external_req_id] += 1
self.ramp_chunk_count[external_req_id] += 1
@@ -535,11 +649,29 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
def cleanup_sender(self, external_req_id: str) -> None:
"""Reclaim sender-side per-request state (keyed by external id).
- Must only be called after the terminal chunk has actually been
- sent (i.e. from ``_send_single_request``), not before.
+ Called after a terminal chunk is sent or when the scheduler aborts the
+ request before a terminal chunk can be produced. In-flight sends are
+ cancelled here and reclaim their own state in ``finally``; cleanup
+ never waits for connector I/O on the scheduler thread.
Idempotent: calling with an already-cleaned or unknown id is safe.
"""
+ with self._sender_state_lock:
+ sender_token = self._sender_tokens.get(external_req_id)
+ if sender_token is None:
+ self._clear_sender_state_locked(external_req_id)
+ return
+ sender_token.cancelled = True
+ if sender_token.in_flight:
+ # The sender's finally block is the first point where no
+ # in-flight code can mutate this request's state. Let it own
+ # cleanup without making the scheduler wait for connector I/O.
+ return
+ self._sender_tokens.pop(external_req_id, None)
+ self._clear_sender_state_locked(external_req_id)
+
+ def _clear_sender_state_locked(self, external_req_id: str) -> None:
+ """Clear sender state while ``_sender_state_lock`` is held."""
self.put_req_chunk.pop(external_req_id, None)
self.request_payload.pop(external_req_id, None)
self.code_prompt_token_ids.pop(external_req_id, None)
@@ -959,6 +1091,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
}
# First pass: collect requests to remove from queues
+ request_ids = set(request_ids)
for req_id in request_ids:
request = requests.get(req_id) if requests else None
if request is None:
@@ -974,7 +1107,18 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
if req_id in self.requests_origin_status and req_id in connector_owned_ids:
request.status = self.requests_origin_status.pop(req_id)
- request_ids = set(request_ids)
+ # An abort can terminate a long-lived native codec stream before it
+ # emits a terminal payload. Reclaim both sides of the adapter so a
+ # later request reusing the same external id starts with empty state.
+ for req_id in request_ids:
+ request = requests.get(req_id) if requests else None
+ external_req_id = (
+ getattr(request, "external_req_id", None)
+ if request is not None
+ else self.request_ids_mapping.get(req_id)
+ ) or req_id
+ self.cleanup_receiver(req_id)
+ self.cleanup_sender(external_req_id)
self.waiting_for_chunk_waiting_requests = deque(
request for request in self.waiting_for_chunk_waiting_requests if request.request_id not in request_ids
diff --git a/vllm_omni/experimental/fullduplex/nemotron_voicechat/__init__.py b/vllm_omni/experimental/fullduplex/nemotron_voicechat/__init__.py
new file mode 100644
index 000000000..ce738cce6
--- /dev/null
+++ b/vllm_omni/experimental/fullduplex/nemotron_voicechat/__init__.py
@@ -0,0 +1,4 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Native Full-Duplex integration for NVIDIA NemotronLabs VoiceChat."""
diff --git a/vllm_omni/experimental/fullduplex/nemotron_voicechat/data_plane.py b/vllm_omni/experimental/fullduplex/nemotron_voicechat/data_plane.py
new file mode 100644
index 000000000..7cc2a43ef
--- /dev/null
+++ b/vllm_omni/experimental/fullduplex/nemotron_voicechat/data_plane.py
@@ -0,0 +1,403 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Output projection for Nemotron VoiceChat's independent channels."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Callable, Iterator, Mapping
+from dataclasses import dataclass, field
+from typing import Any
+from uuid import uuid4
+
+import numpy as np
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.experimental.fullduplex.engine.contracts import (
+ duplex_resource_request_belongs_to_session,
+)
+from vllm_omni.experimental.fullduplex.output import (
+ get_duplex_output_decision,
+)
+
+logger = init_logger(__name__)
+
+EncodeAudio = Callable[[object, int, str, float | None], str | None]
+
+
+@dataclass(frozen=True, slots=True)
+class NemotronVoiceChatDataPlaneContext:
+ epoch: int = 0
+ turn_id: int = 0
+ auto_responds: bool = True
+ response_format: str = "wav"
+ speed: float | None = None
+ modalities: tuple[str, ...] = ("text", "audio")
+
+
+@dataclass(slots=True)
+class _RequestState:
+ text_frames: int = 0
+ audio_frames: int = 0
+ pending_speech_end_frames: list[int] = field(default_factory=list)
+ function_active: bool = False
+ function_tokens: list[int] = field(default_factory=list)
+ function_call_id: str | None = None
+ terminal: bool = False
+
+
+def _coerce_ints(value: object) -> list[int]:
+ if value is None:
+ return []
+ if isinstance(value, torch.Tensor):
+ value = value.detach().cpu().reshape(-1).tolist()
+ elif isinstance(value, np.ndarray):
+ value = value.reshape(-1).tolist()
+ if isinstance(value, int):
+ return [value]
+ if not isinstance(value, list | tuple):
+ return []
+ values: list[int] = []
+ for item in value:
+ if isinstance(item, torch.Tensor):
+ if item.numel() != 1:
+ continue
+ item = item.item()
+ elif isinstance(item, np.generic):
+ item = item.item()
+ try:
+ values.append(int(item))
+ except (TypeError, ValueError):
+ continue
+ return values
+
+
+def _unwrap(output: object) -> tuple[object, object | None, int | None]:
+ stage_id = getattr(output, "stage_id", None)
+ inner = getattr(output, "request_output", None)
+ if inner is not None and inner is not output:
+ output = inner
+ outputs = getattr(output, "outputs", None)
+ completion = outputs[0] if isinstance(outputs, list) and outputs else None
+ if stage_id is None:
+ stage_id = getattr(output, "stage_id", None)
+ return output, completion, int(stage_id) if isinstance(stage_id, int) else None
+
+
+def _multimodal(output: object, completion: object | None) -> dict[str, object]:
+ decision = get_duplex_output_decision(output)
+ metadata = getattr(decision, "metadata", None)
+ if isinstance(metadata, Mapping):
+ return dict(metadata)
+ for candidate in (
+ getattr(output, "multimodal_output", None),
+ getattr(completion, "multimodal_output", None) if completion is not None else None,
+ ):
+ if isinstance(candidate, Mapping):
+ return dict(candidate)
+ return {}
+
+
+def _audio_value(metadata: Mapping[str, object]) -> object | None:
+ value = next((metadata[key] for key in ("audio", "model_outputs", "latent") if key in metadata), None)
+ if isinstance(value, list) and len(value) == 1:
+ return value[0]
+ return value
+
+
+def _sample_rate(metadata: Mapping[str, object]) -> int:
+ value = metadata.get("sr", metadata.get("sample_rate_hz", 22050))
+ if isinstance(value, list) and value:
+ value = value[0]
+ if hasattr(value, "item"):
+ value = value.item()
+ return int(value) if isinstance(value, int | float) else 22050
+
+
+def _audio_samples(audio: object | None) -> int:
+ if audio is None:
+ return 0
+ if isinstance(audio, torch.Tensor):
+ return int(audio.numel())
+ try:
+ return int(np.asarray(audio, dtype=np.float32).size)
+ except (TypeError, ValueError):
+ return 0
+
+
+def _speech_end_event(request_id: str) -> dict[str, object]:
+ return {
+ "stage_role": "tts",
+ "is_listen": False,
+ "data_plane_request_id": request_id,
+ "text": "",
+ "end_of_turn": True,
+ }
+
+
+class NemotronVoiceChatDataPlaneSession:
+ """Join frame-locked text/function outputs with Stage-2 audio."""
+
+ def __init__(self, encode_audio: EncodeAudio) -> None:
+ self._encode_audio = encode_audio
+ self._requests: dict[str, _RequestState] = {}
+ self._tokenizer = None
+ self._tokenizer_ref: str | None = None
+ self._special_ids: dict[str, int] | None = None
+
+ def configure_runtime(self, runtime_config: Mapping[str, object], *, tokenizer: Any | None = None) -> None:
+ """Install the serving-resolved tokenizer contract as the sole truth."""
+ try:
+ special_ids = {
+ "bos": int(runtime_config["nvc_text_bos_id"]),
+ "eos": int(runtime_config["nvc_text_eos_id"]),
+ "pad": int(runtime_config["nvc_text_pad_id"]),
+ "sotc": int(runtime_config["nvc_function_sotc_id"]),
+ "eotc": int(runtime_config["nvc_function_eotc_id"]),
+ }
+ except (KeyError, TypeError, ValueError) as exc:
+ raise ValueError("Nemotron VoiceChat data plane requires resolved special-token ids") from exc
+ tokenizer_ref = runtime_config.get("nvc_tokenizer_ref")
+ if not isinstance(tokenizer_ref, str) or not tokenizer_ref:
+ raise ValueError("Nemotron VoiceChat data plane requires nvc_tokenizer_ref")
+ self._special_ids = special_ids
+ self._tokenizer_ref = tokenizer_ref
+ self._tokenizer = tokenizer
+
+ def begin_request(self, request_id: str) -> None:
+ self._requests.setdefault(request_id, _RequestState()).terminal = False
+
+ def is_terminal(self, request_id: str | None) -> bool:
+ return bool(request_id and self._requests.get(request_id) and self._requests[request_id].terminal)
+
+ def mark_terminal(self, request_id: str) -> None:
+ self._requests.setdefault(request_id, _RequestState()).terminal = True
+
+ def close_stream(self, request_id: str) -> None:
+ self._requests.pop(request_id, None)
+
+ def close_session(self, session_id: str, *, active_request_id: str | None = None) -> None:
+ if active_request_id:
+ self._requests.pop(active_request_id, None)
+ for request_id in tuple(self._requests):
+ if duplex_resource_request_belongs_to_session(request_id, session_id):
+ self._requests.pop(request_id, None)
+
+ def _load_tokenizer(self):
+ if self._tokenizer is None:
+ from transformers import AutoTokenizer
+
+ if self._tokenizer_ref is None:
+ raise RuntimeError("Nemotron VoiceChat data plane was used before runtime configuration")
+ self._tokenizer = AutoTokenizer.from_pretrained(self._tokenizer_ref, trust_remote_code=False)
+ return self._tokenizer
+
+ def _ids(self) -> dict[str, int]:
+ if self._special_ids is None:
+ raise RuntimeError("Nemotron VoiceChat data plane was used before runtime configuration")
+ return self._special_ids
+
+ def _decode(self, token_ids: list[int]) -> str:
+ if not token_ids:
+ return ""
+ return str(self._load_tokenizer().decode(token_ids, skip_special_tokens=True))
+
+ def project(
+ self,
+ result: object,
+ *,
+ context: NemotronVoiceChatDataPlaneContext | None = None,
+ ) -> Iterator[dict[str, object]]:
+ if not isinstance(result, dict):
+ return
+ outputs = result.get("data_plane_outputs")
+ if not isinstance(outputs, list):
+ return
+ for output in outputs:
+ yield from self.project_output(output, context=context)
+
+ def project_output(
+ self,
+ output: object,
+ *,
+ context: NemotronVoiceChatDataPlaneContext | None = None,
+ ) -> Iterator[dict[str, object]]:
+ context = context or NemotronVoiceChatDataPlaneContext()
+ outer = output
+ output, completion, stage_id = _unwrap(output)
+ request_id = getattr(output, "request_id", None) or getattr(outer, "request_id", None)
+ request_id = str(request_id) if request_id is not None else ""
+ state = self._requests.setdefault(request_id, _RequestState())
+ metadata = _multimodal(outer, completion)
+ modalities = {modality.lower() for modality in context.modalities}
+ wants_text = "text" in modalities
+ wants_audio = "audio" in modalities
+ if stage_id == 0 or "nvc_text_token_ids" in metadata:
+ ids = self._ids()
+ text_ids = _coerce_ints(metadata.get("nvc_text_token_ids"))
+ if not text_ids and completion is not None:
+ text_ids = _coerce_ints(
+ getattr(completion, "token_ids", None) or getattr(completion, "cumulative_token_ids", None)
+ )
+ for token_id in text_ids[-1:]:
+ state.text_frames += 1
+ if token_id == ids["eos"]:
+ if state.audio_frames >= state.text_frames:
+ yield _speech_end_event(request_id)
+ else:
+ state.pending_speech_end_frames.append(state.text_frames)
+ elif token_id == ids["pad"]:
+ yield {
+ "stage_role": "llm",
+ "is_listen": True,
+ "model_listen": True,
+ "listen_source": "model_listen",
+ "data_plane_request_id": request_id,
+ "end_of_turn": False,
+ }
+ elif token_id != ids["bos"] and wants_text:
+ text = self._decode([token_id])
+ if text:
+ yield {
+ "stage_role": "llm",
+ "is_listen": False,
+ "data_plane_request_id": request_id,
+ "text": text,
+ "end_of_turn": False,
+ }
+
+ for function_id in _coerce_ints(metadata.get("nvc_function_token"))[-1:]:
+ yield from self._project_function_token(
+ function_id,
+ state=state,
+ request_id=request_id,
+ ids=ids,
+ )
+ return
+
+ audio = _audio_value(metadata)
+ sample_rate = _sample_rate(metadata)
+ sample_count = _audio_samples(audio)
+ encoded = (
+ self._encode_audio(audio, sample_rate, context.response_format, context.speed)
+ if wants_audio and audio is not None
+ else None
+ )
+ audio_chunk_complete = sample_count > 0
+ if audio_chunk_complete:
+ # A native-duplex Stage-2 segment can contain one or more 80 ms
+ # codec frames. Count them against the frame-locked Stage-0 text
+ # channel. The long-lived codec request deliberately keeps
+ # ``finished=False`` across scheduler wakes; a non-empty decoded
+ # chunk, not request lifetime, is the audio-coverage signal.
+ duration_ms = sample_count * 1000 / max(1, sample_rate)
+ state.audio_frames += max(1, round(duration_ms / 80))
+ end_of_turn = bool(state.pending_speech_end_frames) and state.audio_frames >= state.pending_speech_end_frames[0]
+ if encoded:
+ yield {
+ "stage_role": "tts",
+ "is_listen": False,
+ "data_plane_request_id": request_id,
+ "text": "",
+ "audio_data": encoded,
+ "audio_format": context.response_format,
+ "sample_rate_hz": sample_rate,
+ "audio_duration_ms": round(sample_count * 1000 / max(1, sample_rate)),
+ "audio_text_mark": True,
+ "end_of_turn": end_of_turn,
+ }
+ elif end_of_turn:
+ yield _speech_end_event(request_id)
+ if end_of_turn:
+ state.pending_speech_end_frames.pop(0)
+
+ def _project_function_token(
+ self,
+ token_id: int,
+ *,
+ state: _RequestState,
+ request_id: str,
+ ids: Mapping[str, int],
+ ) -> Iterator[dict[str, object]]:
+ if token_id == ids["sotc"]:
+ state.function_active = True
+ state.function_tokens.clear()
+ state.function_call_id = f"call_{uuid4().hex}"
+ return
+ if token_id == ids["eotc"] and state.function_active:
+ state.function_active = False
+ raw = self._decode(state.function_tokens)
+ state.function_tokens.clear()
+ try:
+ calls = self._parse_calls(raw)
+ except ValueError as exc:
+ state.function_call_id = None
+ yield {
+ "stage_role": "function",
+ "data_plane_request_id": request_id,
+ "error_code": "nemotron_function_call_parse_error",
+ "error": str(exc),
+ }
+ return
+ for index, call in enumerate(calls):
+ yield {
+ "stage_role": "function",
+ "data_plane_request_id": request_id,
+ "function_call": True,
+ "call_id": (
+ state.function_call_id if index == 0 and state.function_call_id else f"call_{uuid4().hex}"
+ ),
+ "name": str(call["name"]),
+ "arguments": (
+ call.get("arguments")
+ if isinstance(call.get("arguments"), str)
+ else json.dumps(call.get("arguments", {}), separators=(",", ":"))
+ ),
+ }
+ state.function_call_id = None
+ return
+ if state.function_active and token_id != ids["pad"]:
+ state.function_tokens.append(token_id)
+
+ @staticmethod
+ def _parse_calls(value: str) -> list[dict[str, object]]:
+ text = value.strip()
+ if "" in text:
+ text = text.split("", 1)[1]
+ if "" in text:
+ text = text.split("", 1)[0]
+ try:
+ parsed = json.loads(text)
+ except json.JSONDecodeError as exc:
+ raise ValueError("Nemotron VoiceChat produced malformed function-call JSON") from exc
+ if isinstance(parsed, dict):
+ parsed = [parsed]
+ if not isinstance(parsed, list) or not parsed:
+ raise ValueError("Nemotron VoiceChat function-call payload must be a non-empty object or list")
+ calls: list[dict[str, object]] = []
+ for item in parsed:
+ if not isinstance(item, dict):
+ raise ValueError("Nemotron VoiceChat function-call entries must be objects")
+ name = item.get("name")
+ if not isinstance(name, str) or not name.strip():
+ raise ValueError("Nemotron VoiceChat function-call entries require a non-empty name")
+ arguments = item.get("arguments", {})
+ if isinstance(arguments, str):
+ try:
+ decoded_arguments = json.loads(arguments)
+ except json.JSONDecodeError as exc:
+ raise ValueError("Nemotron VoiceChat function-call arguments contain malformed JSON") from exc
+ if not isinstance(decoded_arguments, dict):
+ raise ValueError("Nemotron VoiceChat function-call arguments must encode a JSON object")
+ elif not isinstance(arguments, dict):
+ raise ValueError("Nemotron VoiceChat function-call arguments must be a JSON object")
+ calls.append({"name": name.strip(), "arguments": arguments})
+ return calls
+
+
+__all__ = [
+ "NemotronVoiceChatDataPlaneContext",
+ "NemotronVoiceChatDataPlaneSession",
+]
diff --git a/vllm_omni/experimental/fullduplex/nemotron_voicechat/input.py b/vllm_omni/experimental/fullduplex/nemotron_voicechat/input.py
new file mode 100644
index 000000000..e69ebf4f4
--- /dev/null
+++ b/vllm_omni/experimental/fullduplex/nemotron_voicechat/input.py
@@ -0,0 +1,225 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""80 ms PCM packetization for Nemotron VoiceChat native duplex."""
+
+from __future__ import annotations
+
+import binascii
+
+import numpy as np
+import pybase64 as base64
+
+NEMOTRON_VOICECHAT_SAMPLE_RATE = 16000
+NEMOTRON_VOICECHAT_FRAME_SAMPLES = 1280
+_SAMPLE_BYTES = 4
+_FRAME_BYTES = NEMOTRON_VOICECHAT_FRAME_SAMPLES * _SAMPLE_BYTES
+
+
+def _require_terminal_tail(buffered_bytes: int) -> None:
+ if buffered_bytes > _FRAME_BYTES:
+ raise ValueError(
+ "Nemotron VoiceChat manual/deferred commit cannot contain more than one 80 ms frame; "
+ "enable extra_body.auto_response for native full-duplex streaming"
+ )
+
+
+def decode_pcm_f32le(payload: object, *, exact_frame: bool = False) -> bytes:
+ if not isinstance(payload, dict):
+ raise ValueError("Nemotron VoiceChat duplex audio payload must be a mapping")
+ if payload.get("format") != "pcm_f32le":
+ raise ValueError("Nemotron VoiceChat duplex audio format must be pcm_f32le")
+ if payload.get("sample_rate_hz") != NEMOTRON_VOICECHAT_SAMPLE_RATE:
+ raise ValueError("Nemotron VoiceChat duplex audio sample_rate_hz must be 16000")
+ encoded = payload.get("audio")
+ if not isinstance(encoded, str):
+ raise ValueError("Nemotron VoiceChat duplex audio must be base64 pcm_f32le")
+ try:
+ raw = base64.b64decode(encoded, validate=True)
+ except (binascii.Error, ValueError) as exc:
+ raise ValueError("Nemotron VoiceChat duplex audio is not valid base64") from exc
+ if len(raw) % _SAMPLE_BYTES:
+ raise ValueError("Nemotron VoiceChat pcm_f32le payload has a partial sample")
+ values = np.frombuffer(raw, dtype=" dict[str, object]:
+ return {
+ "type": "audio",
+ "audio": base64.b64encode(raw).decode("ascii"),
+ "format": "pcm_f32le",
+ "sample_rate_hz": NEMOTRON_VOICECHAT_SAMPLE_RATE,
+ "final": final,
+ }
+
+
+class NemotronVoiceChatPcmAppendReservation:
+ def __init__(
+ self,
+ owner: NemotronVoiceChatPcmAppendBuffer,
+ *,
+ operation_id: str,
+ payload: dict[str, object] | None,
+ raw: bytes,
+ ) -> None:
+ self._owner = owner
+ self.operation_id = operation_id
+ self.payload = payload
+ self.raw = raw
+ self._active = True
+
+ @property
+ def active(self) -> bool:
+ return self._active
+
+ @property
+ def byte_count(self) -> int:
+ return len(self.raw)
+
+ def commit(self) -> None:
+ if not self._active:
+ return
+ self._active = False
+ if self in self._owner._reservations:
+ self._owner._reservations.remove(self)
+
+ def rollback(self) -> None:
+ if not self._active:
+ return
+ try:
+ index = self._owner._reservations.index(self)
+ except ValueError:
+ self._active = False
+ return
+ restore = bytearray()
+ for reservation in self._owner._reservations[index:]:
+ if reservation._active:
+ restore.extend(reservation.raw)
+ reservation._active = False
+ del self._owner._reservations[index:]
+ self._owner._buffer[:0] = restore
+
+
+class NemotronVoiceChatPcmAppendBuffer:
+ """Adapt Realtime packets to the current full-duplex PCM buffer contract.
+
+ The model consumes exactly one 1280-sample frame per scheduler append.
+ Realtime packets may split a frame arbitrarily, but a single packet is
+ capped at one frame so the serving contract never silently drops a second
+ ready frame behind a single-reservation API.
+ """
+
+ def __init__(self) -> None:
+ self._buffer = bytearray()
+ self._reservations: list[NemotronVoiceChatPcmAppendReservation] = []
+
+ @property
+ def pending_byte_count(self) -> int:
+ return len(self._buffer)
+
+ def clear(self) -> None:
+ for reservation in self._reservations:
+ reservation._active = False
+ self._reservations.clear()
+ self._buffer.clear()
+
+ def clear_force_listen(self) -> None:
+ # Nemotron VoiceChat has no serving-side force-listen packet bit.
+ return
+
+ def has_pending(self) -> bool:
+ return bool(self._buffer)
+
+ def has_reserved(self) -> bool:
+ return any(reservation.active for reservation in self._reservations)
+
+ def _reserve_frame(
+ self,
+ *,
+ operation_id: str,
+ final: bool,
+ ) -> NemotronVoiceChatPcmAppendReservation | None:
+ if len(self._buffer) < _FRAME_BYTES:
+ return None
+ raw = bytes(self._buffer[:_FRAME_BYTES])
+ del self._buffer[:_FRAME_BYTES]
+ reservation = NemotronVoiceChatPcmAppendReservation(
+ self,
+ operation_id=operation_id,
+ payload=_frame_payload(raw, final=final),
+ raw=raw,
+ )
+ self._reservations.append(reservation)
+ return reservation
+
+ def prepare_append(
+ self,
+ payload: dict[str, object],
+ *,
+ operation_id: str,
+ chunk_period_ms: int,
+ allow_emit: bool,
+ ) -> NemotronVoiceChatPcmAppendReservation | None:
+ if chunk_period_ms != 80:
+ raise ValueError("Nemotron VoiceChat native duplex chunk_period_ms must be 80")
+ raw = decode_pcm_f32le(payload)
+ if len(raw) > _FRAME_BYTES:
+ raise ValueError("Nemotron VoiceChat Realtime append must contain at most 1280 samples")
+ self._buffer.extend(raw)
+ if not allow_emit:
+ return None
+ return self._reserve_frame(operation_id=operation_id, final=False)
+
+ def prepare_commit(
+ self,
+ *,
+ operation_id: str,
+ chunk_period_ms: int,
+ ) -> NemotronVoiceChatPcmAppendReservation:
+ if chunk_period_ms != 80:
+ raise ValueError("Nemotron VoiceChat native duplex chunk_period_ms must be 80")
+ _require_terminal_tail(len(self._buffer))
+ if not self._buffer:
+ reservation = NemotronVoiceChatPcmAppendReservation(
+ self,
+ operation_id=operation_id,
+ payload=None,
+ raw=b"",
+ )
+ else:
+ raw = bytes(self._buffer)
+ self._buffer.clear()
+ reservation = NemotronVoiceChatPcmAppendReservation(
+ self,
+ operation_id=operation_id,
+ payload=_frame_payload(raw + bytes(_FRAME_BYTES - len(raw)), final=True),
+ raw=raw,
+ )
+ self._reservations.append(reservation)
+ return reservation
+
+ def flush(self, *, chunk_period_ms: int) -> dict[str, object] | None:
+ if chunk_period_ms != 80:
+ raise ValueError("Nemotron VoiceChat native duplex chunk_period_ms must be 80")
+ _require_terminal_tail(len(self._buffer))
+ if not self._buffer:
+ return None
+ raw = bytes(self._buffer)
+ self._buffer.clear()
+ return _frame_payload(raw + bytes(_FRAME_BYTES - len(raw)), final=True)
+
+
+__all__ = [
+ "NEMOTRON_VOICECHAT_FRAME_SAMPLES",
+ "NEMOTRON_VOICECHAT_SAMPLE_RATE",
+ "NemotronVoiceChatPcmAppendBuffer",
+ "NemotronVoiceChatPcmAppendReservation",
+ "decode_pcm_f32le",
+]
diff --git a/vllm_omni/experimental/fullduplex/nemotron_voicechat/runtime.py b/vllm_omni/experimental/fullduplex/nemotron_voicechat/runtime.py
new file mode 100644
index 000000000..99de30aa7
--- /dev/null
+++ b/vllm_omni/experimental/fullduplex/nemotron_voicechat/runtime.py
@@ -0,0 +1,167 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Engine policy for Nemotron VoiceChat's frame-locked duplex timeline."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from vllm.sampling_params import RequestOutputKind, SamplingParams
+
+from vllm_omni.experimental.fullduplex.engine.contracts import (
+ DuplexAppendPlan,
+ DuplexInputMode,
+ DuplexOutputAction,
+ DuplexOutputDecision,
+)
+from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
+from vllm_omni.experimental.fullduplex.nemotron_voicechat.input import (
+ decode_pcm_f32le,
+)
+
+
+def _plain_token_ids(value: object, *, name: str) -> list[int]:
+ if not isinstance(value, list | tuple) or not value:
+ raise ValueError(f"Nemotron VoiceChat runtime requires non-empty {name}")
+ try:
+ return [int(token_id) for token_id in value]
+ except (TypeError, ValueError) as exc:
+ raise ValueError(f"Nemotron VoiceChat {name} must contain token ids") from exc
+
+
+def _positive_int(value: object, *, name: str) -> int:
+ try:
+ parsed = int(value)
+ except (TypeError, ValueError) as exc:
+ raise ValueError(f"Nemotron VoiceChat requires a positive {name}") from exc
+ if parsed <= 0:
+ raise ValueError(f"Nemotron VoiceChat requires a positive {name}")
+ return parsed
+
+
+class NemotronVoiceChatDuplexRuntimeExtension:
+ """Append one 80 ms acoustic frame to one resumable thinker request."""
+
+ def configure_sampling_params(
+ self,
+ *,
+ runtime_config: dict[str, Any],
+ defaults: tuple[object, ...],
+ ) -> tuple[object, ...]:
+ del runtime_config
+ if not defaults:
+ return defaults
+ # Every scheduler segment is one transport wake of a long-lived
+ # request. CUMULATIVE output makes the multimodal output processor
+ # concatenate all earlier Stage-2 waveforms and resend them on every
+ # wake (80, 160, 240, ... ms), causing O(n^2) audio replay. Duplex
+ # consumers require per-wake deltas at every stage.
+ configured = [params.clone() if isinstance(params, SamplingParams) else params for params in defaults]
+ for params in configured:
+ if isinstance(params, SamplingParams):
+ params.output_kind = RequestOutputKind.DELTA
+ stage0 = defaults[0]
+ if isinstance(stage0, SamplingParams):
+ stage0 = stage0.clone()
+ stage0.temperature = 0.0
+ stage0.top_p = 1.0
+ stage0.top_k = 0
+ stage0.max_tokens = 1
+ stage0.ignore_eos = True
+ stage0.output_kind = RequestOutputKind.DELTA
+ configured[0] = stage0
+ return tuple(configured)
+
+ def plan_append(
+ self,
+ *,
+ request_id: str,
+ fence: DuplexFence,
+ session_config: dict[str, Any],
+ runtime_config: dict[str, Any],
+ seq: int,
+ turn_seq: int,
+ mode: DuplexInputMode,
+ payload: object,
+ final: bool,
+ sampling_params: object,
+ ) -> DuplexAppendPlan:
+ del sampling_params, turn_seq
+ if mode is not DuplexInputMode.APPEND_AUDIO_CHUNK:
+ raise ValueError(f"Nemotron VoiceChat does not support duplex input mode {mode.value!r}")
+ decode_pcm_f32le(payload, exact_frame=True)
+ normalized_payload = dict(payload)
+ prompt_ids = _plain_token_ids(
+ runtime_config.get("nvc_prompt_token_ids"),
+ name="nvc_prompt_token_ids",
+ )
+ max_model_len = _positive_int(
+ runtime_config.get("nvc_max_model_len"),
+ name="nvc_max_model_len",
+ )
+ required_model_len = len(prompt_ids) + seq + 1
+ if required_model_len > max_model_len:
+ raise ValueError(
+ "Nemotron VoiceChat native duplex session exceeds the Stage-0 "
+ f"max_model_len: prompt_tokens={len(prompt_ids)} + input_frames={seq} + "
+ f"sampled_token=1 gives {required_model_len} > {max_model_len}; "
+ "start a new session or raise Stage 0 max_model_len"
+ )
+ try:
+ pad_id = int(runtime_config["nvc_text_pad_id"])
+ except (KeyError, TypeError, ValueError) as exc:
+ raise ValueError("Nemotron VoiceChat runtime requires nvc_text_pad_id") from exc
+ scheduler_prompt = prompt_ids + [pad_id] if seq <= 1 else [pad_id]
+ return DuplexAppendPlan(
+ prompt={
+ "prompt_token_ids": scheduler_prompt,
+ "model_intermediate_buffer": {
+ "request_id": request_id,
+ "global_request_id": [fence.session_id],
+ "duplex": {
+ "data_plane": True,
+ "fence": fence,
+ "session_id": fence.session_id,
+ "incarnation": fence.incarnation,
+ "epoch": fence.epoch,
+ "source_input_seq": seq,
+ "seq": seq,
+ "mode": mode.value,
+ "payload": normalized_payload,
+ "final": final,
+ "session_config": dict(session_config),
+ "runtime_config": dict(runtime_config),
+ "scheduler_token_budget": len(scheduler_prompt),
+ },
+ },
+ }
+ )
+
+ def decide_output(
+ self,
+ *,
+ stage_id: int,
+ final_stage_id: int,
+ segment_finished: bool,
+ segment_token_ids: tuple[int, ...],
+ segment_output_metadata: dict[str, Any],
+ output: object,
+ ) -> DuplexOutputDecision | None:
+ del final_stage_id, segment_finished, output
+ if stage_id != 0:
+ return None
+ metadata = dict(segment_output_metadata)
+ metadata["nvc_text_token_ids"] = list(segment_token_ids)
+ # Stage 0 is a client-visible side channel even though Stage 2 is the
+ # configured audio response stage. Mark it explicitly so the shared
+ # collector does not discard it while waiting for the final stage.
+ metadata["duplex_direct_response"] = True
+ return DuplexOutputDecision(
+ action=DuplexOutputAction.DIRECT_RESPONSE,
+ metadata=metadata,
+ final_output_type="text",
+ )
+
+
+__all__ = ["NemotronVoiceChatDuplexRuntimeExtension"]
diff --git a/vllm_omni/experimental/fullduplex/nemotron_voicechat/serving_adapter.py b/vllm_omni/experimental/fullduplex/nemotron_voicechat/serving_adapter.py
new file mode 100644
index 000000000..2505b3766
--- /dev/null
+++ b/vllm_omni/experimental/fullduplex/nemotron_voicechat/serving_adapter.py
@@ -0,0 +1,363 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Realtime serving adapter for Nemotron VoiceChat native duplex."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+from collections.abc import Callable, Mapping
+from copy import deepcopy
+from typing import Any
+
+from vllm_omni.experimental.fullduplex.minicpmo45.session import (
+ MiniCPMO45ServingSessionState,
+)
+from vllm_omni.experimental.fullduplex.nemotron_voicechat.data_plane import (
+ NemotronVoiceChatDataPlaneContext,
+ NemotronVoiceChatDataPlaneSession,
+)
+from vllm_omni.experimental.fullduplex.nemotron_voicechat.input import (
+ NemotronVoiceChatPcmAppendBuffer,
+)
+from vllm_omni.experimental.fullduplex.openai.protocol import DuplexCapabilities
+from vllm_omni.experimental.fullduplex.openai.runtime_adapter import ServingRuntimeConfigError
+
+EncodeAudio = Callable[[object, int, str, float | None], str | None]
+
+_DEFAULT_SYSTEM_PROMPT = (
+ "You are an AI voice assistant developed by NVIDIA. "
+ "Your name is NVIDIA Voice Chat. "
+ "Answer in a spoken, conversational style rather than a written one. "
+ "Do not repeat the same sentence over and over again. "
+ "Start the conversation by greeting the user."
+)
+_PRIVATE_KEYS = frozenset(
+ {
+ "nvc_prompt_token_ids",
+ "nvc_text_bos_id",
+ "nvc_text_eos_id",
+ "nvc_text_pad_id",
+ "nvc_function_sotc_id",
+ "nvc_function_eotc_id",
+ "nvc_function_eotr_id",
+ "nvc_max_model_len",
+ "nvc_tokenizer_ref",
+ "nvc_tools_signature",
+ "nvc_function_response_generation",
+ "nvc_function_response_token_ids",
+ "nvc_function_response_call_id",
+ "nvc_function_response_batches",
+ }
+)
+
+
+def _stt_config(model_config: Any) -> dict[str, Any]:
+ hf_config = getattr(model_config, "hf_config", None)
+ stt_cfg = getattr(hf_config, "stt_cfg", None)
+ if not isinstance(stt_cfg, dict):
+ raise ServingRuntimeConfigError("Nemotron VoiceChat checkpoint STT configuration is unavailable")
+ return stt_cfg
+
+
+def _normalized_tools(config: object) -> tuple[list[dict[str, object]], str]:
+ extra_body = getattr(config, "extra_body", None)
+ raw_tools = extra_body.get("realtime_tools") if isinstance(extra_body, dict) else None
+ if raw_tools is None:
+ return [], "[]"
+ if not isinstance(raw_tools, list):
+ raise ServingRuntimeConfigError("Nemotron VoiceChat tools must be a list")
+ if len(raw_tools) > 5:
+ raise ServingRuntimeConfigError("Nemotron VoiceChat supports at most 5 tools per session")
+
+ normalized: list[dict[str, object]] = []
+ for index, tool in enumerate(raw_tools):
+ if not isinstance(tool, dict):
+ raise ServingRuntimeConfigError(f"Nemotron VoiceChat tool {index} must be an object")
+ function = tool.get("function", tool)
+ if not isinstance(function, dict):
+ raise ServingRuntimeConfigError(f"Nemotron VoiceChat tool {index} has no function definition")
+ definition = {key: value for key, value in function.items() if key != "type"}
+ if not isinstance(definition.get("name"), str) or not str(definition["name"]).strip():
+ raise ServingRuntimeConfigError(f"Nemotron VoiceChat tool {index} requires a name")
+ normalized.append(definition)
+ signature = json.dumps(normalized, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
+ return normalized, signature
+
+
+def _render_tool_prompt(instructions: str, tools: list[dict[str, object]]) -> str:
+ if not tools:
+ return instructions
+ available = json.dumps(tools, ensure_ascii=False, sort_keys=True)
+ return (
+ f"{instructions}\n\n"
+ "You can use the following tools to assist the user if required:"
+ f"\n{available}\n\n"
+ "If you decide to call any tool(s), use the following format:\n"
+ '[{"name": "tool_name1", "arguments": "tool_args1"}, '
+ '{"name": "tool_name2", "arguments": "tool_args2"}]\n\n'
+ "The user will execute tool-calls and return responses from tool(s) in this format:\n"
+ '[{"tool_response1"}, {"tool_response2"}]\n\n'
+ "Based on the tool responses, you can call additional tools if needed, correct tool calls if any "
+ "errors are found, or just respond to the user."
+ )
+
+
+def _render_tool_response(output: str) -> str:
+ """Serialize one Realtime function result using NVIDIA's function channel."""
+ try:
+ value = json.loads(output)
+ except json.JSONDecodeError:
+ value = output
+ payload = json.dumps([value], ensure_ascii=True, separators=(",", ":"))
+ return f"{payload}"
+
+
+def _require_native_full_duplex(config: object) -> None:
+ extra_body = getattr(config, "extra_body", None)
+ enabled = isinstance(extra_body, dict) and (
+ extra_body.get("auto_response") is True or extra_body.get("full_duplex") is True
+ )
+ if not enabled:
+ raise ServingRuntimeConfigError(
+ "Nemotron VoiceChat currently supports model-native full-duplex streaming only; "
+ "set extra_body.auto_response=true",
+ code="unsupported_nemotron_duplex_mode",
+ )
+
+
+def _tokenize_runtime(model_config: Any, instructions: str) -> tuple[dict[str, object], Any]:
+ from transformers import AutoTokenizer
+
+ stt_cfg = _stt_config(model_config)
+ tokenizer_ref = os.environ.get("NEMOTRON_VOICECHAT_LLM_PATH") or stt_cfg.get(
+ "pretrained_llm", "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
+ )
+ tokenizer = AutoTokenizer.from_pretrained(tokenizer_ref, trust_remote_code=False)
+
+ def token(name: str, default: str) -> int:
+ value = tokenizer.convert_tokens_to_ids(stt_cfg.get(name, default))
+ if value is None:
+ raise ServingRuntimeConfigError(f"Nemotron VoiceChat tokenizer does not define {name}")
+ return int(value)
+
+ bos_id = token("bos_token", "")
+ eos_id = token("eos_token", "")
+ pad_id = token("pad_token", "")
+ prompt_ids = [bos_id] + list(tokenizer.encode(instructions, add_special_tokens=False)) + [eos_id]
+ runtime = {
+ "instructions": instructions,
+ "nvc_prompt_token_ids": prompt_ids,
+ "nvc_text_bos_id": bos_id,
+ "nvc_text_eos_id": eos_id,
+ "nvc_text_pad_id": pad_id,
+ "nvc_function_sotc_id": int(tokenizer.convert_tokens_to_ids("")),
+ "nvc_function_eotc_id": int(tokenizer.convert_tokens_to_ids("")),
+ "nvc_function_eotr_id": int(tokenizer.convert_tokens_to_ids("")),
+ "nvc_tokenizer_ref": str(tokenizer_ref),
+ }
+ return runtime, tokenizer
+
+
+class NemotronVoiceChatServingRuntimeAdapter:
+ adapter_id = "nemotron_voicechat"
+ silence_continuation_samples = 1280
+ # The session-owned persistent drain is the sole data-plane output
+ # consumer. Collecting on append would race that drain and can consume a
+ # Stage2 audio output before the realtime projector observes it.
+ collect_outputs_on_append = False
+ clean_response_done_prefix = ""
+ interrupted_tts_prefix = ""
+ private_runtime_config_keys = _PRIVATE_KEYS
+
+ def __init__(self, encode_audio: EncodeAudio) -> None:
+ self.session_states: dict[str, MiniCPMO45ServingSessionState] = {}
+ self.data_plane = NemotronVoiceChatDataPlaneSession(encode_audio)
+ self._tokenizer: Any | None = None
+
+ def create_session_state(self) -> MiniCPMO45ServingSessionState:
+ return MiniCPMO45ServingSessionState(
+ audio_buffer=NemotronVoiceChatPcmAppendBuffer(),
+ )
+
+ def session_state(self, session_id: str) -> MiniCPMO45ServingSessionState:
+ return self.session_states.setdefault(session_id, self.create_session_state())
+
+ def remove_session_state(self, session_id: str) -> None:
+ self.session_states.pop(session_id, None)
+
+ @staticmethod
+ def is_enabled(config: object) -> bool:
+ del config
+ return True
+
+ @staticmethod
+ def capabilities(*, max_sessions: int) -> DuplexCapabilities:
+ supports_multi_session = max_sessions > 1
+ return DuplexCapabilities(
+ supports_model_native_turn_policy=True,
+ supports_external_turn_signal=False,
+ supports_client_commit=True,
+ supports_barge_in=False,
+ supports_playback_ack=True,
+ supports_input_append=True,
+ supports_replace_latest_chunk=False,
+ supports_reencode_context=False,
+ supports_rollback_to_checkpoint=False,
+ supports_turn_commit_only=False,
+ supports_kv_lease=False,
+ supports_core_kv_lease=False,
+ supports_model_internal_state=True,
+ supports_stage_resumption=True,
+ supports_scheduler_native_append=False,
+ supports_core_resumable_request=True,
+ supports_stage_connector_handoff=True,
+ supports_independent_io_streams=True,
+ supports_realtime_endpoint=True,
+ supports_multi_session=supports_multi_session,
+ supports_multi_session_same_replica=False,
+ supports_session_lease=True,
+ supports_session_resume=True,
+ session_admission_mode="engine_managed",
+ supports_audio_truncate=False,
+ requires_model_runner_kv=True,
+ requires_native_stage_role=True,
+ implementation_level="model_native_duplex",
+ adapter_patterns=["scheduler_data_plane"],
+ input_modes=["append_audio_chunk"],
+ signal_sources=["model_native", "client_event"],
+ stage_handoff_transport="scheduler_data_plane",
+ chunk_period_ms=80,
+ target_barge_in_latency_ms=None,
+ )
+
+ @staticmethod
+ def validate_client_extra_body(extra_body: object) -> None:
+ if not isinstance(extra_body, dict):
+ return
+ private = sorted(_PRIVATE_KEYS.intersection(extra_body))
+ if private:
+ raise ServingRuntimeConfigError(
+ "Nemotron VoiceChat runtime configuration is server-owned: " + ", ".join(private)
+ )
+
+ async def prepare_runtime_config(
+ self,
+ config: object,
+ *,
+ model_config: Any,
+ ) -> dict[str, object]:
+ self.validate_client_extra_body(getattr(config, "extra_body", None))
+ _require_native_full_duplex(config)
+ instructions = str(getattr(config, "instructions", None) or _DEFAULT_SYSTEM_PROMPT)
+ tools, tools_signature = _normalized_tools(config)
+ rendered_prompt = _render_tool_prompt(instructions, tools)
+ runtime, tokenizer = await asyncio.to_thread(_tokenize_runtime, model_config, rendered_prompt)
+ max_model_len = getattr(model_config, "max_model_len", None)
+ if not isinstance(max_model_len, int) or max_model_len <= 0:
+ raise ServingRuntimeConfigError("Nemotron VoiceChat requires a positive Stage-0 max_model_len")
+ runtime["nvc_max_model_len"] = max_model_len
+ # Keep raw session values for immutable-in-incarnation update checks;
+ # only the rendered prompt token ids enter Stage-0 KV.
+ runtime["instructions"] = instructions
+ runtime["nvc_tools_signature"] = tools_signature
+ self.data_plane.configure_runtime(runtime, tokenizer=tokenizer)
+ self._tokenizer = tokenizer
+ return runtime
+
+ def runtime_config_for_function_output(
+ self,
+ item: Mapping[str, object],
+ current: Mapping[str, object],
+ ) -> dict[str, object]:
+ """Queue a client tool result for frame-locked function-channel injection."""
+ call_id = item.get("call_id")
+ output = item.get("output")
+ if not isinstance(call_id, str) or not call_id:
+ raise ServingRuntimeConfigError(
+ "function_call_output requires call_id",
+ code="invalid_function_call_output",
+ )
+ if not isinstance(output, str):
+ raise ServingRuntimeConfigError(
+ "function_call_output requires a string output",
+ code="invalid_function_call_output",
+ )
+ if self._tokenizer is None:
+ raise ServingRuntimeConfigError(
+ "Nemotron VoiceChat tokenizer is unavailable for function output",
+ code="invalid_function_call_output",
+ )
+ response_token_ids = list(
+ self._tokenizer.encode(
+ _render_tool_response(output),
+ add_special_tokens=False,
+ )
+ )
+ if not response_token_ids:
+ raise ServingRuntimeConfigError(
+ "function_call_output tokenized to an empty response",
+ code="invalid_function_call_output",
+ )
+ runtime = deepcopy(dict(current))
+ generation = int(runtime.get("nvc_function_response_generation", 0)) + 1
+ token_ids = [int(token_id) for token_id in response_token_ids]
+ batches = runtime.get("nvc_function_response_batches")
+ if not isinstance(batches, list):
+ batches = []
+ batches = [*batches, {"generation": generation, "call_id": call_id, "token_ids": token_ids}]
+ runtime["nvc_function_response_generation"] = generation
+ runtime["nvc_function_response_token_ids"] = token_ids
+ runtime["nvc_function_response_call_id"] = call_id
+ runtime["nvc_function_response_batches"] = batches
+ return runtime
+
+ @staticmethod
+ def runtime_config_for_update(
+ config: object,
+ current: Mapping[str, object],
+ ) -> dict[str, object]:
+ NemotronVoiceChatServingRuntimeAdapter.validate_client_extra_body(getattr(config, "extra_body", None))
+ _require_native_full_duplex(config)
+ runtime = deepcopy(dict(current))
+ instructions = str(getattr(config, "instructions", None) or _DEFAULT_SYSTEM_PROMPT)
+ _, tools_signature = _normalized_tools(config)
+ # The system/tool prompt is already resident in Stage-0 KV. Updating it
+ # without a new incarnation would make the advertised config disagree
+ # with model state, so reject such updates explicitly.
+ if runtime and instructions != runtime.get("instructions"):
+ raise ServingRuntimeConfigError(
+ "Nemotron VoiceChat instructions cannot change inside an active duplex incarnation"
+ )
+ if runtime and tools_signature != runtime.get("nvc_tools_signature", "[]"):
+ raise ServingRuntimeConfigError(
+ "Nemotron VoiceChat tools cannot change inside an active duplex incarnation"
+ )
+ return runtime
+
+ @staticmethod
+ def data_plane_context(
+ *,
+ epoch: int,
+ turn_id: int,
+ active_response_turn_id: int | None,
+ active_response_id: str | None,
+ auto_responds: bool,
+ response_format: str,
+ speed: float | None,
+ modalities: tuple[str, ...],
+ ) -> NemotronVoiceChatDataPlaneContext:
+ del active_response_turn_id, active_response_id
+ return NemotronVoiceChatDataPlaneContext(
+ epoch=epoch,
+ turn_id=turn_id,
+ auto_responds=auto_responds,
+ response_format=response_format,
+ speed=speed,
+ modalities=modalities,
+ )
+
+
+__all__ = ["NemotronVoiceChatServingRuntimeAdapter"]
diff --git a/vllm_omni/experimental/fullduplex/openai/realtime_input.py b/vllm_omni/experimental/fullduplex/openai/realtime_input.py
index e3fe7f3ec..e9799c1b7 100644
--- a/vllm_omni/experimental/fullduplex/openai/realtime_input.py
+++ b/vllm_omni/experimental/fullduplex/openai/realtime_input.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
from __future__ import annotations
import asyncio
@@ -1052,6 +1055,46 @@ class RealtimeInputTranslator:
item_id = str(item["id"])
item_type = item.get("type")
role = item.get("role")
+ if item_type == "function_call_output":
+ call_id = item.get("call_id")
+ output = item.get("output")
+ matching_call = any(
+ known.get("type") == "function_call" and known.get("call_id") == call_id
+ for known in self._conversation_items.values()
+ )
+ duplicate_output = any(
+ known.get("type") == "function_call_output" and known.get("call_id") == call_id
+ for known in self._conversation_items.values()
+ )
+ if not isinstance(call_id, str) or not call_id or not matching_call:
+ await self._send_realtime_payload(
+ self._realtime_error_payload(
+ "invalid_function_call_output",
+ "function_call_output requires the call_id of a completed function call",
+ event_id=event.get("event_id"),
+ param="item.call_id",
+ )
+ )
+ return None
+ if not isinstance(output, str) or duplicate_output:
+ await self._send_realtime_payload(
+ self._realtime_error_payload(
+ "invalid_function_call_output",
+ (
+ "function_call_output requires a string output"
+ if not isinstance(output, str)
+ else f"function_call_output already exists for call_id {call_id}"
+ ),
+ event_id=event.get("event_id"),
+ param="item.output",
+ )
+ )
+ return None
+ # Do not retain the result until SessionRunner has delivered its
+ # runtime update. On success, the resulting
+ # ``conversation.item.created`` event records it in the shared
+ # input/output protocol state; on failure the client may retry the
+ # same call_id instead of being rejected by provisional history.
if item_type != "message" or role in {"assistant", "system"}:
return {
"type": "turn.signal",
diff --git a/vllm_omni/experimental/fullduplex/openai/realtime_output.py b/vllm_omni/experimental/fullduplex/openai/realtime_output.py
index 57c07c852..7729092e3 100644
--- a/vllm_omni/experimental/fullduplex/openai/realtime_output.py
+++ b/vllm_omni/experimental/fullduplex/openai/realtime_output.py
@@ -346,6 +346,67 @@ class RealtimeOutputProjector:
"reason": event.get("finish_reason") or "stop",
},
)
+ if event_type == "function_call.done":
+ call_id = event.get("call_id")
+ name = event.get("name")
+ arguments = event.get("arguments", "")
+ if not isinstance(call_id, str) or not call_id or not isinstance(name, str) or not name:
+ return [{"type": "duplex.function_call.done", "event": event}]
+ if not isinstance(arguments, str):
+ arguments = str(arguments)
+ response_id = f"resp_{uuid4().hex}"
+ item = {
+ "id": f"item_{uuid4().hex}",
+ "object": "realtime.item",
+ "type": "function_call",
+ "status": "completed",
+ "name": name,
+ "call_id": call_id,
+ "arguments": arguments,
+ }
+ self._conversation_items[str(item["id"])] = item
+ response = {
+ "id": response_id,
+ "object": "realtime.response",
+ "status": "completed",
+ "status_details": {"type": "completed", "reason": "completed"},
+ "output": [dict(item)],
+ "metadata": {},
+ }
+ return [
+ {"type": "response.created", "response": {**response, "status": "in_progress", "output": []}},
+ *self._conversation_item_added_events(item),
+ {
+ "type": "response.output_item.added",
+ "response_id": response_id,
+ "output_index": 0,
+ "item": dict(item),
+ },
+ {
+ "type": "response.function_call_arguments.delta",
+ "response_id": response_id,
+ "item_id": item["id"],
+ "output_index": 0,
+ "call_id": call_id,
+ "delta": arguments,
+ },
+ {
+ "type": "response.function_call_arguments.done",
+ "response_id": response_id,
+ "item_id": item["id"],
+ "output_index": 0,
+ "call_id": call_id,
+ "arguments": arguments,
+ },
+ {
+ "type": "response.output_item.done",
+ "response_id": response_id,
+ "output_index": 0,
+ "item": dict(item),
+ },
+ self._conversation_item_done_event(item),
+ {"type": "response.done", "response": response},
+ ]
if event_type == "error":
raw_error = event.get("error")
if isinstance(raw_error, dict):
diff --git a/vllm_omni/experimental/fullduplex/openai/runtime_adapter.py b/vllm_omni/experimental/fullduplex/openai/runtime_adapter.py
index 710b10022..d05b709ff 100644
--- a/vllm_omni/experimental/fullduplex/openai/runtime_adapter.py
+++ b/vllm_omni/experimental/fullduplex/openai/runtime_adapter.py
@@ -124,6 +124,7 @@ class ServingRuntimeAdapter(Protocol):
clean_response_done_prefix: str
interrupted_tts_prefix: str
private_runtime_config_keys: frozenset[str]
+ collect_outputs_on_append: bool
def create_session_state(self) -> ServingRuntimeSessionState: ...
diff --git a/vllm_omni/experimental/fullduplex/openai/runtime_bridge.py b/vllm_omni/experimental/fullduplex/openai/runtime_bridge.py
index 258f512b6..3f1d54b2a 100644
--- a/vllm_omni/experimental/fullduplex/openai/runtime_bridge.py
+++ b/vllm_omni/experimental/fullduplex/openai/runtime_bridge.py
@@ -135,7 +135,9 @@ class NativeRuntimeBridgeMixin:
incarnation=session.incarnation,
)
if self._callable_accepts_keyword(append_input, "collect_outputs"):
- append_kwargs["collect_outputs"] = False
+ append_kwargs["collect_outputs"] = bool(
+ getattr(self._serving_runtime_adapter, "collect_outputs_on_append", False)
+ )
result = await append_input(session.session_id, **append_kwargs)
except Exception as exc:
logger.exception("Failed to append duplex runtime input: %s", exc)
@@ -315,16 +317,21 @@ class NativeRuntimeBridgeMixin:
native.data_plane_task = task
return True
- # One model unit (1 s at 16 kHz) of pcm_f32le silence, matching the
- # official full-duplex behavior where the microphone keeps streaming
- # silence while the assistant speaks; replies span multiple units.
+ # One MiniCPM model unit (1 s at 16 kHz) is the compatibility default.
+ # Other native-duplex adapters may use a different frame-locked unit.
_NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO = base64.b64encode(bytes(16000 * 4)).decode("ascii")
_NATIVE_RESPONSE_MAX_CONTINUATION_UNITS = 8
def _native_silence_unit_payload(self) -> dict[str, object]:
+ samples = int(getattr(self._serving_runtime_adapter, "silence_continuation_samples", 16000))
+ audio = (
+ self._NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO
+ if samples == 16000
+ else base64.b64encode(bytes(samples * 4)).decode("ascii")
+ )
return {
"type": "audio",
- "audio": self._NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO,
+ "audio": audio,
"format": "pcm_f32le",
"sample_rate_hz": 16000,
}
@@ -826,6 +833,18 @@ class NativeRuntimeBridgeMixin:
}
)
return close_reason, True
+ if native_result.get("function_call") is True:
+ await send_json(
+ {
+ "type": "function_call.done",
+ "session_id": session.session_id,
+ "epoch": session.epoch,
+ "call_id": native_result.get("call_id"),
+ "name": native_result.get("name"),
+ "arguments": native_result.get("arguments", ""),
+ }
+ )
+ return close_reason, True
if native_result.get("requires_stage_handoff") is True or native_result.get("requires_tts_stage") is True:
# Stage0 announces a talker handoff before Stage1 has produced
# client-visible audio. A single client input stream may contain
diff --git a/vllm_omni/experimental/fullduplex/openai/serving.py b/vllm_omni/experimental/fullduplex/openai/serving.py
index d7eb31214..516d9708f 100644
--- a/vllm_omni/experimental/fullduplex/openai/serving.py
+++ b/vllm_omni/experimental/fullduplex/openai/serving.py
@@ -1728,6 +1728,10 @@ class OmniDuplexSessionHandler(
session.truncate_history_item(item_id, audio_end_ms=committed_ms)
new_epoch, old_playback = self._advance_barge_in_epoch(session)
if old_request_id is not None:
+ # Epoch-scoped request ids prevent state reuse, but explicitly
+ # release projector/parser cursors so cancelled epochs do not
+ # accumulate until the whole session closes.
+ self._serving_runtime_adapter.data_plane.close_stream(old_request_id)
await self._abort_request_background(
session,
old_request_id,
diff --git a/vllm_omni/experimental/fullduplex/openai/session_runner.py b/vllm_omni/experimental/fullduplex/openai/session_runner.py
index d94dfb43c..23aee61b6 100644
--- a/vllm_omni/experimental/fullduplex/openai/session_runner.py
+++ b/vllm_omni/experimental/fullduplex/openai/session_runner.py
@@ -345,6 +345,13 @@ class DuplexSessionRunnerMixin:
)
precreated_response_id = session.active_response_id if precreate_response else None
+ def _discard_retained_committed_audio() -> None:
+ if (
+ retained_committed_payload is not None
+ and native.committed_audio_payload is retained_committed_payload
+ ):
+ session.release_input_bytes(native.clear_committed_audio())
+
async def _run() -> bool:
nonlocal runtime_closed
try:
@@ -367,8 +374,10 @@ class DuplexSessionRunnerMixin:
and native.committed_audio_payload is retained_committed_payload
):
session.release_input_bytes(native.clear_committed_audio())
- elif pcm_reservation is not None:
- pcm_reservation.rollback()
+ else:
+ if pcm_reservation is not None:
+ pcm_reservation.rollback()
+ _discard_retained_committed_audio()
if (
not append_ok
and precreated_response_id is not None
@@ -406,6 +415,7 @@ class DuplexSessionRunnerMixin:
except Exception as exc:
if pcm_reservation is not None:
pcm_reservation.rollback()
+ _discard_retained_committed_audio()
logger.exception("Native duplex append task failed: %s", exc)
await self._send_runtime_error(emit_event, "runtime_append_task_failed", exc, session=session)
if session.state != DuplexSessionState.CLOSED:
@@ -440,19 +450,34 @@ class DuplexSessionRunnerMixin:
if not predecessor_ok:
if pcm_reservation is not None:
pcm_reservation.rollback()
+ _discard_retained_committed_audio()
return False
if actor.closing or runtime_closed or session.state != DuplexSessionState.OPEN:
if pcm_reservation is not None:
pcm_reservation.rollback()
+ _discard_retained_committed_audio()
return False
if before_append is not None and not before_append():
if pcm_reservation is not None:
pcm_reservation.rollback()
+ _discard_retained_committed_audio()
return True
if pcm_reservation is not None and not pcm_reservation.active:
+ _discard_retained_committed_audio()
return False
return await _run()
+ def _release_cancelled_retained_audio(done: asyncio.Task[bool]) -> None:
+ if done.cancelled():
+ _discard_retained_committed_audio()
+ return
+ try:
+ append_ok = done.result()
+ except Exception:
+ append_ok = False
+ if not append_ok:
+ _discard_retained_committed_audio()
+
predecessor = actor.native_append_tail
if predecessor is not None and predecessor.done():
try:
@@ -465,6 +490,7 @@ class DuplexSessionRunnerMixin:
# is an explicit retry and starts a new chain.
predecessor = None
task = asyncio.create_task(_run_in_wire_order(predecessor))
+ task.add_done_callback(_release_cancelled_retained_audio)
actor.native_append_tail = task
actor.track_append_task(
task,
@@ -929,6 +955,8 @@ class DuplexSessionRunnerMixin:
had_native_append = await actor.cancel_append_tasks(
response_bound_only=event_type in {"response.cancel", "output_audio_buffer.clear"},
)
+ if event_type == "response.cancel":
+ session.release_input_bytes(native.clear_committed_audio())
had_native_stream = native.data_plane_task is not None
cancelled = await self._cancel_active_response(
session,
@@ -1109,6 +1137,53 @@ class DuplexSessionRunnerMixin:
if turn_event == "conversation.item.create":
payload = event.get("payload")
item = payload.get("item") if isinstance(payload, dict) else None
+ item_type = item.get("type") if isinstance(item, dict) else None
+ prepare_function_output = getattr(
+ self._serving_runtime_adapter,
+ "runtime_config_for_function_output",
+ None,
+ )
+ if item_type == "function_call_output" and callable(prepare_function_output):
+ if not await wait_for_native_append_tail():
+ continue
+ try:
+ candidate_runtime_config = prepare_function_output(
+ item,
+ dict(session.runtime_config),
+ )
+ except ServingRuntimeConfigError as exc:
+ await emit_event(
+ {
+ "type": "error",
+ "session_id": session.session_id,
+ "code": exc.code,
+ "error": str(exc),
+ }
+ )
+ continue
+ if not await self._signal_runtime_session(
+ session,
+ "session.update",
+ emit_event,
+ runtime_config=candidate_runtime_config,
+ ):
+ continue
+ session.replace_runtime_config(candidate_runtime_config)
+ await emit_event(
+ {
+ "type": "conversation.item.created",
+ "session_id": session.session_id,
+ "item": item,
+ "created": True,
+ }
+ )
+ await self._maybe_continue_native_response(
+ emit_event,
+ session=session,
+ expected_epoch=session.epoch,
+ expected_model_turn_id=session.turn_id,
+ )
+ continue
message = self._realtime_item_to_history_message(item)
item_id = item.get("id") if isinstance(item, dict) else None
if message is not None:
diff --git a/vllm_omni/experimental/fullduplex/openai/websocket.py b/vllm_omni/experimental/fullduplex/openai/websocket.py
index 18cd11c8d..3b3dc4d5d 100644
--- a/vllm_omni/experimental/fullduplex/openai/websocket.py
+++ b/vllm_omni/experimental/fullduplex/openai/websocket.py
@@ -121,12 +121,15 @@ class DuplexSessionTasks:
]
if not tasks:
return False
+ cancelled_tail = self.native_append_tail if self.native_append_tail in tasks else None
for task in tasks:
task.cancel()
try:
await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), timeout=timeout_s)
except TimeoutError:
pass
+ if cancelled_tail is not None and self.native_append_tail is cancelled_tail:
+ self.native_append_tail = None
return True
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_encoder.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_encoder.py
index b6f8531a8..8400f350c 100644
--- a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_encoder.py
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_encoder.py
@@ -830,6 +830,75 @@ class ConformerEncoder(NeuralModule, StreamingEncoder, Exportable, AccessMixin):
self.streaming_cfg = streaming_cfg
+ def streaming_post_process(self, rets, keep_all_outputs=True):
+ """Apply NeMo's cache-aware output trimming contract."""
+ if len(rets) == 2:
+ return rets[0], rets[1], None, None, None
+
+ encoded, encoded_len, cache_channel, cache_time, cache_channel_len = rets
+ if cache_channel is not None and self.streaming_cfg.last_channel_cache_size >= 0:
+ if self.streaming_cfg.last_channel_cache_size > 0:
+ cache_channel = cache_channel[:, :, -self.streaming_cfg.last_channel_cache_size :, :]
+
+ if self.streaming_cfg.valid_out_len > 0 and (not keep_all_outputs or self.att_context_style == "regular"):
+ encoded = encoded[:, :, : self.streaming_cfg.valid_out_len]
+ encoded_len = torch.clamp(encoded_len, max=self.streaming_cfg.valid_out_len)
+
+ return encoded, encoded_len, cache_channel, cache_time, cache_channel_len
+
+ def get_initial_cache_state(
+ self,
+ batch_size=1,
+ dtype=torch.float32,
+ device=None,
+ max_dim=0,
+ ):
+ """Create the channel/time caches used by cache-aware streaming."""
+ if self.streaming_cfg is None:
+ self.setup_streaming_params()
+ if device is None:
+ device = next(self.parameters()).device
+ create_tensor = torch.randn if max_dim > 0 else torch.zeros
+ cache_last_channel = create_tensor(
+ (
+ len(self.layers),
+ batch_size,
+ self.streaming_cfg.last_channel_cache_size,
+ self.d_model,
+ ),
+ device=device,
+ dtype=dtype,
+ )
+ cache_last_time = create_tensor(
+ (
+ len(self.layers),
+ batch_size,
+ self.d_model,
+ self.conv_context_size[0],
+ ),
+ device=device,
+ dtype=dtype,
+ )
+ if max_dim > 0:
+ cache_last_channel_len = torch.randint(
+ 0,
+ min(max_dim, self.streaming_cfg.last_channel_cache_size),
+ (batch_size,),
+ device=device,
+ dtype=torch.int64,
+ )
+ for i in range(batch_size):
+ cache_last_channel[:, i, cache_last_channel_len[i] :, :] = 0
+ if cache_last_channel_len[i] == 0:
+ cache_last_time[:, i, :, :] = 0
+ else:
+ cache_last_channel_len = torch.zeros(
+ batch_size,
+ device=device,
+ dtype=torch.int64,
+ )
+ return cache_last_channel, cache_last_time, cache_last_channel_len
+
def change_subsampling_conv_chunking_factor(self, subsampling_conv_chunking_factor: int):
"""
Update the conv_chunking_factor (int)
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/compat.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/compat.py
index f1c648acd..e89b67107 100644
--- a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/compat.py
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/compat.py
@@ -259,6 +259,51 @@ class StreamingEncoder(ABC):
"""
pass
+ @abstractmethod
+ def get_initial_cache_state(self, batch_size, dtype, device, max_dim):
+ pass
+
+ def cache_aware_stream_step(
+ self,
+ processed_signal,
+ processed_signal_length=None,
+ cache_last_channel=None,
+ cache_last_time=None,
+ cache_last_channel_len=None,
+ keep_all_outputs=True,
+ drop_extra_pre_encoded=None,
+ ):
+ """Run one incremental encoder step using NeMo's cache contract."""
+ if self.streaming_cfg is None:
+ self.setup_streaming_params()
+ if drop_extra_pre_encoded is not None:
+ prev_drop_extra_pre_encoded = self.streaming_cfg.drop_extra_pre_encoded
+ self.streaming_cfg.drop_extra_pre_encoded = drop_extra_pre_encoded
+ else:
+ prev_drop_extra_pre_encoded = None
+
+ if processed_signal_length is None:
+ processed_signal_length = processed_signal.new_full(
+ (processed_signal.size(0),),
+ processed_signal.size(-1),
+ )
+
+ try:
+ encoder_output = self(
+ audio_signal=processed_signal,
+ length=processed_signal_length,
+ cache_last_channel=cache_last_channel,
+ cache_last_time=cache_last_time,
+ cache_last_channel_len=cache_last_channel_len,
+ )
+ return self.streaming_post_process(
+ encoder_output,
+ keep_all_outputs=keep_all_outputs,
+ )
+ finally:
+ if prev_drop_extra_pre_encoded is not None:
+ self.streaming_cfg.drop_extra_pre_encoded = prev_drop_extra_pre_encoded
+
# ==============================================================================
# CacheAwareStreamingConfig — exact copy of the dataclass from
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_model.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_model.py
index e83e78cb6..2d7d6277f 100644
--- a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_model.py
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_model.py
@@ -935,8 +935,8 @@ class GatedProjectedSumRMSNorm(nn.Module):
audio_emb = audio_emb / self.num_codebooks
# projections run in model dtype (BF16)
- audio_h = self.audio_proj(audio_emb)
- text_h = self.text_proj(text_emb)
+ audio_h = self.audio_proj(audio_emb.to(dtype=self.audio_proj.weight.dtype))
+ text_h = self.text_proj(text_emb.to(dtype=self.text_proj.weight.dtype))
dtype = audio_h.dtype
@@ -1302,7 +1302,7 @@ class RVQEARTTSModel(nn.Module):
pre_bos_mask = bos_mask.cumsum(dim=1) == 0 # [B, T, 1]
# Apply projection to model size
- code_embed = self.embed_code(code_embed)
+ code_embed = self.embed_code(code_embed.to(dtype=self.embed_code.weight.dtype))
# Choose projection
if self.config.get("audio_prompt_encoder_config", None):
@@ -1336,7 +1336,7 @@ class RVQEARTTSModel(nn.Module):
code_embeds = code_embed + bos_mask * self.bos_emb
else: # Inference
- code_embeds = self.embed_code(self.depthsum_embedding(code))
+ code_embeds = self.embed_code(self.depthsum_embedding(code).to(dtype=self.embed_code.weight.dtype))
uncond_dec_flag = torch.zeros(code.size(0), 1, 1, device=code.device, dtype=torch.bool)
if guidance_enabled:
@@ -1591,7 +1591,7 @@ class RVQEARTTSModel(nn.Module):
top_p_or_k_i = top_p_or_k[i] if top_p_or_k is not None else 1.0
noise_scale_i = noise_scale[i] if noise_scale is not None else 1.0
- mog_input_embeds = self.embed_code(self.depthsum_embedding(code))
+ 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)
if guidance_scale_i > 0.0:
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_code2wav.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_code2wav.py
index 3ceaf93de..d53af34e2 100644
--- a/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_code2wav.py
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_code2wav.py
@@ -24,6 +24,7 @@ import torch.nn as nn
from vllm.config import VllmConfig
from vllm.logger import init_logger
+from vllm_omni.model_executor.models.nemotron_voicechat.runtime_info import scalar_bool
from vllm_omni.model_executor.models.output_templates import OmniOutput
logger = init_logger(__name__)
@@ -83,6 +84,7 @@ class NemotronVoiceChatCode2Wav(nn.Module):
# 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.
@@ -147,6 +149,12 @@ class NemotronVoiceChatCode2Wav(nn.Module):
meta = runtime_info.get("meta") if isinstance(runtime_info, Mapping) else None
if isinstance(meta, Mapping) and meta.get("nemotron_voicechat_dummy_profile") is True:
continue
+ codec_streaming = scalar_bool(meta.get("codec_streaming")) if isinstance(meta, Mapping) else False
+ if isinstance(runtime_info, Mapping):
+ codec_streaming = codec_streaming or scalar_bool(runtime_info.get("meta.codec_streaming"))
+ cache_request_id = meta.get("request_id") if isinstance(meta, Mapping) else None
+ if cache_request_id is None and isinstance(runtime_info, Mapping):
+ cache_request_id = runtime_info.get("meta.request_id")
codes = self._codes_from_runtime_info(runtime_info)
if codes is None:
# No input_ids fallback (unlike PersonaPlexCode2Wav): this
@@ -159,6 +167,12 @@ class NemotronVoiceChatCode2Wav(nn.Module):
"code stack (placeholder input_ids are never decoded)."
)
codes = validate_code_stack(codes, self._num_quantizers, self._codebook_size).to(device=device)
+ logger.debug(
+ "Nemotron VoiceChat code2wav input: request=%s shape=%s codec_streaming=%s",
+ cache_request_id,
+ tuple(codes.shape),
+ codec_streaming,
+ )
frames = codes.shape[0]
if frames == 0:
continue
@@ -176,19 +190,42 @@ class NemotronVoiceChatCode2Wav(nn.Module):
if left_context_frames >= frames:
continue # terminal empty chunk: all frames already emitted
lens = torch.tensor([frames], device=device, dtype=torch.long)
+ codec_cache = None
+ if codec_streaming:
+ 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())
# 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)
+ wav, wav_len = self.audio_codec.decode(
+ codes.unsqueeze(0),
+ lens,
+ cache=codec_cache,
+ )
wav = wav.reshape(-1)[: int(wav_len.reshape(-1)[0])]
- if left_context_frames > 0:
+ if left_context_frames > 0 and not codec_streaming:
wav = wav[left_context_frames * self._wav_per_frame :]
audios[i] = wav.detach().to(dtype=torch.float32).cpu()
+ logger.debug(
+ "Nemotron VoiceChat code2wav output: request=%s frames=%s samples=%s",
+ cache_request_id,
+ frames,
+ int(audios[i].numel()),
+ )
return OmniOutput(
text_hidden_states=None,
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
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_talker.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_talker.py
index f4f98ab00..5764afa98 100644
--- a/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_talker.py
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_talker.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Stage-1 talker for NemotronVoiceChat: frame-locked text -> 31-quantizer codes.
``LLM_AR`` stage driven per frame by the vLLM engine, with the vendored NeMo
@@ -46,6 +46,7 @@ from vllm.logger import init_logger
from vllm_omni.model_executor.models.nemotron_voicechat.runtime_info import (
merge_runtime_info,
require_request_id,
+ scalar_bool,
)
from vllm_omni.model_executor.models.output_templates import OmniOutput
@@ -164,6 +165,7 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module):
info_dicts = kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information") or []
codes_list: list[torch.Tensor] = []
prompt_len: int | None = None
+ codec_streaming = False
for info in info_dicts:
if not isinstance(info, dict):
continue
@@ -177,16 +179,32 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module):
reported = info.get("meta.nvc_logical_prompt_len")
if reported is not None:
prompt_len = int(reported)
+ streaming = meta.get("codec_streaming") if isinstance(meta, dict) else None
+ if streaming is None:
+ streaming = info.get("meta.codec_streaming")
+ codec_streaming = codec_streaming or scalar_bool(streaming)
if not codes_list:
return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
- outputs: dict[str, Any] = {"codes": {"audio": torch.cat(codes_list, dim=0)}}
+ codes = torch.cat(codes_list, dim=0)
+ logger.debug(
+ "Nemotron VoiceChat talker emitted codes: shapes=%s combined=%s prompt_len=%s",
+ [tuple(item.shape) for item in codes_list],
+ tuple(codes.shape),
+ prompt_len,
+ )
+ outputs: dict[str, Any] = {"codes": {"audio": codes}}
# The prompt length must ride the WIRE payload (multimodal_outputs):
# the async-chunk dispatcher hands exactly this dict to the
# talker2code2wav producer, while request.additional_information is
# overwritten concurrently by arriving thinker chunks.
+ output_meta: dict[str, Any] = {}
if prompt_len is not None:
# 1-D on purpose: the wire-payload builder indexes element.shape[0].
- outputs["meta"] = {"nvc_logical_prompt_len": torch.tensor([prompt_len], dtype=torch.long)}
+ output_meta["nvc_logical_prompt_len"] = torch.tensor([prompt_len], dtype=torch.long)
+ if codec_streaming:
+ output_meta["codec_streaming"] = torch.tensor([True], dtype=torch.bool)
+ if output_meta:
+ outputs["meta"] = output_meta
return OmniOutput(
text_hidden_states=hidden,
multimodal_outputs=outputs,
@@ -269,8 +287,9 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module):
# 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)
- generation_config = self.tts._get_generation_config(self._guidance_enabled())
- init_inputs.update({"use_cache": True, "past_key_values": None, "guidance_enabled": self._guidance_enabled()})
+ 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)
@@ -281,6 +300,10 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module):
"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
session = {
"timeline": timeline,
"prompt_len": prompt_len,
@@ -288,9 +311,13 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module):
"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,
# Sync mode ships the whole timeline up front via
# nvc_text_timeline; async-chunk mode grows it across chunks and
# the last chunk sets meta.finished.
@@ -333,25 +360,27 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module):
session["tts_initialized"] = True
else:
prev_subword_id = timeline[t - 1].reshape(1, 1)
- current_subword_mask = torch.ones(1, 1, device=timeline.device, dtype=torch.bool)
with torch.inference_mode():
code, past_key_values = self.tts.infer_codes_one_step(
current_subword_id=current_subword_id,
prev_subword_id=prev_subword_id,
- current_subword_mask=current_subword_mask,
+ current_subword_mask=session["current_subword_mask"],
prev_audio_tokens=session["code"],
past_key_values=session["past_key_values"],
- guidance_enabled=self._guidance_enabled(),
+ guidance_enabled=session["guidance_enabled"],
generation_config=session["generation_config"],
ignore_eos_flag_stop=True,
)
session["code"] = code
session["past_key_values"] = past_key_values
session["step"] = t + 1
- # Async mode: exhausting the received timeline only ends the request
- # once the upstream marked its final chunk; otherwise more entries are
- # in flight and the scheduler parks us until they arrive.
- finished = (t + 1) >= int(session["timeline"].numel()) and bool(session["upstream_finished"])
+ # 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"])
+ )
return code.reshape(1, -1).to(torch.long), finished
def preprocess(
@@ -388,21 +417,34 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module):
# chunk each.
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)
- session["codes_rows"].append(codes_row.cpu())
+ 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
if steps_run == 0:
# Zero-progress wake (duplicate/terminal-marker chunk): no TTS step
- # was run, so emit a plain CONTINUE placeholder — never a
- # fabricated frame. The request finishes only when the upstream is
- # exhausted AND the timeline is fully consumed.
- finished = bool(session["upstream_finished"])
+ # 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"])
+ )
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)
if finished:
embeds[:, 0] = 1.0 # stop flag -> compute_logits emits the stop token
@@ -413,7 +455,12 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module):
# chunks and by these model updates, so the async code2wav producer
# reads the prompt length from the per-step multimodal_output
# instead of racing the request-level dict.
- "meta": {"nvc_tts_step": int(session["step"]), "nvc_logical_prompt_len": int(session["prompt_len"])},
+ "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"],
+ },
}
return input_ids, embeds, info_update
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_thinker.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_thinker.py
index 9f237a885..01fbed1f9 100644
--- a/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_thinker.py
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_thinker.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Stage-0 thinker for NemotronVoiceChat: 16 kHz speech -> frame-locked text.
``LLM_AR`` stage combining:
@@ -12,7 +12,8 @@
* the NemotronH (hybrid Mamba) backbone on vLLM paged execution via
``init_vllm_registered_model`` — this is where vLLM earns its keep;
* the separate ``lm_head`` (text channel, sampled by the vLLM engine) and
- ``function_head`` (greedy, maintained internally per request).
+ ``function_head`` (greedy, maintained internally per request, with client
+ tool results injected into the function channel).
Timeline contract (verified against NeMo ``_init_inference``/``_step_zero``/
``_step_inference``; the golden shape test guards it):
@@ -29,13 +30,14 @@ Timeline contract (verified against NeMo ``_init_inference``/``_step_zero``/
example sets ``max_tokens = acoustic_frame_count``. Never set ``min_tokens``:
the tokenizer's eos_token is — the frame-locked PAD/silence
token — and min_tokens masks "EOS" to -inf, forbidding silence entirely.
-* The function channel exists for numerical correctness (fusion weight 2.0)
- even though tool calling is out of scope; the greedy token timeline is
- exposed for debugging via the request info.
+* The function channel is load-bearing for numerical correctness (fusion
+ weight 2.0), emits tool calls, and accepts client tool results while keeping
+ the text channel padded during forced response injection.
"""
from __future__ import annotations
+import copy
import os
from collections.abc import Iterable
from typing import Any
@@ -47,7 +49,10 @@ from vllm.config import VllmConfig
from vllm.logger import init_logger
from vllm.model_executor.models.interfaces import HasInnerState, IsHybrid
from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
+from vllm.v1.outputs import SamplerOutput
+from vllm.v1.sample.metadata import SamplingMetadata
+from vllm_omni.experimental.fullduplex.model_executor import DuplexSamplingRow
from vllm_omni.model_executor.models.nemotron_voicechat.runtime_info import (
merge_runtime_info,
require_request_id,
@@ -59,6 +64,49 @@ logger = init_logger(__name__)
_LM_ARCHITECTURE = "NemotronHForCausalLM"
+def slice_perception_streaming_mel(
+ processed_signal: torch.Tensor,
+ frame_idx: 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
+
+ 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
+ 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]
+ if cache_mel.shape[-1] < pre_encode_size:
+ cache_mel = torch.nn.functional.pad(
+ cache_mel,
+ (pre_encode_size - cache_mel.shape[-1], 0),
+ )
+ mel_chunk = torch.cat([cache_mel, main_chunk], dim=-1)
+ return mel_chunk, int(streaming_cfg.drop_extra_pre_encoded)
+
+
def validate_timeline_fits_max_model_len(
logical_prompt_token_len: int,
acoustic_frame_count: int,
@@ -210,6 +258,14 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState,
from omegaconf import DictConfig
self.perception = AudioPerceptionModule(DictConfig(perception_cfg))
+ # NeMo's streaming wrapper deliberately uses a second, deterministic
+ # preprocessor. The training/offline frontend has dither enabled and
+ # pads the mel time axis; both are incompatible with exact incremental
+ # chunk boundaries.
+ streaming_preprocessor_cfg = copy.deepcopy(self.perception.cfg.preprocessor)
+ streaming_preprocessor_cfg.dither = 0.0
+ streaming_preprocessor_cfg.pad_to = 0
+ self._streaming_preprocessor = self.perception.from_config_dict(streaming_preprocessor_cfg)
# AddFusion weights (the checkpoint uses fuse_method="add" defaults;
# the same duplex_*_weight keys double as the fusion weights).
@@ -235,6 +291,8 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState,
# Omni AR runner contract.
self.have_multimodal_outputs = True
self.has_preprocess = True
+ self.prefer_model_sampler = True
+ self.omni_client_multimodal_output_keys = ("nvc_function_token",)
self.has_postprocess = True # per-step greedy function-channel token
self.requires_full_prefix_cached_hidden_states = False
# The per-step function token is a scalar; no GPU-resident buffers needed.
@@ -251,6 +309,8 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState,
# request_id -> per-request state (audio frames, prefill embeds, counters).
self._sessions: dict[str, dict[str, Any]] = {}
+ self._duplex_sampling_rows: dict[int, str] = {}
+ self._duplex_previous_text_tokens: dict[str, int] = {}
# ------------------------------------------------------------------
# Special token ids (from the checkpoint's token strings).
@@ -306,7 +366,57 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState,
def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput:
if isinstance(model_outputs, OmniOutput):
return model_outputs
- return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs={})
+ multimodal_outputs: dict[str, torch.Tensor] = {}
+ if self._use_function_head and model_outputs.numel():
+ 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
+ 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
+ if isinstance(forced_token, int):
+ function_token = torch.full_like(function_token, forced_token)
+ multimodal_outputs["nvc_function_token"] = function_token
+ return OmniOutput(
+ text_hidden_states=model_outputs,
+ multimodal_outputs=multimodal_outputs,
+ )
+
+ def prepare_duplex_sampling(
+ self,
+ logits: torch.Tensor,
+ sampling_metadata: SamplingMetadata,
+ rows: tuple[DuplexSamplingRow, ...],
+ ) -> None:
+ del logits, sampling_metadata
+ self._duplex_sampling_rows = {row.row_idx: row.request_id for row in rows}
+
+ def sample(
+ self,
+ logits: torch.Tensor,
+ sampling_metadata: SamplingMetadata,
+ ) -> SamplerOutput | None:
+ del sampling_metadata
+ rows = self._duplex_sampling_rows
+ if not rows:
+ return None
+ expected_rows = set(range(int(logits.shape[0])))
+ if set(rows) != expected_rows:
+ raise RuntimeError(
+ "NemotronVoiceChat native-duplex sampling requires every active batch row "
+ "to be a duplex request; mixed duplex/non-duplex batches are unsupported."
+ )
+ sampled = logits.argmax(dim=-1).to(dtype=torch.int32)
+ for row_idx, request_id in rows.items():
+ session = self._sessions.get(request_id)
+ if isinstance(session, dict) and isinstance(session.get("forced_function_token"), int):
+ sampled[row_idx] = int(session["nvc_text_pad_id"])
+ self._duplex_previous_text_tokens[request_id] = int(sampled[row_idx].item())
+ return SamplerOutput(
+ sampled_token_ids=sampled.unsqueeze(-1),
+ logprobs_tensors=None,
+ )
# ------------------------------------------------------------------
# Fusion helpers.
@@ -324,6 +434,263 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState,
fused = fused + self.embed_tokens(func_ids) * self._w_func
return fused
+ @staticmethod
+ def _duplex_info(info: dict[str, Any]) -> dict[str, Any] | None:
+ duplex = info.get("duplex")
+ if isinstance(duplex, dict) and duplex.get("data_plane") is True:
+ return duplex
+ return None
+
+ def _duplex_stable_frame(
+ self,
+ session: dict[str, Any],
+ duplex: dict[str, Any],
+ device: torch.device,
+ ) -> torch.Tensor:
+ try:
+ source_input_seq = int(duplex["source_input_seq"])
+ except (KeyError, TypeError, ValueError) as exc:
+ raise ValueError("Nemotron VoiceChat duplex input lacks source_input_seq") from exc
+ if source_input_seq == session.get("last_input_seq"):
+ return session["duplex_frame"]
+ last_input_seq = int(session.get("last_input_seq", 0))
+ if source_input_seq < last_input_seq:
+ raise ValueError(
+ f"Nemotron VoiceChat duplex input sequence moved backwards: {source_input_seq} < {last_input_seq}"
+ )
+
+ from vllm_omni.experimental.fullduplex.nemotron_voicechat.input import (
+ decode_pcm_f32le,
+ )
+
+ raw = decode_pcm_f32le(duplex.get("payload"), exact_frame=True)
+ frame_audio = torch.from_numpy(np.frombuffer(raw, dtype=" tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
+ device = input_ids.device
+ runtime_config = duplex.get("runtime_config")
+ if not isinstance(runtime_config, dict):
+ raise ValueError("Nemotron VoiceChat duplex input lacks runtime_config")
+ try:
+ source_input_seq = int(duplex["source_input_seq"])
+ pad_id = int(runtime_config["nvc_text_pad_id"])
+ except (KeyError, TypeError, ValueError) as exc:
+ raise ValueError("Nemotron VoiceChat duplex runtime token metadata is incomplete") from exc
+
+ session = self._sessions.get(request_id)
+ if session is None:
+ prompt_ids_raw = runtime_config.get("nvc_prompt_token_ids")
+ if not isinstance(prompt_ids_raw, list | tuple) or not prompt_ids_raw:
+ raise ValueError("Nemotron VoiceChat duplex runtime requires nvc_prompt_token_ids")
+ prompt_ids = torch.as_tensor(
+ [int(token_id) for token_id in prompt_ids_raw],
+ dtype=torch.long,
+ device=device,
+ )
+ session = {
+ "prompt_len": int(prompt_ids.numel()),
+ "func_token": pad_id,
+ "last_input_seq": 0,
+ "nvc_text_pad_id": pad_id,
+ "function_response_generation": 0,
+ "forced_function_tokens": [],
+ }
+ self._sessions[request_id] = session
+ frame = self._duplex_stable_frame(session, duplex, device)
+ timeline = torch.cat(
+ [self.embed_tokens(prompt_ids).to(self._dtype), frame],
+ dim=0,
+ )
+ pad_ids = torch.full(
+ (timeline.shape[0],),
+ pad_id,
+ dtype=torch.long,
+ device=device,
+ )
+ session["prefill_embeds"] = self._fuse(pad_ids, timeline, pad_ids)
+ else:
+ frame = self._duplex_stable_frame(session, duplex, device)
+
+ 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")
+ return (
+ input_ids,
+ prefill[offset : offset + span].to(self._dtype),
+ {
+ "nvc_text_pad_id": pad_id,
+ },
+ )
+
+ if int(input_ids.shape[0]) != 1:
+ raise ValueError("Nemotron VoiceChat duplex continuation must schedule exactly one frame token")
+ previous_function = info.get("nvc_prev_function_token")
+ if previous_function is not None:
+ if isinstance(previous_function, torch.Tensor):
+ previous_function = previous_function.reshape(-1)[-1].item()
+ session["func_token"] = int(previous_function)
+ previous_text = self._duplex_previous_text_tokens.get(request_id)
+ if previous_text is None:
+ raise RuntimeError("Nemotron VoiceChat duplex continuation has no previous sampled text token")
+ text_id = torch.tensor([previous_text], dtype=torch.long, device=device)
+ function_id = torch.tensor(
+ [int(session["func_token"])],
+ dtype=torch.long,
+ device=device,
+ )
+ fused = self._fuse(text_id, frame, function_id)
+ return (
+ input_ids,
+ fused,
+ {
+ "meta": {"nvc_frame": source_input_seq - 1},
+ "nvc_text_pad_id": pad_id,
+ },
+ )
+
+ @staticmethod
+ def _sync_forced_function_response(
+ session: dict[str, Any],
+ runtime_config: dict[str, Any],
+ ) -> None:
+ try:
+ generation = int(runtime_config.get("nvc_function_response_generation", 0))
+ except (TypeError, ValueError) as exc:
+ raise ValueError("Nemotron VoiceChat function-response generation must be an integer") from exc
+ seen = int(session.get("function_response_generation", 0))
+ if generation < seen:
+ raise ValueError("Nemotron VoiceChat function-response generation moved backwards")
+ if generation > seen:
+ raw_batches = runtime_config.get("nvc_function_response_batches")
+ batches = raw_batches if isinstance(raw_batches, list) else []
+ unseen: list[tuple[int, list[int]]] = []
+ for raw_batch in batches:
+ if not isinstance(raw_batch, dict):
+ raise ValueError("Nemotron VoiceChat function response batch must be an object")
+ try:
+ batch_generation = int(raw_batch["generation"])
+ except (KeyError, TypeError, ValueError) as exc:
+ raise ValueError("Nemotron VoiceChat function response batch requires a generation") from exc
+ if batch_generation <= seen:
+ continue
+ raw_tokens = raw_batch.get("token_ids")
+ if not isinstance(raw_tokens, list | tuple) or not raw_tokens:
+ raise ValueError("Nemotron VoiceChat function response requires token ids")
+ unseen.append((batch_generation, [int(token_id) for token_id in raw_tokens]))
+ if not unseen:
+ # Backward-compatible single-response payload used by older
+ # serving adapters and serialized test fixtures.
+ raw_tokens = runtime_config.get("nvc_function_response_token_ids")
+ if not isinstance(raw_tokens, list | tuple) or not raw_tokens:
+ raise ValueError("Nemotron VoiceChat function response requires token ids")
+ unseen = [(generation, [int(token_id) for token_id in raw_tokens])]
+ expected = seen + 1
+ queue = session.setdefault("forced_function_tokens", [])
+ if not isinstance(queue, list):
+ raise ValueError("Nemotron VoiceChat forced function-token queue is invalid")
+ for batch_generation, token_ids in unseen:
+ if batch_generation != expected:
+ raise ValueError("Nemotron VoiceChat function-response generations are not contiguous")
+ queue.extend(token_ids)
+ expected += 1
+ if expected - 1 != generation:
+ raise ValueError("Nemotron VoiceChat function-response batches are incomplete")
+ session["function_response_generation"] = generation
+ queue = session.get("forced_function_tokens")
+ if session.get("forced_function_token") is None and isinstance(queue, list) and queue:
+ session["forced_function_token"] = int(queue[0])
+
# ------------------------------------------------------------------
# Per-request session.
# ------------------------------------------------------------------
@@ -418,6 +785,14 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState,
is_prefill = bool(is_prefill_raw) if isinstance(is_prefill_raw, bool) else span > 1
request_id = require_request_id(info, "thinker")
+ duplex = self._duplex_info(info)
+ if duplex is not None:
+ return self._preprocess_duplex(
+ request_id=request_id,
+ input_ids=input_ids,
+ info=info,
+ duplex=duplex,
+ )
if is_prefill or request_id not in self._sessions:
session = self._sessions.get(request_id) or self._init_session(request_id, info, device)
offset = max(0, int(info.get("_omni_num_computed_tokens", 0) or 0))
@@ -501,9 +876,25 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState,
prompt_len = int(kwargs.get("_omni_prompt_len", 0) or 0)
if prompt_len and computed + int(hidden_states.shape[0]) < prompt_len:
return {} # intermediate prefill chunk: frame 0 not reached yet
- with torch.inference_mode():
- func_logits = self.function_head(hidden_states[-1, :].reshape(1, -1).to(self._dtype))
- token = func_logits.argmax(dim=-1).reshape(()).detach()
+ multimodal_outputs = kwargs.get("multimodal_outputs")
+ token = multimodal_outputs.get("nvc_function_token") if isinstance(multimodal_outputs, dict) else None
+ if not isinstance(token, torch.Tensor) or token.numel() == 0:
+ with torch.inference_mode():
+ func_logits = self.function_head(hidden_states[-1, :].reshape(1, -1).to(self._dtype))
+ token = func_logits.argmax(dim=-1)
+ token = token.reshape(-1)[-1].reshape(()).detach()
+ request_id = kwargs.get("request_id")
+ 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")
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")
@@ -516,6 +907,7 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState,
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)
+ self._duplex_previous_text_tokens.pop(str(request_id), None)
# ------------------------------------------------------------------
# Weights.
@@ -581,6 +973,9 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState,
)
device = self.vllm_config.device_config.device
self.perception.to(device=device, dtype=self._dtype).eval()
+ # Keep the deterministic mel frontend in fp32, matching NeMo's
+ # PerceptionCacheManager. Its output is cast only at the encoder edge.
+ self._streaming_preprocessor.to(device=device, dtype=torch.float32).eval()
self.embed_tokens.to(device=device, dtype=self._dtype)
self.function_head.to(device=device, dtype=self._dtype)
loaded |= {f"perception.{n}" for n in perception_state}
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/pipeline.py b/vllm_omni/model_executor/models/nemotron_voicechat/pipeline.py
index 683228c21..1bb1e19b8 100644
--- a/vllm_omni/model_executor/models/nemotron_voicechat/pipeline.py
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/pipeline.py
@@ -31,6 +31,13 @@ _PROC = "vllm_omni.model_executor.stage_input_processors.nemotron_voicechat"
NEMOTRON_VOICECHAT_PIPELINE = PipelineConfig(
model_type="nemotron_voicechat",
model_arch="NemotronVoiceChatThinkerForConditionalGeneration",
+ duplex_runtime_extension=(
+ "vllm_omni.experimental.fullduplex.nemotron_voicechat.runtime.NemotronVoiceChatDuplexRuntimeExtension"
+ ),
+ duplex_serving_adapter=(
+ "vllm_omni.experimental.fullduplex.nemotron_voicechat.serving_adapter.NemotronVoiceChatServingRuntimeAdapter"
+ ),
+ duplex_control_enabled=True,
# Named after the alias key so the deploy-yaml stem substring-matches the
# checkpoint directory name (NVIDIA-NemotronLabs-VoiceChat-11B) during
# bare-path auto-detection.
diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/runtime_info.py b/vllm_omni/model_executor/models/nemotron_voicechat/runtime_info.py
index 7f5a897d2..9d8508329 100644
--- a/vllm_omni/model_executor/models/nemotron_voicechat/runtime_info.py
+++ b/vllm_omni/model_executor/models/nemotron_voicechat/runtime_info.py
@@ -6,6 +6,27 @@ from __future__ import annotations
from typing import Any
+import torch
+
+
+def scalar_bool(value: object) -> bool:
+ """Normalize a Python or tensor scalar to bool.
+
+ Inter-stage scalar metadata is tensorized while it rides the engine output
+ channel, so model consumers must not rely on object identity with ``True``.
+ Empty tensors mean "unset"; non-scalar tensors are rejected instead of
+ silently choosing an element.
+ """
+ if isinstance(value, torch.Tensor):
+ if value.numel() == 0:
+ return False
+ if value.numel() != 1:
+ raise ValueError(f"Expected scalar boolean metadata, got tensor shape {tuple(value.shape)}.")
+ value = value.item()
+ elif hasattr(value, "item"):
+ value = value.item()
+ return bool(value)
+
def merge_runtime_info(info_dict: dict[str, Any]) -> dict[str, Any]:
"""Flatten ``additional_information`` into the info dict (top-level wins)."""
diff --git a/vllm_omni/model_executor/stage_input_processors/nemotron_voicechat.py b/vllm_omni/model_executor/stage_input_processors/nemotron_voicechat.py
index 48de5a380..918db4ad5 100644
--- a/vllm_omni/model_executor/stage_input_processors/nemotron_voicechat.py
+++ b/vllm_omni/model_executor/stage_input_processors/nemotron_voicechat.py
@@ -11,9 +11,10 @@ loop), so the prompt slots are reconstructed here, and the timeline metadata
rides ``additional_information``.
talker (1) -> code2wav (2): PersonaPlex-style full payload. The talker
-accumulates one 31-quantizer code stack per frame under ``("codes","audio")``;
-the producer ships the prompt-trimmed ``[frames, 31]`` stack once the request
-finishes, and the sync placeholder builder sizes stage-2's dummy prompt.
+accumulates one 31-quantizer code stack per frame under ``("codes","audio")``.
+The synchronous producer ships the prompt-trimmed ``[frames, 31]`` stack once
+the request finishes. Native duplex wakes expose that same cumulative stack,
+so the streaming producer must delta it before feeding code2wav's causal cache.
"""
from __future__ import annotations
@@ -28,6 +29,7 @@ from vllm_omni.data_entry_keys import (
MetaStruct,
OmniPayloadStruct,
)
+from vllm_omni.model_executor.models.nemotron_voicechat.runtime_info import scalar_bool
# Full-sequence tensors are re-shipped whole; the full-payload accumulator must
# replace (not concatenate) them if the producer fires more than once. The AR
@@ -39,6 +41,7 @@ _FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset({"codes", "codes.audio",
# stt pad_token. The thinker also reports it in its latent metadata, which
# takes precedence when present.
_DEFAULT_TEXT_PAD_ID = 12
+_STEP_TOKEN_SNAPSHOT_MISSING = object()
def _info_get(info: Any, key: str) -> Any:
@@ -51,6 +54,27 @@ def _info_get(info: Any, key: str) -> Any:
return None
+def _request_info(request: Any) -> dict[str, Any]:
+ """Return the live runner payload, accepting both transport generations."""
+ for name in (
+ "model_intermediate_buffer",
+ "additional_information",
+ "additional_information_cpu",
+ ):
+ info = getattr(request, name, None)
+ if isinstance(info, dict):
+ return info
+ if info is not None and hasattr(info, "entries"):
+ from vllm_omni.engine.serialization import (
+ deserialize_additional_information,
+ )
+
+ decoded = deserialize_additional_information(info)
+ if isinstance(decoded, dict):
+ return decoded
+ return {}
+
+
def thinker2talker_token_only(
source_outputs: list,
prompt: Any = None,
@@ -118,32 +142,60 @@ def thinker2talker_async_chunk(
the payload self-contained. ``meta.num_processed_tokens`` carries the
logical prompt length for the code2wav producer's prompt-region trim.
"""
- del multimodal_output, kwargs
+ step_snapshot = kwargs.pop("new_token_ids", _STEP_TOKEN_SNAPSHOT_MISSING)
+ has_step_snapshot = step_snapshot is not _STEP_TOKEN_SNAPSHOT_MISSING
+ generated_delta = [int(token_id) for token_id in ((step_snapshot or ()) if has_step_snapshot else ())]
request_id = request.external_req_id
+ request_info = _request_info(request)
+ duplex_info = request_info.get("duplex") if isinstance(request_info, dict) else None
+ is_duplex = isinstance(duplex_info, dict) and duplex_info.get("data_plane") is True
+
+ state = transfer_manager.request_payload.get(request_id)
+ 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 [])
- # vLLM prompt = logical prompt + 1 placeholder (acoustic frame 0).
- logical_prompt_len = max(len(prompt_ids) - 1, 0)
- generated = list(getattr(request, "output_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)
+ previous = list(state.get("nvc_duplex_text_tokens", ())) if isinstance(state, dict) else []
+ if has_step_snapshot:
+ generated = previous + generated_delta
+ elif generated_now[: len(previous)] == previous and len(generated_now) > len(previous):
+ generated = generated_now
+ elif generated_now:
+ generated = previous + [generated_now[-1]]
+ else:
+ generated = previous
+ else:
+ # vLLM prompt = logical prompt + 1 placeholder (acoustic frame 0).
+ logical_prompt_len = max(len(prompt_ids) - 1, 0)
+ generated = generated_now
pad_id = _DEFAULT_TEXT_PAD_ID
- reported_pad = _info_get(getattr(request, "additional_information", None), "nvc_text_pad_id")
+ reported_pad = _info_get(request_info, "nvc_text_pad_id")
if reported_pad is not None:
pad_id = int(reported_pad)
- timeline = [pad_id] * logical_prompt_len + [int(t) for t in generated]
+ timeline = [pad_id] * logical_prompt_len + generated
# Skip no-progress wakeups (same length as the last emitted chunk).
- state = transfer_manager.request_payload.get(request_id)
last_len = int(state.get("nvc_timeline_len", -1)) if isinstance(state, dict) else -1
if len(timeline) <= last_len and not is_finished:
return None
- transfer_manager.request_payload[request_id] = {"nvc_timeline_len": len(timeline)}
+ transfer_manager.request_payload[request_id] = {
+ "nvc_timeline_len": len(timeline),
+ "nvc_duplex_text_tokens": generated if is_duplex else (),
+ }
return OmniPayloadStruct(
- ids=IdsStruct(all=timeline),
+ ids=IdsStruct(all=timeline, prompt=[0]),
meta=MetaStruct(
- finished=torch.tensor(bool(is_finished), dtype=torch.bool),
+ 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,
+ codec_streaming=is_duplex,
+ request_id=request_id if is_duplex else None,
),
)
@@ -155,18 +207,30 @@ def talker2code2wav_async_chunk(
is_finished: bool = False,
**kwargs: Any,
) -> OmniPayloadStruct | None:
- """Streaming producer: ship prompt-trimmed CUMULATIVE code stacks in chunks.
+ """Streaming producer: ship prompt-trimmed code stacks in chunks.
- Follows the NeMo incremental-decode recipe (``decode_one_audio_step``):
- each chunk carries the FULL trimmed code history and
- ``meta.left_context_size`` = frames the codec already emitted, so the
- code2wav stage re-decodes the whole prefix and slices off only the new
- samples — no seams, no crossfade, and the concatenated stream equals the
- prefix-decode of the final stack. Chunk cadence comes from the connector
- config key ``codec_chunk_frames`` (default 13 frames ~= 1 s of audio).
+ The regular path follows NeMo's prefix-decode recipe: each chunk carries
+ the full trimmed history and ``meta.left_context_size`` identifies frames
+ already emitted. Native duplex talker wakes also carry cumulative history,
+ but code2wav keeps causal codec state for that path and therefore must see
+ only rows not observed on an earlier wake. Chunk cadence comes from
+ ``codec_chunk_frames`` (default 13 frames ~= 1 s of audio).
"""
- del kwargs
request_id = request.external_req_id
+ del kwargs
+ request_info = _request_info(request)
+ codec_streaming = scalar_bool(_info_get(request_info, "codec_streaming")) or scalar_bool(
+ _info_get(multimodal_output, "codec_streaming")
+ )
+ state = transfer_manager.request_payload.get(request_id)
+ if isinstance(state, dict):
+ codec_streaming = codec_streaming or bool(state.get("nvc_codec_stream_started", False))
+ # ``is_finished`` describes the current resumable Stage-1 scheduler
+ # segment, not the lifetime of a native-duplex codec stream. Marking each
+ # segment terminal makes the connector clean sender state after its first
+ # chunk; every later one-row wake is then mistaken for a fresh stream and
+ # trimmed as prompt context. Duplex lifetime is closed by session abort.
+ stream_finished = bool(is_finished and not codec_streaming)
codes = None
if isinstance(multimodal_output, dict):
@@ -175,9 +239,19 @@ def talker2code2wav_async_chunk(
if codes is None:
codes = multimodal_output.get("codes.audio")
if not isinstance(codes, torch.Tensor) or codes.numel() == 0:
- if not is_finished:
+ if not stream_finished:
+ # Returning ``None`` at a resumable segment boundary makes the
+ # generic chunk adapter synthesize ``is_segment_finished=True``.
+ # Stage 2 then stops polling after the first 80 ms wake. An
+ # explicit non-terminal empty payload consumes this connector
+ # sequence number without closing the long-lived codec stream.
+ if codec_streaming:
+ return _empty_streaming_payload(request_id=request_id)
return None
- return _empty_finished_payload()
+ return _empty_finished_payload(
+ codec_streaming=codec_streaming,
+ request_id=request_id,
+ )
if codes.ndim == 1:
codes = codes.reshape(1, -1)
@@ -186,15 +260,13 @@ def talker2code2wav_async_chunk(
# request-level additional_information is overwritten both by arriving
# thinker chunks and by talker info updates, so it races. Cache the first
# sighting in the transfer-manager state as a fallback.
- state = transfer_manager.request_payload.get(request_id)
prompt_len = _info_get(multimodal_output if isinstance(multimodal_output, dict) else None, "nvc_logical_prompt_len")
if prompt_len is None and isinstance(multimodal_output, dict):
prompt_len = multimodal_output.get("meta.nvc_logical_prompt_len")
if prompt_len is None:
- info = getattr(request, "additional_information", None)
- prompt_len = _info_get(info, "nvc_logical_prompt_len")
+ prompt_len = _info_get(request_info, "nvc_logical_prompt_len")
if prompt_len is None:
- prompt_len = _info_get(info, "num_processed_tokens")
+ prompt_len = _info_get(request_info, "num_processed_tokens")
if prompt_len is None and isinstance(state, dict):
prompt_len = state.get("nvc_prompt_len")
if prompt_len is None:
@@ -203,41 +275,111 @@ def talker2code2wav_async_chunk(
"('nvc_logical_prompt_len' or async meta.num_processed_tokens); cannot "
"trim the prompt-region codes for streaming code2wav."
)
- trimmed = codes[max(int(prompt_len) - 1, 0) :]
-
frames_sent = int(state.get("nvc_frames_sent", 0)) if isinstance(state, dict) else 0
- new_frames = int(trimmed.shape[0]) - frames_sent
- if new_frames <= 0 and not is_finished:
+ if codec_streaming:
+ # The AR runner exposes its cumulative model_intermediate_buffer on
+ # every wake. Feeding that whole prefix into a stateful causal codec
+ # replays history into its cache (live traces showed lengths
+ # 19,20,20,21,21,...) and corrupts the waveform. Track rows OBSERVED,
+ # separately from rows EMITTED at the chunk cadence, and ship only the
+ # suffix first seen on this wake.
+ full_trimmed = codes[max(int(prompt_len) - 1, 0) :].detach().to(dtype=torch.long, device="cpu")
+ frames_seen = int(state.get("nvc_frames_seen", 0)) if isinstance(state, dict) else 0
+ if int(full_trimmed.shape[0]) < frames_seen:
+ raise ValueError(
+ "NemotronVoiceChat duplex talker code history shrank from "
+ f"{frames_seen} to {int(full_trimmed.shape[0])} frames; cannot "
+ "safely continue the stateful code2wav stream."
+ )
+ trimmed = full_trimmed[frames_seen:]
+ frames_seen = int(full_trimmed.shape[0])
+ pending = state.get("nvc_pending_codes") if isinstance(state, dict) else None
+ if isinstance(pending, torch.Tensor) and pending.numel() > 0:
+ trimmed = torch.cat([pending.reshape(-1, codes.shape[-1]), trimmed], dim=0)
+ new_frames = int(trimmed.shape[0])
+ else:
+ trimmed = codes[max(int(prompt_len) - 1, 0) :]
+ new_frames = int(trimmed.shape[0]) - frames_sent
+ if new_frames <= 0 and not stream_finished:
+ if codec_streaming:
+ transfer_manager.request_payload[request_id] = {
+ "nvc_frames_sent": frames_sent,
+ "nvc_frames_seen": frames_seen,
+ "nvc_prompt_len": int(prompt_len),
+ "nvc_codec_stream_started": True,
+ "nvc_pending_codes": trimmed,
+ }
+ return _empty_streaming_payload(request_id=request_id)
return None
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 {}
chunk_frames = int(cfg.get("codec_chunk_frames", 13))
- if new_frames < chunk_frames and not is_finished:
+ if new_frames < chunk_frames and not stream_finished:
+ if codec_streaming:
+ transfer_manager.request_payload[request_id] = {
+ "nvc_frames_sent": frames_sent,
+ "nvc_frames_seen": frames_seen,
+ "nvc_prompt_len": int(prompt_len),
+ "nvc_codec_stream_started": True,
+ "nvc_pending_codes": trimmed,
+ }
+ return _empty_streaming_payload(request_id=request_id)
return None
- if new_frames <= 0 and is_finished and frames_sent > 0:
+ if new_frames <= 0 and stream_finished and frames_sent > 0:
# Everything already shipped; emit a terminal empty chunk so the
# code2wav request observes meta.finished.
- return _empty_finished_payload()
+ return _empty_finished_payload(
+ codec_streaming=codec_streaming,
+ request_id=request_id,
+ )
transfer_manager.request_payload[request_id] = {
- "nvc_frames_sent": int(trimmed.shape[0]),
+ "nvc_frames_sent": frames_sent + new_frames if codec_streaming else int(trimmed.shape[0]),
+ **({"nvc_frames_seen": frames_seen} if codec_streaming else {}),
"nvc_prompt_len": int(prompt_len),
+ "nvc_codec_stream_started": codec_streaming,
}
+ code_payload = trimmed
return OmniPayloadStruct(
- codes=CodesStruct(audio=trimmed.detach().to(dtype=torch.long, device="cpu")),
+ codes=CodesStruct(audio=code_payload.detach().to(dtype=torch.long, device="cpu")),
meta=MetaStruct(
- left_context_size=frames_sent,
- finished=torch.tensor(bool(is_finished), dtype=torch.bool),
+ left_context_size=0 if codec_streaming else frames_sent,
+ codec_streaming=codec_streaming,
+ request_id=request_id if codec_streaming else None,
+ finished=torch.tensor(stream_finished, dtype=torch.bool),
+ # Native duplex scheduler segments are transport wakes, not codec
+ # stream boundaries. Preserve polling across successive wakes.
+ is_segment_finished=(torch.tensor(False, dtype=torch.bool) if codec_streaming else None),
),
)
-def _empty_finished_payload() -> OmniPayloadStruct:
+def _empty_streaming_payload(*, request_id: str) -> OmniPayloadStruct:
return OmniPayloadStruct(
codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
- meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)),
+ meta=MetaStruct(
+ codec_streaming=True,
+ request_id=request_id,
+ finished=torch.tensor(False, dtype=torch.bool),
+ is_segment_finished=torch.tensor(False, dtype=torch.bool),
+ ),
+ )
+
+
+def _empty_finished_payload(
+ *,
+ codec_streaming: bool = False,
+ request_id: str | None = None,
+) -> OmniPayloadStruct:
+ return OmniPayloadStruct(
+ codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
+ meta=MetaStruct(
+ finished=torch.tensor(True, dtype=torch.bool),
+ codec_streaming=codec_streaming,
+ request_id=request_id if codec_streaming else None,
+ ),
)
diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py
index be084f8f5..029254d6c 100644
--- a/vllm_omni/worker/gpu_ar_model_runner.py
+++ b/vllm_omni/worker/gpu_ar_model_runner.py
@@ -444,9 +444,11 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
def _resolve_pooler_payload_req_ids(self, req_ids_output_copy: list[str]) -> tuple[str, list[str]]:
downstream_req_ids = [rid for rid in req_ids_output_copy if self._request_needs_downstream_stage_payload(rid)]
engine_output_type = (self.vllm_config.model_config.engine_output_type or "").lower()
+ if self._client_multimodal_output_keys():
+ downstream_req_ids = req_ids_output_copy
# Single-stage AR TTS models (e.g. VoxCPM2) finish on this stage but still
# need multimodal payloads for final audio postprocess/output.
- if engine_output_type == "audio" and not downstream_req_ids:
+ elif engine_output_type == "audio" and not downstream_req_ids:
downstream_req_ids = req_ids_output_copy
return engine_output_type, downstream_req_ids
@@ -1555,7 +1557,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
lists are wrapped into tensors, and anything that cannot be safely
converted is dropped.
"""
- if self.vllm_config.model_config.engine_output_type == "text":
+ if self.vllm_config.model_config.engine_output_type == "text" and not self._client_multimodal_output_keys():
return None
if per_req_payloads is None:
return None
@@ -1614,6 +1616,18 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
def _runner_model_omni_flag(self, name: str, default: bool = False) -> bool:
return self._model_omni_flag(getattr(self, "model", None), name, default)
+ def _client_multimodal_output_keys(self) -> tuple[str, ...]:
+ raw = getattr(
+ getattr(self, "model", None),
+ "omni_client_multimodal_output_keys",
+ (),
+ )
+ if not isinstance(raw, tuple) or any(not isinstance(key, str) or not key for key in raw):
+ raise TypeError("omni_client_multimodal_output_keys must be a tuple of non-empty strings")
+ if len(raw) != len(set(raw)):
+ raise ValueError("omni_client_multimodal_output_keys must not contain duplicates")
+ return raw
+
def _model_omni_pooler_payload_include_hidden(self) -> bool:
return self._runner_model_omni_flag("omni_pooler_payload_include_hidden", default=True)
@@ -1921,6 +1935,12 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
# Connector-less stages expose the same payload through the
# orchestrator bridge; non-async stages preserve legacy behavior.
pooler_inter, pooler_client = pooler_output, pooler_output
+ client_output_keys = self._client_multimodal_output_keys()
+ if client_output_keys:
+ allowed = frozenset(client_output_keys)
+ pooler_client = [
+ {key: value for key, value in payload.items() if key in allowed} for payload in pooler_output
+ ]
if pooler_inter and self._should_accumulate_full_payload_output():
with record_function_or_nullcontext("omni_output_builder:accumulate_full_payload_output"):