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"):