From 1a2bd8d03e1133281e7c51e65ea1bdfab39ffc95 Mon Sep 17 00:00:00 2001 From: Elena Rastorgueva Date: Tue, 1 Sep 2026 17:40:22 +0000 Subject: [PATCH 1/6] feat(speechlm2): add streaming VoiceChat inference with native and vLLM-Omni backends StreamingS2SPipeline and the VoiceChat wrapper share one frame loop, with LLM and TTS engines independently selectable. Signed-off-by: Elena Rastorgueva --- docs/source/speechlm2/intro.rst | 43 +- docs/source/speechlm2/models.rst | 2 +- docs/source/speechlm2/streaming_inference.rst | 725 ++++++ .../conf/s2s_streaming.yaml | 163 ++ .../s2s_streaming_infer.py | 122 ++ .../speechlm2/data/duplex_ear_tts_dataset.py | 18 +- .../speechlm2/data/duplex_stt_dataset.py | 26 +- .../collections/speechlm2/data/force_align.py | 18 +- .../collections/speechlm2/data/s2s_dataset.py | 19 +- .../speechlm2/inference/__init__.py | 34 + .../speechlm2/inference/factory/__init__.py | 13 + .../inference/factory/s2s_pipeline_builder.py | 74 + .../inference/model_wrappers/__init__.py | 13 + .../model_wrappers/backend/__init__.py | 37 + .../model_wrappers/backend/eartts.py | 57 + .../inference/model_wrappers/backend/llm.py | 93 + .../backend/pytorch/__init__.py | 13 + .../model_wrappers/backend/pytorch/eartts.py | 178 ++ .../model_wrappers/backend/pytorch/llm.py | 391 ++++ .../model_wrappers/backend/vllm/__init__.py | 39 + .../model_wrappers/backend/vllm/eartts.py | 66 + .../model_wrappers/backend/vllm/llm.py | 78 + .../inference/model_wrappers/capabilities.py | 61 + .../inference/model_wrappers/codec.py | 88 + .../model_wrappers/config_overrides.py | 195 ++ .../inference/model_wrappers/decode_state.py | 202 ++ .../model_wrappers/engine_selection.py | 141 ++ .../nemotron_voicechat_inference_wrapper.py | 1108 ++++++++++ .../model_wrappers/perception_cache.py | 573 +++++ .../inference/model_wrappers/text_sampling.py | 96 + .../speechlm2/inference/pipelines/__init__.py | 13 + .../pipelines/s2s_pipeline_interface.py | 80 + .../pipelines/streaming_s2s_pipeline.py | 855 ++++++++ .../speechlm2/inference/streaming/__init__.py | 13 + .../inference/streaming/framing/__init__.py | 13 + .../streaming/framing/s2s_request_options.py | 67 + .../framing/silence_padded_frame_streamer.py | 75 + .../framing/silence_padded_stream.py | 97 + .../inference/streaming/state/__init__.py | 13 + .../streaming/state/s2s_context_manager.py | 125 ++ .../streaming/state/s2s_streaming_output.py | 230 ++ .../speechlm2/inference/utils/__init__.py | 13 + .../speechlm2/inference/utils/audio_data.py | 178 ++ .../inference/utils/stepprogressbar.py | 74 + .../speechlm2/inference/vllm_omni/__init__.py | 59 + .../inference/vllm_omni/checkpoint.py | 311 +++ .../inference/vllm_omni/deploy/eartts.yaml | 45 + .../vllm_omni/deploy/nemotron_voicechat.yaml | 62 + .../inference/vllm_omni/eartts/__init__.py | 17 + .../vllm_omni/eartts/configuration_eartts.py | 217 ++ .../inference/vllm_omni/eartts/eartts.py | 1945 +++++++++++++++++ .../inference/vllm_omni/eartts/pipeline.py | 63 + .../inference/vllm_omni/eartts/scheduler.py | 451 ++++ .../vllm_omni/nemotron_duplex_h/__init__.py | 19 + .../nemotron_duplex_h/nemotron_duplex_h.py | 771 +++++++ .../vllm_omni/nemotron_duplex_h/sampling.py | 247 +++ .../vllm_omni/nemotron_voicechat/__init__.py | 27 + .../vllm_omni/nemotron_voicechat/pipeline.py | 80 + .../vllm_omni/nemotron_voicechat/scheduler.py | 68 + .../speechlm2/inference/vllm_omni/outputs.py | 128 ++ .../speechlm2/inference/vllm_omni/register.py | 154 ++ .../speechlm2/inference/vllm_omni/runtime.py | 322 +++ .../inference/vllm_omni/scripts/__init__.py | 13 + .../convert_duplex_eartts_checkpoint.py | 357 +++ .../scripts/convert_duplex_stt_checkpoint.py | 343 +++ .../speechlm2/inference/vllm_omni/session.py | 592 +++++ .../speechlm2/models/duplex_ear_tts.py | 12 +- .../speechlm2/models/duplex_stt_model.py | 125 +- .../speechlm2/models/nemotron_voicechat.py | 298 ++- .../speechlm2/modules/ear_tts_model.py | 52 +- nemo/collections/speechlm2/parts/hf_hub.py | 22 + .../speechlm2/parts/logit_boosts.py | 121 + nemo/collections/speechlm2/parts/precision.py | 100 + .../collections/speechlm2/parts/pretrained.py | 49 +- .../collections/speechlm2/parts/text_utils.py | 239 +- .../streaming/duplex_stt_inference.py | 150 +- pyproject.toml | 17 + .../nemo_inference_pipelines/conftest.py | 461 ++++ .../test_config_overrides.py | 193 ++ .../test_engine_selection.py | 77 + ...est_nemotron_voicechat_pipeline_nocrash.py | 231 ++ ...test_nemotron_voicechat_pipeline_parity.py | 323 +++ .../test_nemotron_voicechat_pipeline_vllm.py | 87 + .../test_text_sampling.py | 146 ++ .../test_vllm_omni_checkpoint.py | 112 + .../test_vllm_omni_eartts_cfg.py | 289 +++ ...emotron_voicechat.py => test_voicechat.py} | 245 ++- 87 files changed, 15998 insertions(+), 174 deletions(-) create mode 100644 docs/source/speechlm2/streaming_inference.rst create mode 100644 examples/speechlm2/nemo_inference_pipelines/conf/s2s_streaming.yaml create mode 100644 examples/speechlm2/nemo_inference_pipelines/s2s_streaming_infer.py create mode 100644 nemo/collections/speechlm2/inference/__init__.py create mode 100644 nemo/collections/speechlm2/inference/factory/__init__.py create mode 100644 nemo/collections/speechlm2/inference/factory/s2s_pipeline_builder.py create mode 100644 nemo/collections/speechlm2/inference/model_wrappers/__init__.py create mode 100644 nemo/collections/speechlm2/inference/model_wrappers/backend/__init__.py create mode 100644 nemo/collections/speechlm2/inference/model_wrappers/backend/eartts.py create mode 100644 nemo/collections/speechlm2/inference/model_wrappers/backend/llm.py create mode 100644 nemo/collections/speechlm2/inference/model_wrappers/backend/pytorch/__init__.py create mode 100644 nemo/collections/speechlm2/inference/model_wrappers/backend/pytorch/eartts.py create mode 100644 nemo/collections/speechlm2/inference/model_wrappers/backend/pytorch/llm.py create mode 100644 nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/__init__.py create mode 100644 nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/eartts.py create mode 100644 nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/llm.py create mode 100644 nemo/collections/speechlm2/inference/model_wrappers/capabilities.py create mode 100644 nemo/collections/speechlm2/inference/model_wrappers/codec.py create mode 100644 nemo/collections/speechlm2/inference/model_wrappers/config_overrides.py create mode 100644 nemo/collections/speechlm2/inference/model_wrappers/decode_state.py create mode 100644 nemo/collections/speechlm2/inference/model_wrappers/engine_selection.py create mode 100644 nemo/collections/speechlm2/inference/model_wrappers/nemotron_voicechat_inference_wrapper.py create mode 100644 nemo/collections/speechlm2/inference/model_wrappers/perception_cache.py create mode 100644 nemo/collections/speechlm2/inference/model_wrappers/text_sampling.py create mode 100644 nemo/collections/speechlm2/inference/pipelines/__init__.py create mode 100644 nemo/collections/speechlm2/inference/pipelines/s2s_pipeline_interface.py create mode 100644 nemo/collections/speechlm2/inference/pipelines/streaming_s2s_pipeline.py create mode 100644 nemo/collections/speechlm2/inference/streaming/__init__.py create mode 100644 nemo/collections/speechlm2/inference/streaming/framing/__init__.py create mode 100644 nemo/collections/speechlm2/inference/streaming/framing/s2s_request_options.py create mode 100644 nemo/collections/speechlm2/inference/streaming/framing/silence_padded_frame_streamer.py create mode 100644 nemo/collections/speechlm2/inference/streaming/framing/silence_padded_stream.py create mode 100644 nemo/collections/speechlm2/inference/streaming/state/__init__.py create mode 100644 nemo/collections/speechlm2/inference/streaming/state/s2s_context_manager.py create mode 100644 nemo/collections/speechlm2/inference/streaming/state/s2s_streaming_output.py create mode 100644 nemo/collections/speechlm2/inference/utils/__init__.py create mode 100644 nemo/collections/speechlm2/inference/utils/audio_data.py create mode 100644 nemo/collections/speechlm2/inference/utils/stepprogressbar.py create mode 100644 nemo/collections/speechlm2/inference/vllm_omni/__init__.py create mode 100644 nemo/collections/speechlm2/inference/vllm_omni/checkpoint.py create mode 100644 nemo/collections/speechlm2/inference/vllm_omni/deploy/eartts.yaml create mode 100644 nemo/collections/speechlm2/inference/vllm_omni/deploy/nemotron_voicechat.yaml create mode 100644 nemo/collections/speechlm2/inference/vllm_omni/eartts/__init__.py create mode 100644 nemo/collections/speechlm2/inference/vllm_omni/eartts/configuration_eartts.py create mode 100644 nemo/collections/speechlm2/inference/vllm_omni/eartts/eartts.py create mode 100644 nemo/collections/speechlm2/inference/vllm_omni/eartts/pipeline.py create mode 100644 nemo/collections/speechlm2/inference/vllm_omni/eartts/scheduler.py create mode 100644 nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/__init__.py create mode 100644 nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/nemotron_duplex_h.py create mode 100644 nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/sampling.py create mode 100644 nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/__init__.py create mode 100644 nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/pipeline.py create mode 100644 nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/scheduler.py create mode 100644 nemo/collections/speechlm2/inference/vllm_omni/outputs.py create mode 100644 nemo/collections/speechlm2/inference/vllm_omni/register.py create mode 100644 nemo/collections/speechlm2/inference/vllm_omni/runtime.py create mode 100644 nemo/collections/speechlm2/inference/vllm_omni/scripts/__init__.py create mode 100644 nemo/collections/speechlm2/inference/vllm_omni/scripts/convert_duplex_eartts_checkpoint.py create mode 100644 nemo/collections/speechlm2/inference/vllm_omni/scripts/convert_duplex_stt_checkpoint.py create mode 100644 nemo/collections/speechlm2/inference/vllm_omni/session.py create mode 100644 nemo/collections/speechlm2/parts/logit_boosts.py create mode 100644 tests/collections/speechlm2/nemo_inference_pipelines/conftest.py create mode 100644 tests/collections/speechlm2/nemo_inference_pipelines/test_config_overrides.py create mode 100644 tests/collections/speechlm2/nemo_inference_pipelines/test_engine_selection.py create mode 100644 tests/collections/speechlm2/nemo_inference_pipelines/test_nemotron_voicechat_pipeline_nocrash.py create mode 100644 tests/collections/speechlm2/nemo_inference_pipelines/test_nemotron_voicechat_pipeline_parity.py create mode 100644 tests/collections/speechlm2/nemo_inference_pipelines/test_nemotron_voicechat_pipeline_vllm.py create mode 100644 tests/collections/speechlm2/nemo_inference_pipelines/test_text_sampling.py create mode 100644 tests/collections/speechlm2/nemo_inference_pipelines/test_vllm_omni_checkpoint.py create mode 100644 tests/collections/speechlm2/nemo_inference_pipelines/test_vllm_omni_eartts_cfg.py rename tests/collections/speechlm2/{test_nemotron_voicechat.py => test_voicechat.py} (53%) diff --git a/docs/source/speechlm2/intro.rst b/docs/source/speechlm2/intro.rst index 8aac63f9f12f..1de9c44d7896 100644 --- a/docs/source/speechlm2/intro.rst +++ b/docs/source/speechlm2/intro.rst @@ -245,7 +245,7 @@ You can evaluate and run full-duplex inference using the `NemotronVoiceChat` pip from nemo.collections.audio.parts.utils.transforms import resample import nemo.collections.speechlm2 as slm - model = slm.models.NemotronVoiceChat.from_pretrained("path/to/pretrained_checkpoint").eval() + model = slm.models.NemotronVoiceChat.from_pretrained("nvidia/NVIDIA-NemotronLabs-VoiceChat-11B").eval() # Load user audio prompt audio_path = "path/to/user_audio.wav" @@ -268,7 +268,7 @@ You can evaluate and run full-duplex inference using the `NemotronVoiceChat` pip # Note: If an explicit audio reference is not passed into `offline_inference`, # the model relies on the internal config parameters: - # 1. model.cfg.inference_speaker_name (Highest priority preset, e.g., 'Megan') + # 1. model.cfg.inference_speaker_name (Highest priority preset, e.g., 'Aria') # 2. model.cfg.inference_speaker_reference (Fallback audio file path) # Run full offline inference @@ -285,7 +285,43 @@ You can evaluate and run full-duplex inference using the `NemotronVoiceChat` pip print(f"Agent response: {generated_text}") # generated_speech can now be saved or played (sampled at model.target_sample_rate) - + +NemotronVoiceChat Streaming Inference +************************************* + +For real-time, chunk-by-chunk inference (as opposed to the offline mode shown +above), use the Streaming S2S Pipeline: + +.. code-block:: python + + from nemo.collections.speechlm2.inference import S2SPipelineBuilder + from nemo.collections.speechlm2.inference.model_wrappers.engine_selection import ( + inference_precision_from_cfg, + ) + + with inference_precision_from_cfg(cfg.s2s): + pipeline = S2SPipelineBuilder.build_pipeline(cfg) + try: + outputs = pipeline.run(audio_filepaths, options=options) + finally: + pipeline.shutdown() + +Or from the command line: + +.. code-block:: bash + + python examples/speechlm2/nemo_inference_pipelines/s2s_streaming_infer.py \ + audio_file=/path/to/audio \ + s2s.model_path=nvidia/NVIDIA-NemotronLabs-VoiceChat-11B \ + s2s.speaker_name=Aria \ + s2s.llm_engine_type=native \ + s2s.tts_engine_type=native \ + s2s.system_prompt="You are a helpful assistant." \ + streaming.chunk_size_in_secs=0.24 \ + streaming.buffer_size_in_secs=1.68 + +See :doc:`streaming_inference` for full details on configuration, architecture, +and server integration. Training a Model ---------------- @@ -393,3 +429,4 @@ For more information, see additional sections in the SpeechLM2 docs: datasets configs training_and_scaling + streaming_inference diff --git a/docs/source/speechlm2/models.rst b/docs/source/speechlm2/models.rst index 3ff0b9888ea9..bb5907363e82 100644 --- a/docs/source/speechlm2/models.rst +++ b/docs/source/speechlm2/models.rst @@ -312,7 +312,7 @@ All models in the speechlm2 collection can be instantiated from pretrained check ear_tts_model = slm.models.DuplexEARTTS.from_pretrained("path/to/checkpoint") # Load NemotronVoiceChat (Inference Only) - voicechat_model = slm.models.NemotronVoiceChat.from_pretrained("path/to/checkpoint") + voicechat_model = slm.models.NemotronVoiceChat.from_pretrained("nvidia/NVIDIA-NemotronLabs-VoiceChat-11B") Remote HuggingFace code is disabled by default. If a trusted checkpoint requires custom code, opt in at runtime and pin the repository to a reviewed revision: diff --git a/docs/source/speechlm2/streaming_inference.rst b/docs/source/speechlm2/streaming_inference.rst new file mode 100644 index 000000000000..5a049b3e7955 --- /dev/null +++ b/docs/source/speechlm2/streaming_inference.rst @@ -0,0 +1,725 @@ +Streaming Inference +=================== + +The speechlm2 collection provides a streaming inference pipeline for +NemotronVoiceChat that processes audio chunk by chunk, producing text and speech +output incrementally. The pipeline follows a similar API to the NeMo ASR Inference Pipelines +(see ``nemo.collections.asr.inference``). + +There are two ways to use the pipeline: + +* ``StreamingS2SPipeline.run()`` processes complete audio files. It is used by + ``s2s_streaming_infer.py`` for a single ``.wav`` file, a directory of ``.wav`` + files, or a manifest. +* ``StreamingS2SPipeline.generate_step()`` processes one batch of ``Frame`` + objects and returns incremental outputs for that step. Use it for servers, + microphone connectors, or other live audio sources. + +.. code-block:: text + + File inputs: one or more .wav files + (single path, directory, or manifest) + │ + ▼ + run(audio_filepaths) + │ creates Frame chunks + ▼ + generate_step(frames) + │ + ├─ incremental agent audio + text + └─ incremental user ASR text (when the checkpoint has an ASR head) + +Each audio file passed to ``run()`` is treated as one continuous audio stream. ``run()`` +accumulates the per-step outputs for each stream and writes final audio/text +artifacts when the stream ends. + +The script can append trailing silence so the agent is more likely to finish +speaking before the stream ends. When a manifest contains reference ``text`` +fields, it also reports WER for the recognized user speech. + +Streaming inference is single-stream: ``streaming.batch_size`` must be ``1``. + +Script Call Path +---------------- + +The ``s2s_streaming_infer.py`` script follows this call path: + +.. code-block:: text + + Entry Script s2s_streaming_infer.py + │ + ▼ + Pipeline StreamingS2SPipeline.run() + │ - audio buffering + │ - state management + │ - file I/O + ▼ + Model Wrapper NemotronVoicechatInferenceWrapper + │ - infer_one_step() + │ - perception + │ - model_llm_interface (PyTorchLLM) + │ - model_eartts_interface (PyTorchEarTTS) + │ (replaced by independent AsyncOmni engines + │ when a component is vllm_omni) + │ - codec decode + ▼ + Model NemotronVoiceChat + - DuplexSTTModel + DuplexEARTTS + +(With ``s2s.decode_audio=false``, the model still predicts text and any +checkpoint-provided auxiliary tokens, but skips EarTTS generation and codec +decoding.) + +Quick Start +----------- + +File-Based Inference from a Script +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Call the Python script and pass configuration values as command-line overrides: + +.. code-block:: bash + + python examples/speechlm2/nemo_inference_pipelines/s2s_streaming_infer.py \ + --config-path=examples/speechlm2/nemo_inference_pipelines/conf \ + --config-name=s2s_streaming \ + audio_file=/path/to/audio_or_directory_or_manifest.json \ + output_dir=./generated \ + s2s.model_path=nvidia/NVIDIA-NemotronLabs-VoiceChat-11B \ + s2s.speaker_name=Aria \ + s2s.llm_engine_type=native \ + s2s.tts_engine_type=native \ + s2s.system_prompt="You are a helpful assistant." \ + streaming.chunk_size_in_secs=0.24 \ + streaming.buffer_size_in_secs=1.68 + +This will: + +1. Load the NemotronVoiceChat checkpoint. +2. Stream each audio file through the pipeline in chunks. +3. Save per-stream output files under ``output_dir``: generated ``.wav``, + stereo input+output ``.wav``, ``.txt``, and per-token ``.ctm``. +4. Write ``output_processed.json`` and ``output_raw.json`` summarising the run. + +Public checkpoint +^^^^^^^^^^^^^^^^^ + +The public weights are +`NVIDIA-NemotronLabs-VoiceChat-11B +`_. +Pass the Hugging Face repository ID as ``s2s.model_path`` (shown in Quick Start +above). The first run downloads into the Hugging Face cache; a local directory +works for offline use. The registered speaker name is ``Aria``. + +That checkpoint has a function-token channel and no duplex ASR head. The +pipeline feeds the previous function token back into the next frame and +exposes a decoded copy on the function output; it does not execute tool +calls. User-transcription fields are empty and ASR-based forced turn-taking +is disabled; the model's learned duplex turn-taking remains active. The +bundled RNN-T weights are not loaded. + +Other checkpoints may carry an ASR head, a function head, both, or neither. +The heads are independent. When an ASR head is present, user transcription +and ASR-based forced turn-taking are available. + +Leave both engine keys at ``native`` for PyTorch inference. Either component +can be switched to ``vllm_omni`` after converting a wrapper checkpoint. + +Programmatic Usage +^^^^^^^^^^^^^^^^^^ + +.. code-block:: python + + from nemo.collections.speechlm2.inference import S2SPipelineBuilder + from nemo.collections.speechlm2.inference.model_wrappers.engine_selection import ( + inference_precision_from_cfg, + ) + + with inference_precision_from_cfg(cfg.s2s): + pipeline = S2SPipelineBuilder.build_pipeline(cfg) + try: + outputs = pipeline.run(audio_filepaths, options=options) + finally: + pipeline.shutdown() + + # returns list[S2SStreamingOutput], one per input file + # each element has: .output_text_str, .output_asr_text_str, .audio_filepath, ... + +File Inputs and Manifests +^^^^^^^^^^^^^^^^^^^^^^^^^ + +The ``audio_file`` argument accepted by +``examples/speechlm2/nemo_inference_pipelines/s2s_streaming_infer.py`` may be: + +* A single ``.wav`` file. +* A directory, in which case all ``.wav`` files in that directory are streamed. +* A line-delimited ``.json`` or ``.jsonl`` manifest listing audio files. + +Manifest entries must provide ``audio_filepath`` and may also provide +``system_prompt`` and ``text``: + +.. code-block:: json + + {"audio_filepath": "audio/example.wav", "system_prompt": "You are helpful.", "text": "reference user transcript"} + +The JSON/JSONL manifest accepted by ``s2s_streaming_infer.py`` has this schema. +Its ``text`` field is read only as an optional reference transcript for WER on +the ASR/user side. Generated agent text is produced by the model and written to +``pred_text`` in the output JSON. + +This lightweight streaming inference manifest is distinct from the dataset +manifests used for SpeechLM2 training and offline evaluation. For those dataset +formats, see :doc:`SpeechLM2 datasets `. + +File paths in streaming inference manifests are resolved relative to the +manifest file. Audio from file inputs is converted to mono and +resampled to ``streaming.input_sample_rate`` before it is chunked. + + +Configuration +------------- + +The streaming inference configuration is defined in +``examples/speechlm2/nemo_inference_pipelines/conf/s2s_streaming.yaml``. + +Key configuration groups: + +S2S Model Settings (``s2s``) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. list-table:: + :header-rows: 1 + :widths: 30 15 55 + + * - Parameter + - Default + - Description + * - ``model_path`` + - (required) + - Path to the NemotronVoiceChat HuggingFace checkpoint. + * - ``llm_engine_type`` + - ``native`` + - LLM backend: ``native`` or ``vllm_omni``. + * - ``tts_engine_type`` + - ``native`` + - TTS backend: ``native`` or ``vllm_omni``. Independent of the LLM. + * - ``speaker_name`` + - ``null`` + - Required when ``decode_audio`` is true. Must match a speaker registered + in the checkpoint. Public checkpoints do not support cloning from a + reference wav. + * - ``system_prompt`` + - (required) + - Text injected into the LLM KV cache before audio streaming begins. + * - ``compute_dtype`` + - ``bfloat16`` + - Precision for LLM/embedding layers. + * - ``use_perception_cache`` + - ``true`` + - Cache-aware streaming for the perception encoder. + * - ``use_llm_cache`` + - ``false`` + - Reuse the native LLM KV cache instead of replaying history. NemotronH + requires Transformers 5.13 or newer; leave disabled on older runtimes. + * - ``top_p`` + - ``0.5`` + - Top-p sampling threshold. + * - ``temperature`` + - ``0.3`` + - Sampling temperature. + * - ``repetition_penalty`` + - ``1.1`` + - Repetition penalty applied to previously generated tokens. + * - ``deterministic`` + - ``false`` + - Force deterministic mode (native engine only). + * - ``profile_timing`` + - ``false`` + - Insert ``torch.cuda.synchronize()`` around each stage for accurate + per-stage timing. Disabled by default to avoid GPU stalls. + +Streaming Settings (``streaming``) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. list-table:: + :header-rows: 1 + :widths: 30 15 55 + + * - Parameter + - Default + - Description + * - ``chunk_size_in_secs`` + - (required) + - Audio processed per inference step. Must be a multiple of 0.08 s. + * - ``buffer_size_in_secs`` + - (required) + - Sliding-window size passed to the perception encoder. + * - ``batch_size`` + - ``1`` + - Must be ``1``. Single-stream inference only. + * - ``max_len`` + - ``8192`` + - Maximum number of frames per stream. + +Padding Settings (top-level) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +These append trailing silence after the real input so the agent is more likely +to finish speaking before the stream ends. They are not batch padding. At most +one may be set: + +.. list-table:: + :header-rows: 1 + :widths: 30 15 55 + + * - Parameter + - Default + - Description + * - ``pad_audio_to_sec`` + - ``null`` + - Append trailing silence so each input reaches this duration. + * - ``pad_silence_ratio`` + - ``null`` + - Append trailing silence equal to this fraction of the original duration. + * - ``pad_audio_by_sec`` + - ``null`` + - Append this many extra seconds of trailing silence. + + +Server Integration +------------------ + +Use ``generate_step()`` directly when audio does not come from complete files: +for example, from a microphone, socket, Triton server, or browser UI. The +caller owns the input connector: capture audio, convert it to mono +``streaming.input_sample_rate`` samples, split it into chunks, and pass those +chunks as ``Frame`` objects. + +``generate_step()`` returns one ``GenerateStepOutput`` for each input frame, +containing the audio, agent text, and user ASR text produced by that step. + +.. code-block:: python + + from nemo.collections.asr.inference.streaming.framing.request import Frame + from nemo.collections.speechlm2.inference import S2SRequestOptions + + # 1. Initialize the stream before recording starts. + # Send empty audio because prefill will likely take longer than chunk_size_in_secs. + init_frame = Frame( + samples=torch.empty(0), + stream_id=stream_id, + is_first=True, is_last=False, + options=S2SRequestOptions(system_prompt=prompt, top_p=0.9), + ) + pipeline.generate_step([init_frame]) + # -> the input connector can now start sending audio + + # 2. For each input audio chunk, run one streaming step. + for chunk, is_last in audio_source: + frame = Frame( + samples=chunk, + stream_id=stream_id, + is_first=False, is_last=is_last, + ) + outputs = pipeline.generate_step([frame]) + for out in outputs: + send_to_client(out.audio, out.text, out.asr_text) + +Per-stream options (``system_prompt``, ``top_p``, ``temperature``, +``repetition_penalty``) are attached to the ``is_first`` frame via +``S2SRequestOptions``. Any field left as ``None`` falls back to the +pipeline-level YAML default through ``fill_defaults()``. + +.. _init-and-latency: + +Init and Latency +^^^^^^^^^^^^^^^^ + +When ``generate_step`` sees ``is_first``, it always runs stream +initialization (context creation, KV-cache prefill). If the frame also +carries audio, inference runs immediately after init in the same call. + +For **latency-sensitive** integrations, prefill can take hundreds of +milliseconds or even multiple seconds. Send ``is_first`` with **empty audio**, +wait for the response confirming init is done, and only then start sending real +audio. This prevents input audio from queuing up during the expensive prefill +phase. + +For **batch/offline** usage (``run()``), there is no real-time +constraint. The first frame carries both ``is_first`` and real audio, +so init and first-chunk processing happen in one call with no extra +round-trip. + +The pipeline makes no distinction between these cases — it initializes +on ``is_first`` and processes whatever audio is present. The latency +trade-off is entirely the caller's choice. + + +Architecture +------------ + +File Chunking in ``run()`` +^^^^^^^^^^^^^^^^^^^^^^^^^^ + +For file inputs, ``StreamingS2SPipeline.run()`` uses +``SilencePaddedContinuousBatchedFrameStreamer`` to load the audio paths, +convert them to mono ``streaming.input_sample_rate`` audio, and emit ``Frame`` +chunks. The streamer uses the configured chunk size, batch size, and optional +silence-padding settings. ``run()`` then passes each emitted frame batch to +``generate_step``. Live integrations do not need this helper; they can +construct ``Frame`` objects directly and call ``generate_step``. + +The Core Streaming Loop +^^^^^^^^^^^^^^^^^^^^^^^ + +``StreamingS2SPipeline.run()`` orchestrates the streaming loop, delegating +per-chunk inference to ``generate_step()`` and saving outputs as streams +finish. In simplified pseudocode: + +.. code-block:: python + + # Inside StreamingS2SPipeline.run() (simplified): + self.open_session() + for frames in streamer: + # step_outputs[i] carries GenerateStepOutput.audio / .text / .asr_text + # — the new agent audio and text produced by this chunk. + step_outputs = self.generate_step(frames) + self._finalize_and_save_finished_streams(frames, ...) + self.close_session() + # run() then returns list[S2SStreamingOutput], one per input file + +``run()`` returns a list of finalized ``S2SStreamingOutput`` objects (one per +input audio file) with the accumulated texts, token tensors, and audio +filepaths. + +``run()`` writes outputs as each stream finishes, so results appear on disk +before the full run completes. For each stream: + +* ``.txt`` - agent transcript. +* ``.ctm`` & ``_asr.ctm`` - per-token timing for agent text and ASR text. + Timestamps reflect when the text token was generated by the model. +* ``.wav`` & ``_input_output.wav`` - generated agent audio, plus a + stereo file with input on one channel and output on the other. + + * In the stereo file, the generated-output channel is offset by one chunk so + playback reflects the minimum delay from waiting for a full input chunk + before generating output (Note: actual inference time would add to this in + a real deployment). + * Both audio files are skipped when ``s2s.decode_audio=false``. + +After all streams finish, ``s2s_streaming_infer.py`` also writes two JSON +summaries of the run: ``output_raw.json`` (full token stream including padding +tokens) and ``output_processed.json`` (padding tokens removed for legibility). + +``run()`` loops over chunks of existing audio files, calling ``generate_step()`` +on each; ``generate_step()`` can also be called directly when audio comes from +a non-file source. + + +What Happens Inside One Step +^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. code-block:: text + + generate_step(frames) + │ + ├─ for each frame where is_first=True: + │ │ + │ └─ _init_state(stream_id, options) + │ 1. fill_defaults() ← fill None fields from YAML + │ 2. create_state(options) ← pipeline-level state + │ 3. reset context_manager ← fresh decode-state storage + │ 4. prefill system prompt ← populate LLM KV cache + │ + └─ any frames with audio? + │ + NO → return empty outputs (server prefill-only request) + │ + YES → update per-stream sliding audio buffer + │ + ▼ + generate_step_for_frames() + 1. perception encoder + 2. per-frame LLM loop + 3. per-frame TTS (when decode_audio=true) + 4. codec decode (when decode_audio=true) + 5. state updates + output accumulation + 6. return list[GenerateStepOutput] + +Each call to ``generate_step(frames)`` performs: + +1. **Stream init on** ``is_first`` -- If a frame has ``is_first=True``, the + private ``_init_state()`` method runs: per-stream options are merged with + pipeline defaults (via ``S2SRequestOptions.fill_defaults()``), + a fresh ``S2SStreamingOutput`` is created, the context manager is + allocated, and the LLM KV cache is prefilled with the system prompt and + TTS speaker embedding. This mirrors ASR's ``init_state()`` called inside + ``transcribe_step()``. If the frame carries no audio (zero-length + samples), the method returns after init — this is the recommended + pattern for latency-sensitive deployments (see + :ref:`init-and-latency` above). + +2. **Audio buffer update** -- ``generate_step`` updates each stream's rolling + audio buffer so the model receives the current ``buffer_size_in_secs``-size + window of audio. + +3. **Model inference** via ``infer_one_step(audio_buffer, state)``: + + a. **Perception** -- The audio buffer is encoded by the streaming + FastConformer encoder into frame embeddings. + b. **Per-frame LLM loop** -- For each of the ``num_frames_per_chunk`` + frames, the pipeline builds an input embedding (user audio + + previous-step text/ASR tokens), runs it through the LLM, and obtains + predicted text and ASR tokens. + c. **TTS code generation** -- When ``s2s.decode_audio=true``, predicted text + tokens are fed into the EarTTS model to produce audio codec codes. + d. **Codec decode** -- When ``s2s.decode_audio=true``, the accumulated codes + are decoded into a waveform. + +4. **State updates** -- The per-stream ``StreamingDecodeState`` is updated + with model-side decode state such as generated-token history and caches. + +5. **Output accumulation** -- Decoded audio and text are appended to the + per-stream ``S2SStreamingOutput``. + + +Data Objects +^^^^^^^^^^^^ + +The streaming pipeline uses four data objects. Two are **model-level** +(owned by the model wrapper) and two are **pipeline-level** (owned by +``StreamingS2SPipeline``): + +.. code-block:: text + + Model level (decode_state.py) + ───────────────────────────────────────────────────────────── + StreamingDecodeState created per stream + GPU KV caches, token mutated in-place by infer_one_step() + workspaces, perception destroyed at end-of-stream + cache, codec cache + │ + │ infer_one_step() + ▼ + InferenceStepResult created each step + predicted tokens, text returned to the pipeline + strings, decoded audio consumed immediately + + Pipeline level (streaming_s2s_pipeline.py, s2s_streaming_output.py) + ───────────────────────────────────────────────────────────── + S2SStreamingOutput created per stream + accumulates audio chunks finalized fields (text_with_timestamps, + and text across steps audio_filepath, etc.) filled at end-of-stream + returned by run() + ▲ + │ each step appends + │ + GenerateStepOutput created each step + incremental per-stream returned by generate_step() + audio + text used by server integrations + +**StreamingDecodeState** lives in ``S2SContextManager`` and holds the heavy +GPU tensors (KV caches, perception cache, token workspaces). It is created +by the model wrapper, mutated in-place by ``infer_one_step()``, and +destroyed at end-of-stream. + +**S2SStreamingOutput** lives in the pipeline's ``_state_pool``. During +streaming it accumulates audio chunks and text parts. At end-of-stream the +pipeline populates its finalized fields (``text_with_timestamps``, +``raw_text``, ``audio_filepath``, token tensors) and returns the same +object from ``run()``. + + +Inference Backends +^^^^^^^^^^^^^^^^^^ + +NemotronVoiceChat has two inference components that each need a backend: + +- **LLM** (DuplexSTT backbone) -- takes audio embeddings from the perception + encoder and predicts text tokens plus checkpoint-dependent ASR and function + tokens at each frame. +- **TTS** (EarTTS) -- takes the predicted text token and produces audio codec + codes (RVQ acoustic tokens). + +Two engines can drive those components. ``llm_engine_type`` and +``tts_engine_type`` select them independently; an omitted key defaults to +``native``: + +.. list-table:: + :header-rows: 1 + :widths: 25 25 25 + + * - ``llm_engine_type`` + - ``tts_engine_type`` + - Components + * - ``native`` + - ``native`` + - ``PyTorchLLM`` + ``PyTorchEarTTS`` + * - ``native`` + - ``vllm_omni`` + - ``PyTorchLLM`` + one-stage EarTTS ``AsyncOmni`` + * - ``vllm_omni`` + - ``native`` + - one-stage Nemotron ``AsyncOmni`` + ``PyTorchEarTTS`` + * - ``vllm_omni`` + - ``vllm_omni`` + - independent one-stage Nemotron and EarTTS ``AsyncOmni`` engines + +The two components have one contract each, and each contract has a PyTorch and +a vLLM-Omni implementation: + +.. code-block:: text + + backend/ + llm.py # DuplexLLM ABC: step() -> this frame's tokens + eartts.py # DuplexTTS ABC: step() -> this frame's codes + pytorch/ + llm.py # PyTorchLLM (wraps the DuplexSTT forward pass) + eartts.py # PyTorchEarTTS (wraps DuplexEARTTS.infer_codes_one_step) + vllm/ + llm.py # VllmLLM (drives OmniStreamingSession.step_llm) + eartts.py # VllmEarTTS (drives OmniStreamingSession.step_tts) + +``NemotronVoicechatInferenceWrapper`` selects one implementation per component +at construction and stores them as ``llm_backend`` and ``tts_backend``. Its +frame loop then calls ``step()`` on each without inspecting the engine type, so +all four combinations run the same code path. Perception, the audio codec and +tokenization stay on PyTorch in every combination. + +The contracts hold only ``step()``. Cache creation, prompt prefill and request +abort are PyTorch-only -- vLLM does them inside the engine and the +session -- so they stay on the PyTorch classes, which the wrapper reaches +through ``model_llm_interface`` / ``model_eartts_interface`` during stream +setup. Those two attributes are ``None`` when their component runs on vLLM. + +Text sampling (top-p, repetition penalty, temperature) is shared rather than +duplicated: both backends decode the text head with +``inference.model_wrappers.text_sampling.sample_text_token``. ``PyTorchLLM`` +calls it directly; the vLLM path reaches it through +``SharedTextSamplingLogitsProcessor``. + +Config support by backend +""""""""""""""""""""""""" + +Settings that model code reads off its own config are listed in +``inference.model_wrappers.config_overrides``, which records where each one +lands and which backends honour it. Anything a selected backend ignores is +reported at load time rather than silently doing nothing. + +.. list-table:: + :header-rows: 1 + :widths: 32 12 12 44 + + * - Setting + - native + - vllm_omni + - Notes + * - ``inference_pad/bos/eos_boost`` + - yes + - yes + - Agent text channel. vLLM applies them in the shared sampling hook. + * - ``inference_user_pad/bos/eos_boost`` + - yes + - yes + - ASR channel. Its logits never reach vLLM's sampler, so the converted + Nemotron applies them itself; the wrapper writes the values into + ``nemotron/config.json`` before the engine starts. + * - ``force_turn_taking`` (+ threshold, pad window) + - yes + - yes + - The rewritten text token is fed back explicitly, so Nemotron's history + stays consistent with the text channel. + * - ``inference_force_speech_silence_on_eos`` + - yes + - partial + - Applied inside EarTTS on both paths. The converted EarTTS always + substitutes codec silence on EOS and has no flag, so it cannot honour + ``false``. + * - ``inference_top_p_or_k``, ``inference_noise_scale``, ``inference_guidance_scale`` + - yes + - no + - The vLLM EarTTS takes sampling from the converted checkpoint and + ``vllm_omni_config`` instead. + * - ``deterministic`` + - yes + - no + - Rejected: vLLM's kernels have no deterministic mode. + * - ``use_llm_cache``, ``use_tts_torch_compile``, ``use_tts_subword_cache`` + - yes + - n/a + - Performance knobs whose intent vLLM already meets: it always keeps a + paged KV cache, compiles inside vLLM, and bakes the subword table at + conversion. Setting them warns. + +vLLM-Omni Integration +""""""""""""""""""""" + +Each component selected as ``vllm_omni`` gets the ``Vllm*`` implementation of +its contract, and the corresponding native class is not created. The wrapper +starts only the selected one-stage ``AsyncOmni`` engine or engines: + +- **Nemotron** -- ``NemotronDuplexHForCausalLM``, which consumes the per-step + acoustic embedding and samples a text token plus the checkpoint's optional + ASR or function token. +- **EarTTS** -- ``EarTTSForCausalLM``, which receives each sampled text token + from NeMo and emits one acoustic frame. + +The split keeps the component boundary in NeMo, so either component can be +replaced without changing the other engine. Nemotron settings come from +``inference/vllm_omni/deploy/nemotron_voicechat.yaml`` and EarTTS settings from +``inference/vllm_omni/deploy/eartts.yaml``. Override them independently with +``vllm_omni_config.stage_overrides`` and +``vllm_omni_config.eartts_stage_overrides``. + +Nemotron text sampling uses vLLM's custom logits-processor hook to call the +same PyTorch sampling helper as the native backend. This preserves all-ones +greedy decoding, special-token bypass, top-p, temperature, and repetition +penalty over agent-frame history while leaving vLLM's built-in sampler in +greedy/no-penalty mode. Stochastic runs share the algorithm but are not +guaranteed to produce identical tokens across backends because their worker +processes do not share RNG state. + +EarTTS classifier-free guidance uses two explicit requests in the same engine. +They have independent KV caches, but a custom scheduler advances them in +lockstep. The unconditional request replaces text conditioning with the +checkpoint's ``null_emb``; the MaskGIT sampler applies the native guidance +formula and returns only the conditional stream's codes. Configure it with +``vllm_omni_config.guidance_enabled`` and ``guidance_scale``. + +Perception, the audio codec and tokenization stay on PyTorch in every engine pairing. + +Auxiliary channels +'''''''''''''''''' + +Besides the agent text token, a checkpoint may predict ASR and/or function +tokens per frame. The flags are independent: + +- ``predict_user_text=True`` gives an ASR channel (``asr_head`` plus its own + ``embed_asr_tokens`` table) that transcribes the user. +- ``use_function_head=True`` gives a function channel (``function_head``), + whose feedback is embedded with the *text* ``embed_tokens`` and scaled by + ``duplex_function_channel_weight``. + +The previous frame's auxiliary token is part of the next frame's model input, +so it remains in autoregressive state to reproduce the checkpoint's text and +audio. The pipeline also exposes a decoded copy to clients. Function-channel +text is informational; the pipeline does not execute tool calls. + +Both engines support both channels. On ``vllm_omni`` the converter +(``convert_duplex_stt_checkpoint.py``) records which heads a checkpoint carries +as ``use_asr_head`` / ``use_function_head`` in the wrapper config, because +Nemotron has to decide which modules to build before it sees any weights, and +then feeds each enabled channel back to itself through its +``postprocess`` -> ``preprocess`` buffers. + +On stock vLLM-Omni 0.26, the registered Nemotron stage keeps +``final_output_type="text"`` while using the final multimodal engine-output +path for its optional auxiliary tensor. Text remains on ``RequestOutput`` and +``asr_tokens`` or ``function_tokens`` reaches the caller beside it. This +configuration was validated with the public function-head VoiceChat +checkpoint; ASR-head checkpoints use the same converter and extraction path. + +Client APIs keep the channels distinct: ASR-head checkpoints populate the +existing ASR output, while function-head checkpoints populate the function +output and capability metadata. A missing head always produces an empty +corresponding client field. diff --git a/examples/speechlm2/nemo_inference_pipelines/conf/s2s_streaming.yaml b/examples/speechlm2/nemo_inference_pipelines/conf/s2s_streaming.yaml new file mode 100644 index 000000000000..a305a074bff3 --- /dev/null +++ b/examples/speechlm2/nemo_inference_pipelines/conf/s2s_streaming.yaml @@ -0,0 +1,163 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Runtime args (set in this file or overridden with Hydra command-line arguments) +# audio_file accepts a single .wav file, a directory of .wav files, or a +# line-delimited JSON/JSONL manifest with "audio_filepath" and optional +# "system_prompt"/"text" fields. For this streaming inference script, "text" is +# an optional reference transcript used only for WER. +audio_file: ??? +output_dir: ./generated + +# Input audio padding (set at most one; null = disabled). This appends silence +# after the real input so the agent is more likely to finish speaking before the stream ends. +pad_audio_to_sec: null # Pad each audio to this fixed duration (seconds) +pad_silence_ratio: null # Append silence = ratio * original duration (e.g. 0.2 = 20%) +pad_audio_by_sec: null # Append this fixed number of extra seconds of silence + +pipeline_type: s2s_streaming + +# S2S model block +s2s: + model_path: ??? + decode_audio: true # Whether to conduct the TTS portion of inference, or just the STT portion + speaker_name: null # Required when decode_audio is true. Must match a speaker registered in the checkpoint. + llm_engine_type: native # 'native' or 'vllm_omni' + tts_engine_type: native # 'native' or 'vllm_omni' + + # vllm_omni engine configuration (ignored unless a component is vllm_omni). + # + # On the first run a wrapper directory is built lazily under + # ``$TMPDIR/_vllm_omni_wrapper`` containing + # * ``config.json`` (just ``{"model_type": "nemotron_voicechat"}``) + # * ``nemotron/`` converted NemotronDuplexH checkpoint + # * ``eartts/`` converted EarTTS checkpoint + speaker_latents/ + # It is a full second copy of the weights. Subsequent runs reuse it, so point + # wrapper_dir somewhere persistent and large enough to avoid re-converting. + # + # The wrapper owns one one-stage AsyncOmni engine for each component that + # selects vllm_omni. ``stage_configs_path`` / ``stage_overrides`` configure + # Nemotron; the ``eartts_*`` counterparts configure EarTTS. A ``stage_1`` + # override under ``stage_overrides`` also maps to EarTTS stage 0. + vllm_omni_config: + wrapper_dir: null # null -> $TMPDIR/_vllm_omni_wrapper + stage_configs_path: null # null -> bundled nemotron_voicechat.yaml inside NeMo + eartts_stage_configs_path: null # null -> bundled single-stage eartts.yaml + nemotron_dtype: float32 # dtype for the converted Nemotron LLM checkpoint + eartts_precompute_batch_size: 256 # batch size for baking out the subword lookup table + guidance_enabled: null # null -> converted eartts/config.json enable_guidance + guidance_scale: null # null -> converted eartts/config.json guidance_scale + log_stats: false + stage_init_timeout: 600 # seconds to wait for both stage children to come up + step_timeout: 60.0 # per-step timeout on the synchronous side (seconds) + speaker_name: null # overrides s2s.speaker_name when populated + stage_overrides: null # Nemotron stage 0 overrides + eartts_stage_overrides: null # EarTTS stage 0 overrides + + device: cuda + # ======================== + # Device Configuration + # ======================== + device_id: 0 # GPU device ID + compute_dtype: bfloat16 # Compute precision: 'bfloat16' for Ampere+, + # 'float16' for older GPUs + # 'float32' + # ======================== + # Inference settings + # ======================== + use_perception_cache: true # Enable cache-aware streaming for perception encoder + use_perception_cudagraph: true # Enable CUDA graph-accelerated perception encoder + use_llm_cache: false # Keep the LLM KV cache across steps (native engine only). + # False replays the whole history each step: O(n^2) and slow, + # and is the portable compatibility default. + # NemotronH cache mode requires transformers >= 5.13. + + # TTS speedup flags (default to false; enable to speed up native inference) + use_tts_torch_compile: false # Compile TTS backbone with torch.compile (mode='default') + use_tts_subword_cache: false # Cache CharAwareSubwordEncoder embeddings (skip backbone for repeated tokens) + + # Sampling parameters shared by native and vLLM-Omni. Both temperature=0.0 + # and the all-ones configuration select greedy decoding. + top_p: 0.5 + repetition_penalty: 1.1 + temperature: 0.3 + force_turn_taking: true # Automatically disabled when the checkpoint has no ASR head. + force_turn_taking_threshold: 40 + force_turn_taking_pad_window: 25 + + # ------------------------------------------------------------------ + # Model config overrides + # + # The keys below are read by the model itself rather than by the inference + # wrapper, so they are written into the relevant model config at load time. + # Leaving one null keeps whatever the checkpoint carries. See + # nemo/collections/speechlm2/inference/model_wrappers/config_overrides.py + # for the authoritative list and which backends honour each key; anything a + # selected backend ignores is reported at load time instead of silently + # doing nothing. + # ------------------------------------------------------------------ + + # Inference logit boosts (applied to model logits at inference time) + # User (ASR) side + inference_user_pad_boost: 0.8 # Boost ASR pad logit + inference_user_bos_boost: null # Boost ASR BOS logit + inference_user_eos_boost: null # Boost ASR EOS logit + # Agent (text) side + inference_pad_boost: null # Boost agent text pad logit + inference_bos_boost: null # Boost agent text BOS logit + inference_eos_boost: null # Boost agent text EOS logit + + # EarTTS behaviour. Applied by EarTTS internally, so these affect the TTS + # component selected by tts_engine_type. + inference_force_speech_silence_on_eos: null # null -> checkpoint value (model default: true). + # Substitutes codec silence as the acoustic input of + # the step whose text token is EOS. vllm_omni always + # does this and cannot honour false. + inference_top_p_or_k: null # null -> checkpoint value (model default: 0.8). native TTS only + inference_noise_scale: null # null -> checkpoint value (model default: 0.8). native TTS only + inference_guidance_scale: null # null -> checkpoint value (model default: 0.5). native TTS only; + # the vllm_omni TTS reads vllm_omni_config.guidance_scale, which + # itself falls back to the converted eartts/config.json value. + + system_prompt: ??? + + # ======================== + # Profiling + # ======================== + profile_timing: false # Log per-stage wall-clock times for each inference step + + # ======================== + # Precision & determinism + # ======================== + # These defaults match what was used during model training. + matmul_precision: medium # Matrix multiplication precision: highest, high, medium + allow_tf32: true # Allow TF32 for cuDNN and CUDA matmul (Ampere+ GPUs). + # Set to false for stricter float32 precision. + # Deterministic inference (native engine only). Ensures identical results across + # runs by disabling FlashAttention and forcing deterministic CUDA algorithms. + # Trade-offs: slower inference, might produce worse results than non-deterministic mode, + # since non-deterministic mode was used in training. + deterministic: false + +streaming: + # File inputs are converted to mono and resampled to this rate before chunking. + input_sample_rate: 16000 + # Generated agent audio is saved at this rate. + output_sample_rate: 22050 + batch_size: 1 # Must be 1. Single-stream inference only. + att_context_size: [70,0] # Attention context size: [70,13],[70,6],[70,2],[70,0] + chunk_size_in_secs: ??? # Needs to be multiple of 80ms + buffer_size_in_secs: ??? # Audio buffer size in seconds (larger = more context, better quality) + request_type: frame # Type of request: frame, only frame is supported for cache-aware streaming + max_len: 8192 # Decode-state buffer in 80 ms frames diff --git a/examples/speechlm2/nemo_inference_pipelines/s2s_streaming_infer.py b/examples/speechlm2/nemo_inference_pipelines/s2s_streaming_infer.py new file mode 100644 index 000000000000..0deac21b188b --- /dev/null +++ b/examples/speechlm2/nemo_inference_pipelines/s2s_streaming_infer.py @@ -0,0 +1,122 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +S2S file/manifest streaming inference driver. + +This script loads complete audio files, directories of ``.wav`` files, or +line-delimited JSON/JSONL manifests and streams each input through the +SpeechLM2 ``StreamingS2SPipeline`` chunk by chunk. Live audio integrations +should call ``pipeline.generate_step()`` directly with ``Frame`` objects. + +Usage: + python s2s_streaming_infer.py \ + audio_file=/path/to/audio_or_directory \ + s2s.model_path=nvidia/NVIDIA-NemotronLabs-VoiceChat-11B \ + s2s.speaker_name=Aria \ + streaming.chunk_size_in_secs=0.08 \ + streaming.buffer_size_in_secs=5.6 +""" + +import hydra +from omegaconf import DictConfig + +from nemo.collections.asr.metrics.wer import word_error_rate +from nemo.collections.speechlm2.inference.factory.s2s_pipeline_builder import S2SPipelineBuilder +from nemo.collections.speechlm2.inference.model_wrappers.engine_selection import inference_precision_from_cfg +from nemo.collections.speechlm2.inference.utils.audio_data import ( + calculate_durations_incl_padding, + dump_output_json, + prepare_audio_data, +) +from nemo.collections.speechlm2.inference.utils.stepprogressbar import StepProgressBar +from nemo.collections.speechlm2.parts.text_utils import clean_pred_text +from nemo.utils import logging +from nemo.utils.timers import SimpleTimer + + +@hydra.main(config_path="./conf", config_name="s2s_streaming", version_base=None) +def main(cfg: DictConfig): + default_system_prompt = cfg.get("s2s", {}).get("system_prompt", None) + audio_filepaths, options, ground_truths = prepare_audio_data( + cfg.audio_file, + default_system_prompt=default_system_prompt, + sort_by_duration=False, + ) + logging.info(f"Found {len(audio_filepaths)} audio files to generate") + + # Precision and determinism are torch globals: in effect before any weights + # load, and for the whole run. shutdown() releases any vLLM runtime. + with inference_precision_from_cfg(cfg.s2s): + pipeline = S2SPipelineBuilder.build_pipeline(cfg) + try: + progress_bar = StepProgressBar.from_audio_filepaths( + audio_filepaths, + chunk_size_in_secs=pipeline.chunk_size_in_secs, + pad_audio_to_sec=cfg.get("pad_audio_to_sec"), + pad_silence_ratio=cfg.get("pad_silence_ratio"), + pad_audio_by_sec=cfg.get("pad_audio_by_sec"), + ) + + # Outside the timer: one-time first-pass costs would otherwise dominate + # RTFX on short inputs. + pipeline.warmup() + + timer = SimpleTimer() + timer.start() + outputs = pipeline.run(audio_filepaths, options=options, progress_bar=progress_bar) + timer.stop() + + # Read off the pipeline before it goes out of scope; the reporting below + # only needs these plain values. + special_tokens = pipeline.special_token_strings + finally: + pipeline.shutdown() + + exec_dur = timer.total_sec() + logging.info(f"Generated {len(audio_filepaths)} files in {exec_dur:.2f}s") + + data_dur = sum( + calculate_durations_incl_padding( + audio_filepaths, + cfg.get("pad_audio_to_sec"), + cfg.get("pad_silence_ratio"), + cfg.get("pad_audio_by_sec"), + ) + ) + rtfx = data_dur / exec_dur if exec_dur > 0 else float('inf') + logging.info(f"RTFX: {rtfx:.2f} ({data_dur:.2f}s / {exec_dur:.2f}s)") + + # Compute WER when ground-truth texts are available (micro-average, + # matching the offline eval in speechlm2.parts.metrics.asr_cer_wer) + all_refs, all_hyps = [], [] + for gt, out in zip(ground_truths, outputs): + asr_text = out.asr_text_with_timestamps + if gt and asr_text: + cleaned_gt = clean_pred_text(gt, special_token_strings=special_tokens) + cleaned_pred = clean_pred_text(asr_text, special_token_strings=special_tokens) + if cleaned_gt.strip() and cleaned_pred.strip(): + all_refs.append(cleaned_gt) + all_hyps.append(cleaned_pred) + if all_refs: + wer = word_error_rate(hypotheses=all_hyps, references=all_refs) + logging.info(f"WER: {wer:.4f} ({wer * 100:.2f}%), n={len(all_refs)}") + + output_dir = cfg.get("output_dir", "./generated") + dump_output_json(audio_filepaths, outputs, output_dir, options, ground_truths) + logging.info(f"Transcriptions written to {output_dir}/output_processed.json and {output_dir}/output_raw.json") + + +if __name__ == "__main__": + main() diff --git a/nemo/collections/speechlm2/data/duplex_ear_tts_dataset.py b/nemo/collections/speechlm2/data/duplex_ear_tts_dataset.py index 88cf1d76139d..86f53f1b623d 100644 --- a/nemo/collections/speechlm2/data/duplex_ear_tts_dataset.py +++ b/nemo/collections/speechlm2/data/duplex_ear_tts_dataset.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. import random -import re from copy import deepcopy import torch @@ -26,6 +25,7 @@ from nemo.collections.common.tokenizers import TokenizerSpec from nemo.collections.speechlm2.data.utils import get_pad_id from nemo.collections.speechlm2.parts.precision import fp32_precision +from nemo.collections.speechlm2.parts.text_utils import strip_timestamps from nemo.collections.tts.parts.utils.helpers import get_mask_from_lengths from nemo.utils import logging @@ -160,7 +160,7 @@ def __init__( assert tokenizer.eos is not None, "EOS support in the tokenizer is required for S2S models." def __getitem__(self, cuts: CutSet) -> dict: - cuts = cuts.transform_text(_strip_timestamps) + cuts = cuts.transform_text(strip_timestamps) # ensures fp32 audio load to avoid issues of duration mistakes on fp16 training with fp32_precision(): source_audio, source_audio_lens = collate_audio(cuts.resample(self.source_sample_rate)) @@ -935,20 +935,6 @@ def build_token_channel( return tokens -def _strip_timestamps( - text: str, _TIMESTAMP_PATTERN=re.compile(r"<\|\d+\|>"), _SPACE_PATTERN=re.compile(r"\s+") -) -> str: - """ - Strips timestamp tokens from text, e.g. turns: - '<|0|> Hey <|3|> <|3|> how <|5|> <|7|> are <|8|> <|8|> <|10|> you? <|12|>' - into: - 'Hey how are you?' - """ - # Regexp pattern args are cached compiled patterns (micro-optimization). - text = _TIMESTAMP_PATTERN.sub("", text) # strip timestamp tokens if present - return _SPACE_PATTERN.sub(" ", text).strip() # strip multi-whitespaces - - def sample_audio_segments_repeat( prompt_audio: torch.Tensor, prompt_audio_lens: torch.Tensor, diff --git a/nemo/collections/speechlm2/data/duplex_stt_dataset.py b/nemo/collections/speechlm2/data/duplex_stt_dataset.py index c8044214e8fd..99612f316afb 100644 --- a/nemo/collections/speechlm2/data/duplex_stt_dataset.py +++ b/nemo/collections/speechlm2/data/duplex_stt_dataset.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. import random -import re import torch import torch.utils.data @@ -24,9 +23,9 @@ from nemo.collections.common.data.lhotse.text_adapters import Formattable from nemo.collections.common.tokenizers import TokenizerSpec from nemo.collections.speechlm2.data.force_align import ForceAligner -from nemo.collections.speechlm2.data.s2s_dataset import _strip_timestamps from nemo.collections.speechlm2.data.utils import get_pad_id from nemo.collections.speechlm2.parts.augmentation import AudioAugmenter +from nemo.collections.speechlm2.parts.text_utils import SECONDS_PER_FRAME, TRAINING_TIMESTAMP_RE, strip_timestamps from nemo.utils import logging MCQ_VAL_PROMPT = "Answer the following multiple choice question with an explanation for the answer." @@ -327,7 +326,7 @@ def __getitem__(self, all_cuts: CutSet) -> dict: "source_tokens": source_tokens, "source_token_lens": source_token_lens, "source_texts": [ - " ".join(_strip_timestamps(s.text) for s in cut.supervisions if s.speaker in self.input_roles) + " ".join(strip_timestamps(s.text) for s in cut.supervisions if s.speaker in self.input_roles) for cut in all_cuts_combined ], "target_texts": [ @@ -623,24 +622,23 @@ def _build_token_channel( def _text_to_ids( text: str, tokenizer: TokenizerSpec, - _TIMESTAMP_PATTERN_STR=r"<\|(\d+)\|>", + _TIMESTAMP_PATTERN_RE=TRAINING_TIMESTAMP_RE, available_frames_for_text=None, word_align_position='left', remove_timestamps=False, prepend_word_space=True, ): - if not remove_timestamps and re.compile(_TIMESTAMP_PATTERN_STR).search(text): + if not remove_timestamps and _TIMESTAMP_PATTERN_RE.search(text): text_ids = _text_with_timestamps_to_ids( text, tokenizer, - _TIMESTAMP_PATTERN_STR, + _TIMESTAMP_PATTERN_RE, available_frames_for_text, word_align_position, prepend_word_space=prepend_word_space, ) else: - _TIMESTAMP_PATTERN = re.compile(_TIMESTAMP_PATTERN_STR) - text = _TIMESTAMP_PATTERN.sub("", text) + text = _TIMESTAMP_PATTERN_RE.sub("", text) text = " ".join(text.strip().split()) text_ids = tokenizer.text_to_ids(text) return text_ids @@ -649,7 +647,7 @@ def _text_to_ids( def _text_with_timestamps_to_ids( text: str, tokenizer: TokenizerSpec, - _TIMESTAMP_PATTERN_STR=r"<\|(\d+)\|>", + _TIMESTAMP_PATTERN_RE=TRAINING_TIMESTAMP_RE, available_frames_for_text=None, word_align_position='left', prepend_word_space=True, @@ -657,7 +655,7 @@ def _text_with_timestamps_to_ids( text_ids, start_times, end_times, word_lens = _extract_text_and_time_tokens( text, tokenizer, - _TIMESTAMP_PATTERN_STR, + _TIMESTAMP_PATTERN_RE, prepend_word_space=prepend_word_space, ) text_ids_with_timestamps = _expand_text_with_timestamps_and_word_lengths( @@ -666,7 +664,7 @@ def _text_with_timestamps_to_ids( start_times, end_times, available_frames_for_text, - frame_rate=0.08, + frame_rate=SECONDS_PER_FRAME, pad_id=get_pad_id(tokenizer), word_align_position=word_align_position, ) @@ -674,12 +672,12 @@ def _text_with_timestamps_to_ids( def _extract_text_and_time_tokens( - text, tokenizer: TokenizerSpec, _TIMESTAMP_PATTERN_STR=r"<\|(\d+)\|>", prepend_word_space=True + text, tokenizer: TokenizerSpec, _TIMESTAMP_PATTERN_RE=TRAINING_TIMESTAMP_RE, prepend_word_space=True ): - time_tokens = re.findall(_TIMESTAMP_PATTERN_STR, text) + time_tokens = _TIMESTAMP_PATTERN_RE.findall(text) start_time = [int(time_tokens[i]) for i in range(0, len(time_tokens), 2)] end_time = [int(time_tokens[i]) for i in range(1, len(time_tokens), 2)] - words = re.sub(_TIMESTAMP_PATTERN_STR, '', text).split() + words = _TIMESTAMP_PATTERN_RE.sub('', text).split() text_ids = [] word_lens = [] for i, word in enumerate(words): diff --git a/nemo/collections/speechlm2/data/force_align.py b/nemo/collections/speechlm2/data/force_align.py index 5930bc641cd6..467bac1d72be 100644 --- a/nemo/collections/speechlm2/data/force_align.py +++ b/nemo/collections/speechlm2/data/force_align.py @@ -28,6 +28,7 @@ get_batch_variables, viterbi_decoding, ) +from nemo.collections.speechlm2.parts.text_utils import strip_timestamps class ForceAligner: @@ -150,7 +151,7 @@ def batch_force_align_user_audio(self, cuts: CutSet, source_sample_rate: int = 1 for i, (supervision, cut) in enumerate(zip(user_supervisions, user_cuts)): try: - text = self._strip_timestamps(supervision.text) + text = strip_timestamps(supervision.text) normalized_text = self._normalize_transcript(text) if not normalized_text.strip(): logging.warning(f"Text became empty after normalization: {supervision.text}") @@ -326,18 +327,3 @@ def _convert_alignment_to_timestamped_text( timestamped_words.append(f"<|{start_frame}|> {word} <|{end_frame}|>") return " ".join(timestamped_words) - - def _strip_timestamps(self, text: str) -> str: - """ - Strip timestamp tokens from text. - - Args: - text: Text that may contain timestamp tokens - - Returns: - Text with timestamp tokens removed - """ - text = re.sub(r'<\|[0-9]+\|>', '', text) - text = re.sub(r' +', ' ', text) - - return text.strip() diff --git a/nemo/collections/speechlm2/data/s2s_dataset.py b/nemo/collections/speechlm2/data/s2s_dataset.py index cb89affa0fd8..3835f0934df4 100644 --- a/nemo/collections/speechlm2/data/s2s_dataset.py +++ b/nemo/collections/speechlm2/data/s2s_dataset.py @@ -11,8 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -import re - import torch import torch.utils.data from lhotse import CutSet, Seconds, compute_num_frames @@ -22,6 +20,7 @@ from nemo.collections.common.tokenizers import TokenizerSpec from nemo.collections.speechlm2.data.utils import get_pad_id +from nemo.collections.speechlm2.parts.text_utils import strip_timestamps from nemo.utils import logging @@ -100,7 +99,7 @@ def __init__( assert tokenizer.eos is not None, "EOS support in the tokenizer is required for S2S models." def __getitem__(self, cuts: CutSet) -> dict: - cuts = cuts.transform_text(_strip_timestamps) + cuts = cuts.transform_text(strip_timestamps) source_audio, source_audio_lens = collate_audio(cuts.resample(self.source_sample_rate)) target_audio, target_audio_lens = collate_audio( cuts.resample(self.target_sample_rate, recording_field="target_audio"), recording_field="target_audio" @@ -188,17 +187,3 @@ def build_token_channel( tokens[eospos] = tokenizer.eos return tokens - - -def _strip_timestamps( - text: str, _TIMESTAMP_PATTERN=re.compile(r"<\|\d+\|>"), _SPACE_PATTERN=re.compile(r"\s+") -) -> str: - """ - Strips timestamp tokens from text, e.g. turns: - '<|0|> Hey <|3|> <|3|> how <|5|> <|7|> are <|8|> <|8|> <|10|> you? <|12|>' - into: - 'Hey how are you?' - """ - # Regexp pattern args are cached compiled patterns (micro-optimization). - text = _TIMESTAMP_PATTERN.sub("", text) # strip timestamp tokens if present - return _SPACE_PATTERN.sub(" ", text).strip() # strip multi-whitespaces diff --git a/nemo/collections/speechlm2/inference/__init__.py b/nemo/collections/speechlm2/inference/__init__.py new file mode 100644 index 000000000000..fdc3df7ff497 --- /dev/null +++ b/nemo/collections/speechlm2/inference/__init__.py @@ -0,0 +1,34 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from nemo.collections.speechlm2.inference.factory.s2s_pipeline_builder import S2SPipelineBuilder +from nemo.collections.speechlm2.inference.model_wrappers.capabilities import AuxiliaryOutputCapabilities +from nemo.collections.speechlm2.inference.model_wrappers.decode_state import InferenceStepResult, StreamingDecodeState +from nemo.collections.speechlm2.inference.pipelines.streaming_s2s_pipeline import ( + GenerateStepOutput, + StreamingS2SPipeline, +) +from nemo.collections.speechlm2.inference.streaming.framing.s2s_request_options import S2SRequestOptions +from nemo.collections.speechlm2.inference.streaming.state.s2s_streaming_output import S2SStreamingOutput + +__all__ = [ + 'S2SPipelineBuilder', + 'AuxiliaryOutputCapabilities', + 'InferenceStepResult', + 'StreamingDecodeState', + 'GenerateStepOutput', + 'StreamingS2SPipeline', + 'S2SRequestOptions', + 'S2SStreamingOutput', +] diff --git a/nemo/collections/speechlm2/inference/factory/__init__.py b/nemo/collections/speechlm2/inference/factory/__init__.py new file mode 100644 index 000000000000..9e3fb699d9f6 --- /dev/null +++ b/nemo/collections/speechlm2/inference/factory/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/nemo/collections/speechlm2/inference/factory/s2s_pipeline_builder.py b/nemo/collections/speechlm2/inference/factory/s2s_pipeline_builder.py new file mode 100644 index 000000000000..cc473e5b3a4e --- /dev/null +++ b/nemo/collections/speechlm2/inference/factory/s2s_pipeline_builder.py @@ -0,0 +1,74 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from omegaconf.dictconfig import DictConfig + +from nemo.collections.speechlm2.inference.model_wrappers.nemotron_voicechat_inference_wrapper import ( + NemotronVoicechatInferenceWrapper, +) +from nemo.collections.speechlm2.inference.model_wrappers.engine_selection import precision_matches_cfg +from nemo.collections.speechlm2.inference.pipelines.streaming_s2s_pipeline import StreamingS2SPipeline +from nemo.utils import logging as logger + + +class S2SPipelineBuilder: + """Factory that builds a streaming S2S pipeline.""" + + @classmethod + def build_pipeline(cls, cfg: DictConfig) -> StreamingS2SPipeline: + """ + Build the streaming S2S pipeline based on the config. + + Precision and determinism are torch process globals: they must be in + effect before any weights load and stay in effect for the whole run. + That cannot be owned by an object this call returns, so the caller + scopes it and this call *requires* it, rather than warning afterwards:: + + with inference_precision_from_cfg(cfg.s2s): + pipeline = S2SPipelineBuilder.build_pipeline(cfg) + try: + pipeline.run(audio_filepaths) + finally: + pipeline.shutdown() + + :meth:`StreamingS2SPipeline.shutdown` releases any vLLM-Omni runtime; + it is a no-op for native engines. + + Args: + cfg: (DictConfig) Config + Returns: + Returns StreamingS2SPipeline object + """ + if not precision_matches_cfg(cfg.s2s): + raise RuntimeError( + "This process is not configured the way cfg.s2s asks: allow_tf32, " + "matmul_precision and deterministic are torch process globals that have to be " + "applied before the checkpoint loads, and they are not in effect. Wrap the call:\n" + " from nemo.collections.speechlm2.inference.model_wrappers.engine_selection " + "import inference_precision_from_cfg\n" + " with inference_precision_from_cfg(cfg.s2s):\n" + " pipeline = S2SPipelineBuilder.build_pipeline(cfg)\n" + " try:\n" + " ...\n" + " finally:\n" + " pipeline.shutdown()" + ) + + s2s_model = NemotronVoicechatInferenceWrapper(model_cfg=cfg.s2s) + + logger.info(f"S2S model `{cfg.s2s.model_path}` loaded") + + pipeline = StreamingS2SPipeline(cfg, s2s_model) + logger.info(f"`{type(pipeline).__name__}` pipeline loaded") + return pipeline diff --git a/nemo/collections/speechlm2/inference/model_wrappers/__init__.py b/nemo/collections/speechlm2/inference/model_wrappers/__init__.py new file mode 100644 index 000000000000..9e3fb699d9f6 --- /dev/null +++ b/nemo/collections/speechlm2/inference/model_wrappers/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/nemo/collections/speechlm2/inference/model_wrappers/backend/__init__.py b/nemo/collections/speechlm2/inference/model_wrappers/backend/__init__.py new file mode 100644 index 000000000000..280a18858cea --- /dev/null +++ b/nemo/collections/speechlm2/inference/model_wrappers/backend/__init__.py @@ -0,0 +1,37 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-component inference backends. + +``DuplexLLM`` and ``DuplexTTS`` are the two contracts the VoiceChat wrapper +runs its frame loop against; ``pytorch/`` and ``vllm/`` hold one +implementation of each. The vLLM implementations are not re-exported here so +that importing this package stays free of vLLM. +""" + +from nemo.collections.speechlm2.inference.model_wrappers.backend.eartts import DuplexTTS +from nemo.collections.speechlm2.inference.model_wrappers.backend.llm import DuplexLLM +from nemo.collections.speechlm2.inference.model_wrappers.backend.pytorch.eartts import ( + PyTorchEarTTS, + TTSGenerationResult, +) +from nemo.collections.speechlm2.inference.model_wrappers.backend.pytorch.llm import PyTorchLLM + +__all__ = [ + 'DuplexLLM', + 'DuplexTTS', + 'PyTorchEarTTS', + 'TTSGenerationResult', + 'PyTorchLLM', +] diff --git a/nemo/collections/speechlm2/inference/model_wrappers/backend/eartts.py b/nemo/collections/speechlm2/inference/model_wrappers/backend/eartts.py new file mode 100644 index 000000000000..6ea1f1604006 --- /dev/null +++ b/nemo/collections/speechlm2/inference/model_wrappers/backend/eartts.py @@ -0,0 +1,57 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-frame contract for the VoiceChat TTS (EarTTS) component. + +Implemented by :class:`~nemo.collections.speechlm2.inference.model_wrappers.backend.pytorch.eartts.PyTorchEarTTS` +and :class:`~nemo.collections.speechlm2.inference.model_wrappers.backend.vllm.eartts.VllmEarTTS`. + +Separate from :class:`~nemo.collections.speechlm2.inference.model_wrappers.backend.llm.DuplexLLM` +because the two components are chosen independently, and because an LLM-shaped +call signature never fitted EarTTS: it consumes a text token and emits acoustic +codes. + +Note what is *not* here. ``inference_force_speech_silence_on_eos`` is applied by +both EarTTS implementations internally, on the acoustic input of the step whose +text token is EOS, so it needs no place in this contract. The audio codec runs +natively for every backend and stays in the wrapper. +""" + +from abc import ABC, abstractmethod +from typing import Any + +import torch + + +class DuplexTTS(ABC): + """Turns one committed text token into one frame of acoustic codes.""" + + @abstractmethod + def step( + self, + state: Any, + current_frame_idx: int, + request_id: str, + ) -> torch.Tensor: + """Generate this frame's acoustic codes. + + Reads the committed text token from ``state.gen_text`` rather than + taking it as an argument, so a caller that rewrote it (forced + turn-taking) does not have to remember to pass the new value. + + Returns: + Codes shaped ``(B, T, num_quantizers)``, which is what the native + audio codec consumes. + """ + raise NotImplementedError diff --git a/nemo/collections/speechlm2/inference/model_wrappers/backend/llm.py b/nemo/collections/speechlm2/inference/model_wrappers/backend/llm.py new file mode 100644 index 000000000000..18c75f16babd --- /dev/null +++ b/nemo/collections/speechlm2/inference/model_wrappers/backend/llm.py @@ -0,0 +1,93 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-frame contract for the VoiceChat LLM component. + +Implemented by :class:`~nemo.collections.speechlm2.inference.model_wrappers.backend.pytorch.llm.PyTorchLLM` +and :class:`~nemo.collections.speechlm2.inference.model_wrappers.backend.vllm.llm.VllmLLM`. +The wrapper picks one at construction and then runs the same frame loop, so +nothing above this line inspects which engine is in use. + +The contract is deliberately just :meth:`DuplexLLM.step`. Cache creation, +prompt prefill and request abort are native-only concepts -- vLLM does those +inside the engine and the session -- so they are not part of it; the wrapper +calls them on the native object at stream setup, where it knows it has one. +Adding them here as no-op defaults would advertise methods that only mean +something for one implementation. + +The same rule applies to the *parameters*, not just the method list: anything +only one backend can act on is per-stream state on ``StreamingDecodeState`` +rather than an argument, so the signature does not advertise capabilities an +implementation has to discard. +""" + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any + +import torch + + +@dataclass +class LlmStepResult: + """One frame's LLM output, in the shape both backends can fill. + + A dataclass rather than a dict so the optional channels are discoverable + and callers stop probing with ``"text_logits" in ans``. The auxiliary + tokens are ``None`` when the checkpoint has no such head; the logits are + ``None`` unless the backend has them *and* ``return_debug`` asked (vLLM + keeps its logits inside the engine, so it never fills them). + """ + + predicted_token: torch.Tensor + asr_predicted_token: torch.Tensor | None = None + function_predicted_token: torch.Tensor | None = None + text_logits: torch.Tensor | None = None + asr_logits: torch.Tensor | None = None + function_logits: torch.Tensor | None = None + + +class DuplexLLM(ABC): + """Produces one frame's text (and optional ASR/function) tokens.""" + + @abstractmethod + def step( + self, + frame_embedding: torch.Tensor, + state: Any, + *, + frame_offset: int, + current_frame_idx: int, + has_prompt: bool, + return_debug: bool = False, + sampling_params: dict[str, float] | None = None, + debug_logger: Any = None, + ) -> LlmStepResult: + """Advance one 80 ms frame. + + Args: + frame_embedding: Encoded audio for this frame, shape ``(B, 1, H)``. + state: The stream's ``StreamingDecodeState``. Implementations read + committed history (``gen_text`` and friends) from it and update + their own per-stream fields -- caches, + ``input_embeds_history`` -- in place. + frame_offset: Index of this frame within the current chunk. + current_frame_idx: Index of this frame within the whole stream. + has_prompt: Whether a system prompt is already in the LLM state. + return_debug: Ask for logits in the result when the backend has them. + sampling_params: Per-stream sampling overrides, for backends that + can still apply them at this point. + debug_logger: Receives the per-frame LLM input. + """ + raise NotImplementedError diff --git a/nemo/collections/speechlm2/inference/model_wrappers/backend/pytorch/__init__.py b/nemo/collections/speechlm2/inference/model_wrappers/backend/pytorch/__init__.py new file mode 100644 index 000000000000..9e3fb699d9f6 --- /dev/null +++ b/nemo/collections/speechlm2/inference/model_wrappers/backend/pytorch/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/nemo/collections/speechlm2/inference/model_wrappers/backend/pytorch/eartts.py b/nemo/collections/speechlm2/inference/model_wrappers/backend/pytorch/eartts.py new file mode 100644 index 000000000000..29fc33602b7b --- /dev/null +++ b/nemo/collections/speechlm2/inference/model_wrappers/backend/pytorch/eartts.py @@ -0,0 +1,178 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Native PyTorch backend for the TTS (EarTTS) component of NemotronVoiceChat. + +Wraps the DuplexEARTTS model for direct PyTorch inference, implementing +:class:`~nemo.collections.speechlm2.inference.model_wrappers.backend.eartts.DuplexTTS`; +the vLLM sibling lives in ``backend/vllm/eartts.py``. +""" + +from dataclasses import dataclass, fields +from typing import Any + +import torch + +from nemo.collections.speechlm2.inference.model_wrappers.backend.eartts import DuplexTTS +from nemo.utils import logging + + +@dataclass +class TTSGenerationResult: + """Result from a single TTS generation step (shared by PyTorch and vLLM backends).""" + + codes: torch.Tensor # Generated acoustic tokens + past_key_values: Any # Updated cache (if applicable) + + def __getitem__(self, item: str | int): + """Allows for accessing attributes by key or index.""" + if isinstance(item, str): + return getattr(self, item) + else: + return getattr(self, fields(self)[item].name) + + +class PyTorchEarTTS(DuplexTTS): + """ + Native PyTorch backend for the TTS (EarTTS) component. + + Wraps ``DuplexEARTTS.infer_codes_one_step()`` for per-frame audio code + generation. + """ + + def __init__(self, tts_model, first_context_subword_id=None, generation_config=None): + """ + Args: + tts_model: A ``DuplexEARTTS`` instance (``NemotronVoiceChat.tts_model``). + first_context_subword_id: Context token for frame 0, taken from the + TTS warmup. Set via :meth:`set_warmup_state`. + generation_config: EarTTS generation config from the same warmup. + """ + self.tts_model = tts_model + self.first_context_subword_id = first_context_subword_id + self.generation_config = generation_config + + def set_warmup_state(self, first_context_subword_id, generation_config) -> None: + """Attach the values produced by the wrapper's TTS warmup.""" + self.first_context_subword_id = first_context_subword_id + self.generation_config = generation_config + + def step(self, state: Any, current_frame_idx: int, request_id: str) -> torch.Tensor: + """One native EarTTS step -- see ``DuplexTTS.step``. + + Mutates ``state.tts_code`` and ``state.tts_past_key_values`` in place + and returns a clone of this frame's codes for the caller to decode. + """ + if self.generation_config is None or self.first_context_subword_id is None: + raise RuntimeError("PyTorchEarTTS is missing its warmup state; call set_warmup_state first") + + current_subword_id = state.gen_text[:, current_frame_idx].unsqueeze(-1) + if current_frame_idx == 0: + prev_subword_id = self.first_context_subword_id + else: + prev_subword_id = state.gen_text[:, current_frame_idx - 1].unsqueeze(-1) + + result = self( + { + "current_subword_id": current_subword_id, + "prev_subword_id": prev_subword_id, + "current_subword_mask": state.subword_mask[:, current_frame_idx].unsqueeze(-1), + "prev_audio_tokens": state.tts_code, + "past_key_values": state.tts_past_key_values, + "guidance_enabled": True, + "generation_config": self.generation_config, + "ignore_eos_flag_stop": True, + }, + request_id=request_id, + ) + state.tts_code = result.codes + state.tts_past_key_values = result.past_key_values + return state.tts_code.clone() + + def abort_request(self, request_id: str) -> bool: + """No-op: native PyTorch has no in-flight request to cancel.""" + del request_id + return False + + def __call__(self, inputs: dict, **kwargs) -> TTSGenerationResult: + """ + Run one TTS code-generation step via ``infer_codes_one_step``. + + Args: + inputs: Keyword arguments for ``DuplexEARTTS.infer_codes_one_step`` + (current_subword_id, prev_subword_id, current_subword_mask, + prev_audio_tokens, past_key_values, guidance_enabled, etc.) + + Returns: + TTSGenerationResult with generated codes and updated cache. + """ + codes, cache = self.tts_model.infer_codes_one_step(**inputs) + return TTSGenerationResult(codes=codes, past_key_values=cache) + + def prefill_prompt(self, init_inputs, prompt_token_ids=None, request_id=None, **kwargs): + """Prefill TTS with speaker embedding / warmup inputs. + + Calls ``DuplexEARTTS.tts_model``, the inner EarTTS backbone, directly. + + Args: + init_inputs: Dict of initial TTS inputs from ``get_init_inputs()``. + prompt_token_ids: Unused for native (vLLM-only parameter). + request_id: Unused for native (vLLM-only parameter). + + Returns: + Model outputs (with ``past_key_values`` and ``codes``). + """ + return self.tts_model.tts_model(**init_inputs) + + def compile(self, **kwargs) -> None: + """Apply torch.compile to the TTS backbone if available.""" + tts_backbone = getattr(self.tts_model, 'tts_model', None) + if tts_backbone is not None and hasattr(tts_backbone, 'backbone'): + mode = kwargs.get('mode', 'default') + logging.info(f"Compiling TTS backbone with torch.compile(mode='{mode}')...") + tts_backbone.backbone = torch.compile(tts_backbone.backbone, mode=mode) + logging.info(" TTS backbone compiled") + + def setup_subword_cache(self, cfg) -> None: + """Enable TTS subword embedding cache from config flags.""" + from omegaconf import OmegaConf + + tts_inner = getattr(self.tts_model, 'tts_model', None) + if tts_inner is None or not hasattr(tts_inner, 'config'): + return + if bool(cfg.get("use_tts_subword_cache", False)): + OmegaConf.update(tts_inner.config, "use_tts_subword_cache", True) + logging.info("TTS speedup enabled: use_tts_subword_cache") + embed_subword = getattr(tts_inner, 'embed_subword', None) + if embed_subword is not None and hasattr(embed_subword, 'use_tts_subword_cache'): + embed_subword.use_tts_subword_cache = True + + def to(self, device_or_dtype: torch.device | torch.dtype) -> 'PyTorchEarTTS': + """Move underlying TTS model to device or convert dtype.""" + self.tts_model = self.tts_model.to(device_or_dtype) + return self + + def eval(self) -> 'PyTorchEarTTS': + """Set underlying TTS model to eval mode.""" + self.tts_model.eval() + return self + + @property + def device(self) -> torch.device: + """Get device of the underlying TTS model.""" + try: + return next(self.tts_model.parameters()).device + except StopIteration: + return torch.device('cpu') diff --git a/nemo/collections/speechlm2/inference/model_wrappers/backend/pytorch/llm.py b/nemo/collections/speechlm2/inference/model_wrappers/backend/pytorch/llm.py new file mode 100644 index 000000000000..32747b80747f --- /dev/null +++ b/nemo/collections/speechlm2/inference/model_wrappers/backend/pytorch/llm.py @@ -0,0 +1,391 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Native PyTorch backend for the LLM component of NemotronVoiceChat. + +Wraps the DuplexSTT model (which contains the Nemotron LLM backbone) for +direct PyTorch inference with top-p sampling and repetition penalty support. + +Implements :class:`~nemo.collections.speechlm2.inference.model_wrappers.backend.llm.DuplexLLM`; +the vLLM sibling lives in ``backend/vllm/llm.py``. +""" + +import inspect +import math +from typing import Any + +import torch +from packaging.version import Version + +from nemo.collections.speechlm2.inference.model_wrappers.backend.llm import DuplexLLM, LlmStepResult +from nemo.collections.speechlm2.inference.model_wrappers.text_sampling import sample_text_token +from nemo.collections.speechlm2.parts.text_utils import get_special_token_ids +from nemo.utils import logging + + +def _to_step_result(ans: dict[str, Any]) -> LlmStepResult: + """Project the raw forward-pass dict onto the backend contract. + + ``__call__`` also carries ``cache``, which is per-stream state rather than + a step output, so it is applied to the decode state and dropped here. + """ + return LlmStepResult( + predicted_token=ans["predicted_token"], + asr_predicted_token=ans.get("asr_predicted_token"), + function_predicted_token=ans.get("function_predicted_token"), + text_logits=ans.get("text_logits"), + asr_logits=ans.get("asr_logits"), + function_logits=ans.get("function_logits"), + ) + + +class PyTorchLLM(DuplexLLM): + """ + Native PyTorch backend for the LLM (DuplexSTT) component. + + Wraps the DuplexSTT model's forward pass (``stt_model()``) to produce + text/ASR token predictions. Supports top-p sampling and repetition penalty + through the same :func:`sample_text_token` the vLLM backend reaches via its + logits processor, so the two share one sampling policy. + """ + + def __init__( + self, + model, + special_token_ids: set[int] | None = None, + top_p: float = 1.0, + repetition_penalty: float = 1.0, + temperature: float = 1.0, + use_llm_cache: bool = False, + ): + """ + Initialize with an existing model. + + Args: + model: A :class:`~nemo.collections.speechlm2.models.nemotron_voicechat.NemotronVoiceChat` + (or compatible) model whose ``stt_model`` sub-module drives the LLM channel. + special_token_ids: Set of special token IDs (pad, eos, bos) that should bypass sampling. + These tokens will use greedy decoding and won't be penalized. + If None, auto-extracted from model via + :func:`~nemo.collections.speechlm2.parts.text_utils.get_special_token_ids`. + top_p: Top-p (nucleus) sampling threshold. 1.0 disables it (greedy). Default: 1.0 + repetition_penalty: Penalty for repeated tokens. 1.0 disables it. Default: 1.0 + Recommended value when enabling: 1.2 + temperature: Temperature for sampling. 1.0 = no change, <1.0 = sharper, >1.0 = flatter. + 0.0 = greedy (argmax). Default: 1.0 + use_llm_cache: Keep a KV cache across decode steps instead of replaying the + whole history every step. See :meth:`create_cache`. Default: False + """ + if special_token_ids is None: + try: + stt = model.stt_model + special_token_ids = get_special_token_ids(stt.tokenizer, stt.text_pad_id, model_cfg=stt.cfg) + except AttributeError: + logging.debug("Cannot extract special token IDs: model has no stt_model.tokenizer") + special_token_ids = set() + + if not math.isfinite(temperature): + raise ValueError(f"temperature must be finite, got {temperature}") + if temperature < 0.0: + raise ValueError(f"temperature must be >= 0.0, got {temperature}") + + self.special_token_ids = special_token_ids or set() + self.top_p = top_p + self.repetition_penalty = repetition_penalty + self.temperature = temperature + # Pre-built tensor for special-token filtering in the repetition + # penalty; moved to the logits device on first use. + self._special_ids_tensor = ( + torch.tensor(sorted(self.special_token_ids), dtype=torch.long) if self.special_token_ids else None + ) + + self.model = model + self.use_llm_cache = use_llm_cache + self.cache_key = self._resolve_cache_key() + + logging.debug(f"Special token IDs: {self.special_token_ids}") + + sampling_active = top_p < 1.0 or repetition_penalty != 1.0 or (temperature != 1.0 and temperature != 0.0) + if sampling_active and not self.special_token_ids: + import warnings + + warnings.warn( + "Sampling is enabled but special_token_ids is empty. " + "Could not auto-extract from model.tokenizer. " + "Please provide special_token_ids manually to ensure special tokens use greedy decoding. " + "Otherwise, EOS tokens may be randomly sampled and generation may not stop properly!" + ) + + def _sample_text_token( + self, + logits: torch.Tensor, + generated_tokens: torch.Tensor, + current_step: int, + sampling_params: dict[str, float] | None = None, + ) -> torch.Tensor: + """Sample one text token, honouring per-request overrides. + + Special tokens (pad, BOS, EOS) bypass sampling so generation still + stops reliably when top-p or a repetition penalty is enabled. + """ + params = sampling_params or {} + device = logits.device + if self._special_ids_tensor is not None and self._special_ids_tensor.device != device: + self._special_ids_tensor = self._special_ids_tensor.to(device) + return sample_text_token( + logits, + generated_tokens, + current_step, + top_p=params.get("top_p", self.top_p), + repetition_penalty=params.get("repetition_penalty", self.repetition_penalty), + temperature=params.get("temperature", self.temperature), + special_token_ids=self.special_token_ids, + special_ids_tensor=self._special_ids_tensor, + ) + + def step( + self, + frame_embedding: torch.Tensor, + state: Any, + *, + frame_offset: int, + current_frame_idx: int, + has_prompt: bool, + return_debug: bool = False, + sampling_params: dict[str, float] | None = None, + debug_logger: Any = None, + ) -> LlmStepResult: + """One native forward pass for this frame -- see ``DuplexLLM.step``. + + Builds the duplex input embedding from the committed token history, so + a caller's forced-turn-taking rewrite feeds back here exactly as it does + offline. Uses ``state.llm_cache`` when there is one, otherwise appends + to ``state.input_embeds_history`` and replays it (O(n^2), see + :meth:`create_cache`). + """ + input_emb = self.model.stt_model.build_input_embedding( + frame_embedding, + current_frame_idx, + state.gen_text, + state.gen_asr_text, + state.gen_function, + has_prompt=has_prompt, + ) + if debug_logger is not None: + debug_logger.log_input_embeds(input_emb) + + if state.llm_cache is not None: + ans = self( + input_emb, + cache=state.llm_cache, + cache_position_offset=state.llm_cache_position_offset + frame_offset, + generated_tokens=state.gen_text, + current_step=current_frame_idx, + return_logits=return_debug, + sampling_params=sampling_params, + ) + state.llm_cache = ans["cache"] + return _to_step_result(ans) + + # No cache: this backend owns the replay history, so it extends it here + # rather than handing a list back for the caller to merge. + state.input_embeds_history.append(input_emb) + full_input_embeds = torch.cat(state.input_embeds_history, dim=1) + ans = self( + full_input_embeds, + cache=None, + generated_tokens=state.gen_text, + current_step=current_frame_idx, + return_logits=return_debug, + sampling_params=sampling_params, + ) + return _to_step_result(ans) + + def abort_request(self, request_id: str) -> bool: + """No-op: native PyTorch has no in-flight request to cancel.""" + del request_id + return False + + def _resolve_cache_key(self) -> str: + """Resolve the cache argument from the installed backbone API.""" + configured = str(self.model.stt_model.cfg.get("cache_key", "past_key_values")) + try: + parameters = inspect.signature(self.model.stt_model.llm.forward).parameters + except (TypeError, ValueError): + parameters = {} + + if configured in parameters: + return configured + for candidate in ("past_key_values", "cache_params"): + if candidate in parameters: + if candidate != configured: + logging.info( + f"Using runtime LLM cache key {candidate!r} instead of checkpoint value {configured!r}" + ) + return candidate + return configured + + def create_cache(self): + """Create an LLM KV cache, or None to replay the full history each step. + + ``DynamicCache`` builds its per-layer state from ``config.layer_types``, which + covers NemotronH's hybrid mamba/attention stack as well as plain attention. + Needs transformers >= 5.13, which fixes Mamba2 chunked prefill (huggingface/ + transformers#46741) for any forward with seq_len > 1 and a warm cache. + """ + if not self.use_llm_cache: + logging.info("LLM KV cache disabled: replaying full history each step") + return None + + stt_cfg = self.model.stt_model.cfg + if "Nemotron" in str(stt_cfg.get("pretrained_llm", "")): + import transformers + + installed = Version(transformers.__version__.split("+", 1)[0]) + if installed < Version("5.13.0"): + raise RuntimeError( + "Native NemotronH KV cache requires transformers>=5.13.0 " + f"(installed: {transformers.__version__}). Set use_llm_cache=false " + "or use a compatible runtime." + ) + if self.cache_key != "past_key_values": + raise RuntimeError( + "Installed NemotronH does not expose the expected past_key_values cache API " + f"(resolved cache key: {self.cache_key!r})." + ) + from transformers import DynamicCache + + return DynamicCache(config=self.model.stt_model.llm.config) + + def __call__( + self, + input_embeds: torch.Tensor, + cache: Any | None = None, + cache_position: torch.Tensor | None = None, + cache_position_offset: int | None = None, + generated_tokens: torch.Tensor | None = None, + current_step: int = 0, + return_logits: bool = False, + sampling_params: dict[str, float] | None = None, + **kwargs, + ) -> dict[str, Any]: + """ + Perform inference using the native model. + + Args: + input_embeds: Input embeddings [batch, seq_len, hidden_dim] + cache: Optional DynamicCache for standard transformer models. + cache_position: Optional cache-position tensor for cached decoding. + If not provided, ``cache_position_offset`` is used instead. + cache_position_offset: Optional integer offset; when ``cache_position`` + is None, a single-element tensor ``[cache_position_offset]`` is + built on ``input_embeds.device``. + generated_tokens: Previously generated tokens [batch, num_generated]. + Required for repetition_penalty. If None, creates empty tensor. + current_step: Current decoding step. Used for repetition penalty. + sampling_params: Optional per-request overrides for sampling + (top_p, temperature, repetition_penalty). + **kwargs: Additional arguments passed to the model + + Returns: + Dictionary with 'predicted_token', 'asr_predicted_token', and 'cache' + """ + if cache_position is None and cache_position_offset is not None: + cache_position = torch.tensor([cache_position_offset], device=input_embeds.device) + result = self.model.stt_model( + input_embeds, + cache=cache, + cache_position=cache_position, + cache_key=self.cache_key, + **kwargs, + ) + + if not isinstance(result, dict): + raise TypeError(f"Model returned {type(result)}, expected dict") + + if 'text_logits' not in result: + raise KeyError("Model output must contain 'text_logits' key") + + text_logits = result["text_logits"][:, -1] # [batch, vocab_size] + batch_size = text_logits.shape[0] + + if generated_tokens is None: + gen_tokens = torch.empty(batch_size, 0, device=text_logits.device, dtype=torch.long) + else: + gen_tokens = generated_tokens + + predicted_token = self._sample_text_token( + logits=text_logits, + generated_tokens=gen_tokens, + current_step=current_step, + sampling_params=sampling_params, + ) + + ans = { + "predicted_token": predicted_token, + "asr_predicted_token": None, + "function_predicted_token": None, + "cache": result.get("cache", None), + } + # Auxiliary channels use greedy decoding and are independently optional. + if result.get("asr_logits") is not None: + ans["asr_predicted_token"] = result["asr_logits"][:, -1].argmax(dim=-1) + if result.get("function_logits") is not None: + ans["function_predicted_token"] = result["function_logits"][:, -1].argmax(dim=-1) + if return_logits: + ans["text_logits"] = result["text_logits"] + ans["asr_logits"] = result.get("asr_logits") + ans["function_logits"] = result.get("function_logits") + return ans + + def to(self, device_or_dtype: torch.device | torch.dtype) -> 'PyTorchLLM': + """Move underlying model to device or convert dtype.""" + self.model = self.model.to(device_or_dtype) + return self + + def eval(self) -> 'PyTorchLLM': + """Set underlying model to eval mode.""" + self.model.eval() + return self + + @property + def device(self) -> torch.device: + """Get device of the underlying model.""" + try: + return next(self.model.parameters()).device + except StopIteration: + return torch.device('cpu') + + def prefill_prompt(self, embeddings, cache=None, cache_position=None, **kwargs): + """Prefill the native LLM with prompt embeddings to warm up the KV cache. + + Args: + embeddings: Prompt embeddings [batch, seq_len, hidden_dim]. + cache: KV cache object to update in-place. + cache_position: Position tensor for the prompt tokens. + + Returns: + Dictionary with updated 'cache'. + """ + result = self.model.stt_model( + embeddings, + cache=cache, + cache_position=cache_position, + cache_key=self.cache_key, + **kwargs, + ) + if not isinstance(result, dict): + raise TypeError(f"Model returned {type(result)}, expected dict") + return {"cache": result.get("cache", cache)} diff --git a/nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/__init__.py b/nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/__init__.py new file mode 100644 index 000000000000..bc50bcde6b4a --- /dev/null +++ b/nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/__init__.py @@ -0,0 +1,39 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""vLLM-Omni implementations of the two component contracts. + +Both are thin: the engines are process-scoped and owned by ``OmniRuntime``, +while the request state lives in the per-stream ``OmniStreamingSession`` these +classes read off the decode state. Importing this package does not import +vLLM. +""" + +from typing import Any + + +def require_session(state: Any): + """Return the stream's ``OmniStreamingSession``, or say why there isn't one. + + ``omni_session`` is a declared field on ``StreamingDecodeState``, so this + reads it directly: a missing attribute is a programming error worth an + AttributeError, while ``None`` is the real case worth explaining. + """ + session = state.omni_session + if session is None: + raise RuntimeError( + "A vllm_omni component requires a per-stream OmniStreamingSession; " + "make sure begin_stream(...) ran for this stream before the first frame." + ) + return session diff --git a/nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/eartts.py b/nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/eartts.py new file mode 100644 index 000000000000..fba82d28a3b5 --- /dev/null +++ b/nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/eartts.py @@ -0,0 +1,66 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""vLLM-Omni backend for the TTS (EarTTS) component of NemotronVoiceChat. + +Implements :class:`~nemo.collections.speechlm2.inference.model_wrappers.backend.eartts.DuplexTTS` +against the per-stream ``OmniStreamingSession``; the PyTorch sibling lives in +``backend/pytorch/eartts.py``. +""" + +from typing import Any + +import torch + +from nemo.collections.speechlm2.inference.model_wrappers.backend.eartts import DuplexTTS +from nemo.collections.speechlm2.inference.model_wrappers.backend.vllm import require_session + + +class VllmEarTTS(DuplexTTS): + """Runs EarTTS in a vLLM-Omni engine, one text token per step. + + Stateless itself, like its LLM counterpart: the engine belongs to + ``OmniRuntime`` and the request belongs to the session on the decode state. + With classifier-free guidance enabled, the session's conditional and + unconditional requests are kept in lockstep by the custom scheduler, so one + submission still yields one acoustic frame. + """ + + def __init__(self, device: torch.device): + """ + Args: + device: Device the native audio codec decodes on, so the codes this + backend returns land where the codec expects them. + """ + self.device = device + + def step(self, state: Any, current_frame_idx: int, request_id: str) -> torch.Tensor: + """One EarTTS step -- see ``DuplexTTS.step``. + + ``inference_force_speech_silence_on_eos`` is not applied here: the + converted EarTTS substitutes codec silence itself when the incoming + text token is EOS, matching what DuplexEARTTS does natively. It has no + flag for it, so it cannot honour a ``False`` setting; the wrapper + reports that at load time. + """ + del request_id # The session already owns this stream's request ids. + + session = require_session(state) + text_token = int(state.gen_text[:, current_frame_idx].item()) + session.step_tts(text_token) + audio_chunks = session.drain_audio_codes() + if not audio_chunks: + raise RuntimeError("vLLM EarTTS produced no audio codes for the submitted text token") + # The native codec helpers consume [B, T, num_quantizers]. + return torch.cat(audio_chunks, dim=0).to(self.device, dtype=torch.long).unsqueeze(0) diff --git a/nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/llm.py b/nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/llm.py new file mode 100644 index 000000000000..98679f9644b1 --- /dev/null +++ b/nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/llm.py @@ -0,0 +1,78 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""vLLM-Omni backend for the LLM component of NemotronVoiceChat. + +Implements :class:`~nemo.collections.speechlm2.inference.model_wrappers.backend.llm.DuplexLLM` +against the per-stream ``OmniStreamingSession``; the PyTorch sibling lives in +``backend/pytorch/llm.py``. +""" + +from typing import Any + +import torch + +from nemo.collections.speechlm2.inference.model_wrappers.backend.llm import DuplexLLM, LlmStepResult +from nemo.collections.speechlm2.inference.model_wrappers.backend.vllm import require_session + + +class VllmLLM(DuplexLLM): + """Runs Nemotron in a vLLM-Omni engine, one acoustic frame per step. + + Stateless itself: the engine is process-scoped and owned by + ``OmniRuntime``, and everything request-scoped lives in the session that + the pipeline attached to the decode state at prefill. + """ + + def step( + self, + frame_embedding: torch.Tensor, + state: Any, + *, + frame_offset: int, + current_frame_idx: int, + has_prompt: bool, + return_debug: bool = False, + sampling_params: dict[str, float] | None = None, + debug_logger: Any = None, + ) -> LlmStepResult: + """One Nemotron step -- see ``DuplexLLM.step``. + + Nemotron builds its own duplex input embedding from the acoustic frame, + so there is no ``build_input_embedding`` and no history replay here; + ``frame_offset`` and ``has_prompt`` do not apply. Per-stream sampling + was fixed when the session was created, and logits stay inside the + engine, so ``return_debug`` cannot be honoured either -- the result's + logit fields stay None. + + The previous frame's committed text token is fed back explicitly. That + is what carries a forced-turn-taking rewrite into Nemotron's own + history, the role ``gen_text`` plays for the PyTorch backend. + """ + del frame_offset, has_prompt, sampling_params, return_debug + + session = require_session(state) + if debug_logger is not None: + debug_logger.log_input_embeds(frame_embedding) + + prev_text_token = None + if current_frame_idx > 0: + prev_text_token = int(state.gen_text[0, current_frame_idx - 1].item()) + + tokens = session.step_llm(frame_embedding.reshape(-1), prev_text_token=prev_text_token) + return LlmStepResult( + predicted_token=tokens.text, + asr_predicted_token=tokens.asr, + function_predicted_token=tokens.function, + ) diff --git a/nemo/collections/speechlm2/inference/model_wrappers/capabilities.py b/nemo/collections/speechlm2/inference/model_wrappers/capabilities.py new file mode 100644 index 000000000000..4116c918a26c --- /dev/null +++ b/nemo/collections/speechlm2/inference/model_wrappers/capabilities.py @@ -0,0 +1,61 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Stable capability contract for optional VoiceChat output channels.""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass +from typing import Any + + +@dataclass(frozen=True) +class AuxiliaryOutputCapabilities: + """Which optional VoiceChat output channels this run can produce. + + Both backends surface both auxiliary channels, so a head being present in + the checkpoint is the same thing as its output being available. + """ + + has_asr_head: bool + has_function_head: bool + + def to_dict(self) -> dict[str, bool]: + """Return a JSON-serializable representation with stable field names.""" + return asdict(self) + + +def _head_flag(stt_model: Any, *names: str) -> bool: + """Whether any of *names* marks a head as present on this checkpoint. + + Native checkpoints expose the ASR head as ``predict_user_text``; + converted Nemotron configs use ``use_asr_head``. The attribute is + authoritative when the model defines one; the config is the fallback. + """ + for name in names: + value = getattr(stt_model, name, None) + if value is not None: + return bool(value) + cfg = getattr(stt_model, "cfg", None) + if cfg is None: + return False + return any(bool(cfg.get(name, False)) for name in names) + + +def derive_auxiliary_output_capabilities(stt_model: Any) -> AuxiliaryOutputCapabilities: + """Derive optional-channel capabilities from the checkpoint's heads.""" + return AuxiliaryOutputCapabilities( + has_asr_head=_head_flag(stt_model, "predict_user_text", "use_asr_head"), + has_function_head=_head_flag(stt_model, "use_function_head"), + ) diff --git a/nemo/collections/speechlm2/inference/model_wrappers/codec.py b/nemo/collections/speechlm2/inference/model_wrappers/codec.py new file mode 100644 index 000000000000..aac8f4c92117 --- /dev/null +++ b/nemo/collections/speechlm2/inference/model_wrappers/codec.py @@ -0,0 +1,88 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Streaming decode of EarTTS acoustic codes into a waveform. + +Used for every LLM/TTS engine pairing: the engines emit codes, not audio. +""" + +import torch + +from nemo.collections.speechlm2.parts.precision import fp32_precision +from nemo.utils import logging + + +class AudioCodec: + """Streaming decode of acoustic codes into a waveform.""" + + def __init__(self, tts_model, device: torch.device): + """ + Args: + tts_model: The ``DuplexEARTTS`` that owns ``audio_codec`` and the + control-code table. + device: Device the codec decodes on. + """ + self.tts_model = tts_model + self.device = device + + def create_state(self, max_len: int) -> tuple[torch.Tensor, object]: + """Per-stream codec state: ``(subword_mask, codec_cache)``.""" + from nemo.collections.speechlm2.modules.ear_tts_vae_codec import CausalConv1dCache + + subword_mask = torch.ones((1, max_len), device=self.device, dtype=torch.bool) + return subword_mask, CausalConv1dCache() + + def decode(self, new_codes: list[torch.Tensor], cache) -> torch.Tensor | None: + """Decode this chunk's accumulated codes into a waveform. + + Args: + new_codes: One ``(B, T, num_quantizers)`` tensor per frame. + cache: The stream's ``CausalConv1dCache``, updated in place. + + Returns: + The decoded waveform, or *None* when no codes were produced. + """ + if not new_codes: + return None + + with fp32_precision(), torch.no_grad(): + codes = torch.cat(new_codes, dim=1) + codes = self._replace_control_codes(codes) + code_len = torch.tensor([codes.shape[1]], dtype=torch.long, device=self.device) + decoded_audio, _ = self.tts_model.audio_codec.decode(codes, code_len, cache=cache) + return decoded_audio + + def _replace_control_codes(self, codes: torch.Tensor) -> torch.Tensor: + """Substitute codec silence for the model's control codes. + + Only checkpoints that define a control-code table need this; older ones + have none, and then the codes pass through unchanged. + """ + control_codes = getattr(self.tts_model, "_control_codes", None) + if control_codes is None: + return codes + from nemo.collections.speechlm2.models.duplex_ear_tts import replace_control_speech_codes + + return replace_control_speech_codes( + codes, + control_codes, + getattr(self.tts_model, "codec_silence_tokens", None), + ) + + def log_configuration(self) -> None: + """Record the codec's sample rate and frame rate once, at load time.""" + logging.info( + f"Audio codec ready: target_fps={self.tts_model.target_fps}, " + f"sample_rate={self.tts_model.target_sample_rate}" + ) diff --git a/nemo/collections/speechlm2/inference/model_wrappers/config_overrides.py b/nemo/collections/speechlm2/inference/model_wrappers/config_overrides.py new file mode 100644 index 000000000000..47007043df91 --- /dev/null +++ b/nemo/collections/speechlm2/inference/model_wrappers/config_overrides.py @@ -0,0 +1,195 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Which ``s2s`` config keys each backend honours, and where they land. + +Two kinds of settings arrive in the ``s2s`` config block: + +* **Model-consumed keys** are read by shared model code off its own ``cfg`` + (the logit boosts in ``DuplexSTTModel``, ``force_turn_taking``, the EarTTS + generation config), so the wrapper cannot pass them as arguments. They are + bridged into the relevant model config by :func:`apply_model_cfg_overrides`, + which is also how the streaming path inherits the offline path's behaviour + for the same key. +* **Wrapper-consumed knobs** (``decode_audio``, ``top_p``, ``use_llm_cache``, + ...) are read straight into wrapper attributes. Nothing is bridged for + them, but ignored keys are still reported. + +This module is the single answer to "can I set this, where does it land, and +does the selected backend honour it". :func:`apply_model_cfg_overrides` bridges +the model-consumed keys and warns about everything the selected backends will +ignore. + +``deterministic`` is deliberately not here: it is a process-global torch +setting rather than a model or wrapper key, and has to be settled before any +weights load. See +:func:`~nemo.collections.speechlm2.inference.model_wrappers.engine_selection.reject_unsupported_determinism`. +""" + +from collections.abc import Mapping +from typing import Any + +from omegaconf import OmegaConf + +from nemo.collections.speechlm2.inference.model_wrappers.engine_selection import VLLM_OMNI +from nemo.utils import logging + +# Component that owns the config object each key is read from. The component +# also selects the ``*_engine_type`` that decides whether the key takes effect. +LLM = "llm" # -> model.stt_model.cfg (DuplexSTTModel) +TTS = "tts" # -> model.tts_model.cfg (DuplexEARTTS) + +LLM_KEYS = ( + # Agent (text) and user (ASR) side logit boosts, read by the DuplexSTTModel + # heads; then forced turn-taking, read by _maybe_apply_forced_turn_taking. + "inference_pad_boost", + "inference_bos_boost", + "inference_eos_boost", + "inference_user_pad_boost", + "inference_user_bos_boost", + "inference_user_eos_boost", + "force_turn_taking", + "force_turn_taking_threshold", + "force_turn_taking_pad_window", +) + +TTS_KEYS = ( + # Both backends implement EOS -> codec silence inside the model: + # DuplexEARTTS.infer_codes_one_step (flag-gated, defaults to True) and the + # vLLM EarTTS preprocess (unconditional). See VLLM_FORCES_TRUE. + "inference_force_speech_silence_on_eos", + "inference_top_p_or_k", + "inference_noise_scale", + "inference_guidance_scale", +) + +COMPONENT_OF = {**{key: LLM for key in LLM_KEYS}, **{key: TTS for key in TTS_KEYS}} + +# --- Support tables ------------------------------------------------------- +# Every key not listed below works on both backends. Each entry is +# ``key -> (component, why)``; the component picks which ``*_engine_type`` +# decides, and the reason is quoted verbatim to the user. + +# The run is correct, but the setting does nothing: warn. +VLLM_IGNORES = { + "inference_top_p_or_k": (TTS, "read by DuplexEARTTS._get_generation_config"), + "inference_noise_scale": (TTS, "read by DuplexEARTTS._get_generation_config"), + "inference_guidance_scale": (TTS, "read by DuplexEARTTS._get_generation_config"), + "use_llm_cache": (LLM, "vLLM always keeps a paged KV cache"), + "use_tts_torch_compile": (TTS, "vLLM compiles inside the engine"), + "use_tts_subword_cache": ( + TTS, + "the subword table is baked in at checkpoint conversion, so it is always in effect", + ), +} + +# These boolean flags request an optimization only when enabled. Their false +# values are no-ops, unlike numeric sampling values such as a zero noise scale. +_IGNORED_ENABLE_FLAGS = frozenset({"use_llm_cache", "use_tts_torch_compile", "use_tts_subword_cache"}) + +# vLLM does this unconditionally: it can honour True but not False. Warn only +# when False was asked for. +VLLM_FORCES_TRUE = { + "inference_force_speech_silence_on_eos": ( + TTS, + "the converted EarTTS always substitutes codec silence when the incoming text token is EOS", + ), +} + + +def _selected(component: str, llm_engine_type: str, tts_engine_type: str) -> str: + return str(llm_engine_type if component == LLM else tts_engine_type).lower() + + +def _target_cfg(model, component: str): + """Model config that *component*'s code reads its settings from.""" + if component == LLM: + submodel = model.stt_model + elif component == TTS: + submodel = model.tts_model + else: + raise ValueError(f"Unknown component {component!r}; expected {LLM!r} or {TTS!r}") + return None if submodel is None else submodel.cfg + + +def apply_model_cfg_overrides( + model, + model_cfg: Mapping, + *, + llm_engine_type: str, + tts_engine_type: str, +) -> dict[str, Any]: + """Bridge the model-consumed keys into *model*, then report what is ignored. + + Keys absent from *model_cfg* are left alone, so whatever the checkpoint + carries stays in effect. Keys the selected backend cannot honour are + reported once per component instead of silently doing nothing. + + Returns: + The effective value of every model-consumed key after bridging, for + logging. + """ + for key, component in COMPONENT_OF.items(): + value = model_cfg.get(key, None) + if value is None: + continue + target = _target_cfg(model, component) + if target is None: + logging.warning(f"Ignoring `{key}`: this checkpoint has no {component} component to apply it to.") + continue + OmegaConf.update(target, key, value, force_add=True) + + effective: dict[str, Any] = {} + for key, component in COMPONENT_OF.items(): + target = _target_cfg(model, component) + effective[key] = None if target is None else target.get(key, None) + + # Model-consumed keys are checked post-bridge, so the reported value is the + # one the model will actually read; wrapper knobs come straight from config. + def value_of(key: str) -> Any: + return effective[key] if key in COMPONENT_OF else model_cfg.get(key, None) + + _warn_unsupported(value_of, llm_engine_type=llm_engine_type, tts_engine_type=tts_engine_type) + return effective + + +def _warn_unsupported(value_of, *, llm_engine_type: str, tts_engine_type: str) -> None: + """Report set keys the selected backends ignore, grouped by component.""" + unsupported: dict[str, list[str]] = {} + + for key, (component, why) in VLLM_IGNORES.items(): + if _selected(component, llm_engine_type, tts_engine_type) != VLLM_OMNI: + continue + value = value_of(key) + if value is None or (key in _IGNORED_ENABLE_FLAGS and not value): + continue + unsupported.setdefault(f"{component}_engine_type={VLLM_OMNI}", []).append(f"{key} ({why})") + + for selection, keys in unsupported.items(): + listed = ", ".join(sorted(keys)) + logging.warning( + f"These settings have no effect with {selection}: {listed}. " + "Select the native engine for that component, or remove them from the config " + "so it reflects what the run actually does." + ) + + for key, (component, why) in VLLM_FORCES_TRUE.items(): + if _selected(component, llm_engine_type, tts_engine_type) != VLLM_OMNI: + continue + if value_of(key) is not False: + continue + logging.warning( + f"`{key}=False` is not supported by {component}_engine_type={VLLM_OMNI}: {why}. " + "Expect it to behave as True." + ) diff --git a/nemo/collections/speechlm2/inference/model_wrappers/decode_state.py b/nemo/collections/speechlm2/inference/model_wrappers/decode_state.py new file mode 100644 index 000000000000..1d568a2ab236 --- /dev/null +++ b/nemo/collections/speechlm2/inference/model_wrappers/decode_state.py @@ -0,0 +1,202 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""State and result types for streaming S2S inference. + +These dataclasses define the *model wrapper's* interface contract: + +* :class:`StreamingDecodeState` — mutable per-stream decode state + (KV caches, token workspaces, perception/codec caches). Created by + the wrapper, mutated in-place by ``infer_one_step``, and held between + steps by the pipeline's context manager. + +* :class:`InferenceStepResult` — immutable per-step outputs returned + by ``infer_one_step`` (predicted tokens, text strings, audio). + +Defined here (in ``model_wrappers/``) because the wrapper is the +component that knows what state it needs. The context manager and +pipeline import from here. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +import torch + +from nemo.utils import logging +from nemo.utils.timers import NamedTimer + +if TYPE_CHECKING: + from nemo.collections.speechlm2.inference.model_wrappers.perception_cache import PerceptionCacheState + + +class TimingSummary(NamedTimer): + """Accumulates per-stage wall-clock times across inference steps. + + Extends :class:`~nemo.utils.timers.NamedTimer` (``sync_cuda=True``) + with a :meth:`log_summary` that prints a compact min/mean/max table + at ``logging.info`` level once a stream finishes. + + Usage:: + + timing.start("perception") + # ... run perception ... + timing.stop("perception") + + timing.log_summary(label="Stream 0", chunk_ms=240) + """ + + def __init__(self): + super().__init__(reduction="none", sync_cuda=True) + + def stop(self, name: str = "") -> None: + """Stop the named timer and log its duration at DEBUG level.""" + super().stop(name) + dt_ms = self.timers[name]["dt"][-1] * 1000 + logging.debug(f"[timing] {name}: {dt_ms:.1f}ms") + + def log_summary(self, label: str = "Timing", chunk_ms: float | None = None) -> None: + header = f"{label} timing" + if chunk_ms is not None: + header += f" (chunk={chunk_ms:.0f}ms)" + parts = [] + for name, data in self.timers.items(): + times = data.get("dt", []) + if not times: + continue + mean_ms = sum(times) / len(times) * 1000 + min_ms = min(times) * 1000 + max_ms = max(times) * 1000 + parts.append(f"{name}: mean={mean_ms:.1f}ms min={min_ms:.1f}ms max={max_ms:.1f}ms") + if parts: + logging.info(f"{header}:\n " + "\n ".join(parts)) + + +class NullTimingSummary: + """No-op stand-in for :class:`TimingSummary`.""" + + def start(self, name: str = "") -> None: + pass + + def stop(self, name: str = "") -> None: + pass + + def log_summary(self, label: str = "Timing", chunk_ms: float | None = None) -> None: + pass + + +@dataclass +class StreamingDecodeState: + """Per-stream model-level decode state for streaming S2S inference. + + Holds KV caches, token workspaces, perception cache, and codec state + that persist across inference steps within a single stream. + """ + + frame_idx: int + gen_text: torch.Tensor + gen_asr_text: torch.Tensor | None + gen_function: torch.Tensor | None + input_embeds_history: list[torch.Tensor] + llm_cache: Any # DynamicCache for supported native transformer backbones, otherwise None. + tts_past_key_values: Any + tts_code: torch.Tensor | None + subword_mask: torch.Tensor | None + perception_cache: "PerceptionCacheState" | None = None + tts_codec_cache: Any = None + llm_cache_position_offset: int = 0 + timing: TimingSummary | NullTimingSummary = field(default_factory=NullTimingSummary) + # Present when either component selects vllm_omni. The session may own + # Nemotron, EarTTS, or both. + omni_session: Any = None + + +@dataclass +class InferenceStepResult: + """Output from a single ``infer_one_step`` call. + + State mutations (caches, token workspaces, frame_idx) are applied + in-place on :class:`StreamingDecodeState`. This dataclass carries + only the per-step *outputs* needed by the pipeline. + """ + + predicted_text_tokens: torch.Tensor + asr_predicted_text_tokens: torch.Tensor | None + predicted_text_strs: list[str] + asr_predicted_text_strs: list[str] | None + decoded_audio: torch.Tensor | None = None + debug: dict | None = None + predicted_function_tokens: torch.Tensor | None = None + predicted_function_strs: list[str] | None = None + + +class IntermediateResultLogger: + """Records per-frame debug data (logits, embeddings, indices) during inference. + + Tensors are kept on their original device until :meth:`build_debug_dict` + is called, which performs a single bulk copy to CPU. + """ + + def __init__(self): + self.text_logits: list[torch.Tensor] = [] + self.asr_logits: list[torch.Tensor] = [] + self.input_embeds: list[torch.Tensor] = [] + self.selected_frame_indices: list[int] = [] + + def log_input_embeds(self, emb: torch.Tensor): + self.input_embeds.append(emb.detach()) + + def log_text_logits(self, logits: torch.Tensor): + self.text_logits.append(logits.detach()) + + def log_asr_logits(self, logits: torch.Tensor | None): + if logits is not None: + self.asr_logits.append(logits.detach()) + + def log_selected_frame_index(self, idx: int): + self.selected_frame_indices.append(idx) + + def build_debug_dict( + self, source_encoded: torch.Tensor, gen_text: torch.Tensor, gen_asr_text: torch.Tensor | None + ) -> dict: + return { + "source_encoded": source_encoded.detach().cpu(), + "selected_frame_indices": self.selected_frame_indices, + "input_embeds": torch.cat(self.input_embeds, dim=1).cpu() if self.input_embeds else None, + "gen_text": gen_text.detach().cpu(), + "gen_asr": gen_asr_text.detach().cpu() if gen_asr_text is not None else None, + "text_logits": torch.stack(self.text_logits, dim=1).cpu() if self.text_logits else None, + "asr_logits": torch.stack(self.asr_logits, dim=1).cpu() if self.asr_logits else None, + } + + +class NullIntermediateResultLogger: + """No-op stand-in for :class:`IntermediateResultLogger`.""" + + def log_input_embeds(self, emb): + pass + + def log_text_logits(self, logits): + pass + + def log_asr_logits(self, logits): + pass + + def log_selected_frame_index(self, idx): + pass + + def build_debug_dict(self, *args, **kwargs): + return None diff --git a/nemo/collections/speechlm2/inference/model_wrappers/engine_selection.py b/nemo/collections/speechlm2/inference/model_wrappers/engine_selection.py new file mode 100644 index 000000000000..eae8b78fc2c3 --- /dev/null +++ b/nemo/collections/speechlm2/inference/model_wrappers/engine_selection.py @@ -0,0 +1,141 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""What the ``s2s`` config resolves to, and the process state it implies. + +Kept out of the wrapper so that answering "which engines did this config +select?" does not require constructing a wrapper, loading a checkpoint or +importing torch's model stack. The builder, the wrapper and the tests all read +their answers from here. +""" + +from collections.abc import Mapping + +from nemo.collections.speechlm2.parts.precision import inference_precision, inference_precision_in_effect + +NATIVE = "native" +VLLM_OMNI = "vllm_omni" + +SUPPORTED_S2S_ENGINE_TYPES = frozenset({NATIVE, VLLM_OMNI}) + + +def _component_engine(model_cfg: Mapping, key: str) -> str: + value = model_cfg.get(key, NATIVE) + if value is None: + return NATIVE + return str(value).lower() + + +def resolve_engine_types(model_cfg: Mapping) -> tuple[str, str]: + """Read independent LLM and TTS engines from *model_cfg*. + + Missing or null keys default to ``native``. ``engine_type`` is not a + config key: if it is set, this raises rather than treating it as a + shared default. + """ + leftover = model_cfg.get("engine_type", None) + if leftover is not None: + raise ValueError( + "`engine_type` is not a config key. Set `llm_engine_type` and " + "`tts_engine_type` independently; each must be one of " + f"{sorted(SUPPORTED_S2S_ENGINE_TYPES)}." + ) + llm = _component_engine(model_cfg, "llm_engine_type") + tts = _component_engine(model_cfg, "tts_engine_type") + invalid = { + name: value + for name, value in (("llm_engine_type", llm), ("tts_engine_type", tts)) + if value not in SUPPORTED_S2S_ENGINE_TYPES + } + if invalid: + values = ", ".join(f"{name}={value!r}" for name, value in invalid.items()) + raise ValueError( + f"Unsupported S2S engine selection ({values}); expected one of " f"{sorted(SUPPORTED_S2S_ENGINE_TYPES)}." + ) + return llm, tts + + +def reject_unsupported_determinism(llm_engine_type: str, tts_engine_type: str, deterministic: bool) -> None: + """Raise if ``deterministic`` was asked for alongside a vLLM component. + + vLLM's custom kernels (PagedAttention, FlashAttention) have no + deterministic mode, so this cannot be honoured rather than merely being + slower. Same no-silent-no-op contract as the rest of the config checks. + """ + if not deterministic: + return + vllm_components = [ + name + for name, value in (("llm_engine_type", llm_engine_type), ("tts_engine_type", tts_engine_type)) + if value == VLLM_OMNI + ] + if vllm_components: + raise ValueError( + "`deterministic` is not compatible with vLLM engines because vLLM uses custom " + "CUDA kernels (PagedAttention, FlashAttention) that do not support deterministic mode. " + f"Selected vLLM components: {', '.join(vllm_components)}. " + "Use native engines for deterministic inference." + ) + + +def native_weight_skip_prefixes(llm_engine_type: str, tts_engine_type: str) -> set[str]: + """Checkpoint prefixes not needed by the selected component backends.""" + prefixes = {"stt_model.rnnt_decoder.", "stt_model.rnnt_joint."} + if llm_engine_type == VLLM_OMNI: + prefixes.add("stt_model.llm.") + if tts_engine_type == VLLM_OMNI: + prefixes.add("tts_model.tts_model.") + return prefixes + + +def _precision_settings(model_cfg: Mapping) -> dict: + """The three torch precision switches *model_cfg* asks for.""" + return { + "allow_tf32": bool(model_cfg.get("allow_tf32", True)), + "matmul_precision": str(model_cfg.get("matmul_precision", "medium")), + "deterministic": bool(model_cfg.get("deterministic", False)), + } + + +def precision_matches_cfg(model_cfg: Mapping) -> bool: + """Whether the process is already configured the way *model_cfg* asks. + + What entry points check before loading weights, so a forgotten + :func:`inference_precision_from_cfg` is reported rather than silently + changing the numbers. + """ + return inference_precision_in_effect(**_precision_settings(model_cfg)) + + +def inference_precision_from_cfg(model_cfg: Mapping): + """Scope the process-wide torch settings *model_cfg* implies. + + The switches must be in effect before any weights load and stay on for the + whole run, but they are process globals, so they are restored on exit + rather than left set. That keeps a deterministic run from changing every + later computation in the process. + + The builder requires this scope and does not enter it. Typical caller:: + + with inference_precision_from_cfg(cfg.s2s): + pipeline = S2SPipelineBuilder.build_pipeline(cfg) + try: + pipeline.run(...) + finally: + pipeline.shutdown() + """ + llm_engine_type, tts_engine_type = resolve_engine_types(model_cfg) + settings = _precision_settings(model_cfg) + reject_unsupported_determinism(llm_engine_type, tts_engine_type, settings["deterministic"]) + return inference_precision(**settings) diff --git a/nemo/collections/speechlm2/inference/model_wrappers/nemotron_voicechat_inference_wrapper.py b/nemo/collections/speechlm2/inference/model_wrappers/nemotron_voicechat_inference_wrapper.py new file mode 100644 index 000000000000..f35b7cd208de --- /dev/null +++ b/nemo/collections/speechlm2/inference/model_wrappers/nemotron_voicechat_inference_wrapper.py @@ -0,0 +1,1108 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy +import json +import os +import time + +import torch +from omegaconf import DictConfig, OmegaConf + +from nemo.collections.speechlm2.inference.model_wrappers.backend.llm import LlmStepResult +from nemo.collections.speechlm2.inference.model_wrappers.backend.pytorch.eartts import PyTorchEarTTS +from nemo.collections.speechlm2.inference.model_wrappers.backend.pytorch.llm import PyTorchLLM +from nemo.collections.speechlm2.inference.model_wrappers.backend.vllm.eartts import VllmEarTTS +from nemo.collections.speechlm2.inference.model_wrappers.backend.vllm.llm import VllmLLM +from nemo.collections.speechlm2.inference.model_wrappers.capabilities import ( + AuxiliaryOutputCapabilities, + derive_auxiliary_output_capabilities, +) +from nemo.collections.speechlm2.inference.model_wrappers.config_overrides import apply_model_cfg_overrides +from nemo.collections.speechlm2.inference.model_wrappers.decode_state import ( + InferenceStepResult, + IntermediateResultLogger, + NullIntermediateResultLogger, + NullTimingSummary, + StreamingDecodeState, + TimingSummary, +) +from nemo.collections.speechlm2.inference.model_wrappers.engine_selection import ( + VLLM_OMNI, + native_weight_skip_prefixes, + precision_matches_cfg, + reject_unsupported_determinism, + resolve_engine_types, +) +from nemo.collections.speechlm2.inference.model_wrappers.codec import AudioCodec +from nemo.collections.speechlm2.inference.model_wrappers.perception_cache import ( + PerceptionCacheManager, + PerceptionCacheState, +) +from nemo.collections.speechlm2.models.nemotron_voicechat import NemotronVoiceChat +from nemo.collections.speechlm2.parts.logit_boosts import LogitBoosts +from nemo.collections.speechlm2.parts.text_utils import ( + _decode_tokens_with_specials, + get_special_token_ids, + get_special_token_strings, +) +from nemo.utils import logging, str_to_dtype + +# --- Configuration --- +DEFAULT_DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +# --- Streaming Parameters --- +SAMPLE_RATE = 16000 +FRAME_SIZE_SEC = 0.08 # 80ms per frame +FRAME_SIZE_SAMPLES = int(SAMPLE_RATE * FRAME_SIZE_SEC) # 1280 samples + +TTS_SAMPLE_RATE = 22050 + + +class NemotronVoicechatInferenceWrapper: + """ + Inference wrapper for NemotronVoiceChat models. + Uses a sliding window buffer and processes audio frame by frame. + """ + + def __init__(self, model_cfg: DictConfig): + """ + Initialize the model for realtime streaming inference. + + Args: + model_cfg (DictConfig): Configuration describing the model paths and inference parameters. + """ + if model_cfg is None: + raise ValueError("model_cfg must be provided") + if not isinstance(model_cfg, DictConfig): + model_cfg = OmegaConf.create(model_cfg) + + # Precision and determinism are process globals scoped by the caller's + # `inference_precision_from_cfg`. Validated again here so constructing the + # wrapper directly still rejects an impossible combination. + self.llm_engine_type, self.tts_engine_type = resolve_engine_types(model_cfg) + self._deterministic = bool(model_cfg.get("deterministic", False)) + reject_unsupported_determinism(self.llm_engine_type, self.tts_engine_type, self._deterministic) + if not precision_matches_cfg(model_cfg): + # Direct construction warns: only S2SPipelineBuilder requires the + # precision scope. These torch globals are not applied here. + logging.warning( + "This process is not configured the way model_cfg asks: allow_tf32, " + "matmul_precision and deterministic are torch process globals and are NOT in " + "effect for this model, so sampling and numerics may differ from a configured " + "run. Wrap construction in `inference_precision_from_cfg(model_cfg)`, or use " + "S2SPipelineBuilder.build_pipeline, which requires it." + ) + + self.model_cfg = model_cfg + + self.model_path = model_cfg.get("model_path") + if not self.model_path: + raise ValueError("`model_cfg.model_path` must be provided.") + + self.decode_audio = bool(model_cfg.get("decode_audio", True)) + + if model_cfg.get("speaker_reference"): + raise ValueError( + "s2s.speaker_reference is not supported. Checkpoints only generate with " + "speakers registered at export (register_speaker_dict in " + "examples/speechlm2/nemotron_voicechat_to_hf.py). Set s2s.speaker_name " + "to a registered speaker." + ) + self.speaker_name = model_cfg.get("speaker_name", None) + if self.decode_audio and not self.speaker_name: + raise ValueError( + "`model_cfg.speaker_name` must be provided when decode_audio is enabled. " + "It must match a speaker registered in the checkpoint." + ) + + self.dtype = str_to_dtype(model_cfg.get("compute_dtype", "bfloat16")) + + device = model_cfg.get("device") + device_id = model_cfg.get("device_id") + if device is None: + self.device = DEFAULT_DEVICE + else: + device_str = str(device) + if device_id is not None and device_str.startswith("cuda") and ":" not in device_str: + device_str = f"{device_str}:{device_id}" + self.device = torch.device(device_str) + + logging.info("=" * 70) + logging.info("INITIALIZING REALTIME STREAMING INFERENCE") + logging.info("=" * 70) + logging.info(f"Frame size: {FRAME_SIZE_SEC}s ({FRAME_SIZE_SAMPLES} samples @ {SAMPLE_RATE}Hz)") + logging.info(f"Device: {self.device}") + logging.info(f"Compute dtype: {self.dtype}") + logging.info(f"Decode audio: {self.decode_audio}") + logging.info(f"Engine types: LLM={self.llm_engine_type}, TTS={self.tts_engine_type}") + logging.info( + f"Sampling - top_p: {model_cfg.get('top_p', 0.5)}, repetition_penalty: {model_cfg.get('repetition_penalty', 1.1)}, temperature: {model_cfg.get('temperature', 0.3)}" + ) + logging.info(f"Precision (configured): deterministic={self._deterministic}") + logging.info( + f"Precision (effective): float32_matmul_precision={torch.get_float32_matmul_precision()}, cudnn.allow_tf32={torch.backends.cudnn.allow_tf32}, cuda.matmul.allow_tf32={torch.backends.cuda.matmul.allow_tf32}" + ) + logging.info("=" * 70) + + # Profiling: when True, a TimingSummary (extending NamedTimer with + # sync_cuda=True) is attached to each decode state, recording + # per-stage wall-clock times. Disabled by default to avoid + # unnecessary GPU stalls in production. + self._profile_timing = bool(model_cfg.get("profile_timing", False)) + + # Cached TTS helpers populated during initialization/warmup + self.first_context_subword_id = None + self.generation_config = None + self.first_tts_code_input = None + self.first_tts_past_key_values_input = None + + self.model = None + # Selected per-component backends (DuplexLLM / DuplexTTS). + self.llm_backend = None + self.tts_backend = None + # Codec decode is always PyTorch, for every LLM/TTS engine pair. + self.codec: AudioCodec | None = None + # Native objects only; None when that component runs on vLLM. + self.model_llm_interface = None + self.model_eartts_interface = None + self.tokenizer = None + self.special_token_ids: set[int] = set() + + self._output_capabilities = AuxiliaryOutputCapabilities( + has_asr_head=False, + has_function_head=False, + ) + self.request_id = "streaming_request_0" # For vLLM streaming + + # vLLM-Omni runtime + speaker latent (selected components only). + self.vllm_omni_config = model_cfg.get("vllm_omni_config", None) + self.omni_runtime = None + self.omni_wrapper_dir: str | None = None + self.omni_speaker_latent: torch.Tensor | None = None + self.omni_guidance_enabled = True + self.omni_guidance_scale = 0.5 + + # Sampling parameters (defaults match s2s_streaming.yaml) + self.top_p = float(model_cfg.get("top_p", 0.5)) + self.repetition_penalty = float(model_cfg.get("repetition_penalty", 1.1)) + self.temperature = float(model_cfg.get("temperature", 0.3)) + + # Native LLM KV cache (vLLM engines manage their own cache and ignore this). + self.use_llm_cache = bool(model_cfg.get("use_llm_cache", False)) + + # Perception cache configuration (defaults match s2s_streaming.yaml) + self.use_perception_cache = bool(model_cfg.get("use_perception_cache", True)) + use_perception_cudagraph = bool(model_cfg.get("use_perception_cudagraph", True)) + if use_perception_cudagraph and not self.use_perception_cache: + raise ValueError( + "use_perception_cudagraph requires use_perception_cache to be enabled. " + "Please also set use_perception_cache=True." + ) + self.perception_cache_mgr: PerceptionCacheManager | None = None + self._use_perception_cudagraph = use_perception_cudagraph + + self._initialize_model() + + logging.info("NemotronVoicechatInferenceWrapper initialized successfully.") + + # ``llm_engine_type`` and ``tts_engine_type`` are the only stored selection; + # everything else derives from them so the two cannot drift apart. + # Perception, codec and tokenization stay on PyTorch in every combination. + + @property + def use_vllm_llm(self) -> bool: + return self.llm_engine_type == VLLM_OMNI + + @property + def use_vllm_tts(self) -> bool: + return self.tts_engine_type == VLLM_OMNI + + @property + def use_vllm_omni(self) -> bool: + return self.use_vllm_llm or self.use_vllm_tts + + @property + def output_capabilities(self) -> AuxiliaryOutputCapabilities: + """Optional checkpoint heads this run can produce.""" + return self._output_capabilities + + def _initialize_model(self): + """Initialize the NemotronVoiceChat model from an HF checkpoint.""" + logging.info("Initializing model structure...") + start_model_init = time.time() + + # Tell from_pretrained to skip loading checkpoint weights for + # submodules that vLLM will replace — avoids wasted I/O and memory. + # The streaming pipeline also does not instantiate the auxiliary RNN-T + # decoder some checkpoints bundle; its weights are independent of the + # duplex text/audio path. + skip_prefixes = native_weight_skip_prefixes(self.llm_engine_type, self.tts_engine_type) + + self.model = NemotronVoiceChat.from_pretrained( + self.model_path, + skip_prefixes=skip_prefixes, + ) + logging.info(f"NemotronVoiceChat initialized in {time.time() - start_model_init:.1f}s") + + # Remove skipped submodules (still on meta device / uninitialized) + if self.use_vllm_llm: + del self.model.stt_model.llm + self.model.stt_model.llm = None + if self.use_vllm_tts: + del self.model.tts_model.tts_model + + self.model.to(self.device) + self.model.safe_cast_to(self.dtype) + self.model.eval() + + self.tokenizer = self.model.stt_model.tokenizer + + # Bridge the config keys that shared model code reads off its own cfg. + # See config_overrides for the full set of keys and which + # backends honour each one. + effective_overrides = apply_model_cfg_overrides( + self.model, + self.model_cfg, + llm_engine_type=self.llm_engine_type, + tts_engine_type=self.tts_engine_type, + ) + if self.model.stt_model.cfg.get("force_turn_taking", False) and not self.model.stt_model.predict_user_text: + logging.warning( + "Disabling force_turn_taking because this checkpoint has no ASR head. " + "The model's learned duplex turn-taking remains active." + ) + OmegaConf.update(self.model.stt_model.cfg, "force_turn_taking", False) + effective_overrides["force_turn_taking"] = False + logging.info(f"Effective model config overrides: {effective_overrides}") + + stt = self.model.stt_model + self._output_capabilities = derive_auxiliary_output_capabilities(stt) + logging.info(f"Auxiliary output capabilities: {self._output_capabilities.to_dict()}") + self.special_token_ids = get_special_token_ids( + stt.tokenizer, + stt.text_pad_id, + model_cfg=stt.cfg, + ) + if self.use_vllm_omni: + self._initialize_vllm_omni_backend() + + # One implementation per component, chosen here and nowhere else. The + # PyTorch objects stay reachable as model_*_interface because stream + # setup needs cache creation, prompt prefill and abort, which are not + # on the shared contracts. + if self.use_vllm_llm: + self.llm_backend = VllmLLM() + else: + self.model_llm_interface = PyTorchLLM( + model=self.model, + special_token_ids=self.special_token_ids, + top_p=self.top_p, + repetition_penalty=self.repetition_penalty, + temperature=self.temperature, + use_llm_cache=self.use_llm_cache, + ) + self.llm_backend = self.model_llm_interface + logging.info(f"LLM backend: {type(self.llm_backend).__name__}") + + if self.use_vllm_tts: + self.tts_backend = VllmEarTTS(device=self.device) + else: + self.model_eartts_interface = PyTorchEarTTS(tts_model=self.model.tts_model) + self.tts_backend = self.model_eartts_interface + + # PyTorch TTS-only speedups are delegated to its backend. + if bool(self.model_cfg.get("use_tts_torch_compile", False)): + self.model_eartts_interface.compile() + self.model_eartts_interface.setup_subword_cache(self.model_cfg) + logging.info(f"TTS backend: {type(self.tts_backend).__name__}") + + # Codec decode is always PyTorch, regardless of which TTS engine + # produced the codes. + if hasattr(self.model, "tts_model"): + self.target_fps = self.model.tts_model.target_fps + self.target_sample_rate = self.model.tts_model.target_sample_rate + self.codec = AudioCodec(self.model.tts_model, self.device) + self.codec.log_configuration() + if self.decode_audio and not self.use_vllm_tts: + self._prepare_tts_initial_state() + else: + logging.warning("Warning: TTS model not found in the model") + + # Setup perception cache if enabled + if self.use_perception_cache: + self.perception_cache_mgr = PerceptionCacheManager( + model=self.model, + device=self.device, + dtype=self.dtype, + use_cudagraph=self._use_perception_cudagraph, + ) + if not self.perception_cache_mgr.setup(): + self.use_perception_cache = False + self.perception_cache_mgr = None + + # ------------------------------------------------------------------ + # vLLM-Omni backend + # ------------------------------------------------------------------ + + def _initialize_vllm_omni_backend(self): + """Build the wrapper checkpoint, start the AsyncOmni runtime, and + pre-load the speaker latent. + + The wrapper checkpoint (``config.json`` + ``nemotron/`` + ``eartts/``) + is converted lazily on first use under + ``$TMPDIR/_vllm_omni_wrapper`` and reused afterwards; set + ``vllm_omni_config.wrapper_dir`` to keep it somewhere persistent. + """ + # Deferred because these reach vLLM-Omni, an optional dependency. + from nemo.collections.speechlm2.inference.vllm_omni.checkpoint import ( + EARTTS_SUBDIR, + build_wrapper_checkpoint, + load_speaker_latent, + write_nemotron_inference_overrides, + ) + from nemo.collections.speechlm2.inference.vllm_omni.runtime import OmniRuntime + + cfg = self.vllm_omni_config or {} + + wrapper_dir = build_wrapper_checkpoint( + self.model_path, + wrapper_dir=cfg.get("wrapper_dir", None), + nemotron_dtype=cfg.get("nemotron_dtype", "float32"), + eartts_precompute_batch_size=int(cfg.get("eartts_precompute_batch_size", 256)), + include_nemotron=self.use_vllm_llm, + include_eartts=self.use_vllm_tts, + ) + self.omni_wrapper_dir = wrapper_dir + + if self.use_vllm_llm: + # Must happen before the stage child loads the checkpoint. + user_boosts = LogitBoosts.user_from_cfg(self.model.stt_model.cfg) + write_nemotron_inference_overrides( + wrapper_dir, + { + "inference_user_pad_boost": user_boosts.pad, + "inference_user_bos_boost": user_boosts.bos, + "inference_user_eos_boost": user_boosts.eos, + }, + ) + + self.omni_runtime = OmniRuntime( + wrapper_dir, + stage_configs_path=cfg.get("stage_configs_path", None), + eartts_stage_configs_path=cfg.get("eartts_stage_configs_path", None), + stage_overrides=cfg.get("stage_overrides", None), + eartts_stage_overrides=cfg.get("eartts_stage_overrides", None), + log_stats=bool(cfg.get("log_stats", False)), + stage_init_timeout=int(cfg.get("stage_init_timeout", 600)), + enable_llm=self.use_vllm_llm, + enable_tts=self.use_vllm_tts, + ) + + if self.use_vllm_tts: + eartts_dir = os.path.join(wrapper_dir, EARTTS_SUBDIR) + with open(os.path.join(eartts_dir, "config.json"), encoding="utf-8") as fh: + eartts_config = json.load(fh) + guidance_enabled = cfg.get("guidance_enabled") + if guidance_enabled is None: + guidance_enabled = eartts_config.get("enable_guidance", True) + guidance_scale = cfg.get("guidance_scale") + if guidance_scale is None: + guidance_scale = eartts_config.get("guidance_scale", 0.5) + self.omni_guidance_enabled = bool(guidance_enabled) + self.omni_guidance_scale = float(guidance_scale) + speaker_name = self.speaker_name or cfg.get("speaker_name") + if speaker_name is None: + raise ValueError( + "tts_engine_type='vllm_omni' requires a speaker_name (set " + "s2s.speaker_name or s2s.vllm_omni_config.speaker_name); the " + "speaker latent is read from the converted EarTTS checkpoint." + ) + self.omni_speaker_latent = load_speaker_latent(eartts_dir, speaker_name) + logging.info( + "vllm_omni speaker_latent: name='%s', shape=%s; CFG=%s scale=%s", + speaker_name, + tuple(self.omni_speaker_latent.shape), + self.omni_guidance_enabled, + self.omni_guidance_scale, + ) + + def start_vllm_omni_session( + self, + state: StreamingDecodeState, + system_prompt: str | None, + *, + request_id: str, + sampling_params: dict[str, float] | None = None, + ) -> None: + """Create and attach a per-stream :class:`OmniStreamingSession`. + + Called by the streaming pipeline once it has the system prompt for + the new stream; this replaces native-engine prefill. The session only + enqueues the prefill chunk, so Nemotron does not actually run until + the first :meth:`infer_one_step` call. Per-stream sampling parameters + are fixed when this long-lived vLLM request is created. + + One session class covers all three vLLM combinations: it reads which + components exist off the runtime, which is the same decision that built + the runtime in the first place. + """ + if not self.use_vllm_omni: + return + if self.omni_runtime is None: + raise RuntimeError("vllm_omni backend was not initialized; call _initialize_model first.") + from nemo.collections.speechlm2.inference.vllm_omni.checkpoint import ( + NEMOTRON_SUBDIR, + compute_prefill_len, + ) + from nemo.collections.speechlm2.inference.vllm_omni.session import OmniStreamingSession + + t_prefill = 0 + if self.use_vllm_llm: + nemotron_dir = os.path.join(self.omni_wrapper_dir, NEMOTRON_SUBDIR) + t_prefill = compute_prefill_len(nemotron_dir, system_prompt or "") + + effective_sampling_params = { + "temperature": float(self.temperature), + "top_p": float(self.top_p), + "repetition_penalty": float(self.repetition_penalty), + } + if sampling_params: + effective_sampling_params.update( + {key: float(value) for key, value in sampling_params.items() if key in effective_sampling_params} + ) + + stt = self.model.stt_model + state.omni_session = OmniStreamingSession( + self.omni_runtime, + request_id=request_id, + system_prompt=system_prompt or "", + speaker_latent=self.omni_speaker_latent, + t_prefill=t_prefill, + sampling_params=effective_sampling_params, + special_token_ids=self.special_token_ids, + guidance_enabled=self.omni_guidance_enabled, + guidance_scale=self.omni_guidance_scale, + step_timeout=float((self.vllm_omni_config or {}).get("step_timeout", 60.0)), + profile=self._profile_timing, + # The agent-channel boosts ride along with sampling; the ASR-channel + # ones are applied by the converted model, which reads them from its + # own config at load time. + agent_logit_boosts=LogitBoosts.agent_from_cfg(stt.cfg), + text_token_ids={ + "pad_id": int(stt.text_pad_id), + "bos_id": int(stt.text_bos_id), + "eos_id": int(stt.text_eos_id), + }, + ) + + # ------------------------------------------------------------------ + # Per-stream lifecycle + # ------------------------------------------------------------------ + + def begin_stream( + self, + state: StreamingDecodeState, + system_prompt: str | None, + *, + request_id: str, + sampling_params: dict[str, float] | None = None, + ) -> None: + """Prepare per-stream backend state before the first audio frame. + + The caller does not need to know which components run on vLLM. This + opens a vLLM session for those that do and prefills the system prompt + natively when the LLM is native. A vLLM Nemotron consumes the prompt + inside its own long-lived request, so there is nothing to prefill. + """ + if self.use_vllm_omni: + self.start_vllm_omni_session( + state, + system_prompt or "", + request_id=request_id, + sampling_params=sampling_params, + ) + logging.info(f"vllm_omni: started streaming session (request_id={request_id!r}).") + + if self.use_vllm_llm or not system_prompt: + return + self._prefill_system_prompt_native(state, system_prompt) + + def end_stream(self, state: StreamingDecodeState | None, *, request_id: str) -> None: + """Release per-stream backend state. Idempotent. + + Must run before the decode state is discarded: a vLLM session's + consumer task hangs off the state and would otherwise leak until + process exit, with asyncio reporting a pending destroyed task. + """ + session = state.omni_session if state is not None else None + if session is not None: + state.omni_session = None + self._close_session(session, request_id) + + # The session above covers whichever components run on vLLM; only the + # native backends still hold a request to abort. + if not self.use_vllm_llm: + self.model_llm_interface.abort_request(request_id) + if not self.use_vllm_tts: + self.model_eartts_interface.abort_request(request_id) + + @staticmethod + def _close_session(session, request_id: str) -> None: + """Close a vLLM session, falling back to a hard abort. + + The only teardown step that is allowed to fail: ``finish`` waits for the + engine's consumer task to drain, which can time out or raise if the + engine is already unhealthy. ``abort`` drops the request without + waiting, so it is the correct second attempt rather than a blanket + except. The native aborts above are local bookkeeping and are left to + raise. + """ + try: + session.finish() + return + except Exception as exc: + logging.warning(f"vllm_omni session.finish() failed for request {request_id}: {exc}; aborting instead.") + try: + session.abort() + except Exception as exc: + logging.warning(f"vllm_omni session.abort() also failed for request {request_id}: {exc}") + + def _prefill_system_prompt_native(self, state: StreamingDecodeState, system_prompt: str) -> None: + """Put the system prompt into the native LLM's state for this stream. + + Either warms the KV cache or seeds ``input_embeds_history``, depending + on whether this stream was given a cache. + """ + logging.info("Prefilling system prompt...") + start = time.time() + prompt_embedded, prompt_len = self._prepare_system_prompt_embeddings(system_prompt) + logging.debug(f"Time taken to get prompt embeddings: {time.time() - start:.3f}s") + + if prompt_embedded is None: + logging.warning("System prompt embedding returned None, skipping prefill") + return + + if state.llm_cache is not None: + with torch.no_grad(): + cache_position = torch.arange(prompt_len, device=self.device) + ans = self.model_llm_interface.prefill_prompt( + prompt_embedded, + cache=state.llm_cache, + cache_position=cache_position, + ) + state.llm_cache = ans.get("cache", state.llm_cache) + state.llm_cache_position_offset = prompt_len + logging.info(f"System prompt processed, cache updated ({prompt_len} tokens, offset={prompt_len})") + else: + for t in range(prompt_len): + state.input_embeds_history.append(prompt_embedded[:, t : t + 1, :]) + logging.info(f"Added {prompt_len} prompt embeddings to input_embeds_history") + + def shutdown(self) -> None: + """Tear down the AsyncOmni runtime. Idempotent; no-op for native engines. + + Called by :meth:`StreamingS2SPipeline.shutdown`, so the engine + subprocesses and the runtime's daemon thread go away at a known point + rather than at garbage collection or process exit. + """ + if self.omni_runtime is None: + return + runtime, self.omni_runtime = self.omni_runtime, None + try: + runtime.shutdown() + except Exception as exc: + # Callers invoke this from ``finally`` / server finalize, where + # raising would replace whatever error is already unwinding. + logging.warning(f"OmniRuntime.shutdown raised: {exc!r}") + + def _prepare_system_prompt_embeddings( + self, + system_prompt: str, + ) -> tuple[torch.Tensor | None, int]: + if not system_prompt or not system_prompt.strip(): + return None, 0 + + prompt_token_ids = self._build_prompt_token_ids(system_prompt) + prompt_tokens = torch.tensor(prompt_token_ids, dtype=torch.long, device=self.device).unsqueeze(0) + prompt_embedded = self.model.stt_model.embed_tokens(prompt_tokens).to(dtype=self.dtype) + prompt_len = prompt_tokens.shape[1] + + stt = self.model.stt_model + pad_id = stt.text_pad_id + pad_token = torch.full((1,), fill_value=pad_id, device=self.device, dtype=torch.long) + pad_emb = stt.embed_tokens(pad_token).to(dtype=self.dtype) + bos_emb = stt._get_bos_embedding().to(dtype=self.dtype) + + if prompt_len > 1: + prompt_embedded[:, 1:, :] += pad_emb + if stt.predict_user_text: + pad_asr_emb = stt.embed_asr_tokens(pad_token).to(dtype=self.dtype) + prompt_embedded[:, 1:, :] += pad_asr_emb + + prompt_embedded[:, 0, :] += bos_emb.squeeze(0) + if stt.predict_user_text: + asr_bos_emb = stt._get_asr_bos_embedding().to(dtype=self.dtype) + prompt_embedded[:, 0, :] += asr_bos_emb.squeeze(0) + if stt.use_function_head: + # Match the channel order in DuplexSTTModel.build_input_embedding. + prompt_embedded += pad_emb.expand(1, prompt_len, -1) * stt.cfg.get("duplex_function_channel_weight", 1.0) + + return prompt_embedded, prompt_len + + def _clone_cache(self, cache): + """Deep clone cache structures to ensure complete isolation between streams.""" + if cache is None: + return None + if isinstance(cache, torch.Tensor): + return cache.detach().clone() + if isinstance(cache, (list, tuple)): + return type(cache)(self._clone_cache(x) for x in cache) + if isinstance(cache, dict): + return {k: self._clone_cache(v) for k, v in cache.items()} + if hasattr(cache, "__dict__"): + return copy.deepcopy(cache) + return cache + + def _build_prompt_token_ids(self, system_prompt: str | None) -> list[int]: + if not system_prompt or not system_prompt.strip(): + return [] + return [self.tokenizer.bos_id] + self.tokenizer.text_to_ids(system_prompt) + [self.tokenizer.eos_id] + + def _init_token_buffers(self, max_len: int): + stt_model = self.model.stt_model + gen_text = torch.full((1, max_len), stt_model.text_pad_id, device=self.device, dtype=torch.long) + gen_asr_text = None + if stt_model.predict_user_text: + gen_asr_text = torch.full((1, max_len), stt_model.text_pad_id, device=self.device, dtype=torch.long) + gen_function = None + if stt_model.use_function_head: + gen_function = torch.full((1, max_len), stt_model.text_pad_id, device=self.device, dtype=torch.long) + return gen_text, gen_asr_text, gen_function + + def _prepare_tts_initial_state(self): + if not self.decode_audio: + return + if not hasattr(self.model, "tts_model"): + return + + logging.info("Preparing TTS warmup state...") + + if self.speaker_name not in self.model.tts_model.audio_prompt_latents: + registered = list(self.model.tts_model.audio_prompt_latents.keys()) + raise ValueError( + f"Unknown speaker_name {self.speaker_name!r}. Registered speakers: " + f"{registered or '(none)'}. Register speakers at export with " + "register_speaker_dict, then pass s2s.speaker_name." + ) + logging.info(f"Using registered speaker name: {self.speaker_name}") + + self.model.tts_model.set_init_inputs( + speaker_audio=None, + speaker_audio_lens=None, + speaker_name=self.speaker_name, + ) + init_inputs = self.model.tts_model.get_init_inputs(B=1) + + self.generation_config = self.model.tts_model._get_generation_config(guidance_enabled=True) + init_inputs.update({"use_cache": True, "past_key_values": None, "guidance_enabled": True}) + + with torch.no_grad(): + outputs = self.model_eartts_interface.prefill_prompt( + init_inputs, + prompt_token_ids=None, + request_id="tts_warmup", + ) + self.model_eartts_interface.abort_request("tts_warmup") + + code = init_inputs["code"][:, -1:] + + self.first_context_subword_id = init_inputs["subword_ids"][:, -1].unsqueeze(-1) + self.first_tts_code_input = code.detach().clone() + self.first_tts_past_key_values_input = self._clone_cache(outputs.past_key_values) + # The backend needs the frame-0 context token and generation config to + # run its per-frame step; the wrapper keeps the initial cache/codes + # because those are seeded per stream in create_decode_state. + self.model_eartts_interface.set_warmup_state( + self.first_context_subword_id, + self.generation_config, + ) + + logging.info("TTS warmup state prepared") + + def create_decode_state(self, max_len: int) -> StreamingDecodeState: + gen_text, gen_asr_text, gen_function = self._init_token_buffers(max_len) + + llm_cache = None if self.use_vllm_llm else self.model_llm_interface.create_cache() + # One codec for every backend pairing, so no branch on the TTS engine. + subword_mask, tts_codec_cache = self.codec.create_state(max_len) if self.decode_audio else (None, None) + perception_cache = None + if self.use_perception_cache and self.perception_cache_mgr is not None: + perception_cache = self.perception_cache_mgr.get_initial_state(batch_size=1) + + tts_past_key_values = None + tts_code = None + if self.decode_audio and not self.use_vllm_tts and self.first_tts_code_input is not None: + tts_past_key_values = self._clone_cache(self.first_tts_past_key_values_input) + tts_code = self.first_tts_code_input.detach().clone() + + return StreamingDecodeState( + frame_idx=0, + gen_text=gen_text, + gen_asr_text=gen_asr_text, + gen_function=gen_function, + input_embeds_history=[], + llm_cache=llm_cache, + tts_past_key_values=tts_past_key_values, + tts_code=tts_code, + subword_mask=subword_mask, + perception_cache=perception_cache, + tts_codec_cache=tts_codec_cache, + llm_cache_position_offset=0, + timing=TimingSummary() if self._profile_timing else NullTimingSummary(), + omni_session=None, + ) + + def infer_one_step( + self, + audio_input: torch.Tensor, + num_frames_per_chunk: int, + state: StreamingDecodeState, + *, + request_id: str | None = None, + has_prompt: bool = False, + return_debug: bool = False, + sampling_params: dict[str, float] | None = None, + ) -> InferenceStepResult: + """Run one streaming inference step: perception -> LLM -> TTS -> audio decode. + + All mutable decode state (caches, gen_text, gen_asr_text, code, etc.) is + updated **in-place** on *state*. The returned :class:`InferenceStepResult` + carries only per-step outputs needed by the pipeline. + + Args: + audio_input (torch.Tensor): Raw audio tensor for this chunk, shape ``(1, samples)``. + num_frames_per_chunk (int): Number of 80 ms frames in this chunk. + state (StreamingDecodeState): Mutable decode state (KV caches, token workspaces, etc.). + request_id (str | None): Unique ID for this stream (used by vLLM engines). + has_prompt (bool): Whether the LLM state already contains a prefilled + system prompt. Affects the first-frame embedding (PAD vs BOS). + return_debug (bool): If True, attach per-step debug info to the result. + sampling_params (dict[str, float] | None): Optional per-stream sampling overrides + (``top_p``, ``temperature``, ``repetition_penalty``). + Keys that are absent fall back to the pipeline-level defaults. + """ + effective_request_id = request_id or self.request_id + frame_idx = state.frame_idx + + state.timing.start("total_step") + has_llm_cache = state.llm_cache is not None + B = state.gen_text.shape[0] + if B != 1 and self.use_vllm_omni: + # The session API is per-stream and its steps take scalars. + raise ValueError(f"vllm_omni components only support batch size 1 (got gen_text batch={B}).") + + def per_step_tokens(channel: torch.Tensor | None) -> torch.Tensor | None: + """Pad-filled buffer for one channel, or None if the channel is absent. + + Pad rather than empty: a frame the backend reports no token for + then decodes to nothing instead of to uninitialized memory. + """ + if channel is None: + return None + return torch.full( + (B, num_frames_per_chunk), + self.model.stt_model.text_pad_id, + dtype=state.gen_text.dtype, + device=state.gen_text.device, + ) + + predicted_tokens = per_step_tokens(state.gen_text) + asr_predicted_tokens = per_step_tokens(state.gen_asr_text) + function_predicted_tokens = per_step_tokens(state.gen_function) + + debug_logger = IntermediateResultLogger() if return_debug else NullIntermediateResultLogger() + + # --- Stage 1: Perception --- + state.timing.start("perception") + source_encoded, state.perception_cache = self._run_perception( + audio_input, + frame_idx, + num_frames_per_chunk, + state.perception_cache, + ) + state.timing.stop("perception") + base_frame_index = self._base_frame_index(source_encoded, state, num_frames_per_chunk) + + # --- Stage 2: Per-frame generation loop --- + new_codes_for_decode = [] + for frame_offset in range(num_frames_per_chunk): + current_frame_idx = frame_idx + frame_offset + current_frame_index = min(base_frame_index + frame_offset, source_encoded.shape[1] - 1) + debug_logger.log_selected_frame_index(current_frame_index) + frame_embedding = source_encoded[:, current_frame_index : current_frame_index + 1, :] + + ans = self._run_llm_step( + frame_embedding, + state, + frame_offset=frame_offset, + current_frame_idx=current_frame_idx, + has_prompt=has_prompt, + return_debug=return_debug, + sampling_params=sampling_params, + debug_logger=debug_logger, + ) + + if ans.text_logits is not None: + debug_logger.log_text_logits(ans.text_logits[:, -1]) + if ans.asr_logits is not None: + debug_logger.log_asr_logits(ans.asr_logits[:, -1]) + + state.gen_text[:, current_frame_idx] = ans.predicted_token + if state.gen_asr_text is not None: + asr_token = ans.asr_predicted_token + if asr_token is None and self.output_capabilities.has_asr_head: + raise RuntimeError("Checkpoint has an ASR head but the LLM backend returned no ASR token") + if asr_token is not None: + state.gen_asr_text[:, current_frame_idx] = asr_token + asr_predicted_tokens[:, frame_offset] = asr_token + self.model.stt_model.streaming_inference._maybe_apply_forced_turn_taking( + current_frame_idx, state.gen_text, state.gen_asr_text + ) + if state.gen_function is not None: + function_token = ans.function_predicted_token + if function_token is None and self.output_capabilities.has_function_head: + raise RuntimeError("Checkpoint has a function head but the LLM backend returned no function token") + if function_token is not None: + state.gen_function[:, current_frame_idx] = function_token + function_predicted_tokens[:, frame_offset] = function_token + # Read back rather than reusing ans: forced turn-taking above may + # have rewritten this frame's text token in place. + predicted_tokens[:, frame_offset] = state.gen_text[:, current_frame_idx] + + if self.decode_audio: + new_code = self._run_tts_step( + state, + current_frame_idx, + effective_request_id, + ) + new_codes_for_decode.append(new_code) + + # --- Stage 3: Audio decode --- + # No-op when self.decode_audio is False: _decode_audio returns None immediately. + decoded_audio_new = self._decode_audio(new_codes_for_decode, state, frame_idx, num_frames_per_chunk) + + # --- Stage 4: Token -> string conversion --- + predicted_text_strs = self._tokens_to_strings(predicted_tokens) + asr_predicted_text_strs = ( + self._tokens_to_strings(asr_predicted_tokens) if asr_predicted_tokens is not None else None + ) + predicted_function_strs = ( + self._tokens_to_strings(function_predicted_tokens) if function_predicted_tokens is not None else None + ) + + logging.debug(f"frame {frame_idx}: USER asr: {asr_predicted_text_strs}") + logging.debug(f"frame {frame_idx}: FUNCTION: {predicted_function_strs}") + logging.debug(f"frame {frame_idx}: AGENT txt: {predicted_text_strs}") + + # --- Update remaining state fields --- + # `input_embeds_history` is extended by the native no-cache backend as + # it goes; see PyTorchLLM.step. + if has_llm_cache: + state.llm_cache_position_offset += num_frames_per_chunk + + state.timing.stop("total_step") + + debug = debug_logger.build_debug_dict(source_encoded, state.gen_text, state.gen_asr_text) + + return InferenceStepResult( + predicted_text_tokens=predicted_tokens, + asr_predicted_text_tokens=asr_predicted_tokens, + predicted_text_strs=predicted_text_strs, + asr_predicted_text_strs=asr_predicted_text_strs, + predicted_function_tokens=function_predicted_tokens, + predicted_function_strs=predicted_function_strs, + decoded_audio=decoded_audio_new, + debug=debug, + ) + + # ------------------------------------------------------------------ + # infer_one_step sub-stages + # ------------------------------------------------------------------ + + def _run_llm_step( + self, + frame_embedding: torch.Tensor, + state: StreamingDecodeState, + *, + frame_offset: int, + current_frame_idx: int, + has_prompt: bool, + return_debug: bool, + sampling_params: dict[str, float] | None, + debug_logger, + ) -> LlmStepResult: + """Time one :meth:`DuplexLLM.step` on the selected LLM backend. + + Both backends fill the same :class:`LlmStepResult`, so the caller does + not need to know which one ran; the optional fields are None when the + checkpoint has no such head or the backend cannot expose logits. + """ + state.timing.start("stt_model") + try: + return self.llm_backend.step( + frame_embedding, + state, + frame_offset=frame_offset, + current_frame_idx=current_frame_idx, + has_prompt=has_prompt, + return_debug=return_debug, + sampling_params=sampling_params, + debug_logger=debug_logger, + ) + finally: + state.timing.stop("stt_model") + + def _run_tts_step( + self, + state: StreamingDecodeState, + current_frame_idx: int, + request_id: str, + ) -> torch.Tensor: + """Time one :meth:`DuplexTTS.step` on the selected TTS backend. + + Returns this frame's codes as ``(B, T, num_quantizers)`` for the shared + native codec to decode. Both backends read the committed text token from + ``state.gen_text``, so a forced-turn-taking rewrite reaches TTS without + the caller having to pass it. + + ``inference_force_speech_silence_on_eos`` belongs to the backends, not + here: each applies it internally to the acoustic input of the step whose + text token is EOS -- natively in ``DuplexEARTTS.infer_codes_one_step``, + unconditionally in the vLLM preprocess. + """ + state.timing.start("tts_model") + try: + return self.tts_backend.step(state, current_frame_idx, request_id) + finally: + state.timing.stop("tts_model") + + def _decode_audio( + self, + new_codes_for_decode: list[torch.Tensor], + state: StreamingDecodeState, + frame_idx: int, + num_frames_per_chunk: int, + ) -> torch.Tensor | None: + """Decode accumulated TTS codes into a waveform. + + Returns the decoded audio tensor or *None* when ``decode_audio`` + is disabled or no codes were produced. + """ + if not self.decode_audio or not new_codes_for_decode: + return None + + logging.debug(f"Decoding audio for {frame_idx}-th frame ({num_frames_per_chunk=})") + + state.timing.start("audio_codec") + try: + return self.codec.decode(new_codes_for_decode, state.tts_codec_cache) + finally: + state.timing.stop("audio_codec") + + def _base_frame_index( + self, + source_encoded: torch.Tensor, + state: StreamingDecodeState, + num_frames_per_chunk: int, + ) -> int: + """Index of the first encoded frame belonging to this chunk.""" + if ( + self.use_perception_cache + and state.perception_cache is not None + and state.perception_cache.is_initialized() + ): + # With cache: we get exactly num_frames_per_chunk output frames + return 0 + # Without cache: use the second-to-last encoded frame as the + # "newest" because the model expects 10ms / 80ms / 80ms ... framing + # but we always feed 80ms chunks, so the final frame contains + # silence padding. + newest = source_encoded.shape[1] - 2 + return max(newest - (num_frames_per_chunk - 1), 0) + + def _run_perception( + self, + audio_input: torch.Tensor, + frame_idx: int, + num_frames_per_chunk: int, + perception_cache: PerceptionCacheState | None, + ) -> tuple[torch.Tensor, PerceptionCacheState | None]: + """Run the perception encoder and return (source_encoded, updated_cache).""" + if self.use_perception_cache and perception_cache is not None and perception_cache.is_initialized(): + source_encoded, perception_cache = self.perception_cache_mgr.step( + audio_input=audio_input, + frame_idx=frame_idx, + num_frames_per_chunk=num_frames_per_chunk, + perception_cache=perception_cache, + ) + else: + buffer_len = torch.tensor([audio_input.shape[1]], dtype=torch.long, device=self.device) + source_encoded, _, _ = self.model.stt_model.perception( + input_signal=audio_input, + input_signal_length=buffer_len, + return_encoder_emb=True, + ) + + source_encoded = source_encoded.to(self.dtype) + return source_encoded, perception_cache + + def _tokens_to_strings(self, token_ids: torch.Tensor) -> list[str]: + """Convert a [B, T] tensor of token IDs to a list of strings. + + Uses ``_decode_tokens_with_specials`` so byte-level BPE is decoded + properly (e.g. ``âĢĻ`` -> ``'``) via HF ``convert_tokens_to_string``. + + Leading spaces are preserved in the output: in byte-level BPE, + word-initial tokens carry a space prefix that ``convert_tokens_to_string`` + keeps intact. So callers can concatenate successive chunk strings to + recover properly spaced text. A leading space means "new word"; no + leading space means the token continues the previous word. For + example, three chunks producing ``"Hi"``, ``" how can"``, + ``" I help"`` concatenate to ``"Hi how can I help"`` (not + ``"Hihow canI help"``). + + NOTE: multi-byte UTF-8 characters whose BPE tokens span two frames + will show as replacement chars (U+FFFD) because each frame is decoded + independently. + """ + pad_token_str = self.tokenizer.ids_to_tokens([self.model.stt_model.text_pad_id])[0] + result = [] + for tok_ids_b in token_ids: + toks = self.tokenizer.ids_to_tokens(tok_ids_b.tolist()) + result.append( + _decode_tokens_with_specials( + toks, + self.tokenizer, + pad_token_str=pad_token_str, + keep_pad=False, + ) + ) + return result + + @property + def special_token_strings(self) -> set[str]: + """Token strings that should be stripped from decoded text for clean output.""" + stt = self.model.stt_model + return get_special_token_strings(stt.tokenizer, stt.text_pad_id, model_cfg=stt.cfg) diff --git a/nemo/collections/speechlm2/inference/model_wrappers/perception_cache.py b/nemo/collections/speechlm2/inference/model_wrappers/perception_cache.py new file mode 100644 index 000000000000..9ebc68f79c1c --- /dev/null +++ b/nemo/collections/speechlm2/inference/model_wrappers/perception_cache.py @@ -0,0 +1,573 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Cache-aware perception encoder for streaming S2S inference. + +Provides incremental mel-spectrogram encoding with optional CUDA graph +acceleration, so that only new audio needs to be processed each step +instead of re-encoding the entire buffer. +""" + +import copy +from dataclasses import dataclass + +import torch +from omegaconf import OmegaConf + +from nemo.utils import logging + + +@dataclass +class PerceptionCacheState: + """Cache state for streaming perception inference. + + Holds the cache tensors for the ASR encoder used in the perception module. + This enables cache-aware streaming inference without needing the full audio buffer. + """ + + cache_last_channel: torch.Tensor | None = None + cache_last_time: torch.Tensor | None = None + cache_last_channel_len: torch.Tensor | None = None + + def is_initialized(self) -> bool: + """Check if the cache has been initialized.""" + return None not in [self.cache_last_channel, self.cache_last_time, self.cache_last_channel_len] + + +@dataclass +class PerceptionCUDAGraphState: + """State for CUDA graph-accelerated perception encoder. + + Holds separate graphs for first chunk (different size) and subsequent chunks. + Also holds static buffers for inputs/outputs to enable graph replay. + """ + + # CUDA graphs + graph_first: torch.cuda.CUDAGraph | None = None + graph_subsequent: torch.cuda.CUDAGraph | None = None + + # Static input buffers (for copying data before graph replay) + static_mel_first: torch.Tensor | None = None + static_mel_subsequent: torch.Tensor | None = None + static_mel_len_first: torch.Tensor | None = None + static_mel_len_subsequent: torch.Tensor | None = None + + # Static cache input buffers + static_cache_channel_in: torch.Tensor | None = None + static_cache_time_in: torch.Tensor | None = None + static_cache_channel_len_in: torch.Tensor | None = None + + # Static output buffers (results are written here during replay) + static_encoded_first: torch.Tensor | None = None + static_encoded_subsequent: torch.Tensor | None = None + static_encoded_len_first: torch.Tensor | None = None + static_encoded_len_subsequent: torch.Tensor | None = None + + # Static cache output buffers - SEPARATE for first and subsequent graphs + # (each graph writes to its own output tensors during replay) + static_cache_channel_out_first: torch.Tensor | None = None + static_cache_time_out_first: torch.Tensor | None = None + static_cache_channel_len_out_first: torch.Tensor | None = None + static_cache_channel_out_subsequent: torch.Tensor | None = None + static_cache_time_out_subsequent: torch.Tensor | None = None + static_cache_channel_len_out_subsequent: torch.Tensor | None = None + + def is_captured(self) -> bool: + """Check if graphs have been captured.""" + return self.graph_first is not None and self.graph_subsequent is not None + + +class PerceptionCacheManager: + """Manages cache-aware streaming perception encoding with optional CUDA graphs. + + This class encapsulates all perception cache setup, CUDA graph capture, + and the incremental encoding step. It is created by the inference wrapper + when ``use_perception_cache=True``. + """ + + def __init__(self, model, device: torch.device, dtype: torch.dtype, use_cudagraph: bool = False): + self.model = model + self.device = device + self.dtype = dtype + self.use_cudagraph = use_cudagraph + + self.streaming_cfg = None + self.preprocessor = None + self.subsampling_factor = None + self.input_features = None + self.sampling_frames = None + self.cudagraph_state: PerceptionCUDAGraphState | None = None + + def setup(self) -> bool: + """Setup cache-aware streaming for the perception encoder. + + Returns: + True if setup succeeded, False if the encoder doesn't support streaming. + """ + from nemo.collections.asr.parts.mixins.streaming import StreamingEncoder + + perception = self.model.stt_model.perception + encoder = perception.encoder + + if not isinstance(encoder, StreamingEncoder): + logging.warning("Perception encoder does not support streaming. Disabling perception cache.") + return False + + if encoder.streaming_cfg is None: + encoder.setup_streaming_params() + + self.streaming_cfg = encoder.streaming_cfg + + cfg = copy.deepcopy(perception.cfg) + OmegaConf.set_struct(cfg.preprocessor, False) + cfg.preprocessor.dither = 0.0 + cfg.preprocessor.pad_to = 0 + + self.preprocessor = perception.from_config_dict(cfg.preprocessor) + self.preprocessor.to(self.device) + + self.subsampling_factor = encoder.subsampling_factor + self.input_features = encoder._feat_in + + if hasattr(encoder, "pre_encode") and hasattr(encoder.pre_encode, "get_sampling_frames"): + self.sampling_frames = encoder.pre_encode.get_sampling_frames() + else: + self.sampling_frames = None + + logging.info("Perception cache setup complete:") + logging.info( + f" Streaming config: chunk_size={self.streaming_cfg.chunk_size}, " + f"shift_size={self.streaming_cfg.shift_size}" + ) + logging.info(f" Pre-encode cache size: {self.streaming_cfg.pre_encode_cache_size}") + logging.info(f" Subsampling factor: {self.subsampling_factor}") + + if self.use_cudagraph: + logging.info(" Setting up CUDA graphs for perception encoder...") + self.capture_cudagraphs() + logging.info(" CUDA graphs captured") + + return True + + def capture_cudagraphs(self): + """Capture CUDA graphs for perception encoder with both chunk sizes. + + Note: "chunk" in the streaming encoder config (chunk_size, shift_size, etc.) + follows NeMo's cache-aware streaming encoder API and is measured in + mel-spectrogram time-steps, not audio samples or seconds. + """ + encoder = self.model.stt_model.perception.encoder + perception = self.model.stt_model.perception + streaming_cfg = self.streaming_cfg + + if isinstance(streaming_cfg.chunk_size, list): + chunk_size_first = streaming_cfg.chunk_size[0] + chunk_size_subsequent = streaming_cfg.chunk_size[1] + else: + chunk_size_first = streaming_cfg.chunk_size + chunk_size_subsequent = streaming_cfg.chunk_size + + if isinstance(streaming_cfg.pre_encode_cache_size, list): + pre_encode_cache_first = streaming_cfg.pre_encode_cache_size[0] + pre_encode_cache_subsequent = streaming_cfg.pre_encode_cache_size[1] + else: + pre_encode_cache_first = streaming_cfg.pre_encode_cache_size + pre_encode_cache_subsequent = streaming_cfg.pre_encode_cache_size + + mel_len_first = chunk_size_first + pre_encode_cache_first + mel_len_subsequent = chunk_size_subsequent + pre_encode_cache_subsequent + + logging.info(f" CUDA graph mel lengths: first={mel_len_first}, subsequent={mel_len_subsequent}") + + cache_last_channel, cache_last_time, cache_last_channel_len = encoder.get_initial_cache_state(batch_size=1) + + state = PerceptionCUDAGraphState() + + state.static_mel_first = torch.zeros( + (1, self.input_features, mel_len_first), dtype=torch.float32, device=self.device + ) + state.static_mel_subsequent = torch.zeros( + (1, self.input_features, mel_len_subsequent), dtype=torch.float32, device=self.device + ) + state.static_mel_len_first = torch.tensor([mel_len_first], dtype=torch.long, device=self.device) + state.static_mel_len_subsequent = torch.tensor([mel_len_subsequent], dtype=torch.long, device=self.device) + + if cache_last_channel is not None: + state.static_cache_channel_in = cache_last_channel.clone() + if cache_last_time is not None: + state.static_cache_time_in = cache_last_time.clone() + if cache_last_channel_len is not None: + state.static_cache_channel_len_in = cache_last_channel_len.clone() + + logging.info(" Warming up encoder for CUDA graph capture...") + # PyTorch recommends a few eager warmup iterations before CUDA graph + # capture on a side stream; its example uses three iterations: + # https://pytorch.org/docs/stable/notes/cuda.html#cuda-graphs + warmup_stream = torch.cuda.Stream(device=self.device) + warmup_stream.wait_stream(torch.cuda.current_stream(self.device)) + with torch.cuda.stream(warmup_stream), torch.no_grad(): + for _ in range(3): + _ = encoder.cache_aware_stream_step( + processed_signal=state.static_mel_first, + processed_signal_length=state.static_mel_len_first, + cache_last_channel=( + state.static_cache_channel_in.clone() if state.static_cache_channel_in is not None else None + ), + cache_last_time=( + state.static_cache_time_in.clone() if state.static_cache_time_in is not None else None + ), + cache_last_channel_len=( + state.static_cache_channel_len_in.clone() + if state.static_cache_channel_len_in is not None + else None + ), + keep_all_outputs=True, + drop_extra_pre_encoded=0, + ) + _ = encoder.cache_aware_stream_step( + processed_signal=state.static_mel_subsequent, + processed_signal_length=state.static_mel_len_subsequent, + cache_last_channel=( + state.static_cache_channel_in.clone() if state.static_cache_channel_in is not None else None + ), + cache_last_time=( + state.static_cache_time_in.clone() if state.static_cache_time_in is not None else None + ), + cache_last_channel_len=( + state.static_cache_channel_len_in.clone() + if state.static_cache_channel_len_in is not None + else None + ), + keep_all_outputs=True, + drop_extra_pre_encoded=streaming_cfg.drop_extra_pre_encoded, + ) + torch.cuda.current_stream(self.device).wait_stream(warmup_stream) + + # Capture graph for FIRST chunk + logging.info(f" Capturing CUDA graph for first chunk (mel_len={mel_len_first})...") + state.graph_first = torch.cuda.CUDAGraph() + + if state.static_cache_channel_in is not None: + state.static_cache_channel_in.copy_(cache_last_channel) + if state.static_cache_time_in is not None: + state.static_cache_time_in.copy_(cache_last_time) + if state.static_cache_channel_len_in is not None: + state.static_cache_channel_len_in.copy_(cache_last_channel_len) + + with torch.cuda.graph(state.graph_first): + ( + encoded_first, + encoded_len_first, + cache_channel_out_first, + cache_time_out_first, + cache_channel_len_out_first, + ) = encoder.cache_aware_stream_step( + processed_signal=state.static_mel_first, + processed_signal_length=state.static_mel_len_first, + cache_last_channel=state.static_cache_channel_in, + cache_last_time=state.static_cache_time_in, + cache_last_channel_len=state.static_cache_channel_len_in, + keep_all_outputs=True, + drop_extra_pre_encoded=0, + ) + encoded_adapted_first, _ = perception.modality_adapter( + audio_signal=encoded_first, length=encoded_len_first + ) + encoded_chunk_first = perception.proj(encoded_adapted_first.transpose(1, 2)) + + state.static_encoded_first = encoded_chunk_first + state.static_encoded_len_first = encoded_len_first + state.static_cache_channel_out_first = cache_channel_out_first + state.static_cache_time_out_first = cache_time_out_first + state.static_cache_channel_len_out_first = cache_channel_len_out_first + + # Capture graph for SUBSEQUENT chunks + logging.info(f" Capturing CUDA graph for subsequent chunks (mel_len={mel_len_subsequent})...") + state.graph_subsequent = torch.cuda.CUDAGraph() + + if state.static_cache_channel_in is not None: + state.static_cache_channel_in.copy_(cache_last_channel) + if state.static_cache_time_in is not None: + state.static_cache_time_in.copy_(cache_last_time) + if state.static_cache_channel_len_in is not None: + state.static_cache_channel_len_in.copy_(cache_last_channel_len) + + with torch.cuda.graph(state.graph_subsequent): + ( + encoded_subsequent, + encoded_len_subsequent, + cache_channel_out_subsequent, + cache_time_out_subsequent, + cache_channel_len_out_subsequent, + ) = encoder.cache_aware_stream_step( + processed_signal=state.static_mel_subsequent, + processed_signal_length=state.static_mel_len_subsequent, + cache_last_channel=state.static_cache_channel_in, + cache_last_time=state.static_cache_time_in, + cache_last_channel_len=state.static_cache_channel_len_in, + keep_all_outputs=True, + drop_extra_pre_encoded=streaming_cfg.drop_extra_pre_encoded, + ) + encoded_adapted_subsequent, _ = perception.modality_adapter( + audio_signal=encoded_subsequent, length=encoded_len_subsequent + ) + encoded_chunk_subsequent = perception.proj(encoded_adapted_subsequent.transpose(1, 2)) + + state.static_encoded_subsequent = encoded_chunk_subsequent + state.static_encoded_len_subsequent = encoded_len_subsequent + state.static_cache_channel_out_subsequent = cache_channel_out_subsequent + state.static_cache_time_out_subsequent = cache_time_out_subsequent + state.static_cache_channel_len_out_subsequent = cache_channel_len_out_subsequent + + self.cudagraph_state = state + logging.info(" CUDA graphs captured successfully") + + def get_initial_state(self, batch_size: int = 1) -> PerceptionCacheState: + """Get initial cache state for perception encoder.""" + encoder = self.model.stt_model.perception.encoder + cache_last_channel, cache_last_time, cache_last_channel_len = encoder.get_initial_cache_state( + batch_size=batch_size + ) + + return PerceptionCacheState( + cache_last_channel=cache_last_channel, + cache_last_time=cache_last_time, + cache_last_channel_len=cache_last_channel_len, + ) + + def step( + self, + audio_input: torch.Tensor, + frame_idx: int, + num_frames_per_chunk: int, + perception_cache: PerceptionCacheState, + ) -> tuple[torch.Tensor, PerceptionCacheState]: + """ + Perform cache-aware perception encoding for streaming inference. + + Note: "chunk" in this method (chunk_size, mel_chunk, etc.) follows NeMo's + cache-aware streaming encoder API and is measured in mel-spectrogram time-steps, + not audio samples or seconds. + + This method computes the full mel spectrogram from the audio buffer, then slices + it appropriately based on the frame index. It supports processing multiple + "base steps" in a single call, where each base step processes (lookahead + 1) frames. + + Processing logic per sub-step: + - First sub-step (sub_frame_idx == 0): take first chunk_size_first columns, + prepend zeros for pre_encode_cache + - Subsequent sub-steps (sub_frame_idx > 0): take chunk_size columns starting from + (shift_size_first + (step_number-1)*shift_size), prepend pre_encode_cache_size + columns from mel spec + + The method loops over sub-steps, running the encoder for each and concatenating + the outputs. This allows num_frames_per_chunk to be a multiple of (lookahead + 1). + + Args: + audio_input: Audio buffer tensor [B, T] (full buffer with all samples) + frame_idx: Current frame index in the stream + num_frames_per_chunk: Number of 80ms frames to process. Must be a multiple + of (lookahead + 1), i.e., encoder._cfg.att_context_size[1] + 1 + perception_cache: Current cache state containing encoder caches + + Returns: + Tuple of (encoded_output [B, T_out, D], updated_perception_cache) + where T_out = num_frames_per_chunk (one output frame per input frame) + """ + perception = self.model.stt_model.perception + encoder = perception.encoder + streaming_cfg = self.streaming_cfg + + audio_len = torch.tensor([audio_input.shape[1]], dtype=torch.long, device=self.device) + processed_signal, _ = self.preprocessor( + input_signal=audio_input, + length=audio_len, + ) + + if isinstance(streaming_cfg.chunk_size, list): + chunk_size_first = streaming_cfg.chunk_size[0] + chunk_size = streaming_cfg.chunk_size[1] + else: + chunk_size_first = streaming_cfg.chunk_size + chunk_size = streaming_cfg.chunk_size + + if isinstance(streaming_cfg.shift_size, list): + shift_size_first = streaming_cfg.shift_size[0] + shift_size = streaming_cfg.shift_size[1] + else: + shift_size_first = streaming_cfg.shift_size + shift_size = streaming_cfg.shift_size + + if isinstance(streaming_cfg.pre_encode_cache_size, list): + pre_encode_cache_size_first = streaming_cfg.pre_encode_cache_size[0] + pre_encode_cache_size = streaming_cfg.pre_encode_cache_size[1] + else: + pre_encode_cache_size_first = streaming_cfg.pre_encode_cache_size + pre_encode_cache_size = streaming_cfg.pre_encode_cache_size + + cache_last_channel = perception_cache.cache_last_channel + cache_last_time = perception_cache.cache_last_time + cache_last_channel_len = perception_cache.cache_last_channel_len + + base_step_size = encoder._cfg.att_context_size[1] + 1 + if num_frames_per_chunk % base_step_size != 0: + raise ValueError( + f"num_frames_per_chunk must be a multiple of (lookahead + 1) = {base_step_size}. " + f"Got num_frames_per_chunk={num_frames_per_chunk}" + ) + num_sub_steps = num_frames_per_chunk // base_step_size + + encoded_chunks = [] + + for sub_step in range(num_sub_steps): + sub_frame_idx = frame_idx + (sub_step * base_step_size) + is_first_sub_step = sub_frame_idx == 0 + + if is_first_sub_step: + cur_chunk_size = chunk_size_first + cur_pre_encode_cache_size = pre_encode_cache_size_first + drop_extra_pre_encoded = 0 + + mel_chunk = processed_signal[:, :, :cur_chunk_size] + + if cur_pre_encode_cache_size > 0: + zeros_pad = torch.zeros( + (processed_signal.size(0), self.input_features, cur_pre_encode_cache_size), + device=self.device, + dtype=processed_signal.dtype, + ) + mel_chunk = torch.cat([zeros_pad, mel_chunk], dim=-1) + else: + cur_chunk_size = chunk_size + cur_pre_encode_cache_size = pre_encode_cache_size + drop_extra_pre_encoded = streaming_cfg.drop_extra_pre_encoded + + mel_T = processed_signal.shape[-1] + + step_number = sub_frame_idx // base_step_size + chunk_start = shift_size_first + (step_number - 1) * shift_size + chunk_end = chunk_start + cur_chunk_size + + offset = chunk_size - shift_size_first + if chunk_end > mel_T - offset: + sub_steps_remaining = num_sub_steps - 1 - sub_step + chunk_end = mel_T - offset - sub_steps_remaining * shift_size + chunk_start = chunk_end - cur_chunk_size + + main_chunk = processed_signal[:, :, chunk_start:chunk_end] + + cache_start = max(0, chunk_start - cur_pre_encode_cache_size) + cache_mel = processed_signal[:, :, cache_start:chunk_start] + + if cache_mel.shape[-1] < cur_pre_encode_cache_size: + zeros_pad = torch.zeros( + (cache_mel.size(0), cache_mel.size(1), cur_pre_encode_cache_size - cache_mel.shape[-1]), + device=self.device, + dtype=cache_mel.dtype, + ) + cache_mel = torch.cat([zeros_pad, cache_mel], dim=-1) + + mel_chunk = torch.cat([cache_mel, main_chunk], dim=-1) + + chunk_lengths = torch.tensor([mel_chunk.shape[-1]], dtype=torch.long, device=self.device) + + if self.use_cudagraph and self.cudagraph_state is not None and self.cudagraph_state.is_captured(): + graph_state = self.cudagraph_state + + if is_first_sub_step: + graph_state.static_mel_first.copy_(mel_chunk) + else: + graph_state.static_mel_subsequent.copy_(mel_chunk) + + if graph_state.static_cache_channel_in is not None and cache_last_channel is not None: + graph_state.static_cache_channel_in.copy_(cache_last_channel) + if graph_state.static_cache_time_in is not None and cache_last_time is not None: + graph_state.static_cache_time_in.copy_(cache_last_time) + if graph_state.static_cache_channel_len_in is not None and cache_last_channel_len is not None: + graph_state.static_cache_channel_len_in.copy_(cache_last_channel_len) + + if is_first_sub_step: + graph_state.graph_first.replay() + encoded_chunk = graph_state.static_encoded_first.clone() + cache_last_channel = ( + graph_state.static_cache_channel_out_first.clone() + if graph_state.static_cache_channel_out_first is not None + else None + ) + cache_last_time = ( + graph_state.static_cache_time_out_first.clone() + if graph_state.static_cache_time_out_first is not None + else None + ) + cache_last_channel_len = ( + graph_state.static_cache_channel_len_out_first.clone() + if graph_state.static_cache_channel_len_out_first is not None + else None + ) + else: + graph_state.graph_subsequent.replay() + encoded_chunk = graph_state.static_encoded_subsequent.clone() + cache_last_channel = ( + graph_state.static_cache_channel_out_subsequent.clone() + if graph_state.static_cache_channel_out_subsequent is not None + else None + ) + cache_last_time = ( + graph_state.static_cache_time_out_subsequent.clone() + if graph_state.static_cache_time_out_subsequent is not None + else None + ) + cache_last_channel_len = ( + graph_state.static_cache_channel_len_out_subsequent.clone() + if graph_state.static_cache_channel_len_out_subsequent is not None + else None + ) + + else: + ( + encoded, + encoded_len, + cache_last_channel, + cache_last_time, + cache_last_channel_len, + ) = encoder.cache_aware_stream_step( + processed_signal=mel_chunk, + processed_signal_length=chunk_lengths, + cache_last_channel=cache_last_channel, + cache_last_time=cache_last_time, + cache_last_channel_len=cache_last_channel_len, + keep_all_outputs=True, + drop_extra_pre_encoded=drop_extra_pre_encoded, + ) + + modality_adapter = perception.modality_adapter + encoded_adapted, _ = modality_adapter(audio_signal=encoded, length=encoded_len) + + encoded_chunk = perception.proj(encoded_adapted.transpose(1, 2)) + + encoded_chunks.append(encoded_chunk) + + if len(encoded_chunks) > 1: + encoded_chunk = torch.cat(encoded_chunks, dim=1) + else: + encoded_chunk = encoded_chunks[0] + + new_perception_cache = PerceptionCacheState( + cache_last_channel=cache_last_channel, + cache_last_time=cache_last_time, + cache_last_channel_len=cache_last_channel_len, + ) + + return encoded_chunk, new_perception_cache diff --git a/nemo/collections/speechlm2/inference/model_wrappers/text_sampling.py b/nemo/collections/speechlm2/inference/model_wrappers/text_sampling.py new file mode 100644 index 000000000000..19017471597a --- /dev/null +++ b/nemo/collections/speechlm2/inference/model_wrappers/text_sampling.py @@ -0,0 +1,96 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared text-token sampling for native and vLLM-Omni inference.""" + +import torch + +from nemo.utils import logging + + +def sample_text_token( + logits: torch.Tensor, + generated_tokens: torch.Tensor, + current_step: int, + *, + top_p: float, + repetition_penalty: float, + temperature: float, + special_token_ids: set[int], + special_ids_tensor: torch.Tensor | None = None, +) -> torch.Tensor: + """Sample one token per row with the VoiceChat policy.""" + batch_size, _ = logits.shape + device = logits.device + greedy_tokens = logits.argmax(dim=-1) + + if top_p >= 1.0 and repetition_penalty == 1.0 and temperature in (0.0, 1.0): + return greedy_tokens + if temperature == 0.0: + return greedy_tokens + + sampled_tokens = greedy_tokens.clone() + if special_ids_tensor is not None and special_ids_tensor.device != device: + special_ids_tensor = special_ids_tensor.to(device) + + for batch_idx in range(batch_size): + if greedy_tokens[batch_idx].item() in special_token_ids: + continue + + batch_logits = logits[batch_idx].clone() + if repetition_penalty != 1.0 and current_step > 0: + unique_prev = generated_tokens[batch_idx, :current_step].unique() + if special_ids_tensor is not None: + ids_t = special_ids_tensor + if ids_t.device != unique_prev.device: + ids_t = ids_t.to(unique_prev.device) + unique_prev = unique_prev[~torch.isin(unique_prev, ids_t)] + + if unique_prev.numel() > 0: + if unique_prev.device != batch_logits.device: + unique_prev = unique_prev.to(batch_logits.device) + prev_logits = batch_logits[unique_prev] + batch_logits[unique_prev] = torch.where( + prev_logits > 0, + prev_logits / repetition_penalty, + prev_logits * repetition_penalty, + ) + + if temperature != 1.0: + batch_logits = batch_logits / temperature + + if not torch.isfinite(batch_logits).all(): + logging.warning( + f"sample_text_token: logits contain NaN or inf at step {current_step}, " + f"batch {batch_idx}: nan={batch_logits.isnan().sum().item()}, " + f"inf={batch_logits.isinf().sum().item()}, " + f"min={batch_logits[~batch_logits.isnan()].min().item() if not batch_logits.isnan().all() else 'all_nan'}, " + f"max={batch_logits[~batch_logits.isnan()].max().item() if not batch_logits.isnan().all() else 'all_nan'}" + ) + sampled_tokens[batch_idx] = greedy_tokens[batch_idx] + continue + + if top_p < 1.0: + sorted_logits, sorted_indices = torch.sort(batch_logits, descending=True) + sorted_probs = torch.softmax(sorted_logits, dim=-1) + cumulative_probs = torch.cumsum(sorted_probs, dim=-1) + sorted_indices_to_remove = cumulative_probs > top_p + sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].clone() + sorted_indices_to_remove[0] = False + batch_logits[sorted_indices[sorted_indices_to_remove]] = float("-inf") + + probs = torch.softmax(batch_logits, dim=-1) + sampled_tokens[batch_idx] = torch.multinomial(probs, num_samples=1).item() + + return sampled_tokens diff --git a/nemo/collections/speechlm2/inference/pipelines/__init__.py b/nemo/collections/speechlm2/inference/pipelines/__init__.py new file mode 100644 index 000000000000..9e3fb699d9f6 --- /dev/null +++ b/nemo/collections/speechlm2/inference/pipelines/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/nemo/collections/speechlm2/inference/pipelines/s2s_pipeline_interface.py b/nemo/collections/speechlm2/inference/pipelines/s2s_pipeline_interface.py new file mode 100644 index 000000000000..543b391ed6c8 --- /dev/null +++ b/nemo/collections/speechlm2/inference/pipelines/s2s_pipeline_interface.py @@ -0,0 +1,80 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from abc import ABC, abstractmethod + +from nemo.collections.speechlm2.inference.streaming.framing.s2s_request_options import S2SRequestOptions +from nemo.collections.speechlm2.inference.streaming.state.s2s_streaming_output import S2SStreamingOutput + + +class S2SPipelineInterface(ABC): + """Base class for all streaming S2S pipelines. + + This class is intentionally kept minimal and mirrors the behaviour of + ``BasePipeline`` that is used for streaming ASR pipelines. It + provides an in-memory *state pool* that stores per-stream + :class:`S2SStreamingOutput` objects (accumulated audio, text, and + finalized token fields) required by a concrete pipeline + implementation. Sub-classes are expected to implement + :py:meth:`create_state` to construct a fresh output object. + """ + + def __init__(self) -> None: + self._state_pool: dict[int, S2SStreamingOutput] = {} + + # ------------------------------------------------------------------ + # State helpers + # ------------------------------------------------------------------ + def get_state(self, stream_id: int) -> S2SStreamingOutput | None: + """Return the state object for *stream_id* or *None* if it does not exist.""" + return self._state_pool.get(stream_id, None) + + def delete_state(self, stream_id: int) -> None: + """Delete the state associated with *stream_id* (noop if missing).""" + if stream_id in self._state_pool: + del self._state_pool[stream_id] + + @abstractmethod + def create_state(self, options: S2SRequestOptions | None = None) -> S2SStreamingOutput: + """Create and return a *new*, *empty* state object. + + Args: + options: Per-stream request options (system prompt, sampling + overrides, etc.). Stored on the state so they can be + consulted throughout the stream's lifetime. + """ + raise NotImplementedError + + def get_or_create_state(self, stream_id: int, options: S2SRequestOptions | None = None) -> S2SStreamingOutput: + """Return existing state for *stream_id* or create a new one via :py:meth:`create_state`.""" + if stream_id not in self._state_pool: + self._state_pool[stream_id] = self.create_state(options) + return self._state_pool[stream_id] + + # ------------------------------------------------------------------ + # Session helpers – identical to *BasePipeline* + # ------------------------------------------------------------------ + def reset_session(self) -> None: + """Clear the internal *state pool* – effectively resetting the pipeline.""" + self._state_pool.clear() + + def open_session(self) -> None: + """Alias for :py:meth:`reset_session` to start a fresh streaming session.""" + self.reset_session() + + def close_session(self) -> None: + """Alias for :py:meth:`reset_session` to end the current streaming session.""" + self.reset_session() diff --git a/nemo/collections/speechlm2/inference/pipelines/streaming_s2s_pipeline.py b/nemo/collections/speechlm2/inference/pipelines/streaming_s2s_pipeline.py new file mode 100644 index 000000000000..faa1468dc682 --- /dev/null +++ b/nemo/collections/speechlm2/inference/pipelines/streaming_s2s_pipeline.py @@ -0,0 +1,855 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +import os +import tempfile +import time +from dataclasses import dataclass + +import librosa +import soundfile as sf +import torch +from omegaconf import DictConfig +from torch import Tensor + +from nemo.collections.asr.inference.streaming.buffering.audio_bufferer import BatchedAudioBufferer +from nemo.collections.asr.inference.streaming.framing.request import Frame +from nemo.collections.asr.inference.utils.audio_io import read_audio +from nemo.collections.asr.inference.utils.enums import RequestType +from nemo.collections.speechlm2.inference.model_wrappers.decode_state import NullTimingSummary +from nemo.collections.speechlm2.inference.model_wrappers.nemotron_voicechat_inference_wrapper import ( + NemotronVoicechatInferenceWrapper, +) +from nemo.collections.speechlm2.inference.pipelines.s2s_pipeline_interface import S2SPipelineInterface +from nemo.collections.speechlm2.inference.streaming.framing.s2s_request_options import S2SRequestOptions +from nemo.collections.speechlm2.inference.streaming.framing.silence_padded_frame_streamer import ( + SilencePaddedContinuousBatchedFrameStreamer, +) +from nemo.collections.speechlm2.inference.streaming.state.s2s_context_manager import S2SContextManager +from nemo.collections.speechlm2.inference.streaming.state.s2s_streaming_output import S2SStreamingOutput +from nemo.collections.speechlm2.inference.utils.stepprogressbar import StepProgressBar +from nemo.collections.speechlm2.parts.text_utils import _decode_tokens_with_specials +from nemo.utils import logging + + +@dataclass +class GenerateStepOutput: + """Output of a single :meth:`StreamingS2SPipeline.generate_step` call + for one stream. + + Analogous to :class:`TranscribeStepOutput` in the ASR pipelines, + this carries the **incremental** (new-this-step) audio and text so + that callers don't have to diff against accumulated state. + + The underlying :class:`S2SStreamingOutput` still accumulates + everything for batch/offline use. + """ + + stream_id: int + audio: torch.Tensor + text: str = "" + asr_text: str = "" + function_text: str = "" + + +class StreamingS2SPipeline(S2SPipelineInterface): + """Streaming speech-to-speech pipeline. Single-stream only (``streaming.batch_size`` must be 1).""" + + def __init__(self, cfg: DictConfig, s2s_model: NemotronVoicechatInferenceWrapper): + # ------------------------------------------------------------------ + # Model & device + # ------------------------------------------------------------------ + self.s2s_model = s2s_model + self.device = self.s2s_model.device + self.decode_audio = self.s2s_model.decode_audio + self.collect_debug = False + + # ------------------------------------------------------------------ + # Streaming configuration + # ------------------------------------------------------------------ + # ``.get`` throughout rather than ``getattr``: the fallback below can be + # a plain dict, and ``getattr`` would silently return the default for + # keys that are actually present in it. + self.streaming_cfg = cfg.get("streaming", {}) + self.input_sample_rate = self.streaming_cfg.get("input_sample_rate", 16000) + self.output_sample_rate = self.streaming_cfg.get("output_sample_rate", 22050) + self.batch_size = self.streaming_cfg.get("batch_size", 1) + self.max_len = self.streaming_cfg.get("max_len", 8192) + if self.batch_size != 1: + raise ValueError( + "StreamingS2SPipeline supports a single stream " + "(streaming.batch_size must be 1)." + ) + + # ------------------------------------------------------------------ + # Chunk & buffer sizes + # Terminology: "frame" = 80ms audio unit, "chunk" = 1 or more frames + # A chunk is the amount of audio that is processed per inference step. + # ------------------------------------------------------------------ + self.chunk_size_in_secs = self.streaming_cfg.get("chunk_size_in_secs", 0.08) + # Check if self.chunk_size_in_secs is a multiple of 0.08. + # Because of quirks of floating point arithmetic, the remainder could be either ~0 or ~0.08, + # so we check for both cases. + remainder = self.chunk_size_in_secs % 0.08 + if not (math.isclose(remainder, 0, abs_tol=1e-9) or math.isclose(remainder, 0.08, abs_tol=1e-9)): + raise ValueError(f"Chunk size must be a multiple of 0.08s, but got {self.chunk_size_in_secs}") + + self.num_frames_per_chunk = int(self.chunk_size_in_secs / 0.08) + + # Buffer size determines how much audio is passed to the perception encoder + # Default: 5.68 seconds (71 * 0.08). This is the minimum valid buffer size without the perception cache. + # i.e. att_context_size[0] + att_context_size[1] + 1 frames = 70+0+1 = 71 frames = 5.68 seconds + self.buffer_size_in_secs = self.streaming_cfg.get("buffer_size_in_secs", 71 * 0.08) + + self.att_context_size = self.streaming_cfg.get("att_context_size", [70, 0]) + + # ------------------------------------------------------------------ + # bufferer – reused from ASR utilities + # ------------------------------------------------------------------ + self.bufferer = BatchedAudioBufferer( + sample_rate=self.input_sample_rate, + buffer_size_in_secs=self.buffer_size_in_secs, + ) + + # ------------------------------------------------------------------ + # System prompt & sampling defaults (from YAML s2s block) + # ------------------------------------------------------------------ + s2s_cfg = cfg.get("s2s", {}) + self.system_prompt: str | None = s2s_cfg.get("system_prompt", None) + if self.system_prompt: + logging.info( + f"System prompt configured: {self.system_prompt[:100]}{'...' if len(self.system_prompt) > 100 else ''}" + ) + + self._default_top_p: float | None = s2s_cfg.get("top_p", None) + self._default_temperature: float | None = s2s_cfg.get("temperature", None) + self._default_repetition_penalty: float | None = s2s_cfg.get("repetition_penalty", None) + + # Context manager + self.context_manager = S2SContextManager( + s2s_model=self.s2s_model, + max_len=self.max_len, + ) + + # Output directory for generated files + self.output_dir = cfg.get("output_dir", "./generated") + + # Parse and validate request type early, with a safe default + req_type_cfg = self.streaming_cfg.get("request_type", "frame") + + # Parse and validate the request type; only 'frame' is supported for s2s. + self.request_type = RequestType.from_str(req_type_cfg) + if self.request_type is not RequestType.FRAME: + raise ValueError(f"Request type {self.request_type} is not supported for s2s.") + + self._stream_has_prompt: bool = False + + # ------------------------------------------------------------------ + # Input audio padding (silence appended after real audio) + # ------------------------------------------------------------------ + self.pad_audio_to_sec: float | None = cfg.get("pad_audio_to_sec", None) + self.pad_silence_ratio: float | None = cfg.get("pad_silence_ratio", None) + self.pad_audio_by_sec: float | None = cfg.get("pad_audio_by_sec", None) + if sum(x is not None for x in [self.pad_audio_to_sec, self.pad_silence_ratio, self.pad_audio_by_sec]) > 1: + raise ValueError("Set at most one of: pad_audio_to_sec, pad_silence_ratio, pad_audio_by_sec") + + super().__init__() + + def shutdown(self) -> None: + """Release process-level resources owned by the model. + + Only the vLLM-Omni path has any -- engine subprocesses and a daemon + thread -- but this is safe to call for every backend. Native engines + no-op. Unlike the precision globals this is object-scoped, which is + why it lives on the pipeline rather than the builder. + + Call from a ``finally`` after :meth:`run`, or from a server finalize. + Idempotent. + """ + self.s2s_model.shutdown() + + # ------------------------------------------------------------------ + # State helpers + # ------------------------------------------------------------------ + @property + def special_token_strings(self) -> set[str]: + """Token strings that should be stripped from decoded text for clean output. + + Pass to :func:`~nemo.collections.speechlm2.parts.text_utils.clean_pred_text`. + """ + return self.s2s_model.special_token_strings + + @property + def output_capabilities(self): + """Optional checkpoint heads and active-backend output availability.""" + return self.s2s_model.output_capabilities + + def create_state(self, options: S2SRequestOptions | None = None) -> S2SStreamingOutput: + """Create new empty state with optional per-stream options.""" + return S2SStreamingOutput( + device=self.device, + dtype=self.s2s_model.dtype, + output_sample_rate=self.output_sample_rate, + options=options or S2SRequestOptions(), + capabilities=self.output_capabilities, + ) + + def _init_state(self, stream_id: int, options: S2SRequestOptions | None = None) -> None: + """Initialize a new stream: resolve defaults, create state, create context, prefill. + + This is the S2S equivalent of ASR's ``init_state()`` in ``BasePipeline``. + Called automatically by :meth:`generate_step` when a frame has + ``is_first=True``. + + The method always runs stream initialization (state creation, + context-manager allocation, KV-cache prefill). If the triggering + frame also carries audio, :meth:`generate_step` will process it + immediately after this method returns. For latency-sensitive + deployments (real-time voice chat), callers should send the first + frame with **empty audio** so that prefill completes before the + user starts speaking — this prevents audio from queuing up during + the expensive prefill phase. + """ + if self.get_state(stream_id) is not None or stream_id in self.context_manager.active_stream_ids: + raise RuntimeError( + f"Stream {stream_id} is already active. Send is_last=True to finalize " + f"the existing stream before re-using the same stream_id." + ) + + raw_opts = options or S2SRequestOptions() + if raw_opts.system_prompt is not None and not raw_opts.system_prompt.strip(): + raw_opts = S2SRequestOptions( + system_prompt=None, + top_p=raw_opts.top_p, + temperature=raw_opts.temperature, + repetition_penalty=raw_opts.repetition_penalty, + ) + opts = raw_opts.fill_defaults( + default_system_prompt=self.system_prompt, + default_top_p=self._default_top_p, + default_temperature=self._default_temperature, + default_repetition_penalty=self._default_repetition_penalty, + ) + self.get_or_create_state(stream_id, options=opts) + + # Prefill can take hundreds of ms, or even tens of seconds if the + # prompt is long and the model is not warmed up. + prompt = opts.system_prompt + sampling_params = { + key: getattr(opts, key) + for key in ("top_p", "temperature", "repetition_penalty") + if getattr(opts, key) is not None + } + start_prefill = time.time() + with torch.no_grad(), torch.inference_mode(): + self._prefill_system_prompt( + stream_id, + prompt, + sampling_params=sampling_params or None, + ) + torch.cuda.synchronize() + logging.info(f"_init_state: stream_id={stream_id}, prefill={1000*(time.time()-start_prefill):.1f}ms") + + # Will tell generate_step_for_frames whether the KV cache already contains + # a system prompt, so it can choose the right first-frame embedding + # (PAD tokens if prefilled, BOS tokens if not). Consumed and + # cleared on the first audio frame. + self._stream_has_prompt = bool(prompt) + + def generate_step_for_frames(self, frames: list[Frame], buffers: list[Tensor]) -> list[GenerateStepOutput]: + """Generate speech for audio Frames using a shared ContextManager. + + This is the S2S equivalent of ASR's ``transcribe_step_for_frames`` + in ``BasePipeline``. Like its ASR counterpart, it is never called + directly — :meth:`generate_step` (the public API, analogous to + ``transcribe_step``) handles stream init and then delegates here + for the actual audio processing. + + Stream initialization (state, context, prefill) is always handled + by :meth:`_init_state` *before* this method is called. + """ + if len(frames) == 0: + return [] + + stream_ids = [f.stream_id for f in frames] + eos_flags = [f.is_last for f in frames] + + logging.debug(f"stream_ids={stream_ids} eos_flags={eos_flags}") + + if len(frames) != 1: + raise ValueError("StreamingS2SPipeline supports a single stream (batch_size must be 1)") + + has_prompt = self._stream_has_prompt + self._stream_has_prompt = False + + request_id = self._request_id_for_stream(stream_ids[0]) + + context = self.context_manager.get_context(stream_ids) + + audio_buffer = buffers[0] + if audio_buffer.dim() == 1: + audio_buffer = audio_buffer.unsqueeze(0) + audio_buffer = audio_buffer.to(self.s2s_model.device, dtype=self.s2s_model.dtype) + + # Sampling overrides were resolved by _init_state via fill_defaults + # and stored on state.options. Build the dict for infer_one_step. + pipeline_state = self.get_state(stream_ids[0]) + if pipeline_state is None: + raise RuntimeError( + f"No state initialized for stream {stream_ids[0]}. " + "Clients must send an is_first=True frame before streaming audio." + ) + sampling_params = { + k: getattr(pipeline_state.options, k) + for k in ("top_p", "temperature", "repetition_penalty") + if getattr(pipeline_state.options, k) is not None + } + + result = self.s2s_model.infer_one_step( + audio_input=audio_buffer, + num_frames_per_chunk=self.num_frames_per_chunk, + state=context, + request_id=request_id, + has_prompt=has_prompt, + return_debug=self.collect_debug, + sampling_params=sampling_params or None, + ) + + if self.collect_debug and result.debug is not None: + state = self.get_or_create_state(stream_ids[0]) + state.debug_data.append(result.debug) + + # Persist updated cache & clean finished streams + self.context_manager.update_context(stream_ids, result, self.num_frames_per_chunk) + + # Finalize token tensors and timing from the decode context before it + # is destroyed by reset_streams. + tokenizer = self.s2s_model.tokenizer + pad_id = self.s2s_model.model.stt_model.text_pad_id + timing_by_stream: dict[int, object] = {} + for stream_id, eos_flag in zip(stream_ids, eos_flags): + if eos_flag: + ctx = self.context_manager.get_context_for_stream(stream_id) + if ctx is not None: + state = self.get_or_create_state(stream_id) + state.finalize_tokens( + ctx.gen_text, + ctx.gen_asr_text if self.output_capabilities.has_asr_head else None, + ctx.frame_idx, + tokenizer=tokenizer, + pad_id=pad_id, + gen_function=( + ctx.gen_function if self.output_capabilities.has_function_head else None + ), + ) + timing_by_stream[stream_id] = ctx.timing + + # Close finished streams before reset_streams discards their decode + # context: a vLLM session hangs off that context, and its consumer task + # would leak until process exit. end_stream is idempotent, so any later + # _abort_stream_request for the same stream is a no-op. + for stream_id, eos_flag in zip(stream_ids, eos_flags): + if eos_flag: + self._abort_stream_request(stream_id) + + self.context_manager.reset_streams(stream_ids, eos_flags) + + # Log summary and clean up finished streams + pad_str = tokenizer.ids_to_tokens([pad_id])[0] + for stream_id, eos_flag in zip(stream_ids, eos_flags): + if eos_flag: + state = self.get_state(stream_id) + audio_sec = state._total_audio_samples / self.output_sample_rate if self.output_sample_rate > 0 else 0 + logging.info( + f"Stream {stream_id} finished: {state.token_length or 0} frames, " + f"{audio_sec:.1f}s audio, " + f"agent: {state.output_text_str!r}, user: {state.output_asr_text_str!r}, " + f"function: {state.output_function_text_str!r}" + ) + + # Replace verbose pad token (e.g. '') with '·' for compact logging + compact_agent = state.raw_text.replace(pad_str, "·") + compact_user = (state.raw_asr_text or "").replace(pad_str, "·") + compact_function = (state.raw_function_text or "").replace(pad_str, "·") + logging.info(f"Stream {stream_id} agent (with padding): {compact_agent}") + logging.info(f"Stream {stream_id} user (with padding): {compact_user}") + logging.info(f"Stream {stream_id} function (with padding): {compact_function}") + + # Timing summary (no-op when profile_timing is off) + timing_by_stream.get(stream_id, NullTimingSummary()).log_summary( + label=f"Stream {stream_id}", + chunk_ms=self.chunk_size_in_secs * 1000, + ) + + self.bufferer.rm_bufferer(stream_id) + self._abort_stream_request(stream_id) + + # Split the batch-level InferenceStepResult into per-frame outputs. + # Each frame's incremental audio/text is: + # 1. Appended to the per-stream S2SStreamingOutput accumulator + # (persists across steps; finalized at end-of-stream and + # returned by run()). + # 2. Wrapped in a GenerateStepOutput and returned to the caller + # (used by server integrations to stream partial results + # to clients without diffing accumulated state). + outputs: list[GenerateStepOutput] = [] + for idx, frame in enumerate(frames): + state = self.get_state(frame.stream_id) + audio = result.decoded_audio[idx : idx + 1] if result.decoded_audio is not None else None + text = result.predicted_text_strs[idx] if result.predicted_text_strs else "" + asr_text = result.asr_predicted_text_strs[idx] if result.asr_predicted_text_strs else "" + function_text = result.predicted_function_strs[idx] if result.predicted_function_strs else "" + + state.append_step_output(audio, text=text, asr_text=asr_text, function_text=function_text) + + outputs.append( + GenerateStepOutput( + stream_id=frame.stream_id, + audio=audio if audio is not None else torch.empty(1, 0), + text=text, + asr_text=asr_text, + function_text=function_text, + ) + ) + return outputs + + _WARMUP_FALLBACK_PROMPT = "Mock system prompt for warmup." + # Enough consecutive chunks that the autoregressive steps which trigger + # torch.compile recompilation are all paid during warmup. + _WARMUP_NUM_CHUNKS = 8 + + def warmup(self, system_prompt: str | None = None, num_chunks: int | None = None) -> None: + """Run a throwaway inference cycle to warm up the entire pipeline. + + The very first call through each stage incurs one-time overhead + (e.g. CUDA graph compilation, memory pool allocation, + DynamicCache initialization, torch.compile). Pushing a run of + silence chunks exercises the full path — prefill, perception, LLM + decode, TTS, and codec — so the first real client request is fast. + + Args: + system_prompt: Prompt text to use for warmup. Falls back to + the YAML-configured ``self.system_prompt``, then to a + short fallback string so the LLM prefill path is always + exercised. + num_chunks: How many consecutive chunks to push through. More + than one is needed because each autoregressive step can + trigger a fresh torch.compile. + """ + prompt = system_prompt if system_prompt is not None else self.system_prompt + if not prompt: + prompt = self._WARMUP_FALLBACK_PROMPT + logging.info(f'No system prompt configured — using fallback prompt for warmup: "{prompt}"') + + if num_chunks is None: + num_chunks = self._WARMUP_NUM_CHUNKS + num_chunks = max(1, int(num_chunks)) + + warmup_stream_id = -1 + chunk_samples = int(self.chunk_size_in_secs * self.input_sample_rate) + + logging.info(f"Running pipeline warmup (prefill + {num_chunks} silence chunks)...") + t0 = time.time() + + # Real requests may pass a file path, which decodes through librosa and + # pays a one-time JIT on first use. The silence frames below pass + # samples, so they never reach it; warm it here on a throwaway file + # instead, at double the rate so resampling is covered too. + warmup_rate = self.input_sample_rate * 2 + with tempfile.NamedTemporaryFile(suffix=".wav") as warmup_audio: + sf.write(warmup_audio.name, torch.zeros(warmup_rate // 2).numpy(), warmup_rate) + read_audio(warmup_audio.name, target_sr=self.input_sample_rate) + + for idx in range(num_chunks): + warmup_frame = Frame( + samples=torch.zeros(chunk_samples), + stream_id=warmup_stream_id, + is_first=idx == 0, + is_last=idx == num_chunks - 1, + options=S2SRequestOptions(system_prompt=prompt) if idx == 0 else None, + ) + step_t0 = time.time() + self.generate_step([warmup_frame]) + logging.info(f" warmup chunk {idx + 1}/{num_chunks}: {(time.time() - step_t0) * 1000:.1f}ms") + + # Tear down everything so the engine is clean for real traffic + self.reset_session() + self._stream_has_prompt = False + + logging.info(f"Pipeline warmup complete in {time.time() - t0:.3f}s") + + def generate_step(self, frames: list[Frame]) -> list[GenerateStepOutput]: + """Main streaming API — handles both init and audio processing. + + Mirrors ASR's ``transcribe_step``: on ``is_first`` frames, the + stream is initialized via :meth:`_init_state` (state creation, + context-manager allocation, KV-cache prefill). If the frame also + carries input audio, it is processed in the same call. If there + is no input audio (e.g. a server prefill-only request), the + method returns after init without running inference. + + Returns one :class:`GenerateStepOutput` per input frame carrying + the **incremental** output audio and text produced by this step. + The output audio tensor may be empty when no waveform is produced + (prefill-only frames with no input audio, or when + ``decode_audio=False``). + + For latency-sensitive deployments, send the ``is_first`` frame + with **empty audio** so that the expensive prefill completes + before the user starts speaking. For batch/offline usage the + first frame can carry real audio — init and first-chunk + processing simply happen back-to-back in one call. + """ + # Init phase — like ASR's `if request.is_first: self.init_state(...)` + for frame in frames: + if frame.is_first: + self._init_state(frame.stream_id, frame.options) + + # Decode phase. Prefill-only frames (empty audio) are not passed to + # inference; outputs are stitched back 1:1 with the original *frames*. + non_empty_frames = [f for f in frames if f.samples.numel() > 0] + empty_terminal_frames = [f for f in frames if f.samples.numel() == 0 and f.is_last] + for frame in empty_terminal_frames: + # A prefill-only request may also be the stream's terminal + # request (for example, a client disconnect immediately after + # sequence_start). Tear down model/session state even though no + # decode step ran. + self._abort_stream_request(frame.stream_id) + self.context_manager.reset_streams([frame.stream_id], [True]) + self.bufferer.rm_bufferer(frame.stream_id) + + if not non_empty_frames: + # All frames are prefill-only — nothing to decode. Terminal + # frames were closed above; callers retain the output accumulator + # until they consume and delete it. + return [GenerateStepOutput(stream_id=f.stream_id, audio=torch.empty(1, 0)) for f in frames] + + buffers, left_paddings = self.bufferer.update(non_empty_frames) + # The audio buffer left-pads; strip that so the rest of the pipeline + # sees unpadded samples. + buffers = [b[lp:] for b, lp in zip(buffers, left_paddings)] + with torch.no_grad(), torch.inference_mode(): + step_outputs = self.generate_step_for_frames(non_empty_frames, buffers) + + # Fast path: every frame had audio, so step_outputs is already 1:1. + if len(non_empty_frames) == len(frames): + return step_outputs + + # Prefill-only frames get empty outputs so the returned list stays + # aligned with the input. + output_by_stream: dict[int, GenerateStepOutput] = {o.stream_id: o for o in step_outputs} + return [ + output_by_stream.get( + f.stream_id, + GenerateStepOutput(stream_id=f.stream_id, audio=torch.empty(1, 0)), + ) + for f in frames + ] + + # ------------------------------------------------------------------ + # Finalization helpers + # ------------------------------------------------------------------ + def _finalize_and_save_finished_streams( + self, + frames: list[Frame], + audio_filepaths: list[str], + per_stream_results: dict[int, S2SStreamingOutput], + ) -> None: + """Save output files for streams that ended in this batch. + + Token fields are already populated by :meth:`S2SStreamingOutput.finalize_tokens` + (called in ``generate_step_for_frames`` before the decode context is destroyed). + This method only handles file I/O and memory cleanup. + """ + for frame in frames: + if not frame.is_last: + continue + + stream_id = frame.stream_id + state = self.get_or_create_state(stream_id) + + in_path = audio_filepaths[stream_id] + base = os.path.splitext(os.path.basename(in_path))[0] + + # When decode_audio is False the pipeline produces no waveform, so + # audio_filepath stays None and no wav/stereo files are written. + if self.decode_audio: + state.audio_filepath = self._save_audio_files(state, in_path, base) + self._save_text_files(state, base) + self._save_ctm_files(state, base) + + # Audio has been saved to disk -- drop the (potentially large) + # chunk list so finished streams don't accumulate in memory. + state.clear_audio_buffer() + + per_stream_results[stream_id] = state + self.delete_state(stream_id) + + def _save_audio_files(self, state: S2SStreamingOutput, in_path: str, base: str) -> str | None: + """Save generated mono wav and stereo (input+output) wav. + + Returns the output wav path, or ``None`` if no audio was generated. + """ + generated_audio = state.audio_buffer.detach().cpu().to(torch.float32).flatten() + + if generated_audio.numel() == 0: + return None + + wav_dir = os.path.join(self.output_dir, "wav") + os.makedirs(wav_dir, exist_ok=True) + out_path = os.path.join(wav_dir, f"{base}.wav") + sf.write(out_path, generated_audio.numpy(), self.output_sample_rate) + + # Save a stereo file: input (ch0), output (ch1) + self._save_stereo_file(generated_audio, in_path, base) + return out_path + + def _save_stereo_file(self, generated_audio: torch.Tensor, in_path: str, base: str) -> None: + """Save a stereo wav with input audio on ch0 and generated audio on ch1.""" + stereo_dir = os.path.join(self.output_dir, "stereo") + os.makedirs(stereo_dir, exist_ok=True) + + input_np, _ = librosa.load(in_path, sr=self.output_sample_rate, mono=True) + input_audio = torch.from_numpy(input_np).to(torch.float32) + + # Prepend silence to output channel to account for the one-chunk + # processing delay: the pipeline can't produce output until it has + # received a full input chunk. + delay_samples = int(self.chunk_size_in_secs * self.output_sample_rate) + gen_delayed = torch.cat([torch.zeros(delay_samples), generated_audio]) + + # Pad the shorter channel so both have equal length + max_len = max(input_audio.shape[-1], gen_delayed.shape[-1]) + input_audio = torch.nn.functional.pad(input_audio, (0, max_len - input_audio.shape[-1])) + gen_delayed = torch.nn.functional.pad(gen_delayed, (0, max_len - gen_delayed.shape[-1])) + + stereo = torch.stack([input_audio, gen_delayed], dim=0).T + stereo_path = os.path.join(stereo_dir, f"{base}_input_output.wav") + sf.write(stereo_path, stereo.numpy(), self.output_sample_rate) + + def _save_text_files(self, state: S2SStreamingOutput, base: str) -> None: + """Save agent, ASR, and function text outputs to txt files. + + The agent channel is written even when empty, because an empty agent + turn is a result; the optional channels are only written when the + checkpoint produced them. + """ + txt_dir = os.path.join(self.output_dir, "txt") + os.makedirs(txt_dir, exist_ok=True) + + channels = ( + ("", state.output_text_str, True), + ("_asr", state.output_asr_text_str, False), + ("_function", state.output_function_text_str, False), + ) + for suffix, text, write_when_empty in channels: + if not isinstance(text, str) or not (text or write_when_empty): + continue + path = os.path.join(txt_dir, f"{base}{suffix}.txt") + try: + with open(path, "w", encoding="utf-8") as f: + f.write(text) + except OSError: + logging.warning(f"Failed to write {path}") + + def _save_ctm_files(self, state: S2SStreamingOutput, base: str) -> None: + """Write per-token CTM timing files for agent and ASR channels.""" + if state.token_text is None or state.token_length is None: + return + tokenizer = self.s2s_model.tokenizer + pad_id = self.s2s_model.model.stt_model.text_pad_id + total_frames = state.token_length + total_samples = state._total_audio_samples + self._write_ctm( + base, + state.token_text[0, :total_frames], + total_frames, + total_samples, + tokenizer, + pad_id, + ) + if state.token_asr_text is not None: + self._write_ctm( + base, + state.token_asr_text[0, :total_frames], + total_frames, + total_samples, + tokenizer, + pad_id, + suffix="_asr", + ) + if state.token_function is not None: + self._write_ctm( + base, + state.token_function[0, :total_frames], + total_frames, + total_samples, + tokenizer, + pad_id, + suffix="_function", + ) + + def _write_ctm( + self, + base: str, + token_ids: torch.Tensor, + total_frames: int, + total_audio_samples: int, + tokenizer, + pad_id: int, + suffix: str = "", + ) -> None: + """Write a token-level CTM file derived from a token-ID tensor. + + Each non-pad frame gets one line with evenly-spaced timing based on + the total audio duration divided by the number of frames. + + Args: + suffix: Appended to the filename stem, e.g. ``"_asr"`` produces + ``_asr.ctm``. + """ + if total_frames == 0 or total_audio_samples == 0 or self.output_sample_rate == 0: + return + frame_duration = total_audio_samples / total_frames / self.output_sample_rate + pad_token_str = tokenizer.ids_to_tokens([pad_id])[0] + + ctm_dir = os.path.join(self.output_dir, "ctm") + os.makedirs(ctm_dir, exist_ok=True) + try: + with open(os.path.join(ctm_dir, f"{base}{suffix}.ctm"), "w", encoding="utf-8") as f: + for i in range(total_frames): + tid = int(token_ids[i].item()) + if tid == pad_id: + continue + tok_str = _decode_tokens_with_specials( + tokenizer.ids_to_tokens([tid]), + tokenizer, + pad_token_str=pad_token_str, + keep_pad=False, + ) + if not tok_str: + continue + start = i * frame_duration + f.write(f"{base} A {start:.3f} {frame_duration:.3f} {tok_str}\n") + except OSError: + logging.warning(f"Failed to write CTM for {base}") + + # ------------------------------------------------------------------ + # Session helpers (extend S2SPipelineInterface) + # ------------------------------------------------------------------ + + def reset_session(self) -> None: + """Reset feature buffer and ContextManager together.""" + for stream_id in list(self.context_manager.active_stream_ids): + self._abort_stream_request(stream_id) + self.bufferer.reset() + self.context_manager.reset() + + super().reset_session() # clears state pool + + # ------------------------------------------------------------------ + # Orchestrator – mirrors recognizers' *run* method + # ------------------------------------------------------------------ + def run( + self, + audio_filepaths: list[str], + options: list[S2SRequestOptions] | None = None, + progress_bar: StepProgressBar | None = None, + ) -> list[S2SStreamingOutput]: + """Process audio files through the streaming pipeline, saving outputs to disk. + + Each file is streamed chunk-by-chunk through :meth:`generate_step`. + When a stream finishes, its wav, txt, and CTM files are written + immediately and finalized fields are populated on the + :class:`S2SStreamingOutput`. Returns a list of finalized outputs, + one per input audio file. + + Args: + audio_filepaths: Paths to input audio files. + options: Per-stream request options (system prompt, sampling, etc.). + progress_bar: Optional :class:`StepProgressBar` for per-step + progress with per-stream postfix. + """ + + if options is None: + options = [S2SRequestOptions() for _ in audio_filepaths] + + streamer = SilencePaddedContinuousBatchedFrameStreamer( + n_frames_per_stream=1, + frame_size_in_secs=self.chunk_size_in_secs, + sample_rate=self.input_sample_rate, + batch_size=self.batch_size, + pad_last_frame=True, + pad_to_sec=self.pad_audio_to_sec, + pad_by_sec=self.pad_audio_by_sec, + pad_ratio=self.pad_silence_ratio, + ) + streamer.set_audio_filepaths(audio_filepaths, options) + + os.makedirs(self.output_dir, exist_ok=True) + + per_stream_results: dict[int, S2SStreamingOutput] = {} + + self.open_session() + for frames in streamer: + self.generate_step(frames) + self._finalize_and_save_finished_streams(frames, audio_filepaths, per_stream_results) + + if progress_bar is not None: + for f in frames: + progress_bar.step(f.stream_id) + + if progress_bar is not None: + progress_bar.finish() + + outputs = [per_stream_results.get(idx, self.create_state()) for idx in range(len(audio_filepaths))] + self.close_session() + return outputs + + def _prefill_system_prompt( + self, + stream_id: int, + system_prompt: str | None = None, + sampling_params: dict[str, float] | None = None, + ) -> None: + """Open a new stream on the model, system prompt included. + + Whether that means a native prompt prefill, a vLLM session, or both is + the wrapper's business; the pipeline only announces that a stream is + starting. + + Args: + stream_id: The stream identifier. + system_prompt: The system prompt text for this stream. If *None*, + no prompt is injected. + sampling_params: Per-stream text sampling parameters, for backends + that fix sampling when the stream opens. + """ + self.s2s_model.begin_stream( + self.context_manager.get_context([stream_id]), + system_prompt, + request_id=self._request_id_for_stream(stream_id), + sampling_params=sampling_params, + ) + + def _request_id_for_stream(self, stream_id: int) -> str: + return str(stream_id) + + def _abort_stream_request(self, stream_id: int) -> None: + """Close a stream on the model. Idempotent.""" + self.s2s_model.end_stream( + self.context_manager.get_context_for_stream(stream_id), + request_id=self._request_id_for_stream(stream_id), + ) diff --git a/nemo/collections/speechlm2/inference/streaming/__init__.py b/nemo/collections/speechlm2/inference/streaming/__init__.py new file mode 100644 index 000000000000..9e3fb699d9f6 --- /dev/null +++ b/nemo/collections/speechlm2/inference/streaming/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/nemo/collections/speechlm2/inference/streaming/framing/__init__.py b/nemo/collections/speechlm2/inference/streaming/framing/__init__.py new file mode 100644 index 000000000000..9e3fb699d9f6 --- /dev/null +++ b/nemo/collections/speechlm2/inference/streaming/framing/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/nemo/collections/speechlm2/inference/streaming/framing/s2s_request_options.py b/nemo/collections/speechlm2/inference/streaming/framing/s2s_request_options.py new file mode 100644 index 000000000000..d54a49b1be16 --- /dev/null +++ b/nemo/collections/speechlm2/inference/streaming/framing/s2s_request_options.py @@ -0,0 +1,67 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True, slots=True) +class S2SRequestOptions: + """Immutable per-stream options for S2S inference. + + Attached to the first ``Frame`` of each stream via the ``options`` + field so that the pipeline can read per-stream configuration at the + start of every new audio stream. Frozen so that options cannot be + accidentally modified after the stream is initialised. + + All fields default to ``None``, which means "use the pipeline-level + default". Call :meth:`fill_defaults` to fill ``None`` fields with + pipeline-level values. + """ + + system_prompt: str | None = None + + top_p: float | None = None # (0, 1] + temperature: float | None = None # >= 0 + repetition_penalty: float | None = None # > 0 + + def __post_init__(self) -> None: + if self.top_p is not None and not (0.0 < self.top_p <= 1.0): + raise ValueError(f"top_p must be in (0, 1], got {self.top_p}") + if self.temperature is not None and self.temperature < 0.0: + raise ValueError(f"temperature must be >= 0, got {self.temperature}") + if self.repetition_penalty is not None and self.repetition_penalty <= 0.0: + raise ValueError(f"repetition_penalty must be > 0, got {self.repetition_penalty}") + + @staticmethod + def _with_default(value: Any, default: Any) -> Any: + """Return *value* when it is not ``None``, otherwise *default*.""" + return default if value is None else value + + def fill_defaults( + self, + default_system_prompt: str | None = None, + default_top_p: float | None = None, + default_temperature: float | None = None, + default_repetition_penalty: float | None = None, + ) -> S2SRequestOptions: + """Return a new options instance with ``None`` fields filled from defaults.""" + return S2SRequestOptions( + system_prompt=self._with_default(self.system_prompt, default_system_prompt), + top_p=self._with_default(self.top_p, default_top_p), + temperature=self._with_default(self.temperature, default_temperature), + repetition_penalty=self._with_default(self.repetition_penalty, default_repetition_penalty), + ) diff --git a/nemo/collections/speechlm2/inference/streaming/framing/silence_padded_frame_streamer.py b/nemo/collections/speechlm2/inference/streaming/framing/silence_padded_frame_streamer.py new file mode 100644 index 000000000000..edca88bfd31f --- /dev/null +++ b/nemo/collections/speechlm2/inference/streaming/framing/silence_padded_frame_streamer.py @@ -0,0 +1,75 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from nemo.collections.asr.inference.streaming.framing.mono_stream import MonoStream +from nemo.collections.asr.inference.streaming.framing.multi_stream import ContinuousBatchedFrameStreamer +from nemo.collections.speechlm2.inference.streaming.framing.silence_padded_stream import SilencePaddedStream + + +class SilencePaddedContinuousBatchedFrameStreamer(ContinuousBatchedFrameStreamer): + """``ContinuousBatchedFrameStreamer`` that optionally wraps each + ``MonoStream`` in a :class:`SilencePaddedStream` so extra silence + frames are yielded transparently at the end of each audio file. + + When no padding is configured the behaviour is identical to the base + class. + """ + + def __init__( + self, + *, + pad_to_sec: float | None = None, + pad_by_sec: float | None = None, + pad_ratio: float | None = None, + **kwargs, + ): + super().__init__(**kwargs) + self.pad_to_sec = pad_to_sec + self.pad_by_sec = pad_by_sec + self.pad_ratio = pad_ratio + + @property + def _needs_padding(self) -> bool: + return any(x is not None for x in (self.pad_to_sec, self.pad_by_sec, self.pad_ratio)) + + def add_stream(self) -> None: + if self.stream_id >= self.n_audio_files: + return + + inner = MonoStream( + self.sample_rate, + self.frame_size_in_secs, + stream_id=self.stream_id, + pad_last_frame=self.pad_last_frame, + ) + + if self._needs_padding: + stream = SilencePaddedStream( + inner, + chunk_size_in_secs=self.frame_size_in_secs, + pad_to_sec=self.pad_to_sec, + pad_by_sec=self.pad_by_sec, + pad_ratio=self.pad_ratio, + ) + else: + stream = inner + + audio_filepath = self.audio_filepaths[self.stream_id] + self.sid2filepath[self.stream_id] = audio_filepath + self.elapsed_durations[self.stream_id] = 0.0 + stream.load_audio(audio_filepath, self.options[self.stream_id]) + + self.multi_streamer.add_stream(stream, stream_id=self.stream_id) + self.stream_id += 1 + self.update_progress_bar() diff --git a/nemo/collections/speechlm2/inference/streaming/framing/silence_padded_stream.py b/nemo/collections/speechlm2/inference/streaming/framing/silence_padded_stream.py new file mode 100644 index 000000000000..2110fbe52dc3 --- /dev/null +++ b/nemo/collections/speechlm2/inference/streaming/framing/silence_padded_stream.py @@ -0,0 +1,97 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from nemo.collections.asr.inference.streaming.framing.mono_stream import MonoStream +from nemo.collections.asr.inference.streaming.framing.request import Frame +from nemo.collections.asr.inference.streaming.framing.stream import Stream + + +class SilencePaddedStream(Stream): + """Wraps a ``MonoStream`` and appends silence frames after the real audio + to reach a target duration. + + The pipeline's ``run()`` loop sees a single, longer stream — no frame + mutation or side-channel silence injection is needed. ``MultiStream`` + keeps the stream alive until the final silence frame sets ``is_last=True``. + """ + + def __init__( + self, + inner: MonoStream, + chunk_size_in_secs: float, + pad_to_sec: float | None = None, + pad_by_sec: float | None = None, + pad_ratio: float | None = None, + ): + super().__init__(inner.stream_id) + self.inner = inner + self.chunk_size_in_secs = chunk_size_in_secs + self.pad_to_sec = pad_to_sec + self.pad_by_sec = pad_by_sec + self.pad_ratio = pad_ratio + self._inner_exhausted = False + self._silence_frames_remaining = 0 + + def load_audio(self, audio, options=None): + self.inner.load_audio(audio, options) + audio_secs = self.inner.n_samples / self.inner.rate + remaining = self._padding_secs(audio_secs) + self._silence_frames_remaining = max(1, round(remaining / self.chunk_size_in_secs)) if remaining > 0 else 0 + + def _padding_secs(self, elapsed: float) -> float: + if self.pad_to_sec is not None: + return max(0.0, self.pad_to_sec - elapsed) + if self.pad_ratio is not None: + return elapsed * self.pad_ratio + if self.pad_by_sec is not None: + return self.pad_by_sec + return 0.0 + + def __iter__(self): + self.inner.__iter__() + self._inner_exhausted = False + return self + + def __next__(self) -> list[Frame]: + if not self._inner_exhausted: + frames = next(self.inner) + frame = frames[0] + if frame.is_last and self._silence_frames_remaining > 0: + modified = Frame( + samples=frame.samples, + stream_id=frame.stream_id, + is_first=frame.is_first, + is_last=False, + length=frame.length, + options=frame.options, + ) + self._inner_exhausted = True + return [modified] + return frames + + if self._silence_frames_remaining > 0: + self._silence_frames_remaining -= 1 + return [ + Frame( + samples=torch.zeros(self.inner.frame_size), + stream_id=self.stream_id, + is_first=False, + is_last=(self._silence_frames_remaining == 0), + length=self.inner.frame_size, + ) + ] + + raise StopIteration diff --git a/nemo/collections/speechlm2/inference/streaming/state/__init__.py b/nemo/collections/speechlm2/inference/streaming/state/__init__.py new file mode 100644 index 000000000000..9e3fb699d9f6 --- /dev/null +++ b/nemo/collections/speechlm2/inference/streaming/state/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/nemo/collections/speechlm2/inference/streaming/state/s2s_context_manager.py b/nemo/collections/speechlm2/inference/streaming/state/s2s_context_manager.py new file mode 100644 index 000000000000..d2692adaa58b --- /dev/null +++ b/nemo/collections/speechlm2/inference/streaming/state/s2s_context_manager.py @@ -0,0 +1,125 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from nemo.collections.speechlm2.inference.model_wrappers.decode_state import InferenceStepResult, StreamingDecodeState + + +class S2SContextManager: + """Manages the lifecycle of model-level decode state for S2S streaming inference. + + Each active stream gets a :class:`StreamingDecodeState` that holds LLM KV + caches, TTS KV caches, perception cache, codec cache, token workspaces + (``gen_text``, ``gen_asr_text``), and ``frame_idx``. These are created by + the model wrapper's ``create_decode_state()`` and mutated in-place by + ``infer_one_step()``. One stream at a time (``batch_size`` must be 1). + + This class is kept separate from the pipeline-level output accumulator + (:class:`S2SStreamingOutput`) so that the heavy GPU tensors inside + ``StreamingDecodeState`` can be released as soon as a stream finishes, + independently of how long the caller holds on to the accumulated text + and audio results. + + :meth:`get_context` lazily creates a context on first access, + :meth:`reset_streams` destroys contexts at end-of-stream, and + :meth:`reset` destroys all contexts. A stream ID may be reused + after its context has been destroyed. + """ + + def __init__( + self, + s2s_model, + max_len: int, + ): + self.s2s_model = s2s_model + self.max_len = max_len + self.device = s2s_model.device + self.dtype = s2s_model.dtype + + self._contexts: dict[int, StreamingDecodeState] = {} + + def reset(self) -> None: + """Release all contexts and start fresh.""" + self._contexts.clear() + + @property + def active_stream_ids(self) -> set[int]: + """Stream IDs that currently have an active decode context.""" + return set(self._contexts.keys()) + + def _create_context(self) -> StreamingDecodeState: + """Allocate a fresh context backed by the realtime inference model.""" + if not hasattr(self.s2s_model, "create_decode_state"): + raise RuntimeError("s2s_model must provide create_decode_state(max_len)") + return self.s2s_model.create_decode_state(self.max_len) + + def get_context(self, stream_ids: list[int]) -> StreamingDecodeState: + """Return the decode context for the given stream IDs, creating if needed.""" + if len(stream_ids) == 0: + return self._create_context() + if len(stream_ids) != 1: + raise ValueError("S2SContextManager supports a single stream (batch_size must be 1)") + + stream_id = stream_ids[0] + if stream_id not in self._contexts: + self._contexts[stream_id] = self._create_context() + + return self._contexts[stream_id] + + def get_context_for_stream(self, stream_id: int) -> StreamingDecodeState | None: + """Return the decode context for a single stream, or *None* if absent.""" + return self._contexts.get(stream_id) + + def update_context( + self, + stream_ids: list[int], + step_result: InferenceStepResult, + num_frames: int, + ) -> None: + """Advance frame counter and set subword mask after an inference step. + + All cache and tensor mutations (llm_cache, tts_past_key_values, + tts_code, perception_cache, tts_codec_cache, gen_text, gen_asr_text, + etc.) are already applied in-place on the ``StreamingDecodeState`` by + ``infer_one_step``. This method only bumps ``frame_idx`` and marks + the subword mask for the newly generated frames. + """ + if len(stream_ids) == 0: + return + if len(stream_ids) != 1: + raise ValueError("S2SContextManager supports a single stream (batch_size must be 1)") + + stream_id = stream_ids[0] + context = self._contexts.get(stream_id) + if context is None: + raise RuntimeError(f"Stream {stream_id} is not registered in the context manager") + + start_idx = context.frame_idx + end_idx = start_idx + num_frames + if end_idx > context.gen_text.shape[1]: + raise RuntimeError( + "Context maximum length exceeded. Consider increasing `streaming.max_len` in the configuration." + ) + + context.frame_idx = end_idx + + if context.subword_mask is not None: + context.subword_mask[:, start_idx:end_idx] = True + + def reset_streams(self, stream_ids: list[int], eos_flags: list[bool]) -> None: + """Release contexts for streams that signalled end-of-stream.""" + if len(stream_ids) != len(eos_flags): + raise ValueError("stream_ids and eos_flags must have the same length") + for stream_id, eos_flag in zip(stream_ids, eos_flags): + if eos_flag and stream_id in self._contexts: + del self._contexts[stream_id] diff --git a/nemo/collections/speechlm2/inference/streaming/state/s2s_streaming_output.py b/nemo/collections/speechlm2/inference/streaming/state/s2s_streaming_output.py new file mode 100644 index 000000000000..3dfe24f83e1b --- /dev/null +++ b/nemo/collections/speechlm2/inference/streaming/state/s2s_streaming_output.py @@ -0,0 +1,230 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from dataclasses import dataclass, field + +import torch + +from nemo.collections.speechlm2.inference.model_wrappers.capabilities import AuxiliaryOutputCapabilities +from nemo.collections.speechlm2.inference.streaming.framing.s2s_request_options import S2SRequestOptions +from nemo.collections.speechlm2.parts.text_utils import tokens_to_str + + +@dataclass +class S2SStreamingOutput: + """Pipeline-level output accumulator for a single S2S stream. + + Collects the generated audio chunks, agent text, and ASR (user) text + produced by successive ``generate_step()`` calls, concatenating lazily + to avoid O(n^2) copying. One instance exists per active stream and is + stored in the pipeline's ``_state_pool``. + + At end-of-stream, :meth:`finalize_tokens` snapshots the raw token-ID + tensors from the :class:`StreamingDecodeState` (before it is destroyed) + and decodes them into the *finalized fields* (``text_with_timestamps``, + ``raw_text``, etc.) in a single call. The same object is then returned + by ``run()`` as the final per-stream result. + + This class is **not** the model-level decode state (KV caches, token + workspaces, perception cache). That role belongs to + :class:`~nemo.collections.speechlm2.inference.model_wrappers.decode_state.StreamingDecodeState`, + which lives inside the :class:`S2SContextManager` and is mutated + in-place by the model wrapper's ``infer_one_step()``. + + Typical lifecycle:: + + # Created on is_first frame + output = S2SStreamingOutput(device=..., dtype=..., output_sample_rate=...) + + # Each step appends incremental results + output.append_step_output(audio_chunk, text="", asr_text="Hi") + output.append_step_output(audio_chunk, text="", asr_text="") + output.append_step_output(audio_chunk, text="Hello", asr_text="") + + # At end-of-stream, snapshot token tensors and decode text fields + output.finalize_tokens(ctx.gen_text, ctx.gen_asr_text, ctx.frame_idx, + tokenizer=tok, pad_id=pad_id) + + # Final results read via properties / fields + wav = output.audio_buffer # single concatenated tensor (CPU) + txt = output.output_text_str # joined string + ts = output.text_with_timestamps # decoded with turn-taking markers + """ + + # Required init metadata + device: torch.device + dtype: torch.dtype + output_sample_rate: int + + # Per-stream request options (system prompt, sampling overrides, etc.) + options: S2SRequestOptions = field(default_factory=S2SRequestOptions) + capabilities: AuxiliaryOutputCapabilities | None = None + + # Audio chunks accumulated each step; use the ``audio_buffer`` property + # to get a single concatenated tensor (lazy, O(n) instead of O(n^2)). + _audio_chunks: list[torch.Tensor] = field(default_factory=list, repr=False) + _total_audio_samples: int = field(default=0, repr=False) + + # Text parts accumulated each step; use the ``output_text_str`` / + # ``output_asr_text_str`` properties to get joined strings. + _text_parts: list[str] = field(default_factory=list, repr=False) + _asr_text_parts: list[str] = field(default_factory=list, repr=False) + _function_text_parts: list[str] = field(default_factory=list, repr=False) + + # Per-step debug data (logits, embeddings, etc.) when collect_debug is on. + debug_data: list[dict] = field(default_factory=list, repr=False) + + # -- Finalized fields (populated by finalize_tokens() at end-of-stream) -- + token_text: torch.Tensor | None = None + token_asr_text: torch.Tensor | None = None + token_function: torch.Tensor | None = None + token_length: int | None = None + text_with_timestamps: str | None = None + asr_text_with_timestamps: str | None = None + function_text: str | None = None + raw_text: str | None = None + raw_asr_text: str | None = None + raw_function_text: str | None = None + audio_filepath: str | None = None + + @property + def audio_buffer(self) -> torch.Tensor: + """Concatenated audio from all steps, or loaded from disk after finalization. + + During streaming this concatenates in-memory chunks lazily. + After ``run()`` finishes, the chunks are cleared to save memory + and the audio is re-loaded from ``audio_filepath`` on demand. + """ + if self._audio_chunks: + return torch.cat(self._audio_chunks, dim=-1) + if self.audio_filepath is not None: + import soundfile as sf + + audio_np, _ = sf.read(self.audio_filepath, dtype="float32") + return torch.tensor(audio_np, dtype=torch.float32).unsqueeze(0) + return torch.empty((1, 0), device=self.device, dtype=self.dtype) + + @property + def output_text_str(self) -> str: + """Accumulated agent response text. Joined lazily to avoid O(n^2) copies.""" + return "".join(self._text_parts) + + @property + def output_asr_text_str(self) -> str: + """Accumulated ASR (user) text. Joined lazily to avoid O(n^2) copies.""" + return "".join(self._asr_text_parts) + + @property + def output_function_text_str(self) -> str: + """Accumulated function-channel text. Joined lazily.""" + return "".join(self._function_text_parts) + + def reset(self) -> None: + """Reset all accumulated outputs to initial state.""" + self._audio_chunks.clear() + self._total_audio_samples = 0 + self._text_parts.clear() + self._asr_text_parts.clear() + self._function_text_parts.clear() + self.debug_data.clear() + self.token_text = None + self.token_asr_text = None + self.token_function = None + self.token_length = None + self.text_with_timestamps = None + self.asr_text_with_timestamps = None + self.function_text = None + self.raw_text = None + self.raw_asr_text = None + self.raw_function_text = None + self.audio_filepath = None + + def append_step_output( + self, + audio: torch.Tensor | None, + text: str | None = None, + asr_text: str | None = None, + function_text: str | None = None, + ) -> None: + """Append generated audio and optional text from one inference step.""" + if audio is not None: + if not isinstance(audio, torch.Tensor): + raise TypeError("audio must be a torch.Tensor") + + append_tensor = audio + if append_tensor.dim() > 1: + append_tensor = append_tensor.reshape(1, -1) + elif append_tensor.dim() == 1: + append_tensor = append_tensor.unsqueeze(0) + self._audio_chunks.append(append_tensor.to(self.device, dtype=self.dtype)) + self._total_audio_samples += int(append_tensor.shape[-1]) + + if isinstance(text, str) and text: + self._text_parts.append(text) + + if isinstance(asr_text, str) and asr_text: + self._asr_text_parts.append(asr_text) + + if isinstance(function_text, str) and function_text: + self._function_text_parts.append(function_text) + + def finalize_tokens( + self, + gen_text: torch.Tensor, + gen_asr_text: torch.Tensor | None, + total_frames: int, + tokenizer, + pad_id: int, + gen_function: torch.Tensor | None = None, + ) -> None: + """Snapshot token-ID tensors from the decode context and decode them into text fields. + + Must be called at end-of-stream, before the :class:`StreamingDecodeState` + is destroyed. + """ + self.token_text = gen_text[:, :total_frames].clone().cpu() + self.token_asr_text = gen_asr_text[:, :total_frames].clone().cpu() if gen_asr_text is not None else None + self.token_function = gen_function[:, :total_frames].clone().cpu() if gen_function is not None else None + self.token_length = total_frames + + lengths = torch.tensor([total_frames], dtype=torch.long) + + def _to_str(tokens, **kwargs): + return tokens_to_str(tokens, lengths, tokenizer=tokenizer, pad_id=pad_id, **kwargs)[0] + + self.text_with_timestamps = _to_str(self.token_text, eval_text_turn_taking=True) + self.raw_text = _to_str(self.token_text, keep_pad=True) + if self.token_asr_text is not None: + self.asr_text_with_timestamps = _to_str(self.token_asr_text, eval_text_turn_taking=True) + self.raw_asr_text = _to_str(self.token_asr_text, keep_pad=True) + else: + self.asr_text_with_timestamps = None + self.raw_asr_text = None + if self.token_function is not None: + self.function_text = _to_str(self.token_function) + self.raw_function_text = _to_str(self.token_function, keep_pad=True) + else: + self.function_text = None + self.raw_function_text = None + + def clear_audio_buffer(self) -> None: + """Free in-memory audio chunks. + + Only drops the chunk list; ``_total_audio_samples`` is preserved + so that CTM timing and other metadata remain valid after the + waveform data has been written to disk. + """ + self._audio_chunks.clear() diff --git a/nemo/collections/speechlm2/inference/utils/__init__.py b/nemo/collections/speechlm2/inference/utils/__init__.py new file mode 100644 index 000000000000..9e3fb699d9f6 --- /dev/null +++ b/nemo/collections/speechlm2/inference/utils/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/nemo/collections/speechlm2/inference/utils/audio_data.py b/nemo/collections/speechlm2/inference/utils/audio_data.py new file mode 100644 index 000000000000..2674091de055 --- /dev/null +++ b/nemo/collections/speechlm2/inference/utils/audio_data.py @@ -0,0 +1,178 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Audio data loading and output serialization for S2S inference scripts.""" + +import json +import os + +import soundfile as sf + +from nemo.collections.common.parts.preprocessing.manifest import get_full_path +from nemo.collections.speechlm2.inference.streaming.framing.s2s_request_options import S2SRequestOptions +from nemo.collections.speechlm2.inference.streaming.state.s2s_streaming_output import S2SStreamingOutput + + +def prepare_audio_data( + audio_file: str, + default_system_prompt: str | None = None, + sort_by_duration: bool = True, +) -> tuple[list[str], list[S2SRequestOptions], list[str | None]]: + """Load audio filepaths and per-stream options from a folder, single file, or manifest. + + ``audio_file`` may point to a single ``.wav`` file, a directory of ``.wav`` + files, or a line-delimited ``.json``/``.jsonl`` inference manifest. Each + manifest line must contain ``audio_filepath`` and may also contain + ``system_prompt`` and ``text``:: + + {"audio_filepath": "clip.wav", "text": "...", "system_prompt": "..."} + + If ``system_prompt`` is absent on a line, *default_system_prompt* is used. + The ``text`` field is returned as optional reference transcript for WER on + the ASR/user side. + + Returns: + ``(filepaths, options, ground_truths)`` -- parallel lists of audio paths, + per-stream request options, and ground-truth texts (``None`` when unavailable). + """ + audio_file = audio_file.strip() + if not os.path.isabs(audio_file): + audio_file = os.path.abspath(audio_file) + + options: list[S2SRequestOptions] = [] + ground_truths: list[str | None] = [] + + if os.path.isdir(audio_file): + filepaths = [os.path.join(audio_file, x) for x in os.listdir(audio_file) if x.endswith(".wav")] + options = [S2SRequestOptions(system_prompt=default_system_prompt) for _ in filepaths] + ground_truths = [None] * len(filepaths) + elif audio_file.endswith(".wav"): + filepaths = [audio_file] + options = [S2SRequestOptions(system_prompt=default_system_prompt)] + ground_truths = [None] + elif audio_file.endswith((".json", ".jsonl")): + samples = [] + with open(audio_file, "r") as f: + for line in f.readlines(): + if line.strip(): + samples.append(json.loads(line)) + filepaths = [get_full_path(entry["audio_filepath"], audio_file) for entry in samples] + options = [ + S2SRequestOptions( + system_prompt=entry.get("system_prompt", default_system_prompt), + ) + for entry in samples + ] + ground_truths = [entry.get("text", None) for entry in samples] + else: + raise ValueError(f"audio_file `{audio_file}` needs to be a folder, audio file, or manifest file") + + if sort_by_duration: + durations = [sf.SoundFile(fp).frames for fp in filepaths] + order = sorted(range(len(filepaths)), key=lambda i: durations[i], reverse=True) + filepaths = [filepaths[i] for i in order] + options = [options[i] for i in order] + ground_truths = [ground_truths[i] for i in order] + + return filepaths, options, ground_truths + + +def calculate_durations_incl_padding( + audio_filepaths: list[str], + pad_audio_to_sec: float | None = None, + pad_silence_ratio: float | None = None, + pad_audio_by_sec: float | None = None, +) -> list[float]: + """Return per-file durations in seconds, accounting for silence padding. + + At most one padding argument may be set; when none are set this + returns the raw audio durations. + """ + if sum(x is not None for x in [pad_audio_to_sec, pad_silence_ratio, pad_audio_by_sec]) > 1: + raise ValueError("Set at most one of: pad_audio_to_sec, pad_silence_ratio, pad_audio_by_sec") + durations = [] + for fp in audio_filepaths: + sound = sf.SoundFile(fp) + dur = sound.frames / sound.samplerate + if pad_audio_to_sec is not None: + dur = max(dur, pad_audio_to_sec) + elif pad_silence_ratio is not None: + dur *= 1 + pad_silence_ratio + elif pad_audio_by_sec is not None: + dur += pad_audio_by_sec + durations.append(dur) + return durations + + +def dump_output_json( + audio_filepaths: list[str], + outputs: list[S2SStreamingOutput], + output_dir: str, + options: list[S2SRequestOptions], + ground_truths: list[str | None], +) -> None: + """Dump inference results to output_processed.json and output_raw.json. + + ``output_processed.json`` strips pad and BOS/EOS tokens and annotates + turn boundaries with timestamp annotations: + + * ``<|t|>`` -- turn start (BOS position, in seconds) + * ``<$t$>`` -- turn end (EOS position, in seconds) + + ``output_raw.json`` preserves the full token stream as-is, including: + + * Pad tokens (e.g. ````) for frames with no text output + * Agent turn markers: ```` (BOS) and ```` (EOS) + * User turn markers (e.g. ``^`` for user BOS, ```` for user EOS, + depending on the checkpoint) + + The raw format is useful for debugging token-level model behavior. + """ + output_processed_path = os.path.join(output_dir, "output_processed.json") + output_raw_path = os.path.join(output_dir, "output_raw.json") + + with open(output_processed_path, "w") as f_proc, open(output_raw_path, "w") as f_raw: + for audio_filepath, opts, gt, out in zip(audio_filepaths, options, ground_truths, outputs): + stem = os.path.splitext(os.path.basename(audio_filepath))[0] + pred_audio_path = os.path.join(output_dir, "wav", f"{stem}.wav") + + record_processed = { + "id": stem, + "target_text": "", + "pred_audio": pred_audio_path, + "src_text": gt or "", + "pred_src_text": out.asr_text_with_timestamps or "", + "pred_text": out.text_with_timestamps or "", + "pred_function_text": out.function_text or "", + "output_capabilities": out.capabilities.to_dict() if out.capabilities is not None else {}, + "system_prompt": opts.system_prompt or "", + } + json.dump(record_processed, f_proc, ensure_ascii=False) + f_proc.write("\n") + f_proc.flush() + + record_raw = { + "id": stem, + "target_text": "", + "pred_audio": pred_audio_path, + "src_text": gt or "", + "pred_src_text": out.raw_asr_text or "", + "pred_text": out.raw_text or "", + "pred_function_text": out.raw_function_text or "", + "output_capabilities": out.capabilities.to_dict() if out.capabilities is not None else {}, + "system_prompt": opts.system_prompt or "", + } + json.dump(record_raw, f_raw, ensure_ascii=False) + f_raw.write("\n") + f_raw.flush() diff --git a/nemo/collections/speechlm2/inference/utils/stepprogressbar.py b/nemo/collections/speechlm2/inference/utils/stepprogressbar.py new file mode 100644 index 000000000000..5a7a38de64d4 --- /dev/null +++ b/nemo/collections/speechlm2/inference/utils/stepprogressbar.py @@ -0,0 +1,74 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-inference-step progress bar for S2S streaming pipelines.""" + +from __future__ import annotations + +import math + +from tqdm import tqdm + +from nemo.collections.speechlm2.inference.utils.audio_data import calculate_durations_incl_padding + + +class StepProgressBar: + """Tracks per-step inference progress across one or more streams. + + Each call to :meth:`step` advances the bar by one and updates the + per-stream postfix (e.g. ``stream 2: 45/127``). + + Create via :meth:`from_audio_filepaths`. + """ + + def __init__(self, total_steps: int, steps_per_stream: dict[int, int] | None = None): + self._bar = tqdm(total=total_steps, desc="Inference", unit="step", dynamic_ncols=True) + self._steps_per_stream = steps_per_stream or {} + self._stream_progress: dict[int, int] = {} + + def step(self, stream_id: int) -> None: + """Record one inference step for *stream_id* and advance the bar.""" + self._stream_progress[stream_id] = self._stream_progress.get(stream_id, 0) + 1 + stream_total = self._steps_per_stream.get(stream_id) + if stream_total is not None: + self._bar.set_postfix_str( + f"stream {stream_id}: {self._stream_progress[stream_id]}/{stream_total}", + refresh=False, + ) + self._bar.update(1) + + def finish(self) -> None: + """Close the underlying tqdm bar.""" + self._bar.close() + + @classmethod + def from_audio_filepaths( + cls, + audio_filepaths: list[str], + chunk_size_in_secs: float, + pad_audio_to_sec: float | None = None, + pad_silence_ratio: float | None = None, + pad_audio_by_sec: float | None = None, + ) -> StepProgressBar: + durations = calculate_durations_incl_padding( + audio_filepaths, + pad_audio_to_sec, + pad_silence_ratio, + pad_audio_by_sec, + ) + steps_per_stream = {idx: math.ceil(dur / chunk_size_in_secs) for idx, dur in enumerate(durations)} + return cls( + total_steps=sum(steps_per_stream.values()), + steps_per_stream=steps_per_stream, + ) diff --git a/nemo/collections/speechlm2/inference/vllm_omni/__init__.py b/nemo/collections/speechlm2/inference/vllm_omni/__init__.py new file mode 100644 index 000000000000..472e8728604d --- /dev/null +++ b/nemo/collections/speechlm2/inference/vllm_omni/__init__.py @@ -0,0 +1,59 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""NemotronDuplexH + EarTTS model/pipeline package for vLLM-Omni. + +This package lives inside NeMo and plugs into ``vllm-omni`` at runtime +through :func:`register_nemo_voicechat`. The function is auto-invoked in +every vllm-omni subprocess via the ``vllm_omni.general_plugins`` entry +point declared in NeMo's ``pyproject.toml`` (vllm-omni uses ``spawn`` for +stage children, so the entry-point hook is required — PYTHONPATH alone +is not enough). + +The plugin registers three things: + +* HF config ``"eartts"`` → :class:`EarTTSConfig` +* Model arch ``"NemotronDuplexHForCausalLM"`` → + :class:`nemo.collections.speechlm2.inference.vllm_omni.nemotron_duplex_h.nemotron_duplex_h.NemotronDuplexHForCausalLM` +* Model arch ``"EarTTSForCausalLM"`` → + :class:`nemo.collections.speechlm2.inference.vllm_omni.eartts.eartts.EarTTSForCausalLM` +* One-stage pipelines ``model_type = "nemotron_voicechat"`` and ``"eartts"``. + +Bundled deploy YAMLs for the independent engines live under ``deploy/``. +""" + +from __future__ import annotations + +from pathlib import Path + + +def default_deploy_yaml() -> Path: + """Return the absolute path to the bundled ``nemotron_voicechat.yaml``.""" + return Path(__file__).resolve().parent / "deploy" / "nemotron_voicechat.yaml" + + +def default_eartts_deploy_yaml() -> Path: + """Return the absolute path to the bundled single-stage ``eartts.yaml``.""" + return Path(__file__).resolve().parent / "deploy" / "eartts.yaml" + + +from nemo.collections.speechlm2.inference.vllm_omni.register import ( + register_nemo_voicechat, +) + +__all__ = [ + "default_deploy_yaml", + "default_eartts_deploy_yaml", + "register_nemo_voicechat", +] diff --git a/nemo/collections/speechlm2/inference/vllm_omni/checkpoint.py b/nemo/collections/speechlm2/inference/vllm_omni/checkpoint.py new file mode 100644 index 000000000000..00264d78afae --- /dev/null +++ b/nemo/collections/speechlm2/inference/vllm_omni/checkpoint.py @@ -0,0 +1,311 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Building and reading the vLLM-Omni wrapper checkpoint. + +``AsyncOmni(model=...)`` wants a directory laid out as:: + + / + config.json # {"model_type": "nemotron_voicechat"} + nemotron/ # converted NemotronDuplexH checkpoint + eartts/ # converted EarTTS checkpoint + speaker_latents/ + +Everything that reads or writes that layout lives here: conversion, the source +fingerprint that makes incremental builds safe, the small config patches +applied before an engine starts, and the two loaders that read values back out +of a built wrapper. Nothing here imports vLLM, so it can be exercised without +an engine. +""" + +import hashlib +import json +import os +import shutil +import tempfile +from pathlib import Path +from typing import Any + +import torch + +from nemo.utils import logging + +NEMOTRON_SUBDIR = "nemotron" +EARTTS_SUBDIR = "eartts" +_WRAPPER_CONFIG = {"model_type": "nemotron_voicechat"} +_SOURCE_MANIFEST = ".nemo_source.json" + + +def _checkpoint_fingerprint(model_path: str) -> dict[str, Any]: + """Cheap content identity for safe incremental wrapper construction.""" + root = Path(model_path) + config_path = root / "config.json" + if not config_path.is_file(): + raise FileNotFoundError(f"Checkpoint config not found: {config_path}") + weights = sorted(root.glob("*.safetensors")) + if not weights: + raise FileNotFoundError(f"No safetensors weights found in checkpoint: {root}") + return { + "config_sha256": hashlib.sha256(config_path.read_bytes()).hexdigest(), + "weights": [{"name": path.name, "size": path.stat().st_size} for path in weights], + } + + +def _read_source_manifest(path: str) -> dict[str, Any] | None: + try: + with open(path, encoding="utf-8") as fh: + value = json.load(fh) + return value if isinstance(value, dict) else None + except (FileNotFoundError, json.JSONDecodeError, OSError): + return None + + +def build_wrapper_checkpoint( + model_path: str, + wrapper_dir: str | None = None, + *, + nemotron_dtype: str = "float32", + eartts_precompute_batch_size: int = 256, + include_nemotron: bool = True, + include_eartts: bool = True, +) -> str: + """Build a wrapper checkpoint directory consumed by ``AsyncOmni(model=...)``. + + Layout:: + + / + config.json # {"model_type": "nemotron_voicechat"} + nemotron/ # converted NemotronDuplexH checkpoint + eartts/ # converted EarTTS checkpoint + speaker_latents/ + + Args: + model_path: Path to the source NemotronVoiceChat HF-format checkpoint + directory (``config.json`` + ``model.safetensors``). + wrapper_dir: Where to put the wrapper directory. Defaults to + ``/_vllm_omni_wrapper``, where ```` is + :func:`tempfile.gettempdir` and so honours ``$TMPDIR``. + nemotron_dtype: dtype for the converted Nemotron checkpoint. + eartts_precompute_batch_size: batch size used when baking out the + EarTTS subword-encoder lookup table. + + Returns: + Absolute path to the wrapper directory. If the wrapper directory + already exists and looks complete, the existing one is returned and + nothing is re-converted. + """ + src = os.path.normpath(model_path) + if wrapper_dir is None: + wrapper_dir = os.path.join(tempfile.gettempdir(), os.path.basename(src) + "_vllm_omni_wrapper") + wrapper_dir = os.path.abspath(wrapper_dir) + + nemotron_dir = os.path.join(wrapper_dir, NEMOTRON_SUBDIR) + eartts_dir = os.path.join(wrapper_dir, EARTTS_SUBDIR) + config_path = os.path.join(wrapper_dir, "config.json") + manifest_path = os.path.join(wrapper_dir, _SOURCE_MANIFEST) + source_fingerprint = _checkpoint_fingerprint(src) + + nemotron_ready = ( + os.path.isdir(nemotron_dir) + and os.path.isfile(os.path.join(nemotron_dir, "config.json")) + and os.path.isfile(os.path.join(nemotron_dir, "model.safetensors")) + ) + eartts_ready = ( + os.path.isdir(eartts_dir) + and os.path.isfile(os.path.join(eartts_dir, "config.json")) + and os.path.isfile(os.path.join(eartts_dir, "model.safetensors")) + ) + config_ready = os.path.isfile(config_path) + + if not include_nemotron and not include_eartts: + raise ValueError("At least one vLLM-Omni component must be requested") + + manifest = _read_source_manifest(manifest_path) + if manifest is None: + adding_to_unverified_partial_wrapper = (include_nemotron and not nemotron_ready and eartts_ready) or ( + include_eartts and not eartts_ready and nemotron_ready + ) + if adding_to_unverified_partial_wrapper: + raise ValueError( + f"Cannot safely add a component to wrapper {wrapper_dir}: " + f"{_SOURCE_MANIFEST} is missing, so the existing component's " + "source checkpoint cannot be verified. Use a fresh wrapper_dir." + ) + elif manifest.get("source") != source_fingerprint: + logging.warning( + "Wrapper source checkpoint changed; rebuilding converted components in %s", + wrapper_dir, + ) + for component_dir in (nemotron_dir, eartts_dir): + if os.path.isdir(component_dir): + shutil.rmtree(component_dir) + nemotron_ready = False + eartts_ready = False + manifest = None + else: + if include_nemotron and nemotron_ready and manifest.get("nemotron", {}).get("dtype") != nemotron_dtype: + shutil.rmtree(nemotron_dir) + nemotron_ready = False + if ( + include_eartts + and eartts_ready + and manifest.get("eartts", {}).get("precompute_batch_size") != eartts_precompute_batch_size + ): + shutil.rmtree(eartts_dir) + eartts_ready = False + + if (not include_nemotron or nemotron_ready) and (not include_eartts or eartts_ready) and config_ready: + if manifest is None: + # Wrapper is complete but carries no source manifest. Stamp one + # now: no component is being added, so this cannot mix checkpoints. + logging.warning( + "Adopting vLLM-Omni wrapper without source manifest: %s", + wrapper_dir, + ) + adopted_manifest: dict[str, Any] = {"source": source_fingerprint} + if nemotron_ready: + adopted_manifest["nemotron"] = {"dtype": nemotron_dtype} + if eartts_ready: + adopted_manifest["eartts"] = {"precompute_batch_size": eartts_precompute_batch_size} + with open(manifest_path, "w", encoding="utf-8") as fh: + json.dump(adopted_manifest, fh, indent=2, sort_keys=True) + logging.info(f"Reusing existing vllm-omni wrapper checkpoint at {wrapper_dir}") + return wrapper_dir + + os.makedirs(wrapper_dir, exist_ok=True) + + if include_nemotron and not nemotron_ready: + # Convert the Nemotron LLM with the existing DuplexSTT converter. + # That converter's output (HF NemotronH config + filtered weights) + # is consumed directly by NemotronDuplexHForCausalLM's WeightsMapper. + if os.path.isdir(nemotron_dir): + shutil.rmtree(nemotron_dir) + logging.info(f"Converting Nemotron LLM into {nemotron_dir} ...") + from nemo.collections.speechlm2.inference.vllm_omni.scripts.convert_duplex_stt_checkpoint import ( + convert_to_vllm_format as convert_nemotron, + ) + + convert_nemotron( + checkpoint_path=src, + output_dir=nemotron_dir, + dtype=nemotron_dtype, + ) + + if include_eartts and not eartts_ready: + if os.path.isdir(eartts_dir): + shutil.rmtree(eartts_dir) + logging.info(f"Converting EarTTS into {eartts_dir} ...") + from nemo.collections.speechlm2.inference.vllm_omni.scripts.convert_duplex_eartts_checkpoint import ( + convert_to_vllm_format as convert_eartts, + ) + + convert_eartts( + outdir=eartts_dir, + config=os.path.join(src, "config.json"), + model_path=os.path.join(src, "model.safetensors"), + precompute_batch_size=eartts_precompute_batch_size, + ) + + if not config_ready: + with open(config_path, "w", encoding="utf-8") as fh: + json.dump(_WRAPPER_CONFIG, fh, indent=2) + + completed_manifest: dict[str, Any] = {"source": source_fingerprint} + if os.path.isfile(os.path.join(nemotron_dir, "model.safetensors")): + completed_manifest["nemotron"] = {"dtype": nemotron_dtype} + if os.path.isfile(os.path.join(eartts_dir, "model.safetensors")): + completed_manifest["eartts"] = {"precompute_batch_size": eartts_precompute_batch_size} + with open(manifest_path, "w", encoding="utf-8") as fh: + json.dump(completed_manifest, fh, indent=2, sort_keys=True) + + return wrapper_dir + + +def write_nemotron_inference_overrides(wrapper_dir: str, overrides: dict[str, Any]) -> None: + """Update inference settings in the converted Nemotron ``config.json``. + + ``NemotronDuplexHForCausalLM`` reads some settings off its HF config at + load time -- the user-channel logit boosts, which cannot be delivered per + request because the ASR head's logits never reach vLLM's sampler. Those are + still chosen per run in the inference yaml, so the small JSON is rewritten + here before the stage child starts, rather than re-converting weights. + + Keys whose value is ``None`` are removed, so clearing a boost in the config + clears it in the engine too. + """ + config_path = os.path.join(wrapper_dir, "nemotron", "config.json") + if not os.path.isfile(config_path): + return + with open(config_path, encoding="utf-8") as fh: + config = json.load(fh) + + changed = False + for key, value in overrides.items(): + if value is None: + if config.pop(key, None) is not None: + changed = True + elif config.get(key) != value: + config[key] = value + changed = True + if not changed: + return + + with open(config_path, "w", encoding="utf-8") as fh: + json.dump(config, fh, indent=2) + logging.info(f"Updated Nemotron inference overrides in {config_path}: {overrides}") + + +def load_speaker_latent(eartts_dir: str, speaker_name: str) -> torch.Tensor: + """Load ``/speaker_latents/.pt`` (saved by the + EarTTS converter) and return a contiguous CPU tensor of shape + ``(Tref, hidden_size)``. + """ + latents_dir = os.path.join(eartts_dir, "speaker_latents") + latent_path = os.path.join(latents_dir, f"{speaker_name}.pt") + if not os.path.isfile(latent_path): + available = [] + if os.path.isdir(latents_dir): + available = sorted( + os.path.splitext(name)[0] for name in os.listdir(latents_dir) if name.endswith(".pt") + ) + raise FileNotFoundError( + f"Speaker latent for '{speaker_name}' not found at {latent_path}. " + f"Registered speakers: {available or '(none)'}. " + "Pick a speaker_name present in the EarTTS checkpoint, or re-run the " + "EarTTS converter on a checkpoint that contains the requested " + "audio_prompt_latents." + ) + latent = torch.load(latent_path, weights_only=False) + if isinstance(latent, torch.Tensor) and latent.dim() == 3: + latent = latent[0] + if not isinstance(latent, torch.Tensor) or latent.dim() != 2: + raise ValueError( + f"Expected speaker latent at {latent_path} to be a 2-D tensor [Tref, hidden], " + f"got {type(latent).__name__} with shape " + f"{tuple(latent.shape) if isinstance(latent, torch.Tensor) else 'n/a'}" + ) + return latent.detach().to(torch.float32).cpu().contiguous() + + +def compute_prefill_len(model_dir: str, system_prompt: str) -> int: + """Length of the prefill chunk fed to NemotronDuplexH for a + given system prompt. Mirrors the in-model tokenization: + ``[BOS] + tokenizer.encode(prompt) + [EOS]``. + """ + from transformers import AutoTokenizer + + from nemo.collections.speechlm2.inference.vllm_omni.nemotron_duplex_h.nemotron_duplex_h import ( + NemotronDuplexHForCausalLM, + ) + + tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True) + return NemotronDuplexHForCausalLM.compute_prefix_len(tokenizer, system_prompt) diff --git a/nemo/collections/speechlm2/inference/vllm_omni/deploy/eartts.yaml b/nemo/collections/speechlm2/inference/vllm_omni/deploy/eartts.yaml new file mode 100644 index 000000000000..34750c82eb1d --- /dev/null +++ b/nemo/collections/speechlm2/inference/vllm_omni/deploy/eartts.yaml @@ -0,0 +1,45 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Single-stage EarTTS runtime. The model path passed to AsyncOmni is the +# converted ``eartts/`` directory itself. +async_chunk: false +trust_remote_code: true +enable_prefix_caching: false +enable_chunked_prefill: false +distributed_executor_backend: uni + +stages: + - stage_id: 0 + # One conditional + one unconditional request per VoiceChat stream. + # The no-CFG path uses only one slot. + max_num_seqs: 2 + max_num_batched_tokens: 2048 + max_model_len: 2048 + gpu_memory_utilization: 0.30 + enforce_eager: false + # Keep the synchronous scheduler: paired streaming requests must observe + # the same completed step before their next updates are admitted. + async_scheduling: false + skip_tokenizer_init: true + dtype: float32 + devices: "0" + compilation_config: + cudagraph_mode: PIECEWISE + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 2048 + detokenize: false diff --git a/nemo/collections/speechlm2/inference/vllm_omni/deploy/nemotron_voicechat.yaml b/nemo/collections/speechlm2/inference/vllm_omni/deploy/nemotron_voicechat.yaml new file mode 100644 index 000000000000..536231160e94 --- /dev/null +++ b/nemo/collections/speechlm2/inference/vllm_omni/deploy/nemotron_voicechat.yaml @@ -0,0 +1,62 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Deploy config for the single-stage ``nemotron_voicechat`` NemotronDuplexH +# streaming pipeline. EarTTS is served by ``eartts.yaml`` in a second +# AsyncOmni engine and is coordinated by NeMo. +# +# The model directory passed to ``AsyncOmni(model=...)`` is a small +# user-managed wrapper directory with this layout:: +# +# / +# config.json # {"model_type": "nemotron_voicechat"} +# nemotron/ # directory or symlink → Nemotron ckpt +# eartts/ # loaded separately by the EarTTS runtime +# +# ``model_type = nemotron_voicechat`` dispatches to the one-stage pipeline. +# ``model_subdir`` / ``tokenizer_subdir`` tell the engine to load the nested +# Nemotron checkpoint; see +# ``vllm_omni/engine/stage_init_utils.py:_resolve_model_tokenizer_paths``. +async_chunk: false +trust_remote_code: true +enable_prefix_caching: false +enable_chunked_prefill: false +distributed_executor_backend: uni + +stages: + # Nemotron-Duplex-H (autoregressive text + optional ASR/function channel). + # PIECEWISE compilation only; see the model docstring for why FULL + # cudagraph mode is unsafe with this streaming setup. + - stage_id: 0 + max_num_seqs: 1 + # Prompt tokens plus one token per 80 ms frame. + max_num_batched_tokens: 2048 + max_model_len: 2048 + gpu_memory_utilization: 0.45 + enforce_eager: false + async_scheduling: true + devices: "0" + model_subdir: nemotron + tokenizer_subdir: nemotron + engine_extras: + logits_processors: + - nemo.collections.speechlm2.inference.vllm_omni.nemotron_duplex_h.sampling:SharedTextSamplingLogitsProcessor + compilation_config: + cudagraph_mode: PIECEWISE + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 2048 + detokenize: false diff --git a/nemo/collections/speechlm2/inference/vllm_omni/eartts/__init__.py b/nemo/collections/speechlm2/inference/vllm_omni/eartts/__init__.py new file mode 100644 index 000000000000..b8f12a085699 --- /dev/null +++ b/nemo/collections/speechlm2/inference/vllm_omni/eartts/__init__.py @@ -0,0 +1,17 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from nemo.collections.speechlm2.inference.vllm_omni.eartts.configuration_eartts import EarTTSConfig + +__all__ = ["EarTTSConfig"] diff --git a/nemo/collections/speechlm2/inference/vllm_omni/eartts/configuration_eartts.py b/nemo/collections/speechlm2/inference/vllm_omni/eartts/configuration_eartts.py new file mode 100644 index 000000000000..be85e27cde3b --- /dev/null +++ b/nemo/collections/speechlm2/inference/vllm_omni/eartts/configuration_eartts.py @@ -0,0 +1,217 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""HuggingFace-style configuration for the EarTTS model. + +The configuration mirrors the fields that +:class:`nemo.collections.speechlm2.inference.vllm_omni.eartts.eartts.EarTTSForCausalLM` +reads from ``vllm_config.model_config.hf_config``: + +* Gemma3 backbone fields consumed by ``Gemma3Model`` (``hidden_size``, + ``intermediate_size``, ``num_hidden_layers``, ``num_attention_heads``, + ``num_key_value_heads``, ``head_dim``, ``vocab_size``, + ``max_position_embeddings``, ``query_pre_attn_scalar``, + ``attention_bias``, ``rms_norm_eps``, ``layer_types``, + ``sliding_window``, ``rope_local_base_freq``, ``rope_theta``, + ``rope_scaling``, ``hidden_activation``, ``tie_word_embeddings``, + ``final_logit_softcapping``, ``attn_logits_soft_cap``, + ``use_bidirectional_attention``, ``is_causal``). + +* MaskGIT sampler fields (``num_quantizers``, ``codebook_size``, + ``num_iter``, ``top_p_or_k``, ``noise_scale``, ``exponent``, + ``latent_size``, ``mog_low_rank``, ``mog_num_layers``, + ``mog_num_predictions``, ``mog_min_log_std``, ``mog_eps``). + +* Subword embedding / fusion fields consumed by + :class:`EarTTSInputEmbedding` (``emb_vocab_size``, + ``use_gated_fusion_for_text_audio``, + ``use_audio_prompt_frozen_projection``). The original NeMo model used + a character-aware subword encoder + subword-flag + BOS/EOS additive + embeddings; all of those operations are deterministic per token id + and are baked out at checkpoint-conversion time into a single + ``nn.Embedding`` of size ``(emb_vocab_size, hidden_size)``. + +vLLM's ``patch_rope_parameters`` (transformers-v4 path) auto-populates +``config.rope_parameters`` from ``rope_scaling`` + ``rope_theta`` during +config loading, so the Gemma3 backbone (which expects +``config.rope_parameters``) works without any extra plumbing here. +""" + +from typing import Optional + +from transformers import AutoConfig, PretrainedConfig + + +class EarTTSConfig(PretrainedConfig): + model_type = "eartts" + + def __init__( + self, + # Gemma 3 backbone + hidden_size: int = 1152, + context_hidden_size: int = 1536, + intermediate_size: int = 4608, + num_hidden_layers: int = 28, + num_attention_heads: int = 16, + num_key_value_heads: int = 16, + head_dim: int = 72, + # ``vocab_size`` controls the width of the logits tensor returned + # by ``EarTTSForCausalLM.compute_logits`` — vLLM's sampler and + # ``LogitsProcessor`` size their working buffers from + # ``config.vocab_size``, so it must match the dummy logits the + # model produces. The model emits a 2-class placeholder + # (``[0, -inf]``) so the sampler's argmax always picks index 0; + # the real audio output is the codes tensor exposed via + # ``make_omni_output``. ``2`` is the minimum that keeps vLLM's + # sampler happy. + vocab_size: int = 2, + max_position_embeddings: int = 131072, + # MaskGIT / MoG sampling + num_quantizers: int = 31, + codebook_size: int = 1024, + num_iter: int = 8, + top_p_or_k: float = 0.8, + noise_scale: float = 0.8, + exponent: float = 3.0, + latent_size: int = 512, + mog_low_rank: int = 64, + mog_num_layers: int = 3, + mog_num_predictions: int = 1024, + mog_min_log_std: float = -4.0, + mog_eps: float = 1e-6, + # Classifier-free guidance. The converter exports both fields and + # ``null_emb``; the runtime reads them unless the request overrides + # ``cfg_scale``. + enable_guidance: bool = False, + guidance_scale: float = 0.5, + # Gemma3-specific attributes required by Gemma3Model + query_pre_attn_scalar: float = 256.0, + attention_bias: bool = False, + rms_norm_eps: float = 1e-6, + layer_types: Optional[list] = None, + sliding_window: Optional[int] = 4096, + rope_local_base_freq: float = 10000.0, + # NeMo / EarTTS uses 1M for the global-attention RoPE base. + rope_theta: float = 1000000.0, + rope_scaling: Optional[dict] = None, + hidden_activation: str = "gelu_pytorch_tanh", + tie_word_embeddings: bool = True, + final_logit_softcapping: Optional[float] = None, + attn_logits_soft_cap: Optional[float] = None, + use_bidirectional_attention: bool = False, + is_causal: bool = True, + # Subword encoding. The character-aware subword encoder / + # subword-flag / BOS-EOS embedding tables that NeMo applied at + # runtime are precomputed at checkpoint conversion time into a + # single ``(emb_vocab_size, hidden_size)`` lookup, so only the + # vocab size and the audio-side fusion / projection toggles + # remain as runtime config. + emb_vocab_size: int = 151936, + use_gated_fusion_for_text_audio: bool = True, + use_audio_prompt_frozen_projection: bool = False, + # HF-canonical model dtype (replaces the deprecated + # ``torch_dtype``). Forwarded to ``PretrainedConfig`` so it is + # converted into a real ``torch.dtype`` and exposed as + # ``config.dtype``. + dtype: str = "float32", + # Text-channel specials, copied from the source VoiceChat tokenizer + # at conversion time. Used to pad prefill text and to force codec + # silence when the incoming text token is EOS. + pad_token_id: Optional[int] = None, + eos_token_id: Optional[int] = None, + **kwargs, + ): + # Gemma3 backbone + self.hidden_size = hidden_size + self.context_hidden_size = context_hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + + # MaskGIT / MoG sampling + self.num_quantizers = num_quantizers + self.codebook_size = codebook_size + self.num_iter = num_iter + self.top_p_or_k = top_p_or_k + self.noise_scale = noise_scale + self.exponent = exponent + self.latent_size = latent_size + self.mog_low_rank = mog_low_rank + self.mog_num_layers = mog_num_layers + self.mog_num_predictions = mog_num_predictions + self.mog_min_log_std = mog_min_log_std + self.mog_eps = mog_eps + self.enable_guidance = enable_guidance + self.guidance_scale = guidance_scale + + # Gemma3-specific attributes + self.query_pre_attn_scalar = query_pre_attn_scalar + self.attention_bias = attention_bias + self.rms_norm_eps = rms_norm_eps + # Default all layers to global attention if not specified. + self.layer_types = ( + layer_types if layer_types is not None else ["full_attention"] * num_hidden_layers + ) + self.sliding_window = sliding_window + self.rope_local_base_freq = rope_local_base_freq + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.hidden_activation = hidden_activation + self.final_logit_softcapping = final_logit_softcapping + self.attn_logits_soft_cap = attn_logits_soft_cap + self.use_bidirectional_attention = use_bidirectional_attention + self.is_causal = is_causal + + # Subword encoding (precomputed lookup; see class docstring). + self.emb_vocab_size = emb_vocab_size + self.use_gated_fusion_for_text_audio = use_gated_fusion_for_text_audio + self.use_audio_prompt_frozen_projection = use_audio_prompt_frozen_projection + + # Forward HF-owned fields (``tie_word_embeddings`` and ``dtype``) + # to ``PretrainedConfig`` so they round-trip through + # save/load_pretrained and are visible as ``config.dtype`` / + # ``config.tie_word_embeddings``. Without this, user-supplied + # values silently fall back to PretrainedConfig's defaults. + super().__init__( + tie_word_embeddings=tie_word_embeddings, + dtype=dtype, + pad_token_id=pad_token_id, + eos_token_id=eos_token_id, + **kwargs, + ) + + +def register_eartts_config() -> None: + """Make ``model_type: "eartts"`` resolvable by ``AutoConfig``. + + Idempotent, because it has two callers by necessity and either may run + first: :func:`register_nemo_voicechat` (the plugin entry point) and the + import below. The import-time call is what covers subprocesses that reach + this module without the plugin -- ``StageEngineCoreProc`` unpickles a + config and calls ``AutoConfig.from_pretrained`` -- and mirrors the pattern + used by other vllm-omni custom configs (voxcpm, fish_speech, + mammoth_moda2, ...). + """ + try: + AutoConfig.register(EarTTSConfig.model_type, EarTTSConfig) + except ValueError: + # transformers raises when the model_type is already registered, which + # is the expected outcome for whichever caller runs second. + pass + + +register_eartts_config() diff --git a/nemo/collections/speechlm2/inference/vllm_omni/eartts/eartts.py b/nemo/collections/speechlm2/inference/vllm_omni/eartts/eartts.py new file mode 100644 index 000000000000..d88d9e066313 --- /dev/null +++ b/nemo/collections/speechlm2/inference/vllm_omni/eartts/eartts.py @@ -0,0 +1,1945 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Inference-only EarTTS model definition for vLLM-Omni. + +The model architecture (RMSNorm, MLP, MLPLayer, GatedProjectedSumRMSNorm, +PrecomputedSubwordEmbedding, EarTTSInputEmbedding, MoGHead, +MaskGITSampler, EarTTSModel) matches the PyTorch EarTTS modules. +Classifier-free guidance (CFG) is driven by per-request metadata and keeps +its role, pair, and scale contract in model-owned stable-address buffers. + +The original NeMo model used a character-aware subword encoder (a small +transformer over per-character embeddings) followed by additive +subword-continuation and BOS/EOS flag embeddings to embed text tokens. +Those operations are deterministic per token id, so the checkpoint +converter runs them once over the full vocabulary and stores the result +as a single ``nn.Embedding`` (see :class:`PrecomputedSubwordEmbedding`). + +The outer :class:`EarTTSForCausalLM` exposes the minimal vLLM-Omni +preprocess/postprocess hooks. Per-request inputs are passed via +``additional_information``. + +Inputs (only one mode — streaming text token ids), named after the +categories of :class:`~vllm_omni.data_entry_keys.OmniPayload` so that the +same two keys work whether they arrive on the request or over an +inter-stage connector, which accepts nothing outside that schema: + +* ``embed.voice`` (prefill only): Tensor of shape + ``(Tref, hidden_size)``. The user-supplied speaker latent that + replaces ``embed_code(rvq_sum(acoustic_tokens))`` on every pre-BOS + prefill position. ``Tref`` is also the prefill placeholder length + (the user passes ``prompt_token_ids = [0] * Tref``). +* ``ids.output`` (every decode step): Python ``list[int]`` of the text + tokens the producer has sent most recently. :meth:`preprocess` takes + its **last** entry, so decode step ``k`` consumes ``t_k`` whether the + producer sends one token per step or a growing history. + +There is no whole-utterance text path: callers must always provide token +ids per step via the streaming-text contract above. + +Per-step flow: + +1. ``preprocess`` writes the per-token tensors consumed by + :class:`EarTTSInputEmbedding` — ``acoustic_tokens (BTx31)``, + ``text_tokens (BT)``, ``text_mask (BT)``, ``bos_mask (BT)``, + ``speaker_latent (BT x hidden_size)`` — into the model-owned + static-address buffers at the request's flat-batch offset. Returns + placeholder ``input_ids`` and the ``inputs_embeds`` slice it + received from the runner unchanged (the actual embedding is + computed inside the compiled ``forward``; the buffer's contents + are ignored). + + Prefill is fully derived from ``speaker_latent``: + + * ``acoustic_tokens`` = ``model.sil_tokens`` broadcast to every + prefill position (only the BOS frame's audio embedding actually + contributes to the model output; the rest are replaced by the + speaker latent inside :class:`EarTTSInputEmbedding`). + * ``text_tokens`` = ``[PAD] * (Tref - 1) + [EOS]``. + * ``text_mask`` = ``[0] * (Tref - 2) + [1, 1]``. + * ``bos_mask`` = ``[0] * (Tref - 1) + [1]``. + * ``speaker_latent`` = the user-supplied ``embed.voice`` tensor. + + Decode each step (chooses ``acoustic_tokens`` in this order): + + * ``text_token == EOS`` (``2``) → ``model.sil_tokens``. + * First decode step (``ear_decode_offset == 0``) → the acoustic pad + id (``codebook_size``) broadcast across all quantizers. + * Otherwise → previous-step codes stashed by :meth:`postprocess` + as ``last_acoustic_codes``. + + ``text_tokens = ids.output[-1]``, + ``text_mask = 1``, ``bos_mask = 0``, + ``speaker_latent = 0`` (no replacement on decode). + +2. ``forward`` slices the buffers up to ``num_tokens`` and calls the + compiled :class:`EarTTSModel` (embedding + Gemma3 backbone). The + compiled :class:`EarTTSSamplerModel` (MaskGIT) is invoked + conditionally on decode positions. Generated codes are copied into + a stable ``_out_codes`` buffer for :meth:`make_omni_output`. + +3. ``compute_logits`` returns trivial logits so vLLM's standard + sampler always picks index ``0`` — the actual audio output is the + codes tensor exposed via :meth:`make_omni_output`. + +4. ``postprocess`` stashes the last frame's codes as + ``last_acoustic_codes`` for the next step's :meth:`preprocess`. +""" + +import bisect +import hashlib +import math +from collections.abc import Iterable +from typing import Any, Optional, Union + +import numpy as np +import torch +from torch import nn +from transformers.generation.logits_process import ( + TopKLogitsWarper, + TopPLogitsWarper, +) +from vllm.compilation.backends import set_model_tag +from vllm.compilation.decorators import ( + ignore_torch_compile, + support_torch_compile, +) +from vllm.config import CUDAGraphMode, VllmConfig +from vllm.forward_context import BatchDescriptor, get_forward_context +from vllm.model_executor.models.gemma3 import Gemma3Model +from vllm.model_executor.models.interfaces import SupportsPP +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper +from vllm.sequence import IntermediateTensors + +from vllm_omni.model_executor.models.output_templates import OmniOutput + + + +def _prepare_cfg_sampling_batch( + hidden_states: torch.Tensor, + cfg_enabled: torch.Tensor, + cfg_is_uncond: torch.Tensor, + cfg_pair_id: torch.Tensor, + cfg_scale: torch.Tensor, + valid: torch.Tensor, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + """Order a complete CFG batch and describe its logical pairs. + + All operations stay on-device: CUDA graph replay can therefore change the + role/pair/scale buffers without depending on Python-side state. A batch is + guided only when every valid row has exactly one enabled opposite-role + companion with the same pair id. Padded rows are ignored; the conditional + row's scale is authoritative for each pair. + + Returns ordered hidden states, ordered role/scale tensors, an active-row + mask, each row's partner and conditional representative indices, and the + inverse permutation used to restore the runner's original row order. + """ + batch_size = int(hidden_states.shape[0]) + identity = torch.arange(batch_size, device=hidden_states.device) + + pair_match = ( + valid[:, None] + & valid[None, :] + & cfg_enabled[:, None] + & cfg_enabled[None, :] + & (cfg_pair_id[:, None] == cfg_pair_id[None, :]) + & (cfg_is_uncond[:, None] != cfg_is_uncond[None, :]) + ) + partner_count = pair_match.sum(dim=1) + partner = pair_match.to(torch.long).argmax(dim=1) + complete = valid.any() & ((~valid) | (cfg_enabled & (partner_count == 1))).all() + + # Lexicographic order: valid rows first, then conditional before + # unconditional, with pair ids sorted identically inside both role blocks. + order = torch.argsort(cfg_pair_id, stable=True) + order = order[torch.argsort(cfg_is_uncond[order].to(torch.long), stable=True)] + order = order[torch.argsort((~valid[order]).to(torch.long), stable=True)] + order = torch.where(complete, order, identity) + inverse_order = torch.argsort(order) + + hidden_states = hidden_states[order] + cfg_is_uncond = cfg_is_uncond[order] + cfg_pair_id = cfg_pair_id[order] + cfg_scale = cfg_scale[order] + valid = valid[order] + active = complete & valid + + ordered_pair_match = ( + active[:, None] + & active[None, :] + & (cfg_pair_id[:, None] == cfg_pair_id[None, :]) + & (cfg_is_uncond[:, None] != cfg_is_uncond[None, :]) + ) + partner = ordered_pair_match.to(torch.long).argmax(dim=1) + conditional_rep = torch.where(cfg_is_uncond, partner, identity) + conditional_scale = torch.where( + cfg_is_uncond, + cfg_scale[partner], + cfg_scale, + ) + return ( + hidden_states, + cfg_is_uncond, + conditional_scale, + active, + partner, + conditional_rep, + inverse_order, + ) + + +def _apply_cfg_after_mlp( + x: torch.Tensor, + cfg_is_uncond: torch.Tensor, + cfg_scale: torch.Tensor, + cfg_active: torch.Tensor, + cfg_partner: torch.Tensor, +) -> torch.Tensor: + """Apply EarTTS CFG to MoG MLP outputs before all projections.""" + partner_x = x[cfg_partner] + conditional_x = torch.where(cfg_is_uncond[:, None], partner_x, x) + unconditional_x = torch.where(cfg_is_uncond[:, None], x, partner_x) + guided_x = conditional_x + cfg_scale[:, None].to(x.dtype) * ( + conditional_x - unconditional_x + ) + return torch.where(cfg_active[:, None], guided_x, x) + + +# --------------------------------------------------------------------------- +# Shared EarTTS building blocks, matching the native DuplexEARTTS modules. +# --------------------------------------------------------------------------- + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.zeros(dim)) + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + # Normalize in fp32 and cast back at the end, so low-precision + # activations do not lose the mean-square accumulation. + output = self._norm(x.float()) + # Llama does x.to(float16) * w whilst Gemma3 is (x * w).to(float16) + output = output * (1.0 + self.weight.float()) + return output.type_as(x) + + +class MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + ): + super().__init__() + self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) + self.act_fn = nn.GELU(approximate="tanh") + + def forward(self, x: torch.Tensor) -> torch.Tensor: + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +class MLPLayer(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + eps: float = 1e-6, + ): + super().__init__() + self.pre_norm = RMSNorm(hidden_size, eps=eps) + self.mlp = MLP(hidden_size, intermediate_size) + self.post_norm = RMSNorm(hidden_size, eps=eps) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = self.pre_norm(x) + y = self.mlp(y) + y = self.post_norm(y) + x = x + y + return x + + +class GatedProjectedSumRMSNorm(nn.Module): + def __init__( + self, + audio_dim, + text_dim, + hidden_dim, + final_norm=True, + num_codebooks=31, + init_residual_scale=0.5, + ): + super().__init__() + self.num_codebooks = num_codebooks + + self.audio_proj = nn.Linear(audio_dim, hidden_dim) + self.text_proj = nn.Linear(text_dim, hidden_dim) + + nn.init.normal_(self.audio_proj.weight, mean=0.0, std=0.015) + nn.init.zeros_(self.audio_proj.bias) + nn.init.normal_(self.text_proj.weight, mean=0.0, std=0.015) + nn.init.zeros_(self.text_proj.bias) + + # FP32 gate params + self.gate = nn.Parameter( + torch.zeros(hidden_dim, dtype=torch.float32), requires_grad=False + ) + self.residual_scale = nn.Parameter( + torch.tensor(init_residual_scale, dtype=torch.float32), + requires_grad=False, + ) + + self.final_norm = RMSNorm(hidden_dim) if final_norm else nn.Identity() + + def forward(self, audio_emb, text_emb): + 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) + + dtype = audio_h.dtype + + gate = torch.sigmoid(self.gate) # FP32 + res = torch.sigmoid(self.residual_scale) # FP32 + + h = gate.to(dtype) * audio_h + (1 - gate).to(dtype) * text_h + h = res.to(dtype) * h + h = self.final_norm(h.float()).to(dtype) + + return h + + +class PrecomputedSubwordEmbedding(nn.Module): + """Per-token text embedding lookup baked out at checkpoint-conversion time. + + The original NeMo model embeds text with a character-aware subword + encoder (a small transformer over per-character embeddings) followed + by additive subword-continuation and BOS/EOS flag embeddings. All of + those operations are deterministic per token id, so the converter + runs them once over the full vocabulary and stores the result here + as a single ``nn.Embedding``. + """ + + def __init__(self, vocab_size: int, hidden_size: int): + super().__init__() + self.embed_subwords = nn.Embedding(vocab_size, hidden_size) + + def forward(self, subword_ids: torch.Tensor) -> torch.Tensor: + return self.embed_subwords(subword_ids) + + +class EarTTSInputEmbedding(nn.Module): + """Module that takes text tokens, audio tokens and prepares input + embedding for EarTTS model. + """ + + def __init__(self, config): + super().__init__() + + hidden_size = config.hidden_size + vocab_size = config.emb_vocab_size + + # allows to embed acoustic tokens into a single embeddings + self.rvq_embs = nn.ModuleList( + [ + nn.Embedding(config.codebook_size + 1, config.latent_size) + for _ in range(config.num_quantizers) + ] + ) + self.embed_code = nn.Linear(config.latent_size, hidden_size, bias=False) + # Pre-computed per-token text embedding lookup. Replaces the + # original char-aware subword encoder + subword-flag + BOS/EOS + # additive embeddings; all of those are deterministic per token + # id and are baked into this single table by the checkpoint + # converter. + self.embed_subword = PrecomputedSubwordEmbedding(vocab_size, hidden_size) + self.bos_emb = nn.Parameter(torch.empty(hidden_size)) + # Learned classifier-free text-conditioning embedding. The audio and + # speaker branches remain unchanged for unconditional rows. + self.null_emb = nn.Parameter(torch.empty(hidden_size)) + + self.use_gated_fusion_for_text_audio = config.use_gated_fusion_for_text_audio + if self.use_gated_fusion_for_text_audio: + self.gated_fusion_audio_text = GatedProjectedSumRMSNorm( + hidden_size, hidden_size, hidden_size, config.num_quantizers + ) + + self.use_audio_prompt_frozen_projection = ( + config.use_audio_prompt_frozen_projection + ) + if self.use_audio_prompt_frozen_projection: + self.audio_prompt_projection_W = nn.Parameter( + torch.empty(hidden_size, hidden_size), + requires_grad=False, + ) + + def forward( + self, + acoustic_tokens: torch.Tensor, + text_tokens: torch.Tensor, + text_mask: torch.Tensor, + bos_mask: torch.Tensor, + speaker_latent: Optional[torch.Tensor] = None, + cfg_is_uncond: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """ + Works for context and generation phases to prepare total input + embeddings for EarTTS model. + + Inputs: + acoustic_tokens: (BT x 31) - audio tokens + text_tokens: (BT) - text token to embed + text_mask: (BT) - masks text embeddings for prefill + bos_mask: (BT) - specifies where BOS is applied (first frame of prefill) + speaker_latent: (BT x hidden_size) - external speaker latent. + Non-zero rows replace ``embed_code(rvq_sum(...))`` at + pre-BOS prefill positions; zero rows (decode steps and + the BOS frame) leave ``audio_emb`` untouched. Pass an + all-zero tensor on decode. + + Returns: + embedding of shape (BT x dim) + """ + + # prepare bos emb that is applied to audio embedding + bos_emb = bos_mask.unsqueeze(1) * self.bos_emb # BT x dim + + acoustic_tokens = acoustic_tokens.transpose(0, 1) # 31 x BT + audio_emb = sum( + emb(acoustic_tokens[i]) for i, emb in enumerate(self.rvq_embs) + ) # BT x latent_size + audio_emb = self.embed_code(audio_emb) # BT x hidden_size + + if self.use_audio_prompt_frozen_projection: + if speaker_latent is None: + # No external latent -> derive one from the acoustic + # tokens, matching DuplexEARTTS when no speaker prompt + # is supplied. vLLM-Omni callers always pass a real or + # zero latent, so they take the other branch. + latent_provided = torch.zeros_like(bos_mask).unsqueeze(-1) + latent = torch.nn.functional.linear( + audio_emb, self.audio_prompt_projection_W.T + ) + else: + # ``latent_provided`` is non-zero exactly on the rows + # the user populated with a real speaker latent + # (prefill pre-BOS positions). Decode rows are filled + # with zeros by ``preprocess``, so they read as "not + # provided" here. + latent_provided = ( + speaker_latent.abs().sum(-1, keepdim=True) > 0 + ) # (BT, 1) + latent = speaker_latent + + # Replace only at pre-BOS positions of prefill -- i.e. + # ``bos_mask == 0 AND latent was actually provided``. This + # excludes: + # * the BOS frame (``bos_mask == 1``), where + # ``embed_code(acoustic_tokens)`` of ``sil_tokens`` + # survives (this is the audio_emb the backbone sees on + # the BOS frame). + # * AR decode steps (``latent_provided == False`` because + # ``speaker_latent`` is all zeros). + replace_mask = (bos_mask.unsqueeze(-1) == 0) & latent_provided + audio_emb = torch.where(replace_mask, latent, audio_emb) + + audio_emb = audio_emb + bos_emb + + # Embed text tokens via the pre-computed lookup (subword-flag and + # BOS/EOS additions are baked into the table at conversion time). + # Apply the mask that zeroes this embedding on prefill positions. + text_emb = self.embed_subword(text_tokens) * text_mask.unsqueeze(1) # BT x dim + if cfg_is_uncond is not None: + text_emb = torch.where( + cfg_is_uncond.unsqueeze(1), + self.null_emb.to(text_emb.dtype), + text_emb, + ) + + # prepare total embedding by combining audio and text branches + if self.use_gated_fusion_for_text_audio: + # Gated fusion needs ``audio_emb`` and ``text_emb`` as + # separate inputs (it learns a per-feature gate to mix + # them), which is why neither branch can be folded into a + # single precomputed ``inputs_embeds`` tensor outside the + # compiled forward. + total_emb = self.gated_fusion_audio_text(audio_emb, text_emb) + else: + total_emb = audio_emb + text_emb # BT x dim + return total_emb + + +def gumbel_like(tensor: torch.Tensor, eps: float = 1e-8) -> torch.Tensor: + """ + Generates a tensor of Gumbel noise with the same shape as the input + tensor. Used for the Gumbel-Max trick. + """ + u = torch.rand_like(tensor) + return -torch.log(-torch.log(u + eps) + eps) + + +def batch_matmul(x: torch.Tensor, w: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + """Performs a batched matrix multiplication using PyTorch's native functions. + In NeMo this is implemented as a custom kernel using triton. + + Args: + x: ``[batch_size, d_in]`` + w: ``[num_weights, d_out, d_in]`` + y: ``[batch_size]`` + + Returns: + Tensor of shape ``[batch_size, d_out]``. + """ + return torch.bmm(w[y], x.unsqueeze(2)).squeeze(2) + + +class MoGHead(nn.Module): + """A Mixture of Gaussians (MoG) prediction head. + + This module takes a hidden state and predicts the parameters for a + mixture of Gaussian distributions. It's suitable for modeling + continuous, multi-modal data. + """ + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + out_size: int, + num_layers: int, + num_predictions: int, + low_rank: Optional[int] = 64, + top_p_or_k: Optional[Union[float, int]] = 1.0, + min_log_std: float = -4.0, + eps: float = 1e-6, + ): + super().__init__() + self.out_size = out_size + self.low_rank = low_rank + self.num_predictions = num_predictions + self.min_log_std = min_log_std + self.top_p_or_k = top_p_or_k + + self.logits_processor = ( + TopPLogitsWarper(self.top_p_or_k) + if isinstance(self.top_p_or_k, float) + else ( + TopKLogitsWarper(self.top_p_or_k) + if isinstance(self.top_p_or_k, int) + else None + ) + ) + + self.mlp_stack = nn.Sequential( + *[ + MLPLayer(hidden_size, intermediate_size, eps=eps) + for _ in range(num_layers) + ], + RMSNorm(hidden_size, eps=eps), + ) + + if low_rank is None: + self.proj_logits = nn.Linear(hidden_size, num_predictions, bias=False) + self.proj_mus = nn.Linear( + hidden_size, num_predictions * out_size, bias=False + ) + self.proj_logs = nn.Linear(hidden_size, 1, bias=False) + else: + assert low_rank < out_size + self.proj_logits = nn.Linear(hidden_size, num_predictions, bias=False) + self.proj_mus = nn.Linear( + hidden_size, num_predictions * low_rank, bias=False + ) + self.proj_logs = nn.Linear(hidden_size, 1, bias=False) + self.proj_else = nn.Linear(hidden_size, out_size, bias=False) + self.low_mat = nn.Parameter( + torch.empty(num_predictions, out_size, low_rank) + ) + + def forward( + self, + x: torch.Tensor, + cfg_is_uncond: Optional[torch.Tensor] = None, + cfg_scale: Optional[torch.Tensor] = None, + cfg_active: Optional[torch.Tensor] = None, + cfg_partner: Optional[torch.Tensor] = None, + cfg_conditional_rep: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + bt = x.size(0) + n, d = self.num_predictions, self.low_rank or self.out_size + + x = self.mlp_stack(x) + if ( + cfg_is_uncond is not None + and cfg_scale is not None + and cfg_active is not None + and cfg_partner is not None + ): + # Native EarTTS guidance is applied after the MoG MLP stack and + # before proj_logits/proj_mus/proj_logs/proj_else. + x = _apply_cfg_after_mlp( + x, + cfg_is_uncond=cfg_is_uncond, + cfg_scale=cfg_scale, + cfg_active=cfg_active, + cfg_partner=cfg_partner, + ) + + logits = self.proj_logits(x) + + # Apply top-p or top-k filtering to the mixture logits + if self.logits_processor is not None: + logits = self.logits_processor(None, logits.view(-1, n)).view_as(logits) + + # Sample a mixture component using the Gumbel-Max trick + gumbel = gumbel_like(logits) + if cfg_active is not None and cfg_conditional_rep is not None: + gumbel = torch.where( + cfg_active[:, None], + gumbel[cfg_conditional_rep], + gumbel, + ) + mixture_indices = (nn.functional.log_softmax(logits, dim=-1) + gumbel).argmax( + -1 + ) + + # Select the mean corresponding to the sampled component + mu = batch_matmul( + x.view(bt, -1), + self.proj_mus.weight.detach().view(n, d, -1), + mixture_indices.view(bt), + ).view(bt, d) + if self.proj_mus.bias is not None: + mu += self.proj_mus.bias.detach().view(n, d)[mixture_indices] + + if self.low_rank: + mu = batch_matmul( + mu.view(bt, -1), + self.low_mat.detach().view(n, self.out_size, -1), + mixture_indices.view(bt), + ).view(bt, self.out_size) + mu_res = self.proj_else(x) + else: + mu_res = torch.zeros((bt, d), device=x.device) + + logs = self.proj_logs(x).clamp_min(self.min_log_std) + return mu * torch.exp(logs) + mu_res, logs + + +class MaskGITSampler(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.num_quantizers = self.config.num_quantizers + self.codebook_size = self.config.codebook_size + self.noise_scale = self.config.noise_scale + self.debug_cfg_contract = bool( + getattr(self.config, "debug_cfg_contract", False) + ) + + # pre-compute how many tokens are unmasked at each iteration + rates = np.linspace(0.0, 1.0, self.config.num_iter + 1)[:-1].reshape(-1, 1) + masking_rates = np.power( + 1 - np.power(rates, self.config.exponent), 1 / self.config.exponent + ) + num_maskings = np.ceil(masking_rates * self.num_quantizers).astype(int) + num_maskings_shifted = np.pad( + num_maskings[1:], ((0, 1), (0, 0)), constant_values=0 + ) + sampling_per_step = num_maskings - num_maskings_shifted + sampling_per_step_flat = sampling_per_step.flatten() + # Drop any values at the beginning that are 0 + first_nonzero = np.argmax(sampling_per_step_flat != 0) + self.num_to_sample = sampling_per_step_flat[first_nonzero:].tolist() + + # create layers used for acoustic tokens embedding + self.rvq_embs = nn.Parameter( + torch.empty( + self.config.num_quantizers, + self.config.codebook_size, + self.config.latent_size, + ) + ) + self.embed_code = nn.Linear( + self.config.latent_size, self.config.hidden_size, bias=False + ) + # MoG head for generation (uncompiled part) + self.mog_head = MoGHead( + hidden_size=self.config.hidden_size, + intermediate_size=self.config.intermediate_size, + out_size=self.config.latent_size, + num_layers=self.config.mog_num_layers, + num_predictions=self.config.mog_num_predictions, + low_rank=self.config.mog_low_rank, + top_p_or_k=self.config.top_p_or_k, + min_log_std=self.config.mog_min_log_std, + eps=self.config.mog_eps, + ) + + def _depthsum_embedding(self, code: torch.Tensor) -> torch.Tensor: + """Embeds all codes into a single embedding.""" + embs = nn.functional.pad( + self.rvq_embs, [0, 0, 0, 1] + ) # num_quantizers x (codebook_size + 1) x latent_size + res = nn.functional.embedding(code[0], embs[0]) + for i in range(1, len(embs)): + res = res + nn.functional.embedding(code[i], embs[i]) + return res + + def _depthsum_encoding_step_reshaped( + self, + r: torch.Tensor, # [B*T, hidden_size] + code: torch.Tensor, # [num_quantizers, B*T] + depth_str: int, + k: int, + ) -> torch.Tensor: + """RVQ encoding with reshaped code tensor.""" + for i in range(depth_str, depth_str + k): + # Compute distances: ||emb||² - 2⟨r, emb⟩ + idx_sel = ( + self.rvq_embs[i].pow(2).sum(-1) # [vocab_size] + - 2 * (r @ self.rvq_embs[i].T) # [B*T, vocab_size] + ).argmin(-1) # [B*T] + + # Update residual + emb_i = nn.functional.embedding( + idx_sel, + self.rvq_embs[i], + ) # [B*T, latent_size] + r = r - emb_i + + # Store selected indices + code[i] = idx_sel + + return code + + def forward( + self, + hidden_states: torch.Tensor, + cfg_enabled: Optional[torch.Tensor] = None, + cfg_is_uncond: Optional[torch.Tensor] = None, + cfg_pair_id: Optional[torch.Tensor] = None, + cfg_scale: Optional[torch.Tensor] = None, + valid: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Performs the iterative unmasking process for a single + generation step. + """ + + device = hidden_states.device + batch_size = int(hidden_states.shape[0]) + if cfg_enabled is None: + cfg_enabled = torch.zeros(batch_size, dtype=torch.bool, device=device) + if cfg_is_uncond is None: + cfg_is_uncond = torch.zeros(batch_size, dtype=torch.bool, device=device) + if cfg_pair_id is None: + cfg_pair_id = torch.full( + (batch_size,), + -1, + dtype=torch.long, + device=device, + ) + if cfg_scale is None: + cfg_scale = torch.zeros(batch_size, dtype=torch.float32, device=device) + if valid is None: + valid = torch.ones(batch_size, dtype=torch.bool, device=device) + + ( + hidden_states, + cfg_is_uncond, + cfg_scale, + cfg_active, + cfg_partner, + cfg_conditional_rep, + inverse_order, + ) = _prepare_cfg_sampling_batch( + hidden_states, + cfg_enabled=cfg_enabled, + cfg_is_uncond=cfg_is_uncond, + cfg_pair_id=cfg_pair_id, + cfg_scale=cfg_scale, + valid=valid, + ) + complete_contract = ((~cfg_enabled) | cfg_active).all() + if self.debug_cfg_contract and not bool( + complete_contract.item() + ): + raise RuntimeError( + "Incomplete EarTTS CFG model batch: " + f"enabled={cfg_enabled.tolist()} " + f"is_uncond={cfg_is_uncond.tolist()} " + f"pair_id={cfg_pair_id.tolist()} " + f"valid={valid.tolist()} " + f"active={cfg_active.tolist()}" + ) + + # Initialize the full code tensor + code = ( + torch.zeros( + (self.num_quantizers, hidden_states.shape[0]), + dtype=torch.long, + device=device, + ) + + self.codebook_size + ) + # Iteratively unmask the continuous part of the code + cnt = 0 + for k in self.num_to_sample: + # Prepare input for the MoG head + mog_input_embeds = self.embed_code( + self._depthsum_embedding(code) + ) # (BT x hidden_size) + mog_input_embeds += hidden_states + + mog_mu, mog_logs = self.mog_head( + mog_input_embeds, + cfg_is_uncond=cfg_is_uncond, + cfg_scale=cfg_scale, + cfg_active=cfg_active, + cfg_partner=cfg_partner, + cfg_conditional_rep=cfg_conditional_rep, + ) + normal = torch.randn_like(mog_mu) + normal = torch.where( + cfg_active[:, None], + normal[cfg_conditional_rep], + normal, + ) + z = mog_mu + torch.exp(mog_logs) * normal * self.noise_scale + code = self._depthsum_encoding_step_reshaped(z, code, cnt, k) + # Match PyTorch EarTTS CFG: every MaskGIT iteration feeds the + # conditional code trajectory back into both KV streams. + code = torch.where( + cfg_active.unsqueeze(0), + code[:, cfg_conditional_rep], + code, + ) + + cnt += k + return code.transpose(0, 1)[inverse_order] # BT x num_quantizers + + +@support_torch_compile +class EarTTSModel(nn.Module): + """Embedding preparation + Gemma3 backbone (compiled together). + + MaskGIT sampling lives in :class:`EarTTSSamplerModel` so the iterative + sampler can be skipped on prefill positions while still being CUDA-graph + captured for decode-only batches. See :meth:`EarTTSForCausalLM.forward`. + """ + + def __init__( + self, + *, + vllm_config: VllmConfig, + prefix: str = "", + ): + super().__init__() + config = vllm_config.model_config.hf_config + self.total_emb = EarTTSInputEmbedding(config) + self.backbone = Gemma3Model(vllm_config=vllm_config, prefix=prefix) + + # Per-codebook silence acoustic tokens. Registered as a + # persistent int32 buffer (loaded from the checkpoint under + # ``model.sil_tokens``) rather than nn.Parameter so that vLLM's + # automatic float dtype casting (e.g. ``model.to(bfloat16)``) + # leaves it untouched. + self.register_buffer( + "sil_tokens", + # Zero is a safe dummy-loader default; production checkpoints + # overwrite this persistent buffer in ``load_weights``. + torch.zeros(int(config.num_quantizers), dtype=torch.int32), + persistent=True, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: Optional[IntermediateTensors], + acoustic_tokens: torch.Tensor, + text_tokens: torch.Tensor, + text_mask: torch.Tensor, + bos_mask: torch.Tensor, + speaker_latent: torch.Tensor, + cfg_is_uncond: torch.Tensor, + ) -> torch.Tensor: + """Forward pass through embeddings and backbone transformer. + Returns the backbone's ``hidden_states``. + """ + total_emb = self.total_emb( + acoustic_tokens=acoustic_tokens, + text_tokens=text_tokens, + text_mask=text_mask, + bos_mask=bos_mask, + speaker_latent=speaker_latent, + cfg_is_uncond=cfg_is_uncond, + ) + hidden_states = self.backbone( + input_ids, positions, intermediate_tensors, inputs_embeds=total_emb + ) + return hidden_states + + +@support_torch_compile +class EarTTSSamplerModel(nn.Module): + """MaskGIT sampler in its own compile group. + + Hosting the sampler in a separate ``@support_torch_compile`` module + is what makes it possible for :meth:`EarTTSForCausalLM.forward` to: + + * Capture and replay a CUDA-graph for decode-only batches (where + every position needs sampling). + * Skip the sampler entirely on prefill positions, where the audio + output isn't actually needed. + * Run the sampler on a sliced subset of positions in mixed + prefill+decode batches, with a ``BatchDescriptor`` override so the + sampler's CUDA-graph cache is hit at the padded decode-batch size. + + The :meth:`forward` operates on a stable-address scratch buffer + (:attr:`_sampler_input`) so callers can pass a transient slice + (e.g. ``hidden_states[decode_idx]``) without breaking CUDA-graph + replay. The non-compiled :meth:`sample` wrapper does that copy and + then invokes the compiled :meth:`forward`. + """ + + def __init__( + self, + *, + vllm_config: VllmConfig, + prefix: str = "", + ): + super().__init__() + config = vllm_config.model_config.hf_config + self.sampler = MaskGITSampler(config) + + # Stable-address scratch buffer for the sampler's input. Every + # CUDA-graph replay must read from the same ``data_ptr()``; the + # caller may pass either the full backbone output or a fresh + # ``hidden_states[decode_idx]`` slice, so we copy into this + # buffer (in :meth:`sample`) before invoking :meth:`forward`. + max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens + hidden_size = config.hidden_size + dtype = vllm_config.model_config.dtype + self._sampler_input = torch.zeros( + max_num_tokens, hidden_size, dtype=dtype + ) + self._sampler_cfg_enabled = torch.zeros(max_num_tokens, dtype=torch.bool) + self._sampler_cfg_is_uncond = torch.zeros(max_num_tokens, dtype=torch.bool) + self._sampler_cfg_pair_id = torch.full((max_num_tokens,), -1, dtype=torch.long) + self._sampler_cfg_scale = torch.zeros(max_num_tokens, dtype=torch.float32) + self._sampler_valid = torch.zeros(max_num_tokens, dtype=torch.bool) + + def forward( + self, + hidden_states: torch.Tensor, + cfg_enabled: torch.Tensor, + cfg_is_uncond: torch.Tensor, + cfg_pair_id: torch.Tensor, + cfg_scale: torch.Tensor, + valid: torch.Tensor, + ) -> torch.Tensor: + """Compiled — runs MaskGIT on a (stable-address) hidden buffer.""" + return self.sampler( + hidden_states, + cfg_enabled=cfg_enabled, + cfg_is_uncond=cfg_is_uncond, + cfg_pair_id=cfg_pair_id, + cfg_scale=cfg_scale, + valid=valid, + ) + + def sample( + self, + hidden_states: torch.Tensor, + *, + cfg_enabled: torch.Tensor, + cfg_is_uncond: torch.Tensor, + cfg_pair_id: torch.Tensor, + cfg_scale: torch.Tensor, + valid: torch.Tensor, + ) -> torch.Tensor: + """Non-compiled wrapper — copies into the stable buffer first. + + Mirrors the qwen3-tts code-predictor pattern: transient inputs + are first written into a model-owned static-address buffer so + the captured CUDA-graph for the compiled :meth:`forward` always + reads from the recorded ``data_ptr()``. + """ + seq_len = int(hidden_states.shape[0]) + buf = self._sampler_input[:seq_len] + buf.copy_(hidden_states) + enabled_buf = self._sampler_cfg_enabled[:seq_len] + role_buf = self._sampler_cfg_is_uncond[:seq_len] + pair_buf = self._sampler_cfg_pair_id[:seq_len] + scale_buf = self._sampler_cfg_scale[:seq_len] + valid_buf = self._sampler_valid[:seq_len] + enabled_buf.copy_(cfg_enabled) + role_buf.copy_(cfg_is_uncond) + pair_buf.copy_(cfg_pair_id) + scale_buf.copy_(cfg_scale) + valid_buf.copy_(valid) + return self( + buf, + enabled_buf, + role_buf, + pair_buf, + scale_buf, + valid_buf, + ) + + +# --------------------------------------------------------------------------- +# Outer model — the vLLM-Omni preprocess/postprocess entry point. +# --------------------------------------------------------------------------- + + +# Placeholder token id used to fill the per-step ``input_ids`` returned +# by :meth:`preprocess`. Must be a valid id in ``[0, config.vocab_size)`` +# but is otherwise unused — the actual decode-vs-prefill behaviour is +# driven by the per-token buffers populated in :meth:`preprocess`. +# +# The width of the dummy logits tensor returned by +# :meth:`compute_logits` is taken from ``config.vocab_size`` (see +# :class:`EarTTSConfig`) so vLLM's sampler / ``LogitsProcessor`` and the +# model agree on the logits shape. ``compute_logits`` returns +# ``[0, -inf, ..., -inf]`` so the sampler's argmax always picks index 0 +# regardless of how wide ``vocab_size`` is — the real audio output is +# the codes tensor exposed via :meth:`make_omni_output`. +_DUMMY_TOKEN_ID = 0 + + +@ignore_torch_compile +@support_torch_compile +class EarTTSForCausalLM(nn.Module, SupportsPP): + """EarTTS for vLLM-Omni. + + Inputs (passed via ``additional_information``): + + * ``embed.voice`` (prefill chunk 0 only) — Tensor of shape + ``(Tref, hidden_size)`` carrying the user-supplied speaker + latent. The user must also pass ``prompt_token_ids = [0] * + Tref`` so the prefill placeholder length matches. + * ``ids.output`` — Python ``list[int]`` of the most recently sent + text tokens; preprocess consumes the last entry, so decode step + ``k`` consumes ``t_k``. + * CFG metadata (on every prefill/decode chunk): ``cfg_enabled``, + ``cfg_role`` (``"cond"`` or ``"uncond"``), ``cfg_pair_id``, and + ``cfg_scale``. Unconditional rows replace only text conditioning + with ``model.total_emb.null_emb``. Complete decode pairs are sampled + with native EarTTS guidance and receive identical acoustic codes. + + Per-step flow (see module docstring for details): + + ``preprocess`` populates five model-owned buffers + (:attr:`_acoustic_tokens`, :attr:`_text_tokens`, :attr:`_text_mask`, + :attr:`_bos_mask`, :attr:`_speaker_latent`) at each request's + flat-batch offset. ``forward`` slices them up to ``num_tokens`` and + runs the compiled :class:`EarTTSModel` (embedding + Gemma3 + backbone) for every position, then conditionally invokes the + compiled :class:`EarTTSSamplerModel` (MaskGIT) to produce codes. + The sampler is skipped on prefill positions — see :meth:`forward` + for details. The generated codes (BTx31) are written to + :attr:`_out_codes` and exposed as a multimodal output by + :meth:`make_omni_output` (see there for the key it uses). + ``postprocess`` stashes the final-frame codes under + ``last_acoustic_codes`` for the next decode step's + :meth:`preprocess`. + + Sampler skipping mirrors the qwen3-tts code-predictor pattern: + + * **Profile / dummy run** (``attn_metadata is None``) and + **decode-only batches** (``max_query_len == 1``) run the sampler + on every token so the captured CUDA graph covers all of + ``cudagraph_capture_sizes``. + * **Mixed prefill+decode batches**: only decode-token positions go + through the sampler. The sampler's ``BatchDescriptor`` is + overridden to the padded decode-batch size so the right captured + graph is replayed. + * **Prefill-only batches**: the sampler is skipped entirely. + * Prefill rows of :attr:`_out_codes` are intentionally not + written. ``last_acoustic_codes`` returned by :meth:`postprocess` + after prefill is therefore undefined — :meth:`_preprocess_decode` + seeds the first decode step's acoustic input with the acoustic + pad id (``codebook_size``) so this never matters. + """ + + # ``model.sampler.*`` lands on :attr:`sampler_module` (the MaskGIT compile + # group). Other prefixes (``model.total_emb.``, ``model.backbone.``) match + # the module layout 1:1. + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "model.sampler.": "sampler_module.sampler.", + } + ) + + # Omni preprocess/postprocess hooks (consumed by the gpu model runner). + has_preprocess = True + has_postprocess = True + have_multimodal_outputs = True + + # No ``gpu_resident_buffer_keys``: vLLM-Omni's opt-out from the + # ``model_intermediate_buffer`` D2H round-trip is keyed by + # ``(type_key, qualifier)`` pairs and is only consulted for *nested* + # payload entries, so a model whose payloads are flat -- as both stages + # here are -- cannot express its keys in that form. Declaring flat names + # is not merely inert, it breaks: the runner unpacks every declared key + # as a pair as soon as any nested entry arrives, and 0.24 onwards always + # sends one (``meta``). + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.config = vllm_config.model_config.hf_config + self.vllm_config = vllm_config + + # Embedding + Gemma3 backbone — runs on every position. Built + # under the default ``"backbone"`` model tag (vLLM's compile + # cache key for the main model). We don't wrap this in a + # ``set_model_tag`` block because :func:`set_model_tag` asserts + # the new tag differs from the current one and the default is + # already ``"backbone"``. + self.model = EarTTSModel( + vllm_config=vllm_config, + prefix=prefix, + ) + + # MaskGIT sampler in its own compile group, so it can be invoked + # conditionally (decode positions only, or skipped entirely on + # prefill-only batches) while still being CUDA-graph captured + # for decode-only batches over ``cudagraph_capture_sizes``. The + # ``"sampler"`` tag keys the sampler's compile cache separately + # from the backbone's. + with set_model_tag("sampler"): + self.sampler_module = EarTTSSamplerModel( + vllm_config=vllm_config, + prefix=prefix, + ) + + # Pad ids used by buffers / preprocess. Match the conventions of + # the original EarTTSInputEmbedding: an acoustic token id of + # ``codebook_size`` is the trailing "no audio" pad row in + # ``rvq_embs`` (which has ``codebook_size + 1`` entries). + # How the sampled codes are surfaced from ``make_omni_output``, driven + # by the stage's pipeline-config ``engine_output_type``: + # + # * "audio" (this stage is the final, client-facing one, which is the + # split VoiceChat layout): emit under the ``model_outputs`` key. + # vLLM-Omni's output processor remaps ``model_outputs`` to the + # drainable ``audio`` modality key, so DELTA streaming drains it + # after every step and the client receives one frame per step. + # Any other key is retained across steps *and* concatenated along + # the last dimension (``get_accumulation_strategy`` maps the audio + # modality to ``CONCAT_LAST``), which for a ``T x num_quantizers`` + # code tensor silently widens the per-step frame instead of + # appending to it. + # * otherwise: emit ``audio_codes`` for a downstream stage to consume. + engine_output_type = getattr( + vllm_config.model_config, "engine_output_type", None + ) + self._single_stage_audio = str(engine_output_type or "").lower() == "audio" + + self._num_quantizers: int = int(self.config.num_quantizers) + self._hidden_size: int = int(self.config.hidden_size) + self._acoustic_pad_id: int = int(self.config.codebook_size) + text_pad_id = getattr(self.config, "pad_token_id", None) + eos_token_id = getattr(self.config, "eos_token_id", None) + if text_pad_id is None or eos_token_id is None: + raise ValueError( + "EarTTS config.json must set pad_token_id and eos_token_id from the " + "source VoiceChat tokenizer. Re-run convert_duplex_eartts_checkpoint.py." + ) + self._text_pad_id: int = int(text_pad_id) + self._eos_token_id: int = int(eos_token_id) + + # ── Persistent stable-address buffers ──────────────────────── + # Plain tensor attributes (not nn.Parameter / not register_buffer): + # * AutoWeightsLoader only walks named_parameters() and persistent + # registered buffers, so plain attributes are invisible to it + # (no spurious "missing weight" errors during load_weights). + # * vLLM constructs models inside + # ``with torch.device(device_config.device):`` so a bare + # ``torch.zeros(...)`` here is allocated directly on the GPU. + # * Addresses stay stable across CUDA graph replays as long as + # we never re-assign these names (only do in-place writes via + # copy_/fill_/indexed assignment), which is what the rest of + # this class does. The piecewise CUDAGraphWrapper records + # data_ptr() at capture time and expects the same pointer at + # replay time — that holds with plain tensors. + max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens + model_dtype = vllm_config.model_config.dtype + + self._acoustic_tokens = torch.full( + (max_num_tokens, self._num_quantizers), + self._acoustic_pad_id, + dtype=torch.long, + ) + self._text_tokens = torch.full( + (max_num_tokens,), self._text_pad_id, dtype=torch.long + ) + self._text_mask = torch.zeros(max_num_tokens, dtype=torch.long) + self._bos_mask = torch.zeros(max_num_tokens, dtype=torch.long) + # Speaker latent buffer — model dtype, hidden_size wide. + # Decode rows stay all-zero (which the embedding module reads + # as "latent not provided" so ``audio_emb`` is preserved). + # Prefill rows are populated from the user-supplied tensor. + self._speaker_latent = torch.zeros( + max_num_tokens, self._hidden_size, dtype=model_dtype + ) + # Per-token CFG contract. These plain tensors follow the same + # stable-address rules as the model input buffers above and are copied + # into the sampler's own CUDA-graph scratch buffers before sampling. + self._cfg_enabled = torch.zeros(max_num_tokens, dtype=torch.bool) + self._cfg_is_uncond = torch.zeros(max_num_tokens, dtype=torch.bool) + self._cfg_pair_id = torch.full((max_num_tokens,), -1, dtype=torch.long) + self._cfg_scale = torch.zeros(max_num_tokens, dtype=torch.float32) + # vLLM-Omni 0.26 computes per-request flat-batch slices but does not + # pass ``start``/``end`` into preprocess. The CFG scheduler guarantees + # cond then uncond order, so this cursor reconstructs those slices. + self._preprocess_cursor = 0 + self._out_codes = torch.zeros( + max_num_tokens, self._num_quantizers, dtype=torch.long + ) + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + """Compatibility shim — not actually consumed at runtime since + every forward goes through ``inputs_embeds`` assembled inside + :meth:`forward`. + """ + return self.model.backbone.embed_input_ids(input_ids) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.get_input_embeddings(input_ids) + + @staticmethod + def _unwrap_singleton(value: Any) -> Any: + """Unwrap a possibly list-wrapped scalar (e.g. ``[tensor]``).""" + if isinstance(value, list): + return value[0] if value else None + return value + + @staticmethod + def _payload_get(info_dict: dict[str, Any], category: str, qualifier: str) -> Any: + """Read one ``OmniPayload`` field, e.g. ``("embed", "voice")``. + + Categories arrive as sub-dicts, which is also what lets the runner + keep a prefill-only field alive while a per-step field in another + category is replaced. + """ + sub = info_dict.get(category) + return sub.get(qualifier) if isinstance(sub, dict) else None + + @classmethod + def _cfg_scalar(cls, value: Any) -> Any: + value = cls._unwrap_singleton(value) + if isinstance(value, torch.Tensor): + assert value.numel() == 1, ( + "EarTTS CFG metadata tensors must contain one scalar; " + f"got shape={tuple(value.shape)}." + ) + return value.item() + return value + + @classmethod + def _stable_cfg_pair_id(cls, value: Any) -> int: + """Map request-provided pair ids to a deterministic signed int64.""" + value = cls._cfg_scalar(value) + assert value is not None and not isinstance( + value, bool + ), "EarTTS CFG requires a non-empty ``cfg_pair_id``." + if isinstance(value, int): + assert ( + -(1 << 63) <= value < (1 << 63) + ), f"EarTTS cfg_pair_id={value} does not fit in int64." + return value + if isinstance(value, float): + assert ( + value.is_integer() + ), f"EarTTS cfg_pair_id must be integral or a string; got {value}." + return cls._stable_cfg_pair_id(int(value)) + encoded = str(value).encode("utf-8") + assert encoded, "EarTTS CFG requires a non-empty ``cfg_pair_id``." + return int.from_bytes( + hashlib.blake2b(encoded, digest_size=8).digest(), + "little", + ) & ((1 << 63) - 1) + + def _write_cfg_state( + self, + *, + start: int, + span_len: int, + info_dict: dict[str, Any], + ) -> None: + """Validate one request/chunk's CFG metadata and fill static rows.""" + enabled_value = self._cfg_scalar(info_dict.get("cfg_enabled", False)) + if isinstance(enabled_value, str): + normalized = enabled_value.strip().lower() + assert normalized in { + "true", + "false", + "1", + "0", + }, f"EarTTS cfg_enabled must be boolean; got {enabled_value!r}." + cfg_enabled = normalized in {"true", "1"} + else: + cfg_enabled = bool(enabled_value) + + cfg_is_uncond = False + cfg_pair_id = -1 + scale_value = self._cfg_scalar(info_dict.get("cfg_scale")) + if scale_value is None: + scale_value = getattr(self.config, "guidance_scale", 0.5) + cfg_scale = float(scale_value) + assert math.isfinite( + cfg_scale + ), f"EarTTS cfg_scale must be finite; got {cfg_scale}." + + if cfg_enabled: + role = str(self._cfg_scalar(info_dict.get("cfg_role")) or "").lower() + assert role in {"cond", "uncond"}, ( + "EarTTS CFG requires cfg_role='cond' or 'uncond'; " f"got {role!r}." + ) + cfg_is_uncond = role == "uncond" + cfg_pair_id = self._stable_cfg_pair_id(info_dict.get("cfg_pair_id")) + + end = start + span_len + self._cfg_enabled[start:end].fill_(cfg_enabled) + self._cfg_is_uncond[start:end].fill_(cfg_is_uncond) + self._cfg_pair_id[start:end].fill_(cfg_pair_id) + self._cfg_scale[start:end].fill_(cfg_scale) + + def _validate_speaker_latent(self, value: Any) -> torch.Tensor: + """Assert ``speaker_latent`` has shape ``(Tref, hidden_size)``.""" + x = self._unwrap_singleton(value) + assert isinstance(x, torch.Tensor), ( + f"speaker_latent must be a torch.Tensor; got {type(x).__name__}." + ) + assert x.ndim == 2 and x.shape[1] == self._hidden_size, ( + "speaker_latent must have shape (Tref, hidden_size=" + f"{self._hidden_size}); got {tuple(x.shape)}." + ) + return x.to(dtype=self._speaker_latent.dtype).contiguous() + + def _build_prefill_tensors( + self, + speaker_latent: torch.Tensor, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Build the cached prefill ``(text_tokens, text_mask, bos_mask, + speaker_latent)`` of length ``prefill_len = speaker_latent.shape[0]``. + + Layout: ``text_tokens = [PAD] * (n - 1) + [EOS]``, + ``text_mask = [0] * (n - 2) + [1, 1]``, + ``bos_mask = [0] * (n - 1) + [1]``. Acoustic tokens are not + cached — :meth:`preprocess` broadcasts ``model.sil_tokens`` at + every prefill position. + """ + prefill_len = int(speaker_latent.shape[0]) + assert prefill_len > 0, ( + "speaker_latent must have at least one frame " + f"(got shape={tuple(speaker_latent.shape)})." + ) + + text_tokens = torch.full( + (prefill_len,), self._text_pad_id, dtype=torch.long, device=device + ) + text_tokens[-1] = self._eos_token_id + + text_mask = torch.zeros(prefill_len, dtype=torch.long, device=device) + text_mask[max(0, prefill_len - 2):] = 1 + + bos_mask = torch.zeros(prefill_len, dtype=torch.long, device=device) + bos_mask[-1] = 1 + + speaker_latent = speaker_latent.to( + device=device, dtype=self._speaker_latent.dtype, non_blocking=True + ).contiguous() + + return text_tokens, text_mask, bos_mask, speaker_latent + + # ------------------------------------------------------------------ + # preprocess + # ------------------------------------------------------------------ + + def preprocess( + self, + input_ids: torch.Tensor, + input_embeds: Optional[torch.Tensor], + *, + start: int = 0, + end: int = 0, + **info_dict: Any, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Build per-request ``(input_ids, inputs_embeds)`` for this step. + + Prefill (``span_len > 1``): + On the first prefill chunk, constructs the per-position + prefill tensors of length + ``prefill_len = speaker_latent.shape[0]``: + + * ``text_tokens`` = ``[PAD] * (prefill_len - 1) + [EOS]`` + * ``text_mask`` = ``[0] * (prefill_len - 2) + [1, 1]`` + * ``bos_mask`` = ``[0] * (prefill_len - 1) + [1]`` + * ``acoustic_tokens`` = ``model.sil_tokens`` broadcast at + every position (only the BOS frame's ``audio_emb`` is + actually consumed; the others are replaced by + ``speaker_latent`` inside the embedding module). + * ``speaker_latent`` = the user-supplied tensor. + + ``embed.voice`` is the only required + ``additional_information`` field. Multi-chunk prefill is + tracked by ``ear_prefill_offset``; the cached + ``ear_prefill_speaker_latent`` is sliced into each chunk. + + Decode (``span_len == 1``): + Takes the newest text token from ``ids.output``, which the + producer (a user-driven :class:`StreamingInput` or an + upstream stage in an ``async_chunk`` pipeline such as + ``nemotron_voicechat``) refreshes on every step. + + Acoustic input rules (in order): + * ``text_token == EOS`` → ``model.sil_tokens``. + * First decode (``ear_decode_offset == 0``) → + broadcast acoustic pad id (``codebook_size``). + * Otherwise → ``last_acoustic_codes`` (stashed by + :meth:`postprocess` after the previous step). + + ``text_mask = 1``, ``bos_mask = 0``, + ``speaker_latent = 0``. + """ + # Normalize: some runner paths still pass per-request state + # nested under ``additional_information`` instead of flattened. + nested = info_dict.get("additional_information") + if isinstance(nested, dict): + merged = { + k: v for k, v in info_dict.items() if k != "additional_information" + } + for k, v in nested.items(): + merged.setdefault(k, v) + info_dict = merged + + device = input_ids.device + span_len = int(input_ids.shape[0]) + if span_len <= 0: + base = ( + input_embeds + if input_embeds is not None + else self.embed_input_ids(input_ids) + ) + return input_ids, base, {} + + explicit_start = int(start) + explicit_end = int(end) + if explicit_end > explicit_start: + flat_start = explicit_start + else: + role = str( + self._cfg_scalar(info_dict.get("cfg_role")) + or "" + ).lower() + cfg_enabled = bool( + self._cfg_scalar( + info_dict.get("cfg_enabled", False) + ) + ) + if not cfg_enabled or role == "cond": + self._preprocess_cursor = 0 + flat_start = self._preprocess_cursor + self._preprocess_cursor += span_len + + self._write_cfg_state( + start=flat_start, + span_len=span_len, + info_dict=info_dict, + ) + + if span_len > 1: + return self._preprocess_prefill( + input_ids=input_ids, + input_embeds=input_embeds, + start=flat_start, + span_len=span_len, + device=device, + info_dict=info_dict, + ) + return self._preprocess_decode( + input_ids=input_ids, + input_embeds=input_embeds, + start=flat_start, + device=device, + info_dict=info_dict, + ) + + def _preprocess_prefill( + self, + *, + input_ids: torch.Tensor, + input_embeds: Optional[torch.Tensor], + start: int, + span_len: int, + device: torch.device, + info_dict: dict[str, Any], + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Prefill branch of :meth:`preprocess`. Writes one chunk-slice of + the cached prefill tensors into the static buffers.""" + cached_speaker_latent = info_dict.get("ear_prefill_speaker_latent") + + info_update: dict[str, Any] = {} + if not isinstance(cached_speaker_latent, torch.Tensor): + # First chunk: build & cache prefill tensors from the + # user-supplied speaker latent. + speaker_latent = self._validate_speaker_latent( + self._payload_get(info_dict, "embed", "voice") + ) + ( + cached_text_tokens, + cached_text_mask, + cached_bos_mask, + cached_speaker_latent, + ) = self._build_prefill_tensors(speaker_latent, device=device) + + info_update["ear_prefill_text_tokens"] = cached_text_tokens + info_update["ear_prefill_text_mask"] = cached_text_mask + info_update["ear_prefill_bos_mask"] = cached_bos_mask + info_update["ear_prefill_speaker_latent"] = cached_speaker_latent + info_update["ear_prefill_offset"] = 0 + info_update["ear_decode_offset"] = 0 + else: + cached_text_tokens = info_dict["ear_prefill_text_tokens"] + cached_text_mask = info_dict["ear_prefill_text_mask"] + cached_bos_mask = info_dict["ear_prefill_bos_mask"] + + offset = int(info_dict.get("ear_prefill_offset", 0) or 0) + full_len = int(cached_speaker_latent.shape[0]) + s, e = offset, offset + span_len + assert 0 <= s and e <= full_len, ( + "prefill chunk overshoots cached prefill: offset=" + f"{offset}, span_len={span_len}, prefill_len={full_len}. " + "User must pass prompt_token_ids of length " + "speaker_latent.shape[0]." + ) + + buf_s = start + buf_e = buf_s + span_len + self._text_tokens[buf_s:buf_e].copy_(cached_text_tokens[s:e]) + self._text_mask[buf_s:buf_e].copy_(cached_text_mask[s:e]) + self._bos_mask[buf_s:buf_e].copy_(cached_bos_mask[s:e]) + self._speaker_latent[buf_s:buf_e].copy_(cached_speaker_latent[s:e]) + # Acoustic input is sil_tokens broadcast — only the BOS-frame's + # audio_emb is consumed (the rest get replaced by speaker_latent + # inside EarTTSInputEmbedding). + self._acoustic_tokens[buf_s:buf_e] = self.model.sil_tokens.to( + self._acoustic_tokens.dtype + ) + + info_update["ear_prefill_offset"] = offset + span_len + + # Placeholder input_ids; compiled forward reads the buffers, not these. + input_ids_out = torch.full_like(input_ids, _DUMMY_TOKEN_ID) + return input_ids_out, input_embeds, info_update + + def _preprocess_decode( + self, + *, + input_ids: torch.Tensor, + input_embeds: Optional[torch.Tensor], + start: int, + device: torch.device, + info_dict: dict[str, Any], + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Decode branch of :meth:`preprocess`. Writes one row at ``start``.""" + chunk_text_tokens = self._payload_get(info_dict, "ids", "output") + assert isinstance(chunk_text_tokens, list) and chunk_text_tokens, ( + "EarTTS decode requires a non-empty ``ids.output`` list in " + f"additional_information; got {type(chunk_text_tokens).__name__} " + f"with available keys {sorted(info_dict)}." + ) + + decode_offset = int(info_dict.get("ear_decode_offset", 0) or 0) + # The newest token is the one this step consumes: the producer sends + # exactly one per step, and a producer that resends a history still + # has the current token at the end. + text_token_id = int(chunk_text_tokens[-1]) + + buf_s = start + + # Acoustic input selection: + # * EOS subword → force sil_tokens (return to silence). + # * First decode after prefill → seed with the acoustic pad id + # (codebook_size) broadcast across all quantizers. + # * Otherwise → previous-step predicted codes. + if text_token_id == self._eos_token_id: + self._acoustic_tokens[buf_s].copy_( + self.model.sil_tokens.to(self._acoustic_tokens.dtype) + ) + elif decode_offset == 0: + self._acoustic_tokens[buf_s].fill_(self._acoustic_pad_id) + else: + last_codes = info_dict.get("last_acoustic_codes") + assert isinstance(last_codes, torch.Tensor) and last_codes.numel() > 0, ( + "EarTTS decode (offset > 0) requires " + "``last_acoustic_codes`` from the previous step's " + "postprocess." + ) + ac = ( + last_codes.to(device=device, dtype=torch.long) + .reshape(-1)[: self._num_quantizers] + ) + self._acoustic_tokens[buf_s, : ac.shape[0]].copy_(ac) + + self._text_tokens[buf_s] = text_token_id + self._text_mask[buf_s] = 1 + self._bos_mask[buf_s] = 0 + # Decode never replaces audio_emb with a latent. + self._speaker_latent[buf_s].zero_() + + info_update: dict[str, Any] = {"ear_decode_offset": decode_offset + 1} + return input_ids, input_embeds, info_update + + # ------------------------------------------------------------------ + # forward — runs the compiled embedding + backbone, then the sampler + # only on decode positions (skipping the expensive MaskGIT loop on + # prefill positions). + # ------------------------------------------------------------------ + + def _get_decode_idxs(self): + """Return ``(decode_token_indices, num_requests)`` for sampler dispatch. + + Mirrors the qwen3-tts code-predictor pattern: + + * ``(None, 0)`` → run sampler on every token. Used during + profile / dummy runs (no ``attn_metadata``) and decode-only + batches (``max_query_len == 1``), so the captured CUDA graph + covers all of ``cudagraph_capture_sizes``. + * ``(decode_token_indices, num_requests)`` → run sampler only on + the listed positions. ``decode_token_indices`` is padded up to + the next captured CUDA-graph size (so the sampler's graph + cache is hit) and ``num_requests`` is the unpadded count of + real decode tokens (used to scatter codes back into the right + rows of :attr:`_out_codes`). + """ + ctx = get_forward_context() + attn_metadata = ctx.attn_metadata + if attn_metadata is None: + # Profile / dummy run. Apply sampler everywhere so capture + # covers every cudagraph_capture_sizes value. + return None, 0 + + if isinstance(attn_metadata, dict): + any_layer_meta = next(iter(attn_metadata.values())) + else: + any_layer_meta = attn_metadata + + if any_layer_meta.max_query_len == 1: + # Decode-only batch: every position is a decode position, + # so just run the sampler over the whole flat batch. + return None, 0 + + start_loc = any_layer_meta.query_start_loc + tokens_per_req = start_loc[1:] - start_loc[:-1] + is_decode = (tokens_per_req == 1) + decode_token_indices = start_loc[:-1][is_decode] + + num_requests = decode_token_indices.shape[0] + padded_num_requests = num_requests + if ( + self.vllm_config.compilation_config.cudagraph_mode + != CUDAGraphMode.NONE + ): + sizes = self.vllm_config.compilation_config.cudagraph_capture_sizes + idx = bisect.bisect_left(sizes, num_requests) + if idx < len(sizes): + padded_num_requests = sizes[idx] + if padded_num_requests != num_requests: + decode_token_indices = torch.nn.functional.pad( + decode_token_indices, + (0, padded_num_requests - num_requests), + ) + return decode_token_indices, num_requests + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: Optional[IntermediateTensors] = None, + inputs_embeds: Optional[torch.Tensor] = None, + **_: Any, + ) -> torch.Tensor: + """Run the compiled embedding + backbone over every position, + then run the compiled MaskGIT sampler only on decode positions + (the sampler is skipped on prefill-only batches and on prefill + rows of mixed batches). ``inputs_embeds`` is ignored — the + actual embedding is assembled inside the compiled + :class:`EarTTSInputEmbedding` from the per-token buffers + populated by :meth:`preprocess`. + """ + num_tokens = int(input_ids.shape[0]) + + acoustic_tokens = self._acoustic_tokens[:num_tokens] + text_tokens = self._text_tokens[:num_tokens] + text_mask = self._text_mask[:num_tokens] + bos_mask = self._bos_mask[:num_tokens] + speaker_latent = self._speaker_latent[:num_tokens] + cfg_is_uncond = self._cfg_is_uncond[:num_tokens] + + hidden_states = self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + acoustic_tokens=acoustic_tokens, + text_tokens=text_tokens, + text_mask=text_mask, + bos_mask=bos_mask, + speaker_latent=speaker_latent, + cfg_is_uncond=cfg_is_uncond, + ) + + decode_idx, num_req = self._get_decode_idxs() + if decode_idx is None: + # Dummy/profile run or decode-only batch: sample everywhere. + codes = self.sampler_module.sample( + hidden_states, + cfg_enabled=self._cfg_enabled[:num_tokens], + cfg_is_uncond=self._cfg_is_uncond[:num_tokens], + cfg_pair_id=self._cfg_pair_id[:num_tokens], + cfg_scale=self._cfg_scale[:num_tokens], + valid=torch.ones( + num_tokens, + dtype=torch.bool, + device=hidden_states.device, + ), + ) + self._out_codes[:num_tokens].copy_(codes.to(dtype=torch.long)) + elif num_req > 0: + # Mixed batch: gather decode positions, override the + # BatchDescriptor so the sampler's CUDA-graph cache is hit + # at the padded decode-batch size. + ctx = get_forward_context() + orig_batch_descriptor = ctx.batch_descriptor + ctx.batch_descriptor = BatchDescriptor( + num_tokens=decode_idx.shape[0], + ) + decode_hidden = hidden_states[decode_idx] + sampler_valid = ( + torch.arange( + decode_idx.shape[0], + device=decode_idx.device, + ) + < num_req + ) + codes = self.sampler_module.sample( + decode_hidden, + cfg_enabled=self._cfg_enabled[decode_idx], + cfg_is_uncond=self._cfg_is_uncond[decode_idx], + cfg_pair_id=self._cfg_pair_id[decode_idx], + cfg_scale=self._cfg_scale[decode_idx], + valid=sampler_valid, + ) + ctx.batch_descriptor = orig_batch_descriptor + + valid_dec_idx = decode_idx[:num_req] + self._out_codes[valid_dec_idx] = codes[:num_req].to( + dtype=torch.long + ) + # Prefill-only batch: sampler skipped. ``_out_codes`` rows for + # those positions are not written here on purpose — callers + # must not rely on them; the public per-decode-step contract + # is driven by ``last_acoustic_codes`` from postprocess and + # the seed rules in :meth:`_preprocess_decode`. + + return hidden_states + + # ------------------------------------------------------------------ + # compute_logits — sampler bypass (the real output is ``codes``) + # ------------------------------------------------------------------ + + def compute_logits( + self, + hidden_states: Union[torch.Tensor, OmniOutput], + sampling_metadata: Any = None, + ) -> Optional[torch.Tensor]: + """Return zero logits of width ``config.vocab_size``. + + ``config.vocab_size`` is what vLLM's sampler / ``LogitsProcessor`` + use to size their working buffers, so deriving the width from + the same field guarantees the two agree. The sampled token id + is irrelevant: ``input_ids`` are never consumed by the model + (the per-step decode behaviour is driven by the buffers + populated in :meth:`preprocess`), and the real audio output is + the codes tensor exposed via :meth:`make_omni_output`. + """ + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + if hidden_states is None: + return None + batch_size = hidden_states.shape[0] + return hidden_states.new_zeros(batch_size, int(self.config.vocab_size)) + + # ------------------------------------------------------------------ + # multimodal output plumbing + # ------------------------------------------------------------------ + + def make_omni_output( + self, + model_outputs: Union[torch.Tensor, OmniOutput], + **_: Any, + ) -> OmniOutput: + """Wrap backbone hidden states with the codes generated by the + sampler (BTx31). + + The key depends on whether this stage is the client-facing one; see + ``self._single_stage_audio`` in :meth:`__init__`. ``postprocess`` + accepts either. + """ + if isinstance(model_outputs, OmniOutput): + return model_outputs + + hidden = model_outputs + num_tokens = int(hidden.shape[0]) + audio_codes = self._out_codes[:num_tokens].clone() + key = "model_outputs" if self._single_stage_audio else "audio_codes" + return OmniOutput( + text_hidden_states=hidden, + multimodal_outputs={key: audio_codes}, + ) + + # ------------------------------------------------------------------ + # postprocess — stash last-frame codes for the next decode step + # ------------------------------------------------------------------ + + def postprocess( + self, + hidden_states: torch.Tensor, + multimodal_outputs: Optional[dict[str, Any]] = None, + **_: Any, + ) -> dict[str, Any]: + """Pull the last-frame codes out of the multimodal output (or, + as a fallback, out of :attr:`_out_codes` using the slice's + storage offset) and stash them under ``last_acoustic_codes`` so + the next step's :meth:`preprocess` can use them as the decode + input. + """ + if hidden_states.numel() == 0: + return {} + + mm = multimodal_outputs or {} + audio_codes = mm.get("audio_codes") + if audio_codes is None: + audio_codes = mm.get("model_outputs") + if isinstance(audio_codes, torch.Tensor) and audio_codes.numel() > 0: + # ``hidden_states`` is a slice of the flat batch. Recover + # the request's last position via storage_offset and pick + # the corresponding row from ``audio_codes``. + stride0 = hidden_states.stride(0) or 1 + req_start = hidden_states.storage_offset() // stride0 + last = req_start + hidden_states.shape[0] - 1 + last_codes = audio_codes[last : last + 1].detach() + return {"last_acoustic_codes": last_codes} + + return {} + + # ------------------------------------------------------------------ + # weight loading + # ------------------------------------------------------------------ + + def load_weights( + self, weights: Iterable[tuple[str, torch.Tensor]] + ) -> set[str]: + skip_prefixes: list[str] = [] + if self.config.tie_word_embeddings: + skip_prefixes.append("lm_head.") + + # The Gemma3 backbone keeps a vestigial ``embed_tokens`` layer, + # but this model never consumes ``input_ids`` (every forward + # goes through ``inputs_embeds`` assembled from the audio / text + # buffers in :meth:`preprocess`). We expose a 2-class placeholder + # ``vocab_size`` purely so the vLLM sampler's working buffers + # match the dummy logits returned by :meth:`compute_logits`. + # The checkpoint, however, ships ``embed_tokens.weight`` at the + # original tokenizer vocab size — which trips + # ``VocabParallelEmbedding``'s + # ``loaded_weight.shape[output_dim] == self.org_vocab_size`` + # assertion. Truncate (or pad) the loaded weight to + # ``(config.vocab_size, hidden_size)`` so the assertion passes; + # the surviving rows are never consumed at runtime. + target_vocab = int(self.config.vocab_size) + embed_weight_name = "model.backbone.embed_tokens.weight" + + def _adjusted_weights() -> Iterable[tuple[str, torch.Tensor]]: + for name, w in weights: + if name == embed_weight_name and w.dim() >= 1 and w.shape[0] != target_vocab: + if w.shape[0] >= target_vocab: + yield name, w[:target_vocab].contiguous() + else: + pad = torch.zeros( + target_vocab - w.shape[0], + *w.shape[1:], + dtype=w.dtype, + device=w.device, + ) + yield name, torch.cat([w, pad], dim=0).contiguous() + else: + yield name, w + + # ``AutoWeightsLoader`` only dispatches into child modules and + # ``nn.Parameter``s, so any registered buffer (e.g. + # ``model.sil_tokens``) needs to be routed manually. Resolve the + # buffer names *after* applying ``hf_to_vllm_mapper`` so the + # checkpoint key matches the in-model attribute path. + buffers_dict = dict(self.named_buffers()) + loaded_buffer_names: set[str] = set() + + def _route_buffers( + stream: Iterable[tuple[str, torch.Tensor]], + ) -> Iterable[tuple[str, torch.Tensor]]: + for name, weight in stream: + if name in buffers_dict: + buf = buffers_dict[name] + with torch.no_grad(): + buf.copy_(weight.to(buf.dtype)) + loaded_buffer_names.add(name) + continue + yield name, weight + + # ``hf_to_vllm_mapper`` rewrites ``model.sampler.*`` to + # ``sampler_module.sampler.*`` so the upstream EarTTS checkpoint + # (which still places the MaskGIT sampler under ``model.``) lands + # on the dedicated :attr:`sampler_module` compile group. + loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) + loaded = loader.load_weights( + _route_buffers(_adjusted_weights()), mapper=self.hf_to_vllm_mapper + ) + loaded.update(loaded_buffer_names) + return loaded diff --git a/nemo/collections/speechlm2/inference/vllm_omni/eartts/pipeline.py b/nemo/collections/speechlm2/inference/vllm_omni/eartts/pipeline.py new file mode 100644 index 000000000000..59dfa2adf490 --- /dev/null +++ b/nemo/collections/speechlm2/inference/vllm_omni/eartts/pipeline.py @@ -0,0 +1,63 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Single-stage streaming EarTTS pipeline. + +NeMo submits text tokens directly to this engine after NemotronDuplexH has +produced them. CFG is represented by two explicit requests in the same +engine, giving the conditional and unconditional streams independent vLLM KV +caches while allowing :class:`EarTTSCFGScheduler` to keep them in lockstep. +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_CFG_SCHEDULER = ( + "nemo.collections.speechlm2.inference.vllm_omni." + "eartts.scheduler.EarTTSCFGScheduler" +) + + +EARTTS_PIPELINE = PipelineConfig( + model_type="eartts", + model_arch="EarTTSForCausalLM", + hf_architectures=("EarTTSForCausalLM",), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="eartts", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=True, + final_output_type="audio", + # vLLM-Omni derives the entry-stage ``generate`` task from + # ``owns_tokenizer`` (runtime name: ``is_comprehension``). EarTTS + # consumes token-id placeholders and deploy keeps + # ``skip_tokenizer_init: true``, but this flag must still be true + # for a direct SamplingParams request to pass task validation. + owns_tokenizer=True, + model_arch="EarTTSForCausalLM", + engine_output_type="audio", + retains_state_across_chunks=True, + scheduler_cls=_CFG_SCHEDULER, + sampling_constraints={"detokenize": False}, + ), + ), +) + + +__all__ = ["EARTTS_PIPELINE"] diff --git a/nemo/collections/speechlm2/inference/vllm_omni/eartts/scheduler.py b/nemo/collections/speechlm2/inference/vllm_omni/eartts/scheduler.py new file mode 100644 index 000000000000..3f559dbe3ebd --- /dev/null +++ b/nemo/collections/speechlm2/inference/vllm_omni/eartts/scheduler.py @@ -0,0 +1,451 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Pair-aware synchronous vLLM-Omni scheduler for EarTTS CFG. + +EarTTS classifier-free guidance is represented by two ordinary top-level +requests. Their ``SamplingParams.extra_args`` must contain:: + + { + "cfg_enabled": True, + "cfg_role": "cond" | "uncond", + "cfg_pair_id": "", + "cfg_scale": 0.5, + } + +The model runner is responsible for blending the pair's model outputs. This +scheduler supplies the ordering and lock-step contract required by that +operation. It deliberately subclasses the synchronous Omni AR scheduler: +async placeholder scheduling can put the two members on different token +positions before either result reaches the scheduler. +""" + +from __future__ import annotations + +import math +from typing import Any + +from vllm.logger import init_logger +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.request import Request, RequestStatus, StreamingUpdate +from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler + +logger = init_logger(__name__) + +_COND = "cond" +_UNCOND = "uncond" +_ROLES = (_COND, _UNCOND) + + +def _extra_args(request: Any) -> dict[str, Any]: + sampling_params = getattr(request, "sampling_params", None) + extra_args = getattr(sampling_params, "extra_args", None) + return extra_args if isinstance(extra_args, dict) else {} + + +def _normalized_sampled_tokens(value: Any) -> tuple[int, ...]: + if value is None: + return () + if isinstance(value, int): + return (value,) + if hasattr(value, "tolist"): + value = value.tolist() + if isinstance(value, int): + return (value,) + return tuple(int(token_id) for token_id in value) + + +class EarTTSCFGScheduler(OmniARScheduler): + """Synchronous Omni AR scheduler that keeps EarTTS CFG pairs lock-step.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self._cfg_pairs: dict[str, dict[str, str]] = {} + self._cfg_req_to_pair: dict[str, str] = {} + self._cfg_pair_scales: dict[str, float] = {} + + max_num_seqs = int(getattr(self.scheduler_config, "max_num_seqs", 0) or 0) + if max_num_seqs and max_num_seqs < 2: + raise ValueError("EarTTSCFGScheduler requires max_num_seqs >= 2") + + @staticmethod + def _cfg_metadata(request: Request) -> tuple[str, str, float] | None: + extra_args = _extra_args(request) + if not bool(extra_args.get("cfg_enabled", False)): + return None + + role = extra_args.get("cfg_role") + if role not in _ROLES: + raise ValueError(f"CFG request {request.request_id!r}: cfg_role must be 'cond' or 'uncond', got {role!r}") + + raw_pair_id = extra_args.get("cfg_pair_id") + if raw_pair_id is None or not str(raw_pair_id): + raise ValueError(f"CFG request {request.request_id!r}: cfg_pair_id must be non-empty") + pair_id = str(raw_pair_id) + + raw_scale = extra_args.get("cfg_scale") + if isinstance(raw_scale, bool) or not isinstance(raw_scale, int | float): + raise ValueError(f"CFG request {request.request_id!r}: cfg_scale must be a finite non-negative number") + scale = float(raw_scale) + if not math.isfinite(scale) or scale < 0.0: + raise ValueError(f"CFG request {request.request_id!r}: cfg_scale must be a finite non-negative number") + return pair_id, role, scale + + def add_request(self, request: Request) -> None: + metadata = self._cfg_metadata(request) + if metadata is not None: + pair_id, role, scale = metadata + roles = self._cfg_pairs.get(pair_id, {}) + existing = roles.get(role) + if existing is not None and existing != request.request_id: + raise ValueError(f"CFG pair {pair_id!r} already has {role} request {existing!r}") + existing_scale = self._cfg_pair_scales.get(pair_id) + if existing_scale is not None and existing_scale != scale: + raise ValueError( + f"CFG pair {pair_id!r} has inconsistent cfg_scale values: {existing_scale} and {scale}" + ) + + super().add_request(request) + + if metadata is not None: + pair_id, role, scale = metadata + self._cfg_pairs.setdefault(pair_id, {})[role] = request.request_id + self._cfg_req_to_pair[request.request_id] = pair_id + self._cfg_pair_scales[pair_id] = scale + # Conditional and null prompts need deterministic, identical + # progress; asymmetric prefix-cache hits violate that contract. + if hasattr(request, "skip_reading_prefix_cache"): + request.skip_reading_prefix_cache = True + + def _pair_requests(self, pair_id: str) -> tuple[Request | None, Request | None]: + roles = self._cfg_pairs.get(pair_id, {}) + return self.requests.get(roles.get(_COND, "")), self.requests.get(roles.get(_UNCOND, "")) + + def _drop_pair(self, pair_id: str) -> None: + for request_id in self._cfg_pairs.pop(pair_id, {}).values(): + self._cfg_req_to_pair.pop(request_id, None) + self._cfg_pair_scales.pop(pair_id, None) + + @staticmethod + def _remove_from_queue(queue: Any, requests: list[Request]) -> None: + if not requests: + return + if hasattr(queue, "remove_requests"): + queue.remove_requests(requests) + return + for request in requests: + queue.remove(request) + + @staticmethod + def _prepend_to_queue(queue: Any, requests: list[Request]) -> None: + if not requests: + return + if hasattr(queue, "prepend_requests"): + queue.prepend_requests(requests) + return + if hasattr(queue, "prepend_request"): + for request in reversed(requests): + queue.prepend_request(request) + return + for request in reversed(requests): + queue.insert(0, request) + + @classmethod + def _replace_queue(cls, queue: Any, requests: list[Request]) -> None: + current = list(queue) + cls._remove_from_queue(queue, current) + for request in requests: + queue.add_request(request) if hasattr(queue, "add_request") else queue.append(request) + + def _available_sequence_slots(self) -> int: + capacity = getattr(self, "max_num_running_reqs", None) + if capacity is None: + capacity = getattr(self.scheduler_config, "max_num_seqs", 0) + return max(0, int(capacity or 0) - len(self.running)) + + def _prepare_cfg_waiting(self) -> list[tuple[Any, list[Request]]]: + """Hide unsafe pairs and put admissible pairs first and adjacent.""" + skipped_queue = getattr(self, "skipped_waiting", None) + promoted_held: list[Request] = [] + + # Streaming updates are applied by vLLM while traversing + # ``skipped_waiting``. The first traversal promotes each pair member + # from WAITING_FOR_STREAMING_REQ to WAITING but deliberately skips + # scheduling it (see _try_promote_blocked_waiting_request below). + # Once both members are promoted, move them to the ordinary waiting + # queue together so the admission logic below sees one atomic pair. + if skipped_queue is not None: + skipped_by_id = { + request.request_id: request + for request in list(skipped_queue) + } + for pair_id, roles in self._cfg_pairs.items(): + if set(roles) != set(_ROLES): + continue + cond = skipped_by_id.get(roles[_COND]) + uncond = skipped_by_id.get(roles[_UNCOND]) + if cond is None or uncond is None: + continue + ready = [ + request + for request in (cond, uncond) + if request.status == RequestStatus.WAITING + ] + if len(ready) == 2: + self._remove_from_queue( + skipped_queue, [cond, uncond] + ) + self.waiting.add_request(cond) + self.waiting.add_request(uncond) + elif len(ready) == 1: + self._remove_from_queue(skipped_queue, ready) + promoted_held.extend(ready) + + waiting_items = list(self.waiting) + skipped_items = list(skipped_queue) if skipped_queue is not None else [] + waiting_ids = {request.request_id for request in waiting_items} + skipped_ids = {request.request_id for request in skipped_items} + running_ids = {request.request_id for request in self.running} + + held: list[tuple[Any, list[Request]]] = [] + if promoted_held: + held.append((skipped_queue, promoted_held)) + hold_waiting: set[str] = set() + complete_waiting_pairs: list[str] = [] + + for pair_id, roles in self._cfg_pairs.items(): + pair_ids = {request_id for request_id in roles.values()} + is_complete = set(roles) == set(_ROLES) and all(request_id in self.requests for request_id in pair_ids) + in_waiting = pair_ids & waiting_ids + in_skipped = pair_ids & skipped_ids + in_running = pair_ids & running_ids + + if in_running and (in_waiting or in_skipped): + raise RuntimeError(f"EarTTS CFG pair {pair_id!r} was split across running and waiting queues") + if len(in_running) == 1: + raise RuntimeError(f"EarTTS CFG pair {pair_id!r} has only one running member") + + if not is_complete or len(in_waiting) == 1: + hold_waiting.update(in_waiting) + elif len(in_waiting) == 2: + complete_waiting_pairs.append(pair_id) + + # A pair consumes two sequence slots. Expose only whole pairs to the + # upstream scheduler and put them before ordinary requests so another + # admission cannot consume the second slot between pair members. + admitted_pair_count = self._available_sequence_slots() // 2 + allowed_pairs = set(complete_waiting_pairs[:admitted_pair_count]) + for pair_id in complete_waiting_pairs[admitted_pair_count:]: + hold_waiting.update(self._cfg_pairs[pair_id].values()) + + held_waiting = [request for request in waiting_items if request.request_id in hold_waiting] + self._remove_from_queue(self.waiting, held_waiting) + if held_waiting: + held.append((self.waiting, held_waiting)) + + remaining = list(self.waiting) + ordinary = [request for request in remaining if request.request_id not in self._cfg_req_to_pair] + ordered: list[Request] = [] + for pair_id in complete_waiting_pairs: + if pair_id not in allowed_pairs: + continue + cond, uncond = self._pair_requests(pair_id) + if cond is not None and uncond is not None: + ordered.extend((cond, uncond)) + ordered.extend(ordinary) + if ordered != remaining: + self._replace_queue(self.waiting, ordered) + + actual_ids = [request.request_id for request in self.waiting] + for pair_id in allowed_pairs: + roles = self._cfg_pairs[pair_id] + cond_id, uncond_id = roles[_COND], roles[_UNCOND] + try: + cond_index = actual_ids.index(cond_id) + except ValueError: + continue + if cond_index + 1 >= len(actual_ids) or actual_ids[cond_index + 1] != uncond_id: + raise RuntimeError( + "EarTTSCFGScheduler requires an FCFS-compatible waiting " + f"queue; CFG pair {pair_id!r} could not be made adjacent" + ) + + return held + + def _try_promote_blocked_waiting_request( + self, request: Request + ) -> bool: + promoted = super()._try_promote_blocked_waiting_request( + request + ) + if promoted and request.request_id in self._cfg_req_to_pair: + # The base scheduler would immediately schedule this first + # promoted member. Return False once so it is parked back in + # skipped_waiting; after its peer is promoted, + # _prepare_cfg_waiting moves both to waiting atomically. + return False + return promoted + + def _should_defer_waiting_admission(self) -> bool: + """Install the pair guard after Omni has processed pending inputs. + + ``OmniARScheduler.schedule`` invokes this hook immediately before the + stock vLLM scheduler. Preparing here is important: doing it at the + start of this class's :meth:`schedule` would hide requests from Omni's + chunk/input processing and could leave a streaming pair parked in + ``skipped_waiting`` forever. + """ + self._cfg_decode_ready_before = { + request.request_id + for request in self.running + if self._get_confirmed_num_computed_tokens(request) >= request.num_prompt_tokens + } + self._cfg_held_for_schedule = self._prepare_cfg_waiting() + self._cfg_waiting_before_schedule = {request.request_id for request in self.waiting} + return super()._should_defer_waiting_admission() + + def _restore_held(self, held: list[tuple[Any, list[Request]]]) -> None: + for queue, requests in reversed(held): + self._prepend_to_queue(queue, requests) + + def _equalize_pair_progress(self, scheduler_output: SchedulerOutput) -> None: + scheduled = scheduler_output.num_scheduled_tokens + for pair_id in self._cfg_pairs: + cond, uncond = self._pair_requests(pair_id) + if cond is None or uncond is None: + continue + cond_scheduled = int(scheduled.get(cond.request_id, 0) or 0) + uncond_scheduled = int(scheduled.get(uncond.request_id, 0) or 0) + if not cond_scheduled or not uncond_scheduled: + continue + + target = min(cond.num_computed_tokens, uncond.num_computed_tokens) + feasible = all( + request.num_computed_tokens - target < count + for request, count in ((cond, cond_scheduled), (uncond, uncond_scheduled)) + ) + if not feasible: + continue + + for request, count in ((cond, cond_scheduled), (uncond, uncond_scheduled)): + difference = request.num_computed_tokens - target + if difference <= 0: + continue + request.num_computed_tokens = target + if hasattr(request, "num_in_flight_tokens"): + request.num_in_flight_tokens = max(0, request.num_in_flight_tokens - difference) + scheduler_output.num_scheduled_tokens[request.request_id] = count - difference + scheduler_output.total_num_scheduled_tokens -= difference + + def _assert_atomic_admission(self, scheduler_output: SchedulerOutput, waiting_before: set[str]) -> None: + scheduled = scheduler_output.num_scheduled_tokens + for pair_id, roles in self._cfg_pairs.items(): + pair_ids = {roles.get(_COND), roles.get(_UNCOND)} + pair_ids.discard(None) + if len(pair_ids & waiting_before) != 2: + continue + admitted = {request_id for request_id in pair_ids if scheduled.get(request_id, 0)} + if admitted and admitted != pair_ids: + raise RuntimeError( + f"EarTTS CFG pair {pair_id!r} was not admitted atomically: scheduled={sorted(admitted)}" + ) + + def _assert_complete_decode_pairs(self, scheduler_output: SchedulerOutput, decode_ready_before: set[str]) -> None: + scheduled = scheduler_output.num_scheduled_tokens + for pair_id, roles in self._cfg_pairs.items(): + cond_id = roles.get(_COND) + uncond_id = roles.get(_UNCOND) + if cond_id is None or uncond_id is None: + continue + cond_count = int(scheduled.get(cond_id, 0) or 0) + uncond_count = int(scheduled.get(uncond_id, 0) or 0) + if not cond_count and not uncond_count: + continue + if cond_count != uncond_count: + raise RuntimeError( + f"EarTTS scheduler split CFG decode pair {pair_id!r}: " + f"{cond_id}={cond_count}, {uncond_id}={uncond_count}" + ) + + def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: + self._cfg_held_for_schedule: list[tuple[Any, list[Request]]] = [] + self._cfg_waiting_before_schedule: set[str] = set() + self._cfg_decode_ready_before: set[str] = set() + + scheduler_config = getattr(self, "scheduler_config", None) + original_threshold = getattr(scheduler_config, "long_prefill_token_threshold", None) + if self._cfg_pairs and original_threshold is not None: + budget = int(getattr(self, "max_num_scheduled_tokens", 0) or 0) + if budget: + scheduler_config.long_prefill_token_threshold = max(1, budget // 2) + try: + scheduler_output = super().schedule(throttle_prefills) + finally: + if original_threshold is not None: + scheduler_config.long_prefill_token_threshold = original_threshold + self._restore_held(self._cfg_held_for_schedule) + + self._equalize_pair_progress(scheduler_output) + self._assert_atomic_admission(scheduler_output, self._cfg_waiting_before_schedule) + self._assert_complete_decode_pairs(scheduler_output, self._cfg_decode_ready_before) + return scheduler_output + + def _assert_matching_sampled_tokens(self, scheduler_output: SchedulerOutput, model_runner_output: Any) -> None: + sampled = getattr(model_runner_output, "sampled_token_ids", None) + req_id_to_index = getattr(model_runner_output, "req_id_to_index", None) + if sampled is None or not isinstance(req_id_to_index, dict): + return + + scheduled = scheduler_output.num_scheduled_tokens + for pair_id, roles in self._cfg_pairs.items(): + cond_id = roles.get(_COND) + uncond_id = roles.get(_UNCOND) + if ( + cond_id not in scheduled + or uncond_id not in scheduled + or cond_id not in req_id_to_index + or uncond_id not in req_id_to_index + ): + continue + cond_tokens = _normalized_sampled_tokens(sampled[req_id_to_index[cond_id]]) + uncond_tokens = _normalized_sampled_tokens(sampled[req_id_to_index[uncond_id]]) + if cond_tokens != uncond_tokens: + raise RuntimeError( + f"EarTTS CFG pair {pair_id!r} sampled-token mismatch: " + f"cond={cond_tokens}, uncond={uncond_tokens}" + ) + + def update_from_output(self, scheduler_output: SchedulerOutput, model_runner_output: Any) -> Any: + self._assert_matching_sampled_tokens(scheduler_output, model_runner_output) + # A length stop on a resumable StreamingInput segment is not terminal: + # Omni parks the request until its next chunk. Do not mirror the + # segment's internal free/park operation onto its CFG peer here. + # Explicit abort/final teardown still flows through finish_requests, + # which expands either member to the complete pair below. + return super().update_from_output( + scheduler_output, model_runner_output + ) + + def _update_request_as_session(self, session: Request, update: StreamingUpdate) -> None: + super()._update_request_as_session(session, update) + # Stage 0 receives direct StreamingInput updates. Upstream extends the + # prompt but does not copy this per-chunk payload. Downstream stages + # are intentionally untouched; their connector owns payload delivery. + if self.vllm_config.model_config.stage_id == 0: + additional_information = getattr(update, "additional_information", None) + if additional_information is not None: + session.additional_information = additional_information + + +__all__ = ["EarTTSCFGScheduler"] diff --git a/nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/__init__.py b/nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/__init__.py new file mode 100644 index 000000000000..33bb11b24f6d --- /dev/null +++ b/nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/__init__.py @@ -0,0 +1,19 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from nemo.collections.speechlm2.inference.vllm_omni.nemotron_duplex_h.nemotron_duplex_h import ( + NemotronDuplexHForCausalLM, +) + +__all__ = ["NemotronDuplexHForCausalLM"] diff --git a/nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/nemotron_duplex_h.py b/nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/nemotron_duplex_h.py new file mode 100644 index 000000000000..ee7df96cbaac --- /dev/null +++ b/nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/nemotron_duplex_h.py @@ -0,0 +1,771 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Inference-only NemotronDuplexH model for vLLM-Omni. + +A minimal extension of the upstream :class:`NemotronHForCausalLM` that: + +1. Accepts pre-computed acoustic encoder embeddings per step via + ``acoustic_embedding`` in the per-request payload (one row per + scheduled token). The prefill step receives the *system prompt as + raw text* via ``system_prompt`` in the same payload; the + model's :meth:`preprocess` tokenizes it in-process (using a + HuggingFace tokenizer loaded once in ``__init__`` from the + checkpoint dir) and constructs the prefill combined embedding + itself as + + prompt_embed = embed_tokens([BOS] + text_ids + [EOS]) + + embed_tokens(pad_id) + + embed_asr_tokens(pad_id) + + It then *clears* the buffer entry by returning + ``{"system_prompt": None}`` as its update dict so that subsequent + decode steps fall through to the decode branch. The producer + should still send ``system_prompt=None`` on every decode chunk to + make the intent explicit, but the actual clearing happens + consumer-side because the orchestrator's serialization filters + ``None`` values. +2. Embeds up to two additional per-step token id streams, each fed + **autoregressively from the model itself** via a per-request buffer + that ``postprocess`` keeps populated after every step: + + - ``input_asr_ids`` – the ASR channel (``predict_user_text`` + checkpoints), embedded with its own ``embed_asr_tokens`` table. + - ``input_function_ids`` – the function channel + (``use_function_head`` checkpoints), embedded with the *text* + ``embed_tokens`` table and scaled by + ``duplex_function_channel_weight``, mirroring + ``DuplexSTTModel.build_input_embedding``. + + Which channels exist is checkpoint-dependent. ASR and function can + both be enabled. The converter records whichever heads it found. + +3. Combines the enabled signals into the input embedding fed to the + NemotronH backbone: + + hidden_in = embed_tokens(input_ids) + [+ embed_asr_tokens(input_asr_ids)] + [+ embed_tokens(input_function_ids) * weight] + + acoustic_embedding + +4. Adds a parallel head per enabled channel (``asr_head`` / + ``function_head``) that produces one token at every decoding step. + The head matmul and the ``argmax`` run in :meth:`make_omni_output` + (which the runner invokes *outside* the CUDA-graph wrapper) on the + full-batch ``hidden_states`` returned by :meth:`forward`. The tokens + are exposed under ``OmniOutput.multimodal_outputs["asr_tokens"]`` and + ``["function_tokens"]``, and :meth:`postprocess` stashes the + request's last id of each back into the corresponding buffer so the + next step's :meth:`preprocess` can read it as that channel's + autoregressive input. + + Returning a dict-with-tensor directly from :meth:`forward` is + unsafe under FULL CUDA graphs: ``weak_ref_tensors`` cannot weak-ref + tensors nested inside dicts, and the wrapper coerces ``NamedTuple`` + to a plain ``tuple`` on replay. Routing the multimodal output + through :meth:`make_omni_output` keeps every cudagraph-replayed + value a plain ``Tensor``. + +Text token sampling uses a custom vLLM logits processor which calls the same +PyTorch sampler as the native backend, then forces vLLM's greedy sampler to +the selected token. The auxiliary channels are always greedy. +""" + +from collections.abc import Iterable +from typing import Any + +import torch +from transformers import AutoTokenizer, PreTrainedTokenizerBase +from vllm.config import VllmConfig +from vllm.model_executor.layers.vocab_parallel_embedding import ( + DEFAULT_VOCAB_PADDING_SIZE, + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + maybe_prefix, +) +from vllm.sequence import IntermediateTensors + +from vllm_omni.model_executor.models.output_templates import OmniOutput + +from nemo.collections.speechlm2.parts.logit_boosts import ( + LogitBoosts, + apply_logit_boosts, +) +from nemo.utils import logging as logger + + +def _is_system_prompt_prefill( + system_prompt: Any, + runner_is_prefill: bool, + request_id: str, + prompt_token_cache: dict[str, list[int]], +) -> bool: + """Distinguish initial prompt slices from streaming prompt extensions.""" + has_system_prompt = isinstance(system_prompt, str) and bool( + system_prompt.strip() + ) + return has_system_prompt or ( + runner_is_prefill and request_id in prompt_token_cache + ) + + +def _is_internal_prefill_token( + is_prompt_prefill: bool, + runner_is_prefill: bool, + has_acoustic_embedding: bool, +) -> bool: + """True for vLLM's 1-token generate after the system prompt is in KV. + + The runner labels every streaming extension as prefill because the + prompt grows before that token is computed. After a prompt longer + than the 64-token long-prefill threshold, the last prompt slice pops + the token cache; the engine then schedules one more token (internal + ``t_0``) with ``_omni_is_prefill=True`` and no ``acoustic_embedding``. + A prefill-only ``generate_step`` (empty ``is_first`` frame) does not + queue audio before that happens, so the decode path would assert. + Real client steps always carry an acoustic frame. + """ + return ( + not is_prompt_prefill + and runner_is_prefill + and not has_acoustic_embedding + ) + + +class NemotronDuplexHForCausalLM(NemotronHForCausalLM): + """NemotronH + optional per-step ASR and function token channels.""" + + have_multimodal_outputs = True + has_preprocess = True + has_postprocess = True + + # No ``gpu_resident_buffer_keys``; see the note in ``eartts.py``. The keys + # this stage passes between ``postprocess`` and the next ``preprocess`` are + # flat, which that mechanism cannot express. + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + # NemotronH backbone weights live under + # `stt_model.llm.backbone.*` in the duplex checkpoint and need to + # land under our `model.*`. + "stt_model.llm.backbone": "model", + "stt_model.llm": "model", + "stt_model.embed_tokens": "model.embed_tokens", + "stt_model.embed_asr_tokens": "embed_asr_tokens", + "stt_model.lm_head": "lm_head", + "stt_model.asr_head": "asr_head", + "stt_model.function_head": "function_head", + # Bare-NemotronH naming, kept as a fallback. + "backbone": "model", + }, + orig_to_new_substr={"A_log": "A", "embeddings": "embed_tokens"}, + # Fusing q/k/v into ``qkv_proj`` is done here. This class replaces + # ``NemotronHForCausalLM.hf_to_vllm_mapper`` wholesale, so the stacked + # mapping has to be restated or attention projections fail to load. + orig_to_new_stacked={ + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + }, + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__(vllm_config=vllm_config, prefix=prefix) + + config = vllm_config.model_config.hf_config + + # Missing flags default to ASR on / function off. Fresh conversions + # always write both flags from the checkpoint weights. + self.use_asr_head = bool(getattr(config, "use_asr_head", True)) + self.use_function_head = bool(getattr(config, "use_function_head", False)) + self.function_channel_weight = float(getattr(config, "duplex_function_channel_weight", 1.0)) + + if self.use_asr_head: + self.embed_asr_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + ) + + self.asr_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + padding_size=DEFAULT_VOCAB_PADDING_SIZE, + prefix=maybe_prefix(prefix, "asr_head"), + ) + + # The function channel has no embedding table of its own: its feedback + # token is embedded with the text ``embed_tokens``, exactly as + # ``DuplexSTTModel`` does. + if self.use_function_head: + self.function_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + padding_size=DEFAULT_VOCAB_PADDING_SIZE, + prefix=maybe_prefix(prefix, "function_head"), + ) + + # Tokenizer is used in ``preprocess`` to convert the + # ``additional_information["system_prompt"]`` text into token + # IDs on the prefill chunk. Loaded from the checkpoint dir so + # the vocabulary aligns with ``embed_tokens``. Cached once on + # init to keep the per-step preprocess fast. + model_path = vllm_config.model_config.model + self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + # The runner can split a prompt at its 64-token long-prefill threshold. + # Cache the full tokenization in this model process until every slice + # for a request has been consumed. + self._prompt_token_cache: dict[str, list[int]] = {} + self._seeded_channel_requests: set[str] = set() + self._channel_state: dict[str, dict[str, torch.Tensor]] = {} + # This pipeline is configured with max_num_seqs=1. Keep a fallback + # for vLLM-Omni paths that rewrite the internal request id between + # streaming segments. + self._last_channel_state: dict[str, torch.Tensor] = {} + self._current_is_prompt_prefill = False + + # Special token IDs used to construct the prefill prompt: + # ``[BOS] + text_ids + [EOS]`` and the pad embedding added to + # every prefill position (mirrors the reference STT recipe + # where the BOS / pad embeddings are both ``embed_tokens(pad_id)``). + self.pad_token_id = int(config.pad_token_id) + self.bos_token_id = int(config.bos_token_id) + self.eos_token_id = int(config.eos_token_id) + + # User (ASR) channel boosts, read from the converted config so this + # model applies them exactly as DuplexSTTModel does. The agent-channel + # boosts arrive per request through the shared text sampling hook. + self.user_logit_boosts = LogitBoosts.user_from_cfg(config) + if self.user_logit_boosts: + logger.info( + "NemotronDuplexH user logit boosts: " + f"{self.user_logit_boosts.as_dict()}" + ) + self._last_asr_token = torch.full( + (1,), self.pad_token_id, dtype=torch.long + ) + self._last_function_token = torch.full( + (1,), self.pad_token_id, dtype=torch.long + ) + + # Per-position pad embedding added on every prefill step: the pad id + # embedded once per *enabled* auxiliary channel, shape + # ``(hidden_size,)``. Materialized at the end of + # :meth:`load_weights` because the embedding tables are not + # populated yet in ``__init__``. Registered as a *non-persistent* + # buffer so it follows ``.to(device)`` / dtype casts with the + # rest of the module but is **not** saved in the state_dict + # (it is fully derived from ``embed_tokens`` / + # ``embed_asr_tokens`` which are already saved — duplicating it + # in the checkpoint would just be a footgun). + self.register_buffer("_pad_combined_emb", None, persistent=False) + + # ------------------------------------------------------------------ # + # producer-side helper # + # ------------------------------------------------------------------ # + + @staticmethod + def compute_prefix_len(tokenizer: PreTrainedTokenizerBase, system_prompt: str) -> int: + """Length of the prefill chunk for a given system prompt. + + Mirrors the in-model tokenization done by :meth:`preprocess`: + + [BOS] + tokenizer.encode(system_prompt, add_special_tokens=False) + [EOS] + + The streaming producer needs this number to size the + placeholder ``prompt_token_ids`` it hands vLLM on the prefill + chunk — vLLM schedules off that list's length, while the + actual embedding is constructed inside :meth:`preprocess` + from the ``system_prompt`` string. + + Exposed as a ``@staticmethod`` so callers can compute the + length without instantiating the model (which would download + the full checkpoint). They just need any tokenizer compatible + with the model's vocabulary — typically + ``AutoTokenizer.from_pretrained()``, the + same instance used to decode output tokens. + """ + text_ids = tokenizer.encode(system_prompt, add_special_tokens=False) + return len(text_ids) + 2 # +2 for BOS / EOS wrapped in preprocess + + # ------------------------------------------------------------------ # + # preprocess # + # ------------------------------------------------------------------ # + + def _materialize_pad_combined_emb(self) -> None: + embed_weight = self.model.embed_tokens.weight + device = embed_weight.device + dtype = embed_weight.dtype + pad_tokens = torch.full( + (1,), self.pad_token_id, device=device, dtype=torch.long + ) + pad_emb = self.model.embed_tokens(pad_tokens).to(dtype).squeeze(0) + combined_pad = pad_emb + if self.use_asr_head: + combined_pad = combined_pad + self.embed_asr_tokens( + pad_tokens + ).to(dtype).squeeze(0) + if self.use_function_head: + combined_pad = ( + combined_pad + + pad_emb * self.function_channel_weight + ) + self._pad_combined_emb = combined_pad.detach() + + def _embeds_without_acoustic( + self, input_ids: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Text + pad-channel embeddings, used for prompt slices and internal ``t_0``.""" + if self._pad_combined_emb is None: + self._materialize_pad_combined_emb() + target_dtype = self.model.embed_tokens.weight.dtype + text_emb = self.model.embed_tokens(input_ids).to(target_dtype) + return input_ids, text_emb + self._pad_combined_emb, {"system_prompt": None} + + def preprocess( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None, + **info_dict: Any, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Combine text/asr/speech embeddings into a single per-token vector. + + Three paths: + + * **Prefill construction.** When + ``additional_information["system_prompt"]`` is a non-empty + string, this is the prefill chunk. We tokenize the prompt + in-process as ``[BOS] + tokenizer.encode(prompt) + [EOS]``, + embed it with ``model.embed_tokens``, and add the pad + embedding of each enabled auxiliary channel (which the + reference STT recipe folds in uniformly across every prefill + position): + + prefill_combined = embed_tokens(prompt_token_ids) + + _pad_combined_emb + + The producer-supplied ``input_ids`` for this chunk are + placeholders (their length must match the tokenized prompt + length so vLLM's scheduling sees the right prefill size); + they are returned unchanged so vLLM's bookkeeping is + consistent. We then *clear* the buffer entry by returning + ``{"system_prompt": None}`` in the update dict. Clearing has + to happen here because the orchestrator's serialization + (:func:`vllm_omni.data_entry_keys.serialize_payload`) + silently drops ``None`` values, so the producer cannot + overwrite the buffer with ``None`` via the streaming-input + merge. Decode chunks may carry ``system_prompt=None``, but it + has no effect; the state transition happens only here. + + * **Internal ``t_0`` after a chunked prompt.** vLLM may split the + prompt at its 64-token long-prefill threshold and then + schedule a one-token continuation labeled + ``_omni_is_prefill`` with no ``acoustic_embedding``. That + token is discarded by the session (``output_count <= 1``); + embed it like prefill (text + pad, no user acoustics) so a + prefill-only ``generate_step`` cannot kill the engine. + + * **Decode (single-token step).** Builds the combined embedding + per scheduled token from: + + - ``input_ids`` – per-step text token id (one per + scheduled token; standard vLLM + autoregressive feedback). + - ``input_asr_ids`` – per-step ASR token id, written + back by :meth:`postprocess` on + every step. ASR channel only. + - ``input_function_ids`` – per-step function token id, same + write-back path. Function + channel only. + - ``acoustic_embedding`` – per-step acoustic encoder + embedding, sourced from + ``additional_information``. + + ``input_embeds`` is the runner's pre-allocated scratch buffer + on this path and its contents are ignored. + """ + device = input_ids.device + n = int(input_ids.shape[0]) + + # Prefill vs decode is detected directly on the value of + # ``system_prompt``: a non-empty string means prefill, anything + # else (``None`` / missing / empty) means decode. The buffer + # flips from str → ``None`` inside this method itself (see the + # update dict returned below) because serialization drops + # ``None`` and the producer's "send None on each decode chunk" + # pattern alone is not enough to clear the slot. + system_prompt = info_dict.get("system_prompt") + is_prefill = bool(info_dict.get("_omni_is_prefill", False)) + request_id = str( + info_dict.get("global_request_id") + or info_dict.get("request_id") + or "" + ) + has_system_prompt = isinstance(system_prompt, str) and bool( + system_prompt.strip() + ) + is_prompt_prefill = _is_system_prompt_prefill( + system_prompt, + is_prefill, + request_id, + self._prompt_token_cache, + ) + has_acoustic = isinstance(info_dict.get("acoustic_embedding"), torch.Tensor) + is_internal_t0 = _is_internal_prefill_token( + is_prompt_prefill, is_prefill, has_acoustic + ) + # vLLM labels every one-token streaming extension as prefill because + # the prompt grows before that token is computed. Prompt slices and + # the engine-internal t_0 after them seed auxiliary feedback with PAD; + # client-visible decode steps (which always carry acoustic_embedding) + # do not. + self._current_is_prompt_prefill = is_prompt_prefill or is_internal_t0 + if is_prompt_prefill: + # [BOS] + encode(text, add_special_tokens=False) + [EOS]. + # ``add_special_tokens=False`` keeps full control of which + # specials get wrapped around the text (the underlying HF + # tokenizer would otherwise prepend its own BOS, which may + # or may not equal ``config.bos_token_id``). + if has_system_prompt: + text_ids = self.tokenizer.encode( + system_prompt, add_special_tokens=False + ) + prompt_token_ids = [ + self.bos_token_id, + *text_ids, + self.eos_token_id, + ] + self._prompt_token_cache[request_id] = prompt_token_ids + else: + prompt_token_ids = self._prompt_token_cache[request_id] + prompt_len = len(prompt_token_ids) + expected_prompt_len = info_dict.get("duplex_prompt_len") + if expected_prompt_len is not None: + assert prompt_len == int(expected_prompt_len), ( + f"system_prompt tokenizes to {prompt_len} ids but vLLM " + f"tracks a prompt of length {expected_prompt_len}" + ) + offset = int(info_dict.get("duplex_token_offset", 0) or 0) + end = offset + n + if not (0 <= offset < end <= prompt_len): + # Cache still populated but the runner scheduled past the + # prompt (internal t_0). Same pad-embedding path as below. + self._prompt_token_cache.pop(request_id, None) + return self._embeds_without_acoustic(input_ids) + + prompt_tokens = torch.tensor( + prompt_token_ids[offset:end], + device=device, + dtype=torch.long, + ) + _, prefill_combined, updates = self._embeds_without_acoustic( + prompt_tokens + ) + if end == prompt_len: + self._prompt_token_cache.pop(request_id, None) + return input_ids, prefill_combined, updates + + if is_internal_t0: + return self._embeds_without_acoustic(input_ids) + + combined = self.model.embed_tokens(input_ids) + + request_id = str( + info_dict.get("global_request_id") + or info_dict.get("request_id") + or "" + ) + cached_channel_state = self._channel_state.get( + request_id + ) or self._last_channel_state + for key, value in cached_channel_state.items(): + if not isinstance(info_dict.get(key), torch.Tensor): + info_dict[key] = value + if self.use_asr_head and not isinstance( + info_dict.get("input_asr_ids"), torch.Tensor + ): + info_dict["input_asr_ids"] = self._last_asr_token + if self.use_function_head and not isinstance( + info_dict.get("input_function_ids"), torch.Tensor + ): + info_dict[ + "input_function_ids" + ] = self._last_function_token + if request_id not in self._seeded_channel_requests: + # The initial prompt output is not guaranteed to run postprocess + # before the first direct StreamingInput update on a one-stage + # engine. Seed optional feedback channels exactly as native does; + # subsequent missing state remains an error. + if self.use_asr_head and not isinstance( + info_dict.get("input_asr_ids"), torch.Tensor + ): + info_dict["input_asr_ids"] = torch.full( + (n,), + self.pad_token_id, + device=device, + dtype=torch.long, + ) + if self.use_function_head and not isinstance( + info_dict.get("input_function_ids"), torch.Tensor + ): + info_dict["input_function_ids"] = torch.full( + (n,), + self.pad_token_id, + device=device, + dtype=torch.long, + ) + self._seeded_channel_requests.add(request_id) + + if self.use_asr_head: + asr_ids = self._channel_ids(info_dict, "input_asr_ids", n, device) + combined = combined + self.embed_asr_tokens(asr_ids) + + if self.use_function_head: + function_ids = self._channel_ids(info_dict, "input_function_ids", n, device) + combined = combined + self.model.embed_tokens(function_ids) * self.function_channel_weight + + # Per-step acoustic encoder embedding, sourced from + # ``additional_information["acoustic_embedding"]``. + acoustic = info_dict.get("acoustic_embedding") + assert isinstance(acoustic, torch.Tensor), ( + "acoustic_embedding is required in the per-step payload on every decode step; " + f"got {type(acoustic).__name__} with available keys {sorted(info_dict)}" + ) + acoustic = acoustic.to(device=device, dtype=combined.dtype) + assert acoustic.dim() == 2, f"acoustic_embedding must be 2D, got shape {tuple(acoustic.shape)}" + assert acoustic.shape[0] == n, ( + f"acoustic_embedding length {acoustic.shape[0]} does not match scheduled token count {n}" + ) + combined = combined + acoustic + + return input_ids, combined, {} + + @staticmethod + def _channel_ids(info_dict: dict[str, Any], key: str, n: int, device: torch.device) -> torch.Tensor: + """Read one auxiliary channel's per-step feedback ids from the payload.""" + ids = info_dict.get(key) + assert isinstance(ids, torch.Tensor), ( + f"{key} is required on every decode step but is " + f"{type(ids).__name__}; available keys {sorted(info_dict)}" + ) + ids = ids.to(device=device, dtype=torch.long).reshape(-1) + assert ids.numel() == n, f"{key} length {ids.numel()} does not match scheduled token count {n}" + return ids + + # ------------------------------------------------------------------ # + # postprocess - autoregressive feedback for the auxiliary channels # + # ------------------------------------------------------------------ # + + def postprocess( + self, + hidden_states: torch.Tensor, + multimodal_outputs: dict[str, Any] | None = None, + **info_dict: Any, + ) -> dict[str, Any]: + """Stash this request's last auxiliary tokens as the next step's input. + + ``hidden_states`` is a slice of the full-batch hidden_states tensor, + and each ``multimodal_outputs`` entry is the corresponding full-batch + token tensor produced by :meth:`make_omni_output`. We pick the token + aligned with the last position of this request's slice. + + On the prefill chunk the function channel is seeded with the pad id + instead of the prompt's own prediction, because native starts decoding + with ``gen_function`` still at its ``text_pad_id`` fill value and only + feeds back real function tokens from the second frame onwards. + """ + assert multimodal_outputs + start = hidden_states.storage_offset() // hidden_states.stride(0) + last_idx = start + hidden_states.shape[0] - 1 + # The initial prompt can be split into a 64-token slice plus a + # one-token continuation, so tensor length alone cannot identify it. + is_prefill = self._current_is_prompt_prefill + + def last_token(key: str) -> torch.Tensor: + tokens = multimodal_outputs.get(key) + assert isinstance(tokens, torch.Tensor), f"{key} missing from multimodal_outputs" + return tokens[last_idx : last_idx + 1].detach().to(torch.long) + + updates: dict[str, Any] = {} + if self.use_asr_head: + updates["input_asr_ids"] = last_token("asr_tokens") + if self.use_function_head: + if is_prefill: + updates["input_function_ids"] = torch.full( + (1,), self.pad_token_id, device=hidden_states.device, dtype=torch.long + ) + else: + updates["input_function_ids"] = last_token("function_tokens") + request_id = str( + info_dict.get("global_request_id") + or info_dict.get("request_id") + or "" + ) + if request_id: + channel_state = { + key: value.detach() + for key, value in updates.items() + if isinstance(value, torch.Tensor) + } + self._channel_state[request_id] = channel_state + self._last_channel_state = channel_state + return updates + + # ------------------------------------------------------------------ # + # forward # + # ------------------------------------------------------------------ # + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor | IntermediateTensors: + """Run the backbone and return its hidden states. + + ASR tokens are produced by :meth:`make_omni_output`, which the + runner invokes *outside* the CUDA-graph wrapper. This keeps the + captured graph's output a plain ``Tensor`` (or + ``IntermediateTensors``) — both types that ``weak_ref_tensors`` + handles correctly. Returning a ``NamedTuple`` containing a + ``dict[str, Tensor]`` directly here would corrupt the dict's + tensors on FULL graph replay (the wrapper coerces + ``NamedTuple`` -> plain ``tuple`` and cannot weak-ref tensors + nested in dicts). + + IMPORTANT — cudagraph mode requirement + -------------------------------------- + This model must be run with ``cudagraph_mode="PIECEWISE"`` (or + ``enforce_eager=True``). The streaming-input pattern used here + keeps extending each request's prompt with every audio chunk, + so ``num_computed_tokens < num_prompt_tokens`` is permanently + true and Mamba's metadata builder always classifies the request + as a *prefill* (because + :func:`split_decodes_and_prefills` is called with + ``treat_short_extends_as_decodes=False`` in + ``Mamba2AttentionMetadataBuilder._compute_common_metadata``). + + With FULL cudagraph mode, the persistent + ``state_indices_tensor_d`` buffer is only updated when + ``num_prefills == 0``, so for streaming it stays at the + capture-time dummy value (0) while the FULL decode graph is + still dispatched (the dispatcher only checks ``query_len``). + The captured Mamba kernel then reads slot 0 of ``mamba_cache`` + instead of the real slot, producing garbage hidden states. + PIECEWISE side-steps this because the Mamba layer runs eagerly + and reads the freshly-computed metadata tensor, and the prefill + code path correctly *writes* the chunk into Mamba state on + every step (which is essential — there is no separate "prefill" + phase in this streaming setup). + """ + hidden_states = self.model(input_ids, positions, intermediate_tensors, inputs_embeds) + return hidden_states + + # ------------------------------------------------------------------ # + # make_omni_output - runs eagerly outside the CUDA graph wrapper # + # ------------------------------------------------------------------ # + + def make_omni_output( + self, + model_outputs: torch.Tensor | IntermediateTensors | OmniOutput, + **_: Any, + ) -> OmniOutput: + """Wrap backbone hidden states with the auxiliary channel tokens. + + Invoked by :class:`OmniGPUModelRunner._model_forward` after the + CUDA-graph wrapper has returned, so the auxiliary head matmuls + + ``argmax`` here run eagerly. They operate on the full-batch + ``hidden_states`` tensor in a single GEMM each, so the cost is + negligible relative to the backbone forward. + """ + if isinstance(model_outputs, OmniOutput): + return model_outputs + if isinstance(model_outputs, IntermediateTensors): + return OmniOutput( + text_hidden_states=model_outputs, + intermediate_tensors=model_outputs, + ) + + hidden = model_outputs + multimodal_outputs: dict[str, torch.Tensor] = {} + if self.use_asr_head: + asr_logits = self.logits_processor(self.asr_head, hidden) + # The ASR head's logits never reach vLLM's sampler, so the + # user-channel boosts are applied here rather than in the shared + # text logits processor. Same arithmetic as DuplexSTTModel. + apply_logit_boosts( + asr_logits, + self.user_logit_boosts, + pad_id=self.pad_token_id, + bos_id=self.bos_token_id, + eos_id=self.eos_token_id, + ) + multimodal_outputs["asr_tokens"] = torch.argmax(asr_logits, dim=-1).to(torch.long) + self._last_channel_state["input_asr_ids"] = ( + multimodal_outputs["asr_tokens"][-1:] + .detach() + ) + self._last_asr_token.copy_( + multimodal_outputs["asr_tokens"][-1:] + ) + if self.use_function_head: + function_logits = self.logits_processor(self.function_head, hidden) + multimodal_outputs["function_tokens"] = torch.argmax(function_logits, dim=-1).to(torch.long) + if self._current_is_prompt_prefill: + function_state = torch.full( + (1,), + self.pad_token_id, + device=hidden.device, + dtype=torch.long, + ) + else: + function_state = multimodal_outputs[ + "function_tokens" + ][-1:].detach() + self._last_channel_state[ + "input_function_ids" + ] = function_state + self._last_function_token.copy_(function_state) + + return OmniOutput( + text_hidden_states=hidden, + multimodal_outputs=multimodal_outputs, + ) + + # ------------------------------------------------------------------ # + # weight loading # + # ------------------------------------------------------------------ # + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self, skip_prefixes=["mtp"]) + loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + # Now that the embedding tables are populated, materialize the + # per-prefill pad embedding into the ``_pad_combined_emb`` buffer + # declared in ``__init__``. See the buffer's registration site + # for why this lives here rather than in ``__init__`` (embedding + # tables are empty there) and why it's non-persistent (derived + # from weights that are already saved). + self._materialize_pad_combined_emb() + + return loaded diff --git a/nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/sampling.py b/nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/sampling.py new file mode 100644 index 000000000000..94d90b48c9ff --- /dev/null +++ b/nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/sampling.py @@ -0,0 +1,247 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""VoiceChat text sampling, shared with the PyTorch backend. + +Both backends decode the text head with +:func:`~nemo.collections.speechlm2.inference.model_wrappers.text_sampling.sample_text_token`. +The PyTorch backend calls it directly; vLLM reaches it through the +logits-processor hook implemented here, so the two cannot drift. +""" + +import math +from collections import OrderedDict +from dataclasses import dataclass, field +from typing import Any + +import torch +from vllm import SamplingParams +from vllm.v1.sample.logits_processor import AdapterLogitsProcessor + +from nemo.collections.speechlm2.inference.model_wrappers.text_sampling import sample_text_token +from nemo.collections.speechlm2.parts.logit_boosts import LogitBoosts, apply_logit_boosts + +SHARED_TEXT_SAMPLING_ARG = "nemo_shared_text_sampling" + + +def _sampling_config(params: SamplingParams) -> dict[str, Any] | None: + extra_args = params.extra_args + if not isinstance(extra_args, dict): + return None + value = extra_args.get(SHARED_TEXT_SAMPLING_ARG) + return value if isinstance(value, dict) else None + + +@dataclass +class SharedTextSamplingState: + """Sampling history that survives vLLM streaming segment re-admission.""" + + sample_count: int = 0 + tokens: list[int] = field(default_factory=list) + + +class SharedTextRequestSampler: + """Select with NeMo's sampler, then force vLLM greedy to that token.""" + + def __init__( + self, + *, + top_p: float, + repetition_penalty: float, + temperature: float, + special_token_ids: set[int], + history_skip: int, + state: SharedTextSamplingState | None = None, + boosts: LogitBoosts | None = None, + pad_id: int | None = None, + bos_id: int | None = None, + eos_id: int | None = None, + ) -> None: + self.top_p = top_p + self.repetition_penalty = repetition_penalty + self.temperature = temperature + self.special_token_ids = special_token_ids + self.history_skip = history_skip + self.state = state or SharedTextSamplingState() + self.boosts = boosts or LogitBoosts() + self.pad_id = pad_id + self.bos_id = bos_id + self.eos_id = eos_id + if self.boosts and None in (pad_id, bos_id, eos_id): + raise ValueError("Agent logit boosts require pad_id, bos_id and eos_id") + self._special_ids_tensor = ( + torch.tensor(sorted(special_token_ids), dtype=torch.long) if special_token_ids else None + ) + + def __call__( + self, + output_ids: list[int], + logits: torch.Tensor, + ) -> torch.Tensor: + del output_ids + history = self.state.tokens + generated_tokens = torch.tensor( + history, + device=logits.device, + dtype=torch.long, + ).unsqueeze(0) + if self._special_ids_tensor is not None and self._special_ids_tensor.device != logits.device: + self._special_ids_tensor = self._special_ids_tensor.to(logits.device) + + # Same order as DuplexSTTModel: boost the special tokens, then sample. + apply_logit_boosts( + logits, + self.boosts, + pad_id=self.pad_id, + bos_id=self.bos_id, + eos_id=self.eos_id, + ) + + sampled = sample_text_token( + logits.unsqueeze(0), + generated_tokens, + len(history), + top_p=self.top_p, + repetition_penalty=self.repetition_penalty, + temperature=self.temperature, + special_token_ids=self.special_token_ids, + special_ids_tensor=self._special_ids_tensor, + ) + selected_token = int(sampled[0].item()) + if self.state.sample_count >= self.history_skip: + self.state.tokens.append(selected_token) + self.state.sample_count += 1 + logits.fill_(float("-inf")) + logits[selected_token] = 0.0 + return logits + + +class SharedTextSamplingLogitsProcessor(AdapterLogitsProcessor): + """Batch adapter enabling shared VoiceChat text sampling per request.""" + + def __init__( + self, + vllm_config: Any, + device: torch.device, + is_pin_memory: bool, + ) -> None: + super().__init__(vllm_config, device, is_pin_memory) + max_num_seqs = int(getattr(vllm_config.scheduler_config, "max_num_seqs", 1) or 1) + self._max_history_states = max(1, max_num_seqs) + self._history_states: OrderedDict[str, SharedTextSamplingState] = OrderedDict() + + @classmethod + def validate_params(cls, sampling_params: SamplingParams) -> None: + extra_args = sampling_params.extra_args + if not isinstance(extra_args, dict): + return + raw_config = extra_args.get(SHARED_TEXT_SAMPLING_ARG) + if raw_config is None: + return + if not isinstance(raw_config, dict): + raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG} must be a mapping") + + top_p = raw_config.get("top_p") + temperature = raw_config.get("temperature") + repetition_penalty = raw_config.get("repetition_penalty") + for name, value in ( + ("top_p", top_p), + ("temperature", temperature), + ("repetition_penalty", repetition_penalty), + ): + if isinstance(value, bool) or not isinstance(value, int | float): + raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG}.{name} must be numeric") + if not math.isfinite(float(value)): + raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG}.{name} must be finite") + if not 0.0 < float(top_p) <= 1.0: + raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG}.top_p must be in (0, 1]") + if float(temperature) < 0.0: + raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG}.temperature must be >= 0") + if float(repetition_penalty) <= 0.0: + raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG}.repetition_penalty must be > 0") + + special_ids = raw_config.get("special_token_ids") + if not isinstance(special_ids, list) or any( + isinstance(token_id, bool) or not isinstance(token_id, int) for token_id in special_ids + ): + raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG}.special_token_ids " "must be a list of integers") + history_skip = raw_config.get("history_skip") + if isinstance(history_skip, bool) or not isinstance(history_skip, int) or history_skip < 0: + raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG}.history_skip " "must be a non-negative integer") + history_key = raw_config.get("history_key") + if not isinstance(history_key, str) or not history_key: + raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG}.history_key " "must be a non-empty string") + + boosts = raw_config.get("boosts") + if boosts is None: + return + if not isinstance(boosts, dict): + raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG}.boosts must be a mapping") + for name in ("pad", "bos", "eos"): + value = boosts.get(name) + if value is None: + continue + if isinstance(value, bool) or not isinstance(value, int | float): + raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG}.boosts.{name} must be numeric") + if not math.isfinite(float(value)): + raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG}.boosts.{name} must be finite") + if not any(boosts.get(name) for name in ("pad", "bos", "eos")): + return + for name in ("pad_id", "bos_id", "eos_id"): + token_id = raw_config.get(name) + if isinstance(token_id, bool) or not isinstance(token_id, int) or token_id < 0: + raise ValueError( + f"{SHARED_TEXT_SAMPLING_ARG}.{name} must be a non-negative " "integer when boosts are set" + ) + + def is_argmax_invariant(self) -> bool: + return False + + def new_req_logits_processor( + self, + params: SamplingParams, + ) -> SharedTextRequestSampler | None: + config = _sampling_config(params) + if config is None: + return None + history_key = str(config["history_key"]) + state = self._history_states.get(history_key) + if state is None: + while len(self._history_states) >= self._max_history_states: + self._history_states.popitem(last=False) + state = SharedTextSamplingState() + self._history_states[history_key] = state + else: + self._history_states.move_to_end(history_key) + boosts = LogitBoosts.from_dict(config.get("boosts")) + return SharedTextRequestSampler( + top_p=float(config["top_p"]), + repetition_penalty=float(config["repetition_penalty"]), + temperature=float(config["temperature"]), + special_token_ids=set(config["special_token_ids"]), + history_skip=int(config["history_skip"]), + state=state, + boosts=boosts, + pad_id=config.get("pad_id"), + bos_id=config.get("bos_id"), + eos_id=config.get("eos_id"), + ) + + +__all__ = [ + "SHARED_TEXT_SAMPLING_ARG", + "SharedTextRequestSampler", + "SharedTextSamplingState", + "SharedTextSamplingLogitsProcessor", +] diff --git a/nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/__init__.py b/nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/__init__.py new file mode 100644 index 000000000000..374ca75eab75 --- /dev/null +++ b/nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/__init__.py @@ -0,0 +1,27 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Single-stage ``nemotron_voicechat`` NemotronDuplexH Omni pipeline. + +EarTTS is registered as a separate one-stage pipeline; NeMo coordinates the +two engines. +""" + +from nemo.collections.speechlm2.inference.vllm_omni.nemotron_voicechat.pipeline import ( + NEMOTRON_VOICECHAT_PIPELINE, +) + +__all__ = [ + "NEMOTRON_VOICECHAT_PIPELINE", +] diff --git a/nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/pipeline.py b/nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/pipeline.py new file mode 100644 index 000000000000..0be71e9110cb --- /dev/null +++ b/nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/pipeline.py @@ -0,0 +1,80 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Single-stage ``nemotron_voicechat`` Omni pipeline. + +The VoiceChat wrapper intentionally runs NemotronDuplexH and EarTTS as +independent one-stage engines. This pipeline is the Nemotron half: it +consumes one acoustic encoder embedding per :class:`StreamingInput` update +and emits the text token plus whichever auxiliary channel (ASR or function) +the checkpoint contains. NeMo forwards the text token to the separate +``eartts`` pipeline. + +The pipeline is registered against ``model_type = "nemotron_voicechat"``, +which the component checkpoint does not report natively, so the converted +wrapper directory remains the model root:: + + / + config.json # {"model_type": "nemotron_voicechat"} + nemotron/ # directory or symlink → Nemotron ckpt + eartts/ # directory or symlink → EarTTS ckpt + +Only ``nemotron/`` is loaded by this pipeline. ``eartts/`` is passed +directly to a second :class:`AsyncOmni` instance. The bundled deploy YAML at +``nemo/collections/speechlm2/inference/vllm_omni/deploy/nemotron_voicechat.yaml`` +points this stage at its component via ``model_subdir`` / ``tokenizer_subdir``. +""" + +from __future__ import annotations + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_SCHED_ASYNC = ( + "nemo.collections.speechlm2.inference.vllm_omni." "nemotron_voicechat.scheduler.NemotronVoicechatARAsyncScheduler" +) + + +NEMOTRON_VOICECHAT_PIPELINE = PipelineConfig( + model_type="nemotron_voicechat", + model_arch="NemotronDuplexHForCausalLM", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="nemotron", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=True, + final_output_type="text", + owns_tokenizer=True, + model_arch="NemotronDuplexHForCausalLM", + # Stock vLLM-Omni 0.26 only includes single-stage AR requests in + # the client multimodal pooler payload for the "audio" engine + # output path. The final output remains text, so sampled text + # tokens stay on RequestOutput while OmniOutput.multimodal_outputs + # carries the optional ASR/function token beside it. + engine_output_type="audio", + scheduler_cls=_SCHED_ASYNC, + sampling_constraints={"detokenize": False}, + ), + ), +) + + +__all__ = [ + "NEMOTRON_VOICECHAT_PIPELINE", +] diff --git a/nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/scheduler.py b/nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/scheduler.py new file mode 100644 index 000000000000..8d4dcbf21ced --- /dev/null +++ b/nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/scheduler.py @@ -0,0 +1,68 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Streaming scheduler for the one-stage Nemotron VoiceChat engine. + +The only deviation from vLLM-Omni's async AR scheduler is forwarding each +direct ``StreamingInput`` chunk's ``additional_information`` payload onto the +session. This carries the current acoustic embedding into the model runner +without modifying the installed vLLM-Omni package. +""" + +from __future__ import annotations + +from vllm.v1.request import Request, StreamingUpdate +from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler + + +class NemotronVoicechatSchedulerMixin: + """Forward direct per-chunk payloads on the one-stage session.""" + + def _update_request_as_session(self, session: Request, update: StreamingUpdate) -> None: + super()._update_request_as_session(session, update) + + # Forward the chunk's payload onto the session, which is the courier + # that carries it to ``OmniNewRequestData`` and from there into the + # runner's ``model_intermediate_buffer``. Upstream propagates + # ``model_intermediate_buffer`` itself but not ``additional_information``, + # and this pipeline has to use the latter: the per-chunk payload is a + # tensor (``acoustic_embedding``), and ``model_intermediate_buffer`` is + # typed ``dict[str, Any]`` on the request, so vLLM's msgpack decoder has + # no declared type to rebuild a tensor from and would hand the model a + # ``[dtype, shape, bytes]`` list instead. ``additional_information`` is + # the transport with an explicit tensor encoding, which is why upstream's + # own duplex example keeps ``model_intermediate_buffer`` to plain lists. + # + # Replace rather than merge: this field is a per-chunk message, and + # accumulating whole payloads across chunks would keep stale + # prefill-only keys alive. ``None`` means "this chunk omitted the + # field" rather than "clear the session", so placeholder chunks do not + # drop the initial request's state. The runner does the actual merge + # into the cached buffer, one sub-key at a time. Only stage 0 does this: + # in a downstream stage the chunk transfer adapter is the sole writer of + # the payload, so upstream returns early there. + if self.vllm_config.model_config.stage_id == 0: + new_info = getattr(update, "additional_information", None) + if new_info is not None: + session.additional_information = new_info + + +class NemotronVoicechatARAsyncScheduler(NemotronVoicechatSchedulerMixin, OmniARAsyncScheduler): + """Default: matches upstream's ``async_scheduling=True`` for LLM_AR stages.""" + + +__all__ = [ + "NemotronVoicechatARAsyncScheduler", + "NemotronVoicechatSchedulerMixin", +] diff --git a/nemo/collections/speechlm2/inference/vllm_omni/outputs.py b/nemo/collections/speechlm2/inference/vllm_omni/outputs.py new file mode 100644 index 000000000000..df63105b108f --- /dev/null +++ b/nemo/collections/speechlm2/inference/vllm_omni/outputs.py @@ -0,0 +1,128 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Reading vLLM-Omni stage outputs. + +The shape of an ``OmniStageOutput`` is not part of any contract we control: +the multimodal payload may sit on the stage output or nested on its +completion, and the key it arrives under depends on whether the stage is +client-facing. These readers are the one place that tolerates that, so the +version-sensitivity is contained rather than spread through the session. +""" + +from typing import Any, NamedTuple + +import torch + + +class StepTokens(NamedTuple): + """Per-frame tokens sampled by Nemotron for one acoustic frame. + + ``asr`` and ``function`` are *None* when the checkpoint has no such + channel. ASR and function channels are independently optional. + """ + + text: int + asr: int | None = None + function: int | None = None + + +def _multimodal_output(stage_output: Any, req_out: Any) -> Any: + """Return a stage's multimodal payload, whichever level carries it. + + vLLM-Omni attaches the payload to the ``MultimodalCompletionOutput`` in + ``request_output.outputs[0]`` and also lifts it onto the stage output + itself, so check both (mirroring + ``vllm_omni.metrics.utils.first_multimodal_output``). + + Stock vLLM-Omni 0.26 surfaces EarTTS audio codes here. The registered + Nemotron pipeline also uses the final multimodal engine-output route so + optional ASR/function tensors accompany its text ``RequestOutput``. + """ + mm = getattr(stage_output, "multimodal_output", None) + if mm: + return mm + outputs = getattr(req_out, "outputs", None) or () + for completion in outputs: + nested = getattr(completion, "multimodal_output", None) + if nested: + return nested + return {} + + +def _audio_codes(mm: Any) -> Any: + """Return this step's EarTTS acoustic codes from a multimodal payload. + + ``EarTTSForCausalLM.make_omni_output`` publishes them under + ``model_outputs``, which vLLM-Omni's output processor remaps to the + drainable ``audio`` modality key: in DELTA mode that key is emptied after + every step, so each payload carries only the frames computed this step. + Keys other than the modality's own are retained across steps and merged + with :class:`TensorAccumulationStrategy` ``CONCAT_LAST`` for audio, which + widens a ``T x num_quantizers`` frame instead of appending to it — so + ``audio_codes`` is read last, only for a stage that is not client-facing. + """ + if mm is None: + return None + for key in ("audio", "model_outputs", "audio_codes"): + value = mm.get(key) + if value is not None: + return value + return None + + +def _step_delta(value: Any, finished: bool, *, skip_finished: bool = True): + """Mirror of ``_step_delta`` in the vllm-omni example: pull the + new-this-step multimodal chunk from an :class:`OmniStageOutput`'s + ``multimodal_output`` value (which may be a tensor, a list of tensors, + ``None``, or absent). + + ``skip_finished`` drops a terminal duplicate. The split streaming + requests use one-token segments, so callers pass ``False`` and separately + skip each request's prefill output. + """ + if finished and skip_finished: + return None + if isinstance(value, torch.Tensor): + return value if value.numel() > 0 else None + if isinstance(value, list) and value: + last = value[-1] + return last if isinstance(last, torch.Tensor) and last.numel() > 0 else None + return None + + +def _step_tokens(stage_output: Any) -> StepTokens: + """Extract text and optional auxiliary tokens from one Nemotron output. + + The text token remains on the stock vLLM ``RequestOutput`` even when the + stage uses the multimodal engine-output path. Auxiliary tensors may be + lifted onto ``OmniStageOutput`` or remain nested on its completion. + """ + req_out = stage_output.request_output + mm = _multimodal_output(stage_output, req_out) + finished = bool(getattr(req_out, "finished", False)) + if req_out and req_out.outputs and req_out.outputs[0].token_ids: + text_tok = int(req_out.outputs[0].token_ids[-1]) + else: + text_tok = 0 + + def last_token(key: str) -> int | None: + delta = _step_delta(mm.get(key), finished, skip_finished=False) + return int(delta[-1].item()) if delta is not None else None + + return StepTokens( + text=text_tok, + asr=last_token("asr_tokens"), + function=last_token("function_tokens"), + ) diff --git a/nemo/collections/speechlm2/inference/vllm_omni/register.py b/nemo/collections/speechlm2/inference/vllm_omni/register.py new file mode 100644 index 000000000000..e48eaa644329 --- /dev/null +++ b/nemo/collections/speechlm2/inference/vllm_omni/register.py @@ -0,0 +1,154 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Runtime registration of the NemotronDuplexH + EarTTS models and the +``nemotron_voicechat`` pipeline with vLLM and vLLM-Omni. + +Call :func:`register_nemo_voicechat` once, before constructing +``vllm_omni.AsyncOmni`` / ``vllm_omni.Omni``. It is idempotent. + +NeMo's ``pyproject.toml`` exposes this function as an entry point in **both** +the ``vllm.general_plugins`` and ``vllm_omni.general_plugins`` groups, so a +plugin loader invokes it automatically in every process (orchestrator + each +spawned ``StageEngineCoreProc`` child + workers). vllm-omni uses +``multiprocessing`` with start method ``spawn`` for stage children, so spawned +processes do NOT inherit Python state from the parent — registering in the +parent alone is not enough, which is why a plugin entry point is the correct +hook. vLLM's group is the one that matters in the parent: vllm-omni resolves +the pipeline for ``model_type`` while constructing ``AsyncOmniEngine``, before +it loads its own plugin group. + +Being in vLLM's group means this also runs in processes that have no vllm-omni +installed at all (an ordinary ``vllm serve``), so the vllm-omni half is +optional and skipped when the import fails. + +Three things get registered: + +1. ``EarTTSConfig`` with ``transformers.AutoConfig`` (so ``AutoConfig.from_pretrained`` + resolves ``model_type = "eartts"``) and with vLLM's ``_CONFIG_REGISTRY``. +2. Model architectures (``NemotronDuplexHForCausalLM``, ``EarTTSForCausalLM``) + with both ``vllm.model_executor.models.ModelRegistry`` and + ``vllm_omni.model_executor.models.OmniModelRegistry``. The two registries + serve different lookups inside vLLM-Omni so both have to know about the + new arches. +3. The one-stage :data:`NEMOTRON_VOICECHAT_PIPELINE` and + :data:`EARTTS_PIPELINE` with ``vllm_omni.config.register_pipeline``. +""" + +from __future__ import annotations + +import logging + +logger = logging.getLogger(__name__) + + +_PKG = "nemo.collections.speechlm2.inference.vllm_omni" + + +_ARCH_MAP: dict[str, tuple[str, str]] = { + # arch_name -> (module_path, class_name) + "NemotronDuplexHForCausalLM": ( + f"{_PKG}.nemotron_duplex_h.nemotron_duplex_h", + "NemotronDuplexHForCausalLM", + ), + "EarTTSForCausalLM": ( + f"{_PKG}.eartts.eartts", + "EarTTSForCausalLM", + ), +} + + +_registered = False + + +def register_nemo_voicechat() -> None: + """Register the NemotronDuplexH + EarTTS models, ``EarTTSConfig``, + and the ``nemotron_voicechat`` pipeline with vLLM / vLLM-Omni. + + Safe to call multiple times. + """ + global _registered + if _registered: + return + + _register_hf_configs() + _register_model_archs() + omni = _register_omni() + + _registered = True + logger.info( + "nemo_voicechat: registered NemotronDuplexH + EarTTS%s.", + " + split nemotron_voicechat/eartts pipelines" + if omni + else " (vllm-omni absent, pipelines not registered)", + ) + + +def _register_hf_configs() -> None: + from nemo.collections.speechlm2.inference.vllm_omni.eartts.configuration_eartts import ( + EarTTSConfig, + register_eartts_config, + ) + + register_eartts_config() + + try: + from vllm.transformers_utils.config import _CONFIG_REGISTRY + except ImportError: + _CONFIG_REGISTRY = None + + if _CONFIG_REGISTRY is not None and EarTTSConfig.model_type not in _CONFIG_REGISTRY: + _CONFIG_REGISTRY[EarTTSConfig.model_type] = EarTTSConfig + + +def _register_model_archs() -> None: + # vLLM's public model registry — needed for ``ModelRegistry.is_*`` + # checks and for the arch → module resolution used outside of + # OmniModelConfig. + from vllm.model_executor.models import ModelRegistry + + supported_archs = ModelRegistry.get_supported_archs() + for arch, (module_path, class_name) in _ARCH_MAP.items(): + if arch not in supported_archs: + ModelRegistry.register_model(arch, f"{module_path}:{class_name}") + + +def _register_omni() -> bool: + """Register with vLLM-Omni, if it is installed. Returns whether it was.""" + try: + # OmniModelRegistry is the mirror registry ``OmniModelConfig.registry`` + # returns, used to load the model class for each pipeline stage. + from vllm_omni.config import register_pipeline + from vllm_omni.model_executor.models import OmniModelRegistry + except ImportError: + return False + + omni_supported = OmniModelRegistry.get_supported_archs() + for arch, (module_path, class_name) in _ARCH_MAP.items(): + if arch not in omni_supported: + OmniModelRegistry.register_model(arch, f"{module_path}:{class_name}") + + from nemo.collections.speechlm2.inference.vllm_omni.nemotron_voicechat.pipeline import ( + NEMOTRON_VOICECHAT_PIPELINE, + ) + from nemo.collections.speechlm2.inference.vllm_omni.eartts.pipeline import ( + EARTTS_PIPELINE, + ) + + register_pipeline(NEMOTRON_VOICECHAT_PIPELINE) + register_pipeline(EARTTS_PIPELINE) + return True + + +__all__ = ["register_nemo_voicechat"] diff --git a/nemo/collections/speechlm2/inference/vllm_omni/runtime.py b/nemo/collections/speechlm2/inference/vllm_omni/runtime.py new file mode 100644 index 000000000000..6db1da775e97 --- /dev/null +++ b/nemo/collections/speechlm2/inference/vllm_omni/runtime.py @@ -0,0 +1,322 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Process-scoped vLLM-Omni engines on one background asyncio loop. + +``AsyncOmni`` is asynchronous and its engines are expensive, so exactly one +:class:`OmniRuntime` is built per process and shared by every stream. It owns +a daemon thread running a dedicated event loop, the independent Nemotron and +EarTTS engines, and the deploy-YAML overrides they were started with. + +Nemotron and EarTTS get separate one-stage engines so NeMo can hand tokens +between them and give EarTTS a classifier-free-guidance companion request +without duplicating the much larger Nemotron request. + +Which components exist is decided here, at construction, and read off the +runtime afterwards (``llm_engine``/``tts_engine`` being None) -- callers do not +pass backend flags around. +""" + +import asyncio +import logging as stdlib_logging +import os +import tempfile +import threading +from pathlib import Path +from typing import Any + +import yaml + +from nemo.collections.speechlm2.inference.vllm_omni import default_deploy_yaml, default_eartts_deploy_yaml +from nemo.collections.speechlm2.inference.vllm_omni.checkpoint import EARTTS_SUBDIR +from nemo.utils import logging + + +class _ExpectedJanusShutdownFilter(stdlib_logging.Filter): + """Hide only vLLM-Omni's expected output-queue shutdown traceback.""" + + def filter(self, record: stdlib_logging.LogRecord) -> bool: + if "[AsyncOmni] final_output_loop failed." not in record.getMessage(): + return True + exc = record.exc_info[1] if record.exc_info else None + return not ( + exc is not None + and exc.__class__.__module__.startswith("janus") + and exc.__class__.__name__ + in { + "ShutDown", + "SyncQueueShutDown", + "AsyncQueueShutDown", + } + ) + + +class OmniRuntime: + """Long-lived split AsyncOmni engines + one background asyncio loop. + + Constructed once by the inference wrapper and shared across streams. + Nemotron and EarTTS have independent one-stage engines so NeMo can hand + tokens between them and create an EarTTS CFG companion without duplicating + the much larger Nemotron request. + """ + + def __init__( + self, + wrapper_dir: str, + *, + stage_configs_path: str | None = None, + eartts_stage_configs_path: str | None = None, + stage_overrides: dict | None = None, + eartts_stage_overrides: dict | None = None, + log_stats: bool = False, + stage_init_timeout: int = 600, + enable_llm: bool = True, + enable_tts: bool = True, + ) -> None: + if not enable_llm and not enable_tts: + raise ValueError("OmniRuntime requires at least one enabled component") + self.enable_llm = bool(enable_llm) + self.enable_tts = bool(enable_tts) + + llm_yaml = Path(stage_configs_path) if stage_configs_path else default_deploy_yaml() + tts_yaml = Path(eartts_stage_configs_path) if eartts_stage_configs_path else default_eartts_deploy_yaml() + required_yamls = [] + if self.enable_llm: + required_yamls.append(llm_yaml) + if self.enable_tts: + required_yamls.append(tts_yaml) + for deploy_yaml in required_yamls: + if not deploy_yaml.is_file(): + raise FileNotFoundError(f"Deploy YAML not found: {deploy_yaml}") + + # Accept single-pipeline override keys as well: ``stage_0`` addresses + # the Nemotron engine and ``stage_1`` the EarTTS engine's stage 0. + llm_overrides, legacy_tts_overrides = self._split_stage_overrides(stage_overrides) + if eartts_stage_overrides is None: + eartts_stage_overrides = legacy_tts_overrides + self._llm_stage_yaml_path = ( + self._maybe_write_overridden_yaml(llm_yaml, llm_overrides, prefix="nemotron_") if self.enable_llm else None + ) + self._tts_stage_yaml_path = ( + self._maybe_write_overridden_yaml(tts_yaml, eartts_stage_overrides, prefix="eartts_") + if self.enable_tts + else None + ) + self._wrapper_dir = wrapper_dir + self._eartts_dir = os.path.join(wrapper_dir, EARTTS_SUBDIR) + self._shutdown = False + + # Start the background loop in a daemon thread first; ``AsyncOmni`` + # is constructed *on* that loop (its ``__init__`` allocates + # ``asyncio.Condition`` / ``asyncio.Queue`` and the orchestrator + # binds them to the current event loop, so the engine must be + # built from inside that loop's thread). + self._loop = asyncio.new_event_loop() + self._ready_evt = threading.Event() + self._thread = threading.Thread( + target=self._loop_runner, + name="OmniRuntimeLoop", + daemon=True, + ) + self._thread.start() + self._ready_evt.wait() + + # Register in this process before constructing the engine: + # ``AsyncOmniEngine.__init__`` resolves ``model_type`` before loading + # plugin groups. An unregistered model type selects the default diffusion + # pipeline, which expects ``model_index.json``. Entry points register the + # same pipeline in spawned stage processes. + from vllm_omni import AsyncOmni + + from nemo.collections.speechlm2.inference.vllm_omni.register import register_nemo_voicechat + + register_nemo_voicechat() + + logging.info( + "Creating split AsyncOmni engines from wrapper=%s (Nemotron) and %s (EarTTS) ...", + wrapper_dir, + self._eartts_dir, + ) + + async def _build_engines() -> tuple[Any | None, Any | None]: + llm_engine = None + tts_engine = None + try: + if self.enable_llm: + llm_engine = AsyncOmni( + model=wrapper_dir, + stage_configs_path=str(self._llm_stage_yaml_path), + log_stats=log_stats, + stage_init_timeout=stage_init_timeout, + ) + if self.enable_tts: + tts_engine = AsyncOmni( + model=self._eartts_dir, + stage_configs_path=str(self._tts_stage_yaml_path), + log_stats=log_stats, + stage_init_timeout=stage_init_timeout, + ) + return llm_engine, tts_engine + except BaseException: + if llm_engine is not None: + llm_engine.shutdown() + raise + + fut = asyncio.run_coroutine_threadsafe(_build_engines(), self._loop) + self.llm_engine, self.tts_engine = fut.result() + logging.info( + "Split AsyncOmni ready (Nemotron=%s, EarTTS=%s)", + f"{self.llm_engine.num_stages} stage" if self.llm_engine is not None else "native", + f"{self.tts_engine.num_stages} stage" if self.tts_engine is not None else "native", + ) + + # ------------------------------------------------------------------ # + # YAML override # + # ------------------------------------------------------------------ # + + @staticmethod + def _split_stage_overrides( + stage_overrides: dict | None, + ) -> tuple[dict | None, dict | None]: + if not stage_overrides: + return None, None + common = dict(stage_overrides.get("common", {}) or {}) + llm: dict[str, Any] = {} + tts: dict[str, Any] = {} + if common: + llm["common"] = common + tts["common"] = common + if stage_overrides.get("stage_0"): + llm["stage_0"] = dict(stage_overrides["stage_0"]) + if stage_overrides.get("stage_1"): + tts["stage_0"] = dict(stage_overrides["stage_1"]) + return llm or None, tts or None + + @staticmethod + def _maybe_write_overridden_yaml( + deploy_yaml: Path, + stage_overrides: dict | None, + *, + prefix: str, + ) -> Path: + """Apply per-stage overrides to the deploy YAML, write to a tmp file. + + ``stage_overrides`` shape:: + + { + "common": {}, + "stage_0": {}, + "stage_1": {}, + } + + Returns the path that ``AsyncOmni`` should load; the original YAML + is returned untouched when no overrides are supplied. + """ + if not stage_overrides: + return deploy_yaml + + with open(deploy_yaml, encoding="utf-8") as fh: + cfg = yaml.safe_load(fh) + + common = stage_overrides.get("common", {}) or {} + per_stage = {int(k.split("_", 1)[1]): v for k, v in stage_overrides.items() if k.startswith("stage_") and v} + + for stage in cfg.get("stages", []): + for key, value in common.items(): + stage[key] = value + sid = int(stage.get("stage_id", -1)) + for key, value in per_stage.get(sid, {}).items(): + stage[key] = value + + tmp = tempfile.NamedTemporaryFile( + mode="w", + suffix=".yaml", + prefix=prefix, + delete=False, + ) + yaml.dump(cfg, tmp, default_flow_style=False, sort_keys=False) + tmp.close() + logging.info(f"Wrote overridden stage config to {tmp.name}") + return Path(tmp.name) + + # ------------------------------------------------------------------ # + # Background loop # + # ------------------------------------------------------------------ # + + def _loop_runner(self) -> None: + asyncio.set_event_loop(self._loop) + self._ready_evt.set() + try: + self._loop.run_forever() + finally: + # ``run_forever`` returns when ``loop.stop()`` is called from + # ``shutdown``. Tear down any pending tasks before closing. + try: + pending = asyncio.all_tasks(self._loop) + for task in pending: + task.cancel() + except RuntimeError: + pass + try: + self._loop.close() + except Exception: + pass + + def submit(self, coro): + """Schedule a coroutine on the background loop, return the concurrent ``Future``.""" + return asyncio.run_coroutine_threadsafe(coro, self._loop) + + def shutdown(self) -> None: + """Stop both engines and the background loop.""" + if self._shutdown: + return + self._shutdown = True + + async def _shutdown_engines() -> None: + for name in ("tts_engine", "llm_engine"): + engine = getattr(self, name, None) + if engine is None: + continue + try: + final_output_task = getattr(engine, "final_output_task", None) + if final_output_task is not None and not final_output_task.done(): + final_output_task.cancel() + await asyncio.gather( + final_output_task, + return_exceptions=True, + ) + engine.final_output_task = None + engine.shutdown() + except Exception as exc: + logging.warning(f"{name}.shutdown() raised: {exc!r}") + + shutdown_filter = _ExpectedJanusShutdownFilter() + async_omni_logger = stdlib_logging.getLogger("vllm_omni.entrypoints.async_omni") + async_omni_logger.addFilter(shutdown_filter) + root_handlers = list(stdlib_logging.getLogger().handlers) + for handler in root_handlers: + handler.addFilter(shutdown_filter) + try: + self.submit(_shutdown_engines()).result(timeout=120) + except Exception as exc: + logging.warning(f"Split AsyncOmni shutdown raised: {exc!r}") + finally: + async_omni_logger.removeFilter(shutdown_filter) + for handler in root_handlers: + handler.removeFilter(shutdown_filter) + try: + self._loop.call_soon_threadsafe(self._loop.stop) + except Exception: + pass + self._thread.join(timeout=10) diff --git a/nemo/collections/speechlm2/inference/vllm_omni/scripts/__init__.py b/nemo/collections/speechlm2/inference/vllm_omni/scripts/__init__.py new file mode 100644 index 000000000000..9e3fb699d9f6 --- /dev/null +++ b/nemo/collections/speechlm2/inference/vllm_omni/scripts/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/nemo/collections/speechlm2/inference/vllm_omni/scripts/convert_duplex_eartts_checkpoint.py b/nemo/collections/speechlm2/inference/vllm_omni/scripts/convert_duplex_eartts_checkpoint.py new file mode 100644 index 000000000000..6e7cce134690 --- /dev/null +++ b/nemo/collections/speechlm2/inference/vllm_omni/scripts/convert_duplex_eartts_checkpoint.py @@ -0,0 +1,357 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Convert the DuplexEARTTS component of a NemotronVoiceChat checkpoint to vLLM format. + +The converter expects the HuggingFace-format NemotronVoiceChat checkpoint layout: +``config.json`` contains ``model.speech_generation`` and ``model.stt`` entries, +and ``model.safetensors`` contains nested ``tts_model.tts_model.*`` weights. + +The character-aware subword encoder (``embed_subword``) is collapsed into a +single pre-computed lookup table mapping ``token_id -> hidden_size`` embedding. +The character/transformer weights of that encoder are dropped, since the lookup +fully captures their deterministic per-token output (including the additive +subword-flag and BOS/EOS contributions). +""" + +import argparse +import json +import os +import tqdm + +import torch +from omegaconf import DictConfig, OmegaConf +from safetensors.torch import load_file, save_file +from transformers import AutoConfig + +from nemo.collections.speechlm2.models.duplex_ear_tts import DuplexEARTTS +from nemo.utils import logging + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--config", type=str, required=True) + parser.add_argument("--model", type=str, required=True) + parser.add_argument("--outdir", type=str, required=True) + parser.add_argument( + "--precompute-batch-size", + type=int, + default=256, + help="Batch size for pre-computing per-token embeddings.", + ) + return parser.parse_args() + + +def _precompute_subword_embeddings(model: DuplexEARTTS, batch_size: int) -> torch.Tensor: + """Run ``embed_subword`` over the entire vocabulary to bake out a lookup table. + + The character-aware subword encoder is fully deterministic per token id + (it takes ``subword_ids`` only and adds id-conditioned flag/BOS-EOS + embeddings). Running it once per id and storing the result lets vLLM + replace the whole encoder with a single ``nn.Embedding`` lookup. + + Returns: + Tensor of shape ``[vocab_size, hidden_size]`` matching the dtype of the + encoder's parameters. + """ + embed_subword = model.tts_model.embed_subword + embed_subword.eval() + + # Run the precomputation on GPU when available; the encoder is small but + # the vocabulary loop is long, so this is a meaningful speedup. Only the + # subword encoder is moved (not the full model) to keep peak memory low. + device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") + embed_subword.to(device) + logging.info(f"Precomputing subword embeddings on {device}") + + dtype = next(embed_subword.parameters()).dtype + + subword_ids_map = embed_subword.subword_id_to_char_ids + vocab_size = max(int(k) for k in subword_ids_map.keys()) + 1 + hidden_size = embed_subword.proj_embedding.out_features + + table = torch.zeros((vocab_size, hidden_size), dtype=dtype, device=device) + + with torch.no_grad(): + for start in tqdm.tqdm(range(0, vocab_size, batch_size), desc="Precomputing subword embeddings"): + end = min(start + batch_size, vocab_size) + ids = torch.arange(start, end, dtype=torch.long, device=device).unsqueeze(0) + mask = torch.ones_like(ids, dtype=torch.bool) + embeds = embed_subword(ids, mask) + table[start:end] = embeds.squeeze(0).to(dtype) + + return table.cpu() + + +def convert_to_vllm_format(outdir: str, config: str, model_path: str, precompute_batch_size: int = 256) -> None: + """Convert DuplexEARTTS weights from a NemotronVoiceChat HF checkpoint for vLLM. + + Args: + outdir: Directory where the vLLM-compatible checkpoint will be written. + config: Path to the NemotronVoiceChat ``config.json`` file. + model_path: Path to the NemotronVoiceChat ``model.safetensors`` file. + precompute_batch_size: Batch size used while running the subword encoder + once per token id to construct the lookup table. + """ + os.makedirs(outdir, exist_ok=True) + + with open(config, "r") as f: + full_config = json.load(f) + + # This converter builds a real DuplexEARTTS from the checkpoint's own + # config, so it needs the same config normalization that + # NemotronVoiceChat.from_pretrained applies to the published VoiceChat + # release. Both helpers are no-ops for any other checkpoint. + from nemo.collections.speechlm2.models.nemotron_voicechat import ( + _apply_nemotron_labs_voicechat_release_config_shim, + _is_nemotron_labs_voicechat_release, + ) + + if _is_nemotron_labs_voicechat_release(os.path.dirname(os.path.abspath(config)), full_config): + logging.info("Applying NemotronLabs VoiceChat release shim before building DuplexEARTTS") + _apply_nemotron_labs_voicechat_release_config_shim(full_config) + + config_dict = full_config["model"]["speech_generation"] + cfg = DictConfig(config_dict) + # Inference-only overrides, so the model built here matches the settings + # DuplexEARTTS.generate() uses when the embeddings are precomputed below. + cfg.model.tts_config.use_unshifthed_prompt = True + cfg.data.add_audio_prompt_after_description = True + cfg.model.tts_config.use_unshifthed_prompt = True + cfg.model.subword_mask_exactly_as_eartts = False + cfg.model.context_hidden_mask_exactly_as_eartts = False + cfg.model.tts_config.disable_eos_prediction = True + cfg.model.inference_force_speech_silence_on_eos = True + cfg.model.use_word_sep_tokenizer = False + cfg.model.num_delay_speech_tokens = 0 + cfg.data.source_sample_rate = 22050 + cfg.data.target_sample_rate = 22050 + cfg.model.pretrained_model = None + + model = DuplexEARTTS(OmegaConf.to_container(cfg, resolve=True)).eval() + hidden_size = cfg.model.tts_config.backbone_config.hidden_size + + # Load the HuggingFace-format NemotronVoiceChat safetensors checkpoint. + raw_weights = load_file(model_path) + # The checkpoint is wrapped by an outer module (NemotronVoiceChat) whose TTS + # attribute is also called ``tts_model``. Strip a single ``tts_model.`` prefix + # to land in the DuplexEARTTS state-dict namespace. + weights = {k[len("tts_model.") :]: v for k, v in raw_weights.items() if k.startswith("tts_model.")} + + # Load the real weights into the DuplexEARTTS model so that running + # ``embed_subword`` produces the trained per-token outputs (otherwise we + # would just bake out random init values). + missing, unexpected = model.load_state_dict(weights, strict=False) + # Some keys (e.g. the unused language model / audio codec heads) may be + # missing or unexpected; that is fine for the embedding sub-tree we care + # about. Surface the diagnostics anyway. + if missing: + logging.info(f"load_state_dict missing keys (expected for unused submodules): {len(missing)}") + if unexpected: + logging.info(f"load_state_dict unexpected keys: {len(unexpected)}") + + # Pre-compute the subword lookup table once per token id. This collapses + # the entire char-aware encoder (char embedding + transformer + projection + # + subword/BOS-EOS flag adds) into a single ``nn.Embedding`` lookup that + # vLLM can use directly. + precomputed_subword_emb = _precompute_subword_embeddings(model, precompute_batch_size) + vocab_size, _ = precomputed_subword_emb.shape + + # Codec silence tokens are produced once at training time by encoding a + # zero waveform with the audio codec and picking the most common frame. + # Bake the resulting per-codebook ids into the vLLM checkpoint so the + # runtime does not need to load / run the codec to know what "silence" + # looks like (used e.g. when forcing silence on EOS). + codec_silence_tokens = model.codec_silence_tokens.detach().clone().cpu().to(torch.int32) + + # Strip the ``tts_model.`` prefix so the renaming below operates on + # RVQEARTTSModel state-dict keys. + weights = {k[len("tts_model.") :]: v for k, v in weights.items() if k.startswith("tts_model.")} + + # duplicate weights for rvq embeddings and embed code + rvq_embs_weight = weights["rvq_embs"].clone() # 31 x codebook_size x latent_size + rvq_embs_weight_pad = torch.nn.functional.pad( + rvq_embs_weight, [0, 0, 0, 1] + ) # 31 x (codebook_size + 1) x latent_size + embed_code_weight = weights["embed_code.weight"].clone() # latent_size x hidden_size + + # ====================== + # embedding module weights + bos_emb = weights["bos_emb"] + null_emb = weights["null_emb"] + + embedding_module_weights = {} + embedding_module_weights["bos_emb"] = bos_emb + # CFG-only (unconditional text branch). Always exported so a wrapper + # works with guidance on or off without reconverting. + embedding_module_weights["null_emb"] = null_emb + + # Single pre-computed lookup replacing the entire char-aware encoder. + embedding_module_weights["embed_subword.embed_subwords.weight"] = precomputed_subword_emb + + # Keep gated fusion + audio prompt projection: these depend on runtime + # tensors, not on token id, so they cannot be pre-computed. + for key, weight in weights.items(): + if key.startswith("gated_fusion_audio_text."): + embedding_module_weights[key] = weight + if "audio_prompt_projection_W" in weights: + embedding_module_weights["audio_prompt_projection_W"] = weights["audio_prompt_projection_W"] + + for i in range(rvq_embs_weight_pad.shape[0]): + embedding_module_weights[f"rvq_embs.{i}.weight"] = rvq_embs_weight_pad[i] + embedding_module_weights["embed_code.weight"] = embed_code_weight + embedding_module_weights = {f"total_emb.{k}": v for k, v in embedding_module_weights.items()} + + # ====================== + # gemma backbone weights + backbone_module_weights = {k: v for k, v in weights.items() if k.startswith("backbone.")} + backbone_module_weights["backbone.embed_tokens.weight"] = torch.randn( + 1, hidden_size, dtype=bos_emb.dtype, device=bos_emb.device + ) + + # ====================== + # sampler weights + used_keys = ("rvq_embs", "embed_code", "mog_head") + sampler_weights = {"sampler." + k: v for k, v in weights.items() if k.startswith(used_keys)} + + # combine embedding module and backbone module weights + weights = {**embedding_module_weights, **backbone_module_weights, **sampler_weights} + weights = {"model." + k: v for k, v in weights.items()} + + # Top-level silence token buffer (int32 tensor of shape [num_quantizers]). + # Stored under ``model.sil_tokens`` so the vLLM model can register it as a + # plain buffer at the top of its module tree. + weights["model.sil_tokens"] = codec_silence_tokens + + # save weights + safetensors_path = os.path.join(outdir, "model.safetensors") + save_file(weights, safetensors_path) + logging.info("Saved weights for vllm model") + weight_map = {name: "model.safetensors" for name in weights.keys()} + index = { + "metadata": {"total_size": sum(w.numel() * w.element_size() for w in weights.values())}, + "weight_map": weight_map, + } + index_path = os.path.join(outdir, "model.safetensors.index.json") + with open(index_path, "w") as f: + json.dump(index, f, indent=2) + logging.info("Saved model index") + + # save config.json + flat_config = {"architectures": ["EarTTSForCausalLM"], "model_type": "eartts"} + # not using vocab size of the backbone model, but need 2 for dummy sampling to work + flat_config["vocab_size"] = 2 + + # Parse backbone config exactly as NeMo does to get all defaults from transformers + backbone_type = cfg.model.tts_config.get("backbone_type", None) + backbone_config_dict = ( + OmegaConf.to_container(cfg.model.tts_config.backbone_config, resolve=True) + if cfg.model.tts_config.get("backbone_config") + else {} + ) + + # Create AutoConfig the same way NeMo does - this fills in all defaults + parsed_backbone_config = AutoConfig.for_model(backbone_type, **backbone_config_dict) + + # Store the backbone type for vllm to use + flat_config["backbone_type"] = backbone_type + + # Forward all backbone configs from the parsed AutoConfig (includes defaults) + for key in [ + "hidden_size", + "intermediate_size", + "num_hidden_layers", + "num_attention_heads", + "num_key_value_heads", + "head_dim", + "max_position_embeddings", + "rope_theta", + "rope_local_base_freq", + "sliding_window", + "layer_types", + ]: + if hasattr(parsed_backbone_config, key): + value = getattr(parsed_backbone_config, key) + # convert to list if it's a tuple or other iterable (except str) + if hasattr(value, '__iter__') and not isinstance(value, (str, dict)): + value = list(value) + flat_config[key] = value + # forward overall configs + for key in ["latent_size", "codebook_size", "num_quantizers", "exponent"]: + flat_config[key] = cfg.model.tts_config[key] + # forward mog head configs + for key in ["num_layers", "low_rank", "num_predictions", "min_log_std", "eps"]: + flat_config[f"mog_{key}"] = cfg.model.tts_config.mog_head_config[key] + + # forward inference configs (with name mapping for vLLM model) + # MaskGIT unmasking iterations, fixed at 8 to match DuplexEARTTS inference. + flat_config["num_iter"] = 8 + flat_config["noise_scale"] = cfg.model.get("inference_noise_scale", 0.8) + flat_config["top_p_or_k"] = cfg.model.get("inference_top_p_or_k", 0.8) + + # Classifier-free guidance. ``guidance_scale`` is the checkpoint default; + # individual requests may override it through ``additional_information["cfg_scale"]``. + flat_config["guidance_scale"] = cfg.model.get("inference_guidance_scale", 0.5) + flat_config["enable_guidance"] = True + + # Text-channel specials from the source tokenizer, not the Gemma backbone. + flat_config["pad_token_id"] = int(model.text_pad_id) + flat_config["eos_token_id"] = int(model.text_eos_id) + + # Embedding module configuration. The char-aware encoder is gone; vLLM only + # needs to know the size of the pre-computed lookup table. + flat_config["emb_vocab_size"] = vocab_size + + flat_config["use_gated_fusion_for_text_audio"] = cfg.model.tts_config.use_gated_fusion_for_text_audio + flat_config["use_audio_prompt_frozen_projection"] = cfg.model.tts_config.use_audio_prompt_frozen_projection + + # configuring custom inputs/outputs + flat_config["custom_input_specs"] = [ + { + "name": "acoustic_tokens", + "dim": flat_config["num_quantizers"], + "dtype": "int32", + }, + {"name": "text_tokens", "dtype": "int32"}, + {"name": "text_mask"}, + {"name": "bos_mask"}, + {"name": "speaker_latent", "dim": flat_config["hidden_size"]}, + ] + flat_config["custom_outputs"] = ["acoustic_tokens"] + + with open(os.path.join(outdir, "config.json"), "w") as f: + json.dump(flat_config, f, indent=2) + logging.info("Saved vllm config") + + # Extract and save pre-computed speaker latents (audio_prompt_latents.*) + # from the NeMo checkpoint so they can be used at inference time. + speaker_latents_dir = os.path.join(outdir, "speaker_latents") + found_latents = False + for key, tensor in raw_weights.items(): + if "audio_prompt_latents." in key: + speaker_name = key.split("audio_prompt_latents.")[-1] + os.makedirs(speaker_latents_dir, exist_ok=True) + latent_path = os.path.join(speaker_latents_dir, f"{speaker_name}.pt") + torch.save(tensor, latent_path) + logging.info(f"Saved speaker latent '{speaker_name}' to {latent_path} (shape={tensor.shape})") + found_latents = True + if not found_latents: + logging.warning( + "No audio_prompt_latents found in checkpoint. " "speaker_name will not work unless latents are added." + ) + + +if __name__ == "__main__": + args = parse_args() + convert_to_vllm_format(args.outdir, args.config, args.model, args.precompute_batch_size) diff --git a/nemo/collections/speechlm2/inference/vllm_omni/scripts/convert_duplex_stt_checkpoint.py b/nemo/collections/speechlm2/inference/vllm_omni/scripts/convert_duplex_stt_checkpoint.py new file mode 100644 index 000000000000..07e2bb2769ac --- /dev/null +++ b/nemo/collections/speechlm2/inference/vllm_omni/scripts/convert_duplex_stt_checkpoint.py @@ -0,0 +1,343 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Convert the DuplexSTT component of a NemotronVoiceChat checkpoint to vLLM format. + +This script extracts weights from a HuggingFace-format NemotronVoiceChat +checkpoint with tensors such as: +- stt_model.llm.layers.* +- stt_model.lm_head.* +- stt_model.asr_head.* +- stt_model.embed_asr_tokens.* +- stt_model.function_head.* +- stt_model.embed_tokens.* + +And converts them to a HuggingFace layout that can be loaded by vLLM with the +custom WeightsMapper defined in nemotron_duplex_h.py. + +Which auxiliary channels a checkpoint carries varies: + +- ``predict_user_text=True`` gives ``asr_head`` + ``embed_asr_tokens`` +- ``use_function_head=True`` gives ``function_head`` (reusing ``embed_tokens`` + for its feedback, so it has no embedding table of its own) + +The converter records whichever heads are present as +``use_asr_head`` / ``use_function_head``, because ``NemotronDuplexHForCausalLM`` +has to decide which modules to build *before* it sees any weights. +""" + +import argparse +import json +import os +from pathlib import Path +import torch +from safetensors.torch import load_file, save_file +from transformers import AutoConfig, AutoTokenizer +from nemo.utils import logging + + +def load_checkpoint(checkpoint_path: str) -> dict[str, torch.Tensor]: + """ + Load a NemotronVoiceChat checkpoint state dict. + + Args: + checkpoint_path: Path to a checkpoint directory, safetensors file, or PyTorch checkpoint file. + + Returns: + Dictionary of tensor names to tensors + """ + if os.path.isdir(checkpoint_path): + checkpoint_path = os.path.join(checkpoint_path, "model.safetensors") + + if checkpoint_path.endswith('.safetensors'): + logging.info(f"Loading safetensors from {checkpoint_path}") + return load_file(checkpoint_path) + else: + logging.info(f"Loading PyTorch checkpoint from {checkpoint_path}") + ckpt = torch.load(checkpoint_path, map_location='cpu') + # Handle different checkpoint formats + if 'state_dict' in ckpt: + return ckpt['state_dict'] + elif 'model' in ckpt: + return ckpt['model'] + else: + return ckpt + + +def filter_tensors(state_dict: dict[str, torch.Tensor], prefixes_to_keep: list[str]) -> dict[str, torch.Tensor]: + """ + Filter tensors to keep only those with specified prefixes. + + Args: + state_dict: Full state dictionary + prefixes_to_keep: List of prefixes to keep (e.g., ["stt_model.llm", "stt_model.asr_head"]) + + Returns: + Filtered state dictionary + """ + filtered_dict = {} + for name, tensor in state_dict.items(): + if any(name.startswith(prefix) for prefix in prefixes_to_keep): + filtered_dict[name] = tensor + logging.debug(f"Keeping: {name} with shape {tensor.shape}") + else: + logging.debug(f"Skipping: {name}") + + logging.info(f"Total tensors kept: {len(filtered_dict)}") + return filtered_dict + + +def _apply_source_special_tokens(base_config, tokenizer, source_config: dict | None) -> None: + """Match the converted tokenizer/config to the VoiceChat channel tokens. + + VoiceChat training overrides the LLM-backbone tokenizer specials + (typically ```` for EOS and ```` for PAD). Keeping the + backbone's original EOS/PAD ids corrupts system-prompt prefill even + though text tokenization itself appears valid. Copy whatever the source + VoiceChat config actually used. + """ + try: + model_config = source_config["model"]["stt"]["model"] + except (KeyError, TypeError): + return + + overrides = model_config.get("override_tokens", {}) or {} + special_tokens = { + name: overrides.get(name) or model_config.get(name) + for name in ("bos_token", "eos_token", "pad_token") + } + special_tokens = {name: token for name, token in special_tokens.items() if token} + if not special_tokens: + return + + vocabulary = tokenizer.get_vocab() + missing = [token for token in special_tokens.values() if token not in vocabulary] + if missing: + raise ValueError( + "VoiceChat special tokens must already exist in the backbone vocabulary; " + f"missing={missing}" + ) + added = tokenizer.add_special_tokens(special_tokens) + if added: + raise ValueError( + "VoiceChat special-token overrides unexpectedly expanded the vocabulary; " + f"added={added}" + ) + + for name, token in special_tokens.items(): + token_id = int(tokenizer.convert_tokens_to_ids(token)) + setattr(base_config, f"{name}_id", token_id) + logging.info( + "VoiceChat special-token IDs: bos=%s eos=%s pad=%s", + getattr(base_config, "bos_token_id", None), + getattr(base_config, "eos_token_id", None), + getattr(base_config, "pad_token_id", None), + ) + + +def convert_to_vllm_format( + checkpoint_path: str, + output_dir: str, + config_path: str | None = None, + pretrained_llm: str | None = None, + tensors_to_keep: list[str] | None = None, + dtype: str = "float32", +) -> None: + """ + Convert the DuplexSTT component to vLLM-compatible HuggingFace format. + + Args: + checkpoint_path: Path to the NeMo checkpoint (.safetensors or .pt) + output_dir: Directory to save the converted checkpoint + config_path: Path to config.json (if None, will look in same dir as checkpoint) + pretrained_llm: HuggingFace model name to get base config from + tensors_to_keep: List of tensor prefixes to keep (default: all stt_model.* tensors) + dtype: Data type for tensors ("float32", "float16", "bfloat16") + """ + # Default prefixes to keep. The auxiliary-channel entries are only present + # in some checkpoints; absent ones simply match nothing. + if tensors_to_keep is None: + tensors_to_keep = [ + "stt_model.llm", + "stt_model.lm_head", + "stt_model.asr_head", + "stt_model.embed_asr_tokens", + "stt_model.function_head", + "stt_model.embed_tokens", + ] + + # Load config to get pretrained_llm if not provided + if config_path is None: + ckpt_dir = checkpoint_path if os.path.isdir(checkpoint_path) else os.path.dirname(checkpoint_path) + config_path = os.path.join(ckpt_dir, "config.json") + + config = None + if os.path.exists(config_path): + logging.info(f"Loading config from {config_path}") + with open(config_path, "r") as f: + config = json.load(f) + + try: + pretrained_llm = config["model"]["stt"]["model"]["pretrained_llm"] + logging.info(f"Found pretrained_llm in config: {pretrained_llm}") + except KeyError: + if pretrained_llm is None: + raise ValueError("Could not find pretrained_llm in config and none provided via argument") + else: + if pretrained_llm is None: + raise ValueError(f"Config file not found at {config_path} and pretrained_llm not provided") + + # Create output directory + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + + # Load base config from pretrained model + logging.info(f"Loading base config from {pretrained_llm}") + base_config = AutoConfig.from_pretrained(pretrained_llm, trust_remote_code=True) + + # Load tokenizer from pretrained model + logging.info(f"Loading tokenizer from {pretrained_llm}") + tokenizer = AutoTokenizer.from_pretrained(pretrained_llm, trust_remote_code=True) + _apply_source_special_tokens(base_config, tokenizer, config) + + # Load checkpoint + logging.info(f"Loading checkpoint from {checkpoint_path}") + state_dict = load_checkpoint(checkpoint_path) + + # Filter tensors + logging.info(f"Filtering tensors to keep prefixes: {tensors_to_keep}") + filtered_state_dict = filter_tensors(state_dict, tensors_to_keep) + + if len(filtered_state_dict) == 0: + raise ValueError( + f"No tensors found with prefixes {tensors_to_keep}. " + f"Available prefixes: {set(k.split('.')[0] for k in state_dict.keys())}" + ) + + # Record which auxiliary channels this checkpoint actually carries, so the + # vLLM model builds exactly those modules. Detected from the weights that + # made it through the filter rather than from the source config, so a + # narrowed --tensors-to-keep stays consistent with what gets saved. + has_asr_head = any(name.startswith("stt_model.asr_head") for name in filtered_state_dict) + has_function_head = any(name.startswith("stt_model.function_head") for name in filtered_state_dict) + + if has_asr_head and not any(name.startswith("stt_model.embed_asr_tokens") for name in filtered_state_dict): + raise ValueError( + "Checkpoint has stt_model.asr_head but no stt_model.embed_asr_tokens; " + "the ASR channel needs both (the head to predict the token and the " + "embedding table to feed it back on the next step)." + ) + + # The function channel scales its feedback embedding by this weight, matching + # DuplexSTTModel.build_input_embedding. + function_channel_weight = 1.0 + if has_function_head and config is not None: + try: + function_channel_weight = float(config["model"]["stt"]["model"].get("duplex_function_channel_weight", 1.0)) + except (KeyError, TypeError): + logging.warning("Could not read duplex_function_channel_weight from source config; defaulting to 1.0") + + custom_outputs = ["text_logits"] + if has_asr_head: + custom_outputs += ["asr_tokens", "asr_logits"] + if has_function_head: + custom_outputs += ["function_tokens", "function_logits"] + + base_config.update( + { + "custom_input_specs": [{"name": "combined_embeds", "dtype": dtype, "dim": base_config.hidden_size}], + "custom_outputs": custom_outputs, + "use_asr_head": has_asr_head, + "use_function_head": has_function_head, + "duplex_function_channel_weight": function_channel_weight, + } + ) + logging.info( + f"Auxiliary channels: asr_head={has_asr_head}, function_head={has_function_head} " + f"(function_channel_weight={function_channel_weight})" + ) + + # Save tensors + output_model_path = output_path / "model.safetensors" + logging.info(f"Saving tensors to {output_model_path}") + save_file(filtered_state_dict, str(output_model_path)) + + # Save config + output_config_path = output_path / "config.json" + logging.info(f"Saving config to {output_config_path}") + base_config.save_pretrained(str(output_path)) + + # Save tokenizer + logging.info(f"Saving tokenizer to {output_path}") + tokenizer.save_pretrained(str(output_path)) + + logging.info(f"Conversion completed successfully! Output saved to: {output_path}") + + +def main(): + parser = argparse.ArgumentParser(description="Convert NeMo STT checkpoint to HuggingFace format for vLLM") + parser.add_argument( + "--checkpoint", + type=str, + required=True, + help="Path to NeMo checkpoint file (.safetensors or .pt/.pth)", + ) + parser.add_argument( + "--output-dir", + type=str, + required=True, + help="Directory to save converted checkpoint", + ) + parser.add_argument( + "--config", + type=str, + default=None, + help="Path to config.json (default: same directory as checkpoint)", + ) + parser.add_argument( + "--pretrained-llm", + type=str, + default=None, + help="HuggingFace model name to use as base (default: read from config)", + ) + parser.add_argument( + "--tensors-to-keep", + type=str, + nargs="+", + default=None, + help="Tensor prefixes to keep (default: all stt_model.* backbone llm related tensors)", + ) + parser.add_argument( + "--dtype", + type=str, + default="float32", + choices=["float32", "float16", "bfloat16", "fp32", "fp16", "bf16"], + help="Target dtype for tensors (default: float32)", + ) + + args = parser.parse_args() + + convert_to_vllm_format( + checkpoint_path=args.checkpoint, + output_dir=args.output_dir, + config_path=args.config, + pretrained_llm=args.pretrained_llm, + tensors_to_keep=args.tensors_to_keep, + dtype=args.dtype, + ) + + +if __name__ == "__main__": + main() diff --git a/nemo/collections/speechlm2/inference/vllm_omni/session.py b/nemo/collections/speechlm2/inference/vllm_omni/session.py new file mode 100644 index 000000000000..e30c3689a755 --- /dev/null +++ b/nemo/collections/speechlm2/inference/vllm_omni/session.py @@ -0,0 +1,592 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Synchronous per-stream bridge onto the vLLM-Omni engines. + +The S2S wrapper drives inference from a synchronous PyTorch loop (perception +-> per-frame text + ASR -> audio codec decode), while ``AsyncOmni.generate`` +is an ``async for``. Spawning an event loop per chunk would re-pay the +request-init cost every 80 ms and break vllm-omni's session semantics, so a +:class:`OmniStreamingSession` runs consumer tasks on the runtime's shared loop +and hands frames back across two synchronous queues. + +Per-step protocol: + +1. **Prefill** -- Nemotron receives the system prompt, EarTTS the speaker + latent. With CFG enabled EarTTS receives a conditional and an unconditional + prefill, each with its own KV cache. Nemotron's first token ``t_0`` is fed + back internally rather than exposed. +2. **Decode step k** (``k >= 1``) -- ``prompt_token_ids = [t_{k-1}]``, + ``additional_information.acoustic_embedding = ac_emb[k-1]``, producing + ``t_k`` plus whichever auxiliary channel the checkpoint carries. + +The components are stepped independently: :meth:`OmniStreamingSession.step_llm` +returns Nemotron's tokens and :meth:`OmniStreamingSession.step_tts` submits a +text token, so the caller can rewrite that token (forced turn-taking) in +between and both TTS backends see the same value. + +The synchronous side is single-threaded: only one step may be in flight at a +time. :meth:`OmniStreamingSession.finish` closes the request cleanly and +:meth:`OmniStreamingSession.abort` drops it. +""" + +import asyncio +import threading +import time +import uuid +from queue import Queue +from typing import Any + +import torch + +from nemo.collections.speechlm2.inference.vllm_omni.outputs import ( + StepTokens, + _audio_codes, + _multimodal_output, + _step_delta, + _step_tokens, +) +from nemo.collections.speechlm2.inference.vllm_omni.runtime import OmniRuntime +from nemo.collections.speechlm2.parts.logit_boosts import LogitBoosts +from nemo.utils import logging + + +class _Sentinel: + """Marker placed on the sync output queues to signal end-of-stream or an error.""" + + __slots__ = ("exc",) + + def __init__(self, exc: BaseException | None = None): + self.exc = exc + + +_END_OF_STREAM = _Sentinel() + + +class OmniStreamingSession: + """One split Nemotron/EarTTS streaming request. + + The two components are driven independently: :meth:`step_llm` submits an + acoustic frame and returns Nemotron's tokens, :meth:`step_tts` submits a + text token and produces one acoustic frame. A session owning both exposes + both, which lets the caller act on the text token in between rather than + having EarTTS consume it inside the session. With CFG enabled, conditional + and unconditional EarTTS requests run in the same engine and the custom + scheduler keeps their independent KV-cache streams in lockstep. + """ + + def __init__( + self, + runtime: OmniRuntime, + request_id: str, + system_prompt: str = "", + speaker_latent: torch.Tensor | None = None, + t_prefill: int = 0, + *, + sampling_params: dict | None = None, + special_token_ids: set[int] | None = None, + guidance_enabled: bool = True, + guidance_scale: float = 0.5, + step_timeout: float = 60.0, + profile: bool = False, + agent_logit_boosts: "LogitBoosts | None" = None, + text_token_ids: dict[str, int] | None = None, + ) -> None: + self._has_llm = runtime.llm_engine is not None + self._has_tts = runtime.tts_engine is not None + if not self._has_llm and not self._has_tts: + raise ValueError("OmniStreamingSession requires at least one vLLM component") + if self._has_tts and (speaker_latent is None or speaker_latent.numel() == 0): + raise ValueError("speaker_latent is required for OmniStreamingSession") + if self._has_llm and t_prefill <= 0: + raise ValueError(f"t_prefill must be > 0 (got {t_prefill})") + + self._runtime = runtime + self.request_id = request_id + self._llm_request_id = f"{request_id}:nemotron" + self._sampling_history_key = f"{self._llm_request_id}:{uuid.uuid4().hex}" + self._cfg_pair_id = f"{request_id}:eartts" + self._tts_cond_request_id = f"{self._cfg_pair_id}:cond" + self._tts_uncond_request_id = f"{self._cfg_pair_id}:uncond" + self._step_timeout = step_timeout + self._system_prompt = system_prompt + self._speaker_latent = speaker_latent.detach().cpu().contiguous() if speaker_latent is not None else None + self._t_prefill = int(t_prefill) + self._sampling_overrides = dict(sampling_params or {}) + self._special_token_ids = tuple(sorted(special_token_ids or ())) + self._guidance_enabled = bool(guidance_enabled) + self._guidance_scale = float(guidance_scale) + self._agent_logit_boosts = agent_logit_boosts or LogitBoosts() + self._text_token_ids = dict(text_token_ids or {}) + if self._has_llm and self._agent_logit_boosts: + missing = {"pad_id", "bos_id", "eos_id"} - set(self._text_token_ids) + if missing: + raise ValueError(f"agent_logit_boosts requires text_token_ids {sorted(missing)}") + + # Separate completion queues per component: a session may have both, + # and step_llm()/step_tts() must not consume each other's items. + self._text_out_q: "Queue[StepTokens | _Sentinel]" = Queue() + self._tts_done_q: "Queue[StepTokens | _Sentinel]" = Queue() + self._audio_buf_lock = threading.Lock() + self._audio_buf: list[torch.Tensor] = [] + self._closed = False + self._error: BaseException | None = None + self._loop = runtime._loop + self._queues_ready = threading.Event() + + self._input_q: asyncio.Queue | None = None + self._llm_internal_q: asyncio.Queue | None = None + self._tts_cond_input_q: asyncio.Queue | None = None + self._tts_uncond_input_q: asyncio.Queue | None = None + self._pending_tokens_q: asyncio.Queue | None = None + self._uncond_audio_q: asyncio.Queue | None = None + + self._prof: dict[str, list[float]] | None = {} if profile else None + self._t_put = 0.0 + self._t_yield = 0.0 + self._t_llm = 0.0 + self._t_done = 0.0 + + self._consumer_future = runtime.submit(self._run_consumer()) + if not self._queues_ready.wait(timeout=30): + self._consumer_future.cancel() + raise TimeoutError("Timed out creating split vLLM-Omni session queues") + + def _rec(self, name: str, dt_s: float) -> None: + if self._prof is not None: + self._prof.setdefault(name, []).append(dt_s * 1000.0) + + def _rec_ms(self, name: str, dt_ms: float) -> None: + if self._prof is not None: + self._prof.setdefault(name, []).append(float(dt_ms)) + + def log_timing_summary(self) -> None: + if not self._prof: + return + parts = [] + for name, values in self._prof.items(): + mean = sum(values) / len(values) + parts.append( + f"{name}: mean={mean:.1f}ms min={min(values):.1f}ms " f"max={max(values):.1f}ms n={len(values)}" + ) + logging.info(f"OmniStreamingSession {self.request_id} per-frame timing:\n " + "\n ".join(parts)) + + def _cfg_payload(self, role: str) -> dict[str, Any]: + return { + "cfg_enabled": self._guidance_enabled, + "cfg_role": role, + "cfg_pair_id": self._cfg_pair_id, + "cfg_scale": self._guidance_scale, + } + + def _record_stage_metrics(self, prefix: str, stage_output: Any) -> None: + if self._prof is None: + return + for key, value in (getattr(stage_output, "stage_durations", None) or {}).items(): + self._rec_ms(f"{prefix}.{key}", value) + + async def _run_consumer(self) -> None: + tasks: list[asyncio.Task] = [] + try: + self._input_q = asyncio.Queue() + self._llm_internal_q = asyncio.Queue() + self._tts_cond_input_q = asyncio.Queue() + self._tts_uncond_input_q = asyncio.Queue() + self._pending_tokens_q = asyncio.Queue() + self._uncond_audio_q = asyncio.Queue() + self._queues_ready.set() + + if self._has_llm: + tasks.append(asyncio.create_task(self._consume_llm())) + if self._has_tts: + tasks.append(asyncio.create_task(self._consume_tts("cond", self._tts_cond_input_q))) + if self._has_tts and self._guidance_enabled: + tasks.append(asyncio.create_task(self._consume_tts("uncond", self._tts_uncond_input_q))) + await asyncio.gather(*tasks) + except asyncio.CancelledError: + raise + except BaseException as exc: + self._error = exc + # Both queues, so a caller blocked in either step never hangs. + self._text_out_q.put(_Sentinel(exc)) + self._tts_done_q.put(_Sentinel(exc)) + raise + finally: + for task in tasks: + if not task.done(): + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + self._queues_ready.set() + self._text_out_q.put(_END_OF_STREAM) + self._tts_done_q.put(_END_OF_STREAM) + self.log_timing_summary() + + async def _consume_llm(self) -> None: + from vllm import SamplingParams + from vllm.engine.protocol import StreamingInput + from vllm.sampling_params import RequestOutputKind + + from nemo.collections.speechlm2.inference.vllm_omni.nemotron_duplex_h.sampling import SHARED_TEXT_SAMPLING_ARG + + shared_sampling = { + "temperature": float(self._sampling_overrides.get("temperature", 1.0)), + "top_p": float(self._sampling_overrides.get("top_p", 1.0)), + "repetition_penalty": float(self._sampling_overrides.get("repetition_penalty", 1.0)), + "special_token_ids": list(self._special_token_ids), + # The first vLLM output is the internal prefill token t_0. The + # repetition history starts with the first client-visible frame. + "history_skip": 1, + "history_key": self._sampling_history_key, + # Agent-channel boosts, applied before sampling exactly as + # DuplexSTTModel does. The user-channel ones cannot travel this way + # because the ASR head's logits never reach vLLM's sampler; the + # model applies those itself. + "boosts": self._agent_logit_boosts.as_dict(), + **self._text_token_ids, + } + params = SamplingParams( + temperature=0.0, + top_p=1.0, + repetition_penalty=1.0, + max_tokens=1, + detokenize=False, + ignore_eos=True, + output_kind=RequestOutputKind.DELTA, + extra_args={SHARED_TEXT_SAMPLING_ARG: shared_sampling}, + ) + + async def inputs(): + yield StreamingInput( + prompt={ + "prompt_token_ids": [0] * self._t_prefill, + "additional_information": { + "system_prompt": self._system_prompt, + }, + }, + sampling_params=params, + ) + while True: + submission = await self._input_q.get() + if submission is None: + return + acoustic, committed_text = submission + # Always drain the internal queue to stay one output per input, + # then let the caller's committed token win. That is how the + # caller's forced-turn-taking rewrite reaches Nemotron's own + # history, matching what native feedback does through gen_text. + prev_tokens = await self._llm_internal_q.get() + if committed_text is not None: + prev_tokens = prev_tokens._replace(text=int(committed_text)) + additional_information: dict[str, Any] = { + "system_prompt": None, + "acoustic_embedding": acoustic, + } + if prev_tokens.asr is not None: + additional_information["input_asr_ids"] = torch.tensor([prev_tokens.asr], dtype=torch.long) + if prev_tokens.function is not None: + additional_information["input_function_ids"] = torch.tensor( + [prev_tokens.function], dtype=torch.long + ) + self._t_yield = time.perf_counter() + yield StreamingInput( + prompt={ + "prompt_token_ids": [int(prev_tokens.text)], + "additional_information": additional_information, + }, + sampling_params=params, + ) + + output_count = 0 + try: + async for stage_output in self._runtime.llm_engine.generate( + inputs(), + sampling_params_list=[params], + request_id=self._llm_request_id, + ): + now = time.perf_counter() + self._record_stage_metrics("llm", stage_output) + current_tokens = _step_tokens(stage_output) + await self._llm_internal_q.put(current_tokens) + + output_count += 1 + if output_count <= 1: + continue + + self._rec("pull", self._t_yield - self._t_put) + self._rec("llm_engine", now - self._t_yield) + self._t_llm = now + # The token is returned to the caller, never forwarded to + # EarTTS from here. The caller owns what happens in between + # (forced turn-taking rewrites the text token) and submits it + # with step_tts, so both TTS backends see the same token. + self._t_done = now + self._text_out_q.put(current_tokens) + finally: + # Safety net rather than the normal path: finish() closes the TTS + # inputs itself. This covers Nemotron ending first, which would + # otherwise leave the EarTTS consumer waiting and stall the + # gather() in _run_consumer. + if self._has_tts: + await self._tts_cond_input_q.put(None) + if self._guidance_enabled: + await self._tts_uncond_input_q.put(None) + + async def _consume_tts(self, role: str, input_q: asyncio.Queue) -> None: + from vllm import SamplingParams + from vllm.engine.protocol import StreamingInput + from vllm.sampling_params import RequestOutputKind + + extra_args = self._cfg_payload(role) if self._guidance_enabled else {} + params = SamplingParams( + temperature=0.0, + top_p=1.0, + max_tokens=1, + detokenize=False, + ignore_eos=True, + output_kind=RequestOutputKind.DELTA, + extra_args=extra_args, + ) + request_id = self._tts_cond_request_id if role == "cond" else self._tts_uncond_request_id + + async def inputs(): + prefill_info = { + "embed": {"voice": self._speaker_latent.clone()}, + **self._cfg_payload(role), + } + if not self._guidance_enabled: + prefill_info["cfg_enabled"] = False + yield StreamingInput( + prompt={ + "prompt_token_ids": [0] * int(self._speaker_latent.shape[0]), + "additional_information": prefill_info, + }, + sampling_params=params, + ) + while True: + text_tok = await input_q.get() + if text_tok is None: + return + yield StreamingInput( + prompt={ + "prompt_token_ids": [0], + "additional_information": { + "ids": {"output": [int(text_tok)]}, + **self._cfg_payload(role), + }, + }, + sampling_params=params, + ) + + output_count = 0 + async for stage_output in self._runtime.tts_engine.generate( + inputs(), + sampling_params_list=[params], + request_id=request_id, + ): + now = time.perf_counter() + req_out = stage_output.request_output + mm = _multimodal_output(stage_output, req_out) + finished = bool(getattr(req_out, "finished", False)) + self._record_stage_metrics(f"tts_{role}", stage_output) + output_count += 1 + if output_count <= 1: + continue + audio = _step_delta(_audio_codes(mm), finished, skip_finished=False) + if audio is None or audio.ndim != 2 or audio.shape[0] < 1: + continue + # EarTTS emits exactly one acoustic frame per streaming update, so + # keep only the newest row: a prefill update covers the whole + # speaker latent, and a non-drained key would arrive cumulative. + audio = audio[-1:] + audio = audio.detach().cpu().to(torch.long) + if role == "uncond": + await self._uncond_audio_q.put(audio) + continue + + # A final empty update can race after all text tokens have been + # consumed. It has no corresponding wrapper step and must not + # synthesize another frame. + if self._pending_tokens_q.empty(): + continue + tokens = await self._pending_tokens_q.get() + if self._guidance_enabled: + uncond_audio = await self._uncond_audio_q.get() + if not torch.equal(audio, uncond_audio): + raise RuntimeError( + "EarTTS CFG pair produced divergent client-visible " + "codes: " + f"cond_shape={tuple(audio.shape)} " + f"uncond_shape={tuple(uncond_audio.shape)} " + f"cond={audio.tolist()} " + f"uncond={uncond_audio.tolist()}" + ) + with self._audio_buf_lock: + self._audio_buf.append(audio) + if self._has_llm: + self._rec("tts_after_llm", now - self._t_llm) + else: + self._rec("tts_engine", now - self._t_put) + self._t_done = now + self._tts_done_q.put(tokens) + + def step_llm( + self, + acoustic_embedding: torch.Tensor, + *, + prev_text_token: int | None = None, + ) -> StepTokens: + """Submit one acoustic frame to Nemotron and return its tokens. + + Returns as soon as Nemotron has produced the frame's tokens. EarTTS is + not driven from here even when this session owns both components: the + caller submits the (possibly rewritten) text token with + :meth:`step_tts`. + + Args: + acoustic_embedding: This frame's encoded audio. + prev_text_token: Text token to feed back as the previous step's + output, letting a caller that rewrote it (forced turn-taking) + keep Nemotron's history consistent with its own. ``None`` + keeps whatever Nemotron last produced, which is what the first + frame after prefill needs since its predecessor is the + engine-internal prefill token. + """ + if not self._has_llm: + raise RuntimeError("This vLLM-Omni session has no Nemotron component") + if self._closed: + raise RuntimeError(f"OmniStreamingSession {self.request_id} is closed") + ac_emb = acoustic_embedding.detach().cpu().contiguous() + if ac_emb.dim() == 1: + ac_emb = ac_emb.unsqueeze(0) + elif ac_emb.dim() == 3: + ac_emb = ac_emb.reshape(-1, ac_emb.shape[-1]) + if ac_emb.dim() != 2: + raise ValueError( + "acoustic_embedding must be shapeable to 2D [n, hidden], " f"got {tuple(acoustic_embedding.shape)}" + ) + ac_emb = ac_emb.to(torch.float32) + + self._t_put = time.perf_counter() + asyncio.run_coroutine_threadsafe( + self._input_q.put((ac_emb, prev_text_token)), + self._loop, + ).result() + self._rec("put", time.perf_counter() - self._t_put) + + item = self._text_out_q.get(timeout=self._step_timeout) + returned = time.perf_counter() + if not isinstance(item, _Sentinel): + self._rec("deliver", returned - self._t_done) + self._rec("step_total", returned - self._t_put) + if isinstance(item, _Sentinel): + if item.exc is not None: + raise RuntimeError(f"OmniStreamingSession {self.request_id} consumer raised") from item.exc + raise RuntimeError(f"OmniStreamingSession {self.request_id} ended before producing a token") + return item + + def step_tts(self, text_token: int) -> None: + """Submit one text token to EarTTS and wait for its acoustic frame. + + Valid whether or not this session also owns Nemotron; the resulting + codes are collected by :meth:`drain_audio_codes`. + """ + if not self._has_tts: + raise RuntimeError("This vLLM-Omni session has no EarTTS component") + if self._closed: + raise RuntimeError(f"OmniStreamingSession {self.request_id} is closed") + + tokens = StepTokens(int(text_token)) + self._t_put = time.perf_counter() + + async def _put_tts_input() -> None: + await self._pending_tokens_q.put(tokens) + await self._tts_cond_input_q.put(tokens.text) + if self._guidance_enabled: + await self._tts_uncond_input_q.put(tokens.text) + + asyncio.run_coroutine_threadsafe(_put_tts_input(), self._loop).result() + item = self._tts_done_q.get(timeout=self._step_timeout) + if isinstance(item, _Sentinel): + if item.exc is not None: + raise RuntimeError(f"OmniStreamingSession {self.request_id} consumer raised") from item.exc + raise RuntimeError(f"OmniStreamingSession {self.request_id} ended before producing audio") + + def drain_audio_codes(self) -> list[torch.Tensor]: + with self._audio_buf_lock: + out = self._audio_buf + self._audio_buf = [] + return out + + def _abort_engine_requests(self) -> None: + requests = ( + (self._runtime.llm_engine, self._llm_request_id), + (self._runtime.tts_engine, self._tts_cond_request_id), + (self._runtime.tts_engine, self._tts_uncond_request_id), + ) + for engine, request_id in requests: + if engine is None: + continue + if request_id == self._tts_uncond_request_id and not self._guidance_enabled: + continue + try: + abort_result = engine.abort(request_id) + if asyncio.iscoroutine(abort_result): + asyncio.run_coroutine_threadsafe(abort_result, self._loop).result(timeout=5) + except Exception as exc: + logging.debug(f"AsyncOmni.abort({request_id}) raised: {exc!r}") + + def finish(self, *, drain_remaining_audio_s: float = 0.0) -> None: + if self._closed: + return + self._closed = True + try: + + async def _close_inputs() -> None: + # Close every input this session owns. A session with both + # components drives them independently, so closing only one + # would leave the other consumer waiting forever. + if self._has_llm: + await self._input_q.put(None) + if self._has_tts: + await self._tts_cond_input_q.put(None) + if self._guidance_enabled: + await self._tts_uncond_input_q.put(None) + + asyncio.run_coroutine_threadsafe(_close_inputs(), self._loop).result(timeout=5) + except Exception: + pass + try: + self._consumer_future.result(timeout=max(drain_remaining_audio_s, 1.0)) + except Exception as exc: + logging.debug(f"OmniStreamingSession {self.request_id} consumer ended with: {exc!r}") + self._abort_engine_requests() + self._consumer_future.cancel() + try: + self._consumer_future.result(timeout=5) + except Exception: + pass + for queue in (self._text_out_q, self._tts_done_q): + try: + while True: + queue.get_nowait() + except Exception: + pass + + def abort(self) -> None: + if self._closed: + return + self._closed = True + self._abort_engine_requests() + self._consumer_future.cancel() diff --git a/nemo/collections/speechlm2/models/duplex_ear_tts.py b/nemo/collections/speechlm2/models/duplex_ear_tts.py index 02e89080ae3a..4c2e6ee6e27b 100644 --- a/nemo/collections/speechlm2/models/duplex_ear_tts.py +++ b/nemo/collections/speechlm2/models/duplex_ear_tts.py @@ -112,8 +112,16 @@ def __init__(self, cfg: dict) -> None: # compute samples per frame self.source_samples_per_frame = int(self.source_sample_rate * cfg.data.frame_length) - # get codec silence tokens - codec_silence_tokens = self.get_codec_silence_frame() + # Get codec silence tokens (skip when codec has random weights — the + # buffer will be overwritten from the checkpoint. We skip to save time: + # self.get_codec_silence_frame() is relatively slow as it works by running the codec + # encoder on a silence waveform and then picking the most common frame + # from the output). + if self.cfg.get('pretrained_codec_model', None) is not None: + codec_silence_tokens = self.get_codec_silence_frame() + else: + num_q = self.tts_model.config.num_quantizers + codec_silence_tokens = torch.zeros(num_q, dtype=torch.long) self.register_buffer("codec_silence_tokens", codec_silence_tokens) # cached for quicker audio decoding diff --git a/nemo/collections/speechlm2/models/duplex_stt_model.py b/nemo/collections/speechlm2/models/duplex_stt_model.py index ded8a89330b6..492718ff154f 100644 --- a/nemo/collections/speechlm2/models/duplex_stt_model.py +++ b/nemo/collections/speechlm2/models/duplex_stt_model.py @@ -13,7 +13,6 @@ # limitations under the License. import copy import os -import re import torch from lightning import LightningModule @@ -35,6 +34,7 @@ from nemo.collections.speechlm2.data.utils import get_pad_id from nemo.collections.speechlm2.parts.hf_hub import HFHubMixin from nemo.collections.speechlm2.parts.label_prep import maybe_prepend_prompt_tokens, prepare_text_and_asr_labels +from nemo.collections.speechlm2.parts.logit_boosts import LogitBoosts, apply_logit_boosts from nemo.collections.speechlm2.parts.lora import maybe_install_lora from nemo.collections.speechlm2.parts.metrics.bleu import BLEU from nemo.collections.speechlm2.parts.metrics.empty_text import EmptyTextMetric @@ -45,22 +45,24 @@ from nemo.collections.speechlm2.parts.pretrained import ( load_pretrained_hf, maybe_load_pretrained_models, + resolve_pretrained_config, set_model_dict_for_partial_init, setup_speech_encoder, ) +from nemo.collections.speechlm2.parts.text_utils import strip_timestamps from nemo.collections.speechlm2.streaming.duplex_stt_inference import DuplexSTTStreamingInference from nemo.core.neural_types import AudioSignal, LabelsType, LengthsType, NeuralType from nemo.utils import logging -def maybe_rename_llm_kwargs_for_nemotron(kwargs: dict, model_cfg) -> dict: +def maybe_rename_llm_kwargs_for_nemotron(kwargs: dict, model_cfg, cache_key: str | None = None) -> dict: """This is required because Nemotron models have a different signature than other HF models.""" if 'Nemotron' not in model_cfg.pretrained_llm: return kwargs cache = kwargs.pop("past_key_values") if cache is not None: - cache_key = model_cfg.get("cache_key", "past_key_values") - kwargs[cache_key] = cache + resolved_cache_key = cache_key or model_cfg.get("cache_key", "cache_params") + kwargs[resolved_cache_key] = cache return kwargs @@ -79,15 +81,18 @@ def __init__(self, cfg: dict) -> None: self.predict_user_text = self.cfg.get("predict_user_text", False) + pretrained_weights, tokenizer_path = resolve_pretrained_config(self.cfg) + # Load LLM first llm = load_pretrained_hf( self.cfg.pretrained_llm, - pretrained_weights=self.cfg.pretrained_weights, + pretrained_weights=pretrained_weights, trust_remote_code=self.cfg.get("trust_remote_code", False), + use_meta_device=self.cfg.get("use_meta_device", False), ).train() # Initialize tokenizer with optional special tokens from config - tokenizer_src = self.cfg.get("tokenizer_path", None) or self.cfg.pretrained_llm + tokenizer_src = self.cfg.get("tokenizer_path", None) or tokenizer_path self.tokenizer = AutoTokenizer( tokenizer_src, use_fast=True, @@ -110,10 +115,18 @@ def __init__(self, cfg: dict) -> None: self.asr_head = copy.deepcopy(self.lm_head) self.embed_asr_tokens = copy.deepcopy(self.embed_tokens) + # Some VoiceChat checkpoints have a separate function-token output + # channel. Even when the caller does not execute tools, its predicted + # token is part of the next frame's input and must therefore remain in + # the autoregressive state to preserve text/audio behavior. + self.use_function_head = self.cfg.get("use_function_head", False) + if self.use_function_head: + self.function_head = copy.deepcopy(self.lm_head) + maybe_install_lora(self) # Load the pretrained streaming ASR model - setup_speech_encoder(self, pretrained_weights=self.cfg.pretrained_weights) + setup_speech_encoder(self, pretrained_weights=pretrained_weights) maybe_load_pretrained_models(self) @@ -168,12 +181,16 @@ def forward( self, input_embeds: Tensor, cache=None, + cache_position=None, + cache_key: str | None = None, ) -> dict[str, Tensor]: """ Text prediction only (audio_loss_weight=0). """ kwargs = dict(inputs_embeds=input_embeds, past_key_values=cache, use_cache=cache is not None, return_dict=True) - kwargs = maybe_rename_llm_kwargs_for_nemotron(kwargs, self.cfg) + kwargs = maybe_rename_llm_kwargs_for_nemotron(kwargs, self.cfg, cache_key=cache_key) + if cache_position is not None: + kwargs["cache_position"] = cache_position out = self.llm(**kwargs) B, T = input_embeds.shape[:2] @@ -184,22 +201,26 @@ def forward( asr_in = out['last_hidden_state'] asr_logits = self.asr_head(asr_in) # (B, T, asr_vocab_size) + function_logits = None + if self.use_function_head: + function_logits = self.function_head(out['last_hidden_state']) + if not self.training: - if self.cfg.get("inference_pad_boost", None): - text_logits[:, :, self.text_pad_id] += self.cfg.inference_pad_boost - if self.cfg.get("inference_bos_boost", None): - text_logits[:, :, self.text_bos_id] += self.cfg.inference_bos_boost - if self.cfg.get("inference_eos_boost", None): - text_logits[:, :, self.text_eos_id] += self.cfg.inference_eos_boost + token_ids = dict(pad_id=self.text_pad_id, bos_id=self.text_bos_id, eos_id=self.text_eos_id) + apply_logit_boosts(text_logits, LogitBoosts.agent_from_cfg(self.cfg), **token_ids) + if self.predict_user_text: + apply_logit_boosts(asr_logits, LogitBoosts.user_from_cfg(self.cfg), **token_ids) ans = {"text_logits": text_logits} if self.predict_user_text: ans["asr_logits"] = asr_logits + if self.use_function_head: + ans["function_logits"] = function_logits if cache is not None: if 'Nemotron' in self.cfg.pretrained_llm: - cache_key = self.cfg.get("cache_key", "cache_params") - ans["cache"] = getattr(out, cache_key, out.get(cache_key)) + resolved_cache_key = cache_key or self.cfg.get("cache_key", "cache_params") + ans["cache"] = getattr(out, resolved_cache_key, out.get(resolved_cache_key)) else: ans["cache"] = out["past_key_values"] @@ -507,8 +528,7 @@ def validation_step(self, batch: dict, batch_idx: int): prompt_token_lens=prompt_token_lens, ) - # Strip timestamps for metrics - text_clean = [re.sub(r"<[\|$].*?[\|$]>", "", s).strip() for s in results["text"]] + text_clean = [strip_timestamps(s) for s in results["text"]] # Agent text metrics self.bleu.update(name=name, refs=dataset_batch["target_texts"], hyps=text_clean) @@ -571,6 +591,68 @@ def _get_asr_bos_embedding(self) -> torch.Tensor: input_embeds = self.embed_asr_tokens(text_bos) return input_embeds + def build_input_embedding( + self, + frame_embedding: torch.Tensor, + current_frame_idx: int, + gen_text: torch.Tensor, + gen_asr_text: torch.Tensor | None, + gen_function: torch.Tensor | None = None, + has_prompt: bool = False, + ) -> torch.Tensor: + """Compose the LLM input embedding for a single streaming frame. + + Combines the perception embedding (user channel) with the text / + ASR channel embeddings from the previous step. At frame 0 this + is either BOS (no prompt) or pad (after prompt). + + The arithmetic order must match offline inference exactly + (floating-point addition is not associative). For t > 0 the text + and ASR embeddings are summed first, then added to the perception + embedding. For t == 0 the sequential ``+=`` pattern matches the + offline path. + """ + emb = frame_embedding.clone() + emb *= self.cfg.get("duplex_user_channel_weight", 1.0) + + if current_frame_idx == 0 and not has_prompt: + emb += self._get_bos_embedding() * self.cfg.get("duplex_text_channel_weight", 1.0) + if self.predict_user_text: + emb += self._get_asr_bos_embedding() * self.cfg.get("duplex_asr_text_weight", 1.0) + + elif current_frame_idx == 0 and has_prompt: + pad_token = torch.full((1,), fill_value=self.text_pad_id, device=self.device, dtype=torch.long) + emb += self.embed_tokens(pad_token).to(dtype=emb.dtype) + if self.predict_user_text: + emb += self.embed_asr_tokens(pad_token).to(dtype=emb.dtype) + + else: + prev = current_frame_idx - 1 + last_token_emb = self.embed_tokens(gen_text[:, prev]) * self.cfg.get("duplex_text_channel_weight", 1.0) + if self.predict_user_text: + last_asr_token_emb = self.embed_asr_tokens(gen_asr_text[:, prev]) * self.cfg.get( + "duplex_asr_text_weight", 1.0 + ) + emb += last_token_emb + last_asr_token_emb + else: + emb += last_token_emb + + if self.use_function_head: + if gen_function is None: + raise ValueError("gen_function is required when use_function_head=True") + function_idx = 0 if current_frame_idx == 0 else current_frame_idx - 1 + function_emb = self.embed_tokens(gen_function[:, function_idx]) + emb += function_emb * self.cfg.get("duplex_function_channel_weight", 1.0) + + # The reference branch below sums the agent, audio, ASR and function + # channels in a different floating-point order. NeMo keeps the + # ordering above for every checkpoint, so results can differ in the + # last bits from that implementation. + # Reference: + # https://github.com/NVIDIA-NeMo/Speech/blob/14c77efb8110ee46eebdc50a3b15ee6d2c1a3878/nemo/collections/speechlm2/parts/fusion.py#L72-L106 + + return emb + def backward(self, *args, **kwargs): with loss_parallel(): super().backward(*args, **kwargs) @@ -695,7 +777,10 @@ def configure_model(self) -> None: logging.warning(f"Both config and fallback methods failed: {fallback_e}") logging.warning("Skipping tensor parallel configuration for this attention layer") - for m in (self.lm_head,): + output_heads = [self.lm_head] + if self.use_function_head: + output_heads.append(self.function_head) + for m in output_heads: parallelize_module( m, tp_mesh, @@ -717,6 +802,8 @@ def configure_model(self) -> None: self.embed_tokens = fully_shard(self.embed_tokens, **fsdp_config) self.llm = fully_shard(self.llm, **fsdp_config) self.lm_head = fully_shard(self.lm_head, **fsdp_config) + if self.use_function_head: + self.function_head = fully_shard(self.function_head, **fsdp_config) self.perception = fully_shard(self.perception, **fsdp_config) if self.predict_user_text: self.asr_head = fully_shard(self.asr_head, **fsdp_config) diff --git a/nemo/collections/speechlm2/models/nemotron_voicechat.py b/nemo/collections/speechlm2/models/nemotron_voicechat.py index d48b3f46fe73..94ee2247eb97 100644 --- a/nemo/collections/speechlm2/models/nemotron_voicechat.py +++ b/nemo/collections/speechlm2/models/nemotron_voicechat.py @@ -12,10 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. import gc +import json import os +import warnings from pathlib import Path -from typing import Optional, Union - import torch from huggingface_hub import CONFIG_NAME from lightning import LightningModule @@ -35,6 +35,97 @@ from nemo.utils import logging +_NEMOTRON_LABS_VOICECHAT_RELEASE_ID = "nvidia/NVIDIA-NemotronLabs-VoiceChat-11B" + + +def _is_nemotron_labs_voicechat_release(model_id: str | Path, cfg: dict | None = None) -> bool: + """Recognize the public repo ID, local directory name, or release config.""" + model_id = str(model_id).rstrip("/") + if ( + model_id == _NEMOTRON_LABS_VOICECHAT_RELEASE_ID + or Path(model_id).name == Path(_NEMOTRON_LABS_VOICECHAT_RELEASE_ID).name + ): + return True + + # ``hf download --local-dir`` permits arbitrary directory names. In that + # case identify the release by its unique bundled RNN-T metadata plus the + # function/no-ASR channel configuration. + try: + stt_cfg = cfg["model"]["stt"]["model"] + except (KeyError, TypeError): + return False + return ( + "_rnnt_merge_info" in cfg + and stt_cfg.get("pretrained_llm") == "nvidia/NVIDIA-Nemotron-Nano-9B-v2" + and stt_cfg.get("use_function_head") is True + and stt_cfg.get("predict_user_text") is False + ) + + +def _apply_nemotron_labs_voicechat_release_config_shim(cfg: dict) -> dict: + """Adapt the experimental NemotronLabs VoiceChat release to current NeMo. + + This is intentionally a narrow backward-compatibility shim for: + + Checkpoint: + https://huggingface.co/nvidia/NVIDIA-NemotronLabs-VoiceChat-11B + Reference code: + https://github.com/NVIDIA-NeMo/Speech/tree/nemotron-labs-voicechat + + That checkpoint's config was written for the reference branch above, not + against NeMo's checkpoint schema, so do not copy these fields into new + exports. The shim translates only the known config/API differences and + leaves the checkpoint unchanged on disk. + + The release stores STT data and experiment settings next to + ``model.stt.model``, whereas ``DuplexSTTModel`` receives only the inner + model dict, so copy the two runtime fields it requires when absent. + Exports that already match NeMo's schema are left untouched. + """ + try: + stt_section = cfg["model"]["stt"] + stt_model_cfg = stt_section["model"] + except (KeyError, TypeError): + return cfg + + if "source_sample_rate" not in stt_model_cfg: + source_sample_rate = stt_section.get("data", {}).get("source_sample_rate") + if source_sample_rate is None: + source_sample_rate = cfg.get("data", {}).get("source_sample_rate") + if source_sample_rate is not None: + stt_model_cfg["source_sample_rate"] = source_sample_rate + logging.info("NemotronLabs release shim: mapped STT source_sample_rate") + + if "validation_save_path" not in stt_model_cfg: + validation_save_path = stt_section.get("exp_manager", {}).get("explicit_log_dir") + if validation_save_path is None: + validation_save_path = cfg.get("exp_manager", {}).get("explicit_log_dir", "") + stt_model_cfg["validation_save_path"] = validation_save_path + logging.info("NemotronLabs release shim: mapped STT validation_save_path") + + # The released branch used the old Transformers NemotronH attribute names + # (backbone/embeddings). Current Transformers exposes the causal-LM + # backbone as model while retaining embeddings on the inner module. + if "llm_attr_name" not in stt_model_cfg and "base_model_name" in stt_model_cfg: + stt_model_cfg["llm_attr_name"] = "model" + logging.info("NemotronLabs release shim: mapped NemotronH backbone to 'model'") + if "embed_tokens_attr_name" not in stt_model_cfg and "embed_tokens_name" in stt_model_cfg: + stt_model_cfg["embed_tokens_attr_name"] = stt_model_cfg["embed_tokens_name"] + logging.info( + "NemotronLabs release shim: mapped token embedding attribute to " + f"{stt_model_cfg['embed_tokens_attr_name']!r}" + ) + + try: + cas_cfg = cfg["model"]["speech_generation"]["model"]["tts_config"]["cas_config"] + if cas_cfg.pop("pretrained_tokenizer_name", None) is not None: + logging.info("NemotronLabs release shim: removed legacy CAS pretrained_tokenizer_name") + except (KeyError, TypeError): + pass + + return cfg + + class NemotronVoiceChat(LightningModule, HFHubMixin): """ NemotronVoiceChat: End-to-End Duplex Speech-to-Speech Model. @@ -120,8 +211,10 @@ def __init__(self, cfg: dict) -> None: # Load Duplex TTS model self.tts_model = DuplexEARTTS(OmegaConf.to_container(self.cfg.speech_generation, resolve=True)) - # reset silence tokens to avoid inference issues - self.tts_model.codec_silence_tokens = self.tts_model.get_codec_silence_frame() + # reset silence tokens to avoid inference issues (skip when codec + # has random weights — the buffer will be loaded from checkpoint) + if self.tts_model.cfg.get('pretrained_codec_model', None) is not None: + self.tts_model.codec_silence_tokens = self.tts_model.get_codec_silence_frame() self.target_fps = self.tts_model.target_fps # compute source fps self.source_fps = self.source_sample_rate / ( @@ -131,6 +224,58 @@ def __init__(self, cfg: dict) -> None: self._use_fsdp = False self._use_tp = False + _DEFAULT_KEEP_FP32 = frozenset({"stt_model.perception", "tts_model"}) + + def safe_cast_to(self, dtype: torch.dtype, keep_fp32: set[str] | None = None) -> "NemotronVoiceChat": + """Cast model to dtype, keeping specified submodules in float32. + + Args: + dtype: Target dtype for castable modules. + keep_fp32: Dotted submodule name prefixes to keep in float32. + Defaults to {"stt_model.perception", "tts_model"}. + """ + if keep_fp32 is None: + keep_fp32 = self._DEFAULT_KEEP_FP32 + for name, param in self.named_parameters(): + if param.is_floating_point() and not any(name == p or name.startswith(p + ".") for p in keep_fp32): + param.data = param.data.to(dtype=dtype) + for name, buf in self.named_buffers(): + if buf.is_floating_point() and not any(name == p or name.startswith(p + ".") for p in keep_fp32): + buf.data = buf.data.to(dtype=dtype) + return self + + def save_pretrained( + self, + save_directory: str | Path, + **kwargs, + ) -> str | None: + """Save model and export LLM artifacts for offline inference. + + Tokenizer is exported by HFHubMixin._save_pretrained. + This override adds perception config and pretrained_weights=False to config.json. + """ + result = super().save_pretrained(save_directory, **kwargs) + + # Save full perception config and pretrained_weights=False at the top level + # of config.json so that resolve_pretrained_config() can skip pretrained + # ASR/LLM downloads. + try: + config_path = Path(save_directory) / "config.json" + if config_path.exists(): + with open(config_path) as f: + config = json.load(f) + # Tell from_pretrained not to re-download original child models + # (ASR encoder, LLM, codec) — weights come from model.safetensors. + config["pretrained_weights"] = False + config["perception"] = OmegaConf.to_container(self.stt_model.cfg.perception, resolve=True) + with open(config_path, "w") as f: + json.dump(config, f, indent=2) + logging.info(f"Saved perception config to {config_path}") + except Exception as e: + warnings.warn(f"Failed to save perception config: {e}") + + return result + def init_from_model_from_ckpt(self, checkpoint_path): if checkpoint_path is not None: checkpoint_state = torch.load(checkpoint_path, map_location='cpu')['state_dict'] @@ -144,11 +289,11 @@ def _from_pretrained( cls, *, model_id: str, - revision: Optional[str], - cache_dir: Optional[Union[str, Path]], + revision: str | None, + cache_dir: str | Path | None, force_download: bool, local_files_only: bool, - token: Union[str, bool, None], + token: str | bool | None, map_location: str = "cpu", strict: bool = False, **model_kwargs, @@ -157,7 +302,20 @@ def _from_pretrained( Load Pytorch pretrained weights and return the loaded model. Wrapper over PyTorchModelHubMixin that auto-handles config and uses our custom memory-efficient safetensors streaming loader to prevent OOM. + + Extra kwargs (passed via ``from_pretrained(..., key=val)``): + + - ``skip_prefixes`` (set[str] | None): Parameter-name prefixes whose + weights should be skipped during checkpoint loading. The loader + will neither materialize meta-device tensors nor read safetensors + data for keys matching these prefixes — avoiding wasted memory and + I/O for components that the caller will replace (e.g. with vLLM + engines). The caller is responsible for cleaning up or replacing + the corresponding submodules after loading. + Example: ``{"stt_model.llm.", "tts_model.tts_model."}`` """ + skip_prefixes = model_kwargs.pop("skip_prefixes", None) + # Fetch the Config resolved_config_file = cached_file( model_id, @@ -175,10 +333,32 @@ def _from_pretrained( raise RuntimeError(f"Missing {CONFIG_NAME} file for {model_id=}") model_kwargs['cfg'] = OmegaConf.to_container(OmegaConf.load(resolved_config_file)) + if _is_nemotron_labs_voicechat_release(model_id, model_kwargs['cfg']): + logging.info(f"Applying compatibility shim for {_NEMOTRON_LABS_VOICECHAT_RELEASE_ID}") + _apply_nemotron_labs_voicechat_release_config_shim(model_kwargs['cfg']) # Skip loading child module weights natively model_kwargs['cfg']['pretrained_weights'] = False + # Propagate pretrained_weights=False into nested configs so child + # modules skip downloading pretrained ASR, LLM, and codec models. + cfg = model_kwargs['cfg'] + try: + stt_model_cfg = cfg['model']['stt']['model'] + stt_model_cfg['pretrained_weights'] = False + stt_model_cfg['use_meta_device'] = True + if 'perception' in cfg: + stt_model_cfg['perception'] = cfg['perception'] + logging.info("Injected saved perception config into STT model config") + except (KeyError, TypeError): + logging.warning("Could not propagate pretrained_weights=False into nested STT config") + try: + tts_model_cfg = cfg['model']['speech_generation']['model'] + tts_model_cfg['pretrained_model'] = None + tts_model_cfg['pretrained_codec_model'] = None + except (KeyError, TypeError): + logging.warning("Could not nullify pretrained TTS/codec paths in nested TTS config") + # Instantiate the empty model skeleton model = cls(model_kwargs['cfg']) @@ -198,21 +378,62 @@ def _from_pretrained( if resolved_weights_file is None: raise RuntimeError(f"Missing model.safetensors file for {model_id=}") - # Stream the weights safely using your custom memory-efficient loader! + # Stream the weights from safetensors ckpt_dir = os.path.dirname(resolved_weights_file) - model.init_from_safetensors_ckpt(ckpt_dir) + model.init_from_safetensors_ckpt(ckpt_dir, skip_prefixes=skip_prefixes) return model - def init_from_safetensors_ckpt(self, ckpt_path, prefix=""): + def init_from_safetensors_ckpt(self, ckpt_path, prefix="", skip_prefixes: set[str] | None = None): """ Memory-efficient streaming safetensors loader with dynamic audio_prompt_latents recreation support. + Uses ``torch.nn.Module.to_empty()`` to materialize any meta-device + tensors before streaming weights from the checkpoint. + Safe for large models and distributed training (if called before DDP/FSDP wrap). + + Args: + ckpt_path: Directory containing ``model.safetensors``. + prefix: Optional prefix prepended to checkpoint keys when + matching against model parameter names. + skip_prefixes: If provided, parameter-name prefixes to skip. + Matching checkpoint keys will not be read from disk, and + any matching meta-device parameters will not be + materialized to CPU. Skipped meta-device parameters + will not trigger the post-load safety check. Use this + to avoid wasted memory and I/O for submodules the caller + intends to replace (e.g. with vLLM engines). """ + skip_prefixes = set(skip_prefixes) if skip_prefixes else set() + + def _should_skip(name: str) -> bool: + return any(name.startswith(p) for p in skip_prefixes) + + # Materialize meta-device tensors into real (uninitialized) CPU tensors + # so that the streaming copy_() loop below can use target.data.copy_(). + # Only targets tensors actually on the meta device — modules that + # were constructed with real weights are left untouched. + for name, param in list(self.named_parameters()): + if param.is_meta and not _should_skip(name): + parts = name.split(".") + module = self + for part in parts[:-1]: + module = getattr(module, part) + module._parameters[parts[-1]] = torch.nn.Parameter( + torch.empty_like(param, device="cpu"), requires_grad=param.requires_grad + ) + for name, buf in list(self.named_buffers()): + if buf.is_meta and not _should_skip(name): + parts = name.split(".") + module = self + for part in parts[:-1]: + module = getattr(module, part) + module._buffers[parts[-1]] = torch.empty_like(buf, device="cpu") loaded_keys = [] + skipped_keys = [] missing_keys = [] # Build fast lookup tables once @@ -227,6 +448,10 @@ def init_from_safetensors_ckpt(self, ckpt_path, prefix=""): for key in f.keys(): + if _should_skip(key): + skipped_keys.append(key) + continue + try: tensor = f.get_tensor(key) except Exception as e: @@ -242,24 +467,20 @@ def init_from_safetensors_ckpt(self, ckpt_path, prefix=""): if prefix + key in param_dict: target = param_dict[prefix + key] - if target.shape != tensor.shape: - logging.warning(f"Shape mismatch for {key}: " f"model {target.shape} vs ckpt {tensor.shape}") + logging.warning(f"Shape mismatch for {key}: model {target.shape} vs ckpt {tensor.shape}") else: target.data.copy_(tensor) - loaded_keys.append(key) elif prefix + key in buffer_dict: target = buffer_dict[prefix + key] - if target.shape != tensor.shape: logging.warning( - f"Buffer shape mismatch for {key}: " f"model {target.shape} vs ckpt {tensor.shape}" + f"Buffer shape mismatch for {key}: model {target.shape} vs ckpt {tensor.shape}" ) else: target.data.copy_(tensor) - loaded_keys.append(key) else: @@ -272,9 +493,23 @@ def init_from_safetensors_ckpt(self, ckpt_path, prefix=""): logging.info(f"Loaded {len(loaded_keys)} tensors from pretrained model") + if skipped_keys: + logging.info(f"Skipped {len(skipped_keys)} tensors matching skip_prefixes {skip_prefixes}") + if missing_keys: logging.warning(f"{len(missing_keys)} keys in checkpoint not found in model") + # Fail if any *parameters* are still on meta device — those genuinely + # need weights from the checkpoint and their absence is an error. + # Parameters covered by skip_prefixes are excluded: the caller + # is responsible for replacing or deleting those submodules. + meta_params = [n for n, p in self.named_parameters() if p.is_meta and not _should_skip(n)] + if meta_params: + raise RuntimeError( + f"{len(meta_params)} parameters still on meta device after checkpoint load " + f"(missing from checkpoint; showing first 20): {meta_params[:20]}" + ) + gc.collect() def training_step(self, batch: dict, batch_idx: int): @@ -431,6 +666,7 @@ def offline_inference( incremental_audio_decoding: bool = False, generation_config: dict = None, guidance_enabled: bool = True, + return_logits: bool = False, ) -> dict[str, torch.Tensor]: """ Runs full offline duplex speech-to-speech inference. @@ -479,6 +715,12 @@ def offline_inference( guidance_enabled (bool, optional): Enables classifier-free guidance. + return_logits (bool, optional): + When True, collect per-step text and ASR logits and + include them in the returned dict as ``"text_logits"`` + (B, T, V_text) and ``"asr_logits"`` (B, T, V_asr). + Useful for parity testing against incremental inference. + Returns: dict[str, torch.Tensor]: @@ -502,6 +744,12 @@ def offline_inference( Tensor (B,) — waveform lengths in samples (if decode_audio=True). + • "text_logits" (only when return_logits=True): + Tensor (B, T, V_text) — per-step text head logits. + + • "asr_logits" (only when return_logits=True): + Tensor (B, T, V_asr) — per-step ASR head logits. + Notes: • Uses streaming inference backend of DuplexSTTModel. • Uses autoregressive codec generation from DuplexEARTTS. @@ -519,6 +767,10 @@ def offline_inference( B = inference_state["B"] T = inference_state["T"] + if return_logits: + _text_logits = [ans["text_logits"][:, -1].detach()] + _asr_logits = [ans["asr_logits"][:, -1].detach()] if "asr_logits" in ans else [] + # if speaker_name is provided uses it, if not uses the speaker_audio provided, if speaker_audio is None load it from inference_speaker_reference if speaker_audio is None: speaker_name = self.cfg.get("inference_speaker_name", None) @@ -566,7 +818,12 @@ def offline_inference( # Autoregressive loop for t in range(1, T): # do one step inference on Duplex STT model - _ = self.stt_model.streaming_inference._step_inference(t, inference_state, ans) + ans = self.stt_model.streaming_inference._step_inference(t, inference_state, ans) + + if return_logits: + _text_logits.append(ans["text_logits"][:, -1].detach()) + if "asr_logits" in ans: + _asr_logits.append(ans["asr_logits"][:, -1].detach()) # do one step inference on Duplex TTS model # current subword id is always seem @@ -603,7 +860,7 @@ def offline_inference( audio_pred = torch.cat([audio_pred, audio_pred_i], dim=1) audio_pred_len += audio_pred_i_len - logging.info(f"Autoregressive inference step: {t} of {T} !") + logging.debug(f"Autoregressive inference step: {t} of {T} !") # Trim back to local length if padded if self._use_fsdp and T > inference_state["T_local"]: @@ -623,6 +880,11 @@ def offline_inference( ans["audio"] = audio_pred.squeeze(1) ans["audio_len"] = audio_pred_len + if return_logits: + ans["text_logits"] = torch.stack(_text_logits, dim=1) + if _asr_logits: + ans["asr_logits"] = torch.stack(_asr_logits, dim=1) + return ans def load_state_dict(self, state_dict, strict: bool = True): diff --git a/nemo/collections/speechlm2/modules/ear_tts_model.py b/nemo/collections/speechlm2/modules/ear_tts_model.py index 9ea1b238620a..431346cb34cf 100644 --- a/nemo/collections/speechlm2/modules/ear_tts_model.py +++ b/nemo/collections/speechlm2/modules/ear_tts_model.py @@ -808,6 +808,8 @@ def __init__( if self.use_bos_eos_emb: self.bos_eos_emb = BOSEOSEmbedding(tokenizer, self.hidden_size) + self.use_tts_subword_cache = False + def prepare_inputs(self, subword_ids: Tensor, padding_mask: Tensor) -> tuple[Tensor, Tensor]: """ Converts a batch of subword IDs into a padded batch of character IDs. @@ -842,6 +844,10 @@ def forward(self, subword_ids: Tensor, subword_mask: Tensor | None = None) -> Te """ Performs the forward pass to get character-aware subword embeddings. + When use_tts_subword_cache is True and the module is in eval mode, a + per-subword-ID cache skips the expensive char encoding + backbone + + pooling path for previously seen tokens. + Args: subword_ids (Tensor): A tensor of subword IDs. Shape: `[batch, seq_len]`. subword_mask (Tensor | None): A boolean mask for padding. Defaults to None. @@ -852,6 +858,19 @@ def forward(self, subword_ids: Tensor, subword_mask: Tensor | None = None) -> Te if subword_mask is None: subword_mask = torch.ones_like(subword_ids, dtype=torch.bool) + # Inference cache: return cached embeddings if all valid IDs have been seen + if not self.training and self.use_tts_subword_cache: + if not hasattr(self, '_inference_cache'): + self._inference_cache = {} + valid_ids = torch.masked_select(subword_ids, subword_mask).tolist() + if all(sid in self._inference_cache for sid in valid_ids): + cached = torch.stack([self._inference_cache[sid] for sid in valid_ids]) + out = torch.zeros( + subword_ids.shape + (cached.size(-1),), device=subword_ids.device, dtype=cached.dtype + ) + out[subword_mask] = cached + return out + # 1. Convert subword IDs to character IDs char_ids, char_lengths = self.prepare_inputs(subword_ids, subword_mask) @@ -882,6 +901,13 @@ def forward(self, subword_ids: Tensor, subword_mask: Tensor | None = None) -> Te if self.use_bos_eos_emb: subword_embeds = self.bos_eos_emb(subword_embeds, subword_ids) + # Cache results for future lookups + if not self.training and self.use_tts_subword_cache: + valid_ids = torch.masked_select(subword_ids, subword_mask).tolist() + valid_embeds = subword_embeds[subword_mask].detach() + for idx, sid in enumerate(valid_ids): + self._inference_cache[sid] = valid_embeds[idx] + return subword_embeds @@ -1050,15 +1076,13 @@ def depthsum_embedding(self, code: Tensor) -> Tensor: ret: [b, t, h] """ b, t, d = code.size() - _, v, h = self.rvq_embs.size() - device = code.device - - ret = torch.zeros((b, t, h), device=device, dtype=self.rvq_embs.dtype) embs = F.pad(self.rvq_embs, [0, 0, 0, 1]) - for i in range(d): - emb = embs[i] - ret = ret + F.embedding(code[..., i], emb) - return ret + v_padded = embs.shape[1] + offsets = torch.arange(d, device=code.device).view(1, 1, d) * v_padded + flat_indices = (code + offsets).reshape(b * t * d) + flat_embs = embs.reshape(d * v_padded, -1) + gathered = F.embedding(flat_indices, flat_embs).reshape(b, t, d, -1) + return gathered.sum(dim=2) def prepare_training_inputs(self, code: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: """Prepares masked and dropped-out versions of the code for training.""" @@ -1542,11 +1566,13 @@ def generate_step( num_maskings = torch.ceil(masking_rates * self.config.num_quantizers).long() ks = num_maskings - F.pad(num_maskings[1:], [0, 0, 0, 1]) + ks_list = ks.squeeze(-1).tolist() # 4. Iteratively unmask the continuous part of the code cnt = 0 - for i, k in enumerate(ks): - if torch.all(k == 0): + for i, k_val in enumerate(ks_list): + k_val = int(k_val) + if k_val == 0: continue # Prepare input for the MoG head @@ -1556,7 +1582,7 @@ def generate_step( mog_input_embeds = self.embed_code(self.depthsum_embedding(code)) if self.config.random_target_masking: - mog_input_embeds += self.embed_target_mask(cnt + k - 1) + mog_input_embeds += self.embed_target_mask(cnt + k_val - 1) if guidance_scale_i > 0.0: mog_input_embeds = torch.cat( [mog_input_embeds + hidden_states, mog_input_embeds + uncond_hidden_states], 0 @@ -1570,8 +1596,8 @@ def generate_step( top_p_or_k=top_p_or_k_i, ) z = mog_mu + torch.exp(mog_logs) * torch.randn_like(mog_mu) * noise_scale_i - code = depthsum_encoding_step(self.rvq_embs, z, code, cnt, k[0].item()) - cnt += k[0].item() + code = depthsum_encoding_step(self.rvq_embs, z, code, cnt, k_val) + cnt += k_val return code, lm_logits, eos_flag def load_state_dict(self, state_dict, strict: bool = True): diff --git a/nemo/collections/speechlm2/parts/hf_hub.py b/nemo/collections/speechlm2/parts/hf_hub.py index 7bb66841a1ed..68d69a63a6fb 100644 --- a/nemo/collections/speechlm2/parts/hf_hub.py +++ b/nemo/collections/speechlm2/parts/hf_hub.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import warnings from pathlib import Path from typing import Any, Dict, Optional, Union @@ -19,6 +20,8 @@ from omegaconf import DictConfig, OmegaConf from transformers.utils import cached_file +from nemo.utils import logging + SAFETENSORS_SINGLE_FILE = "model.safetensors" LLM_BACKBONE_DIR = "llm_backbone" @@ -127,6 +130,25 @@ def _from_pretrained( cached_file_kwargs=_cached_file_kwargs, ) + def _save_pretrained(self, save_directory: Path) -> None: + """Save model weights and export tokenizer for offline inference.""" + super()._save_pretrained(save_directory) + + tokenizer = getattr(self, "tokenizer", None) + if tokenizer is None: + stt = getattr(self, "stt_model", None) + if stt is not None: + tokenizer = getattr(stt, "tokenizer", None) + if tokenizer is not None: + try: + llm_dir = Path(save_directory) / "llm_artifacts" + llm_dir.mkdir(parents=True, exist_ok=True) + inner = getattr(tokenizer, "tokenizer", tokenizer) + inner.save_pretrained(str(llm_dir)) + logging.info(f"Saved LLM tokenizer to {llm_dir}") + except Exception as e: + warnings.warn(f"Failed to save LLM tokenizer: {e}. Inference will fall back to downloading from HF.") + def save_pretrained( self, save_directory: Union[str, Path], diff --git a/nemo/collections/speechlm2/parts/logit_boosts.py b/nemo/collections/speechlm2/parts/logit_boosts.py new file mode 100644 index 000000000000..2fd202683f5b --- /dev/null +++ b/nemo/collections/speechlm2/parts/logit_boosts.py @@ -0,0 +1,121 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Inference-time special-token logit boosts for the duplex text channels. + +Both the PyTorch and vLLM runtimes bias the pad/BOS/EOS logits of the agent +text channel and the user (ASR) channel before picking a token. Keeping the +values and the arithmetic here means the two cannot disagree about what +``inference_pad_boost`` and friends do. + +The boosts are read from config rather than passed as arguments because the +models that consume them are also used offline, where there is no inference +wrapper to thread them through. See +``nemo/collections/speechlm2/inference/model_wrappers/config_overrides.py`` +for how the streaming pipeline sets them. +""" + +from dataclasses import dataclass +from typing import Any + +import torch + +AGENT_BOOST_KEYS = ("inference_pad_boost", "inference_bos_boost", "inference_eos_boost") +USER_BOOST_KEYS = ( + "inference_user_pad_boost", + "inference_user_bos_boost", + "inference_user_eos_boost", +) + + +@dataclass(frozen=True) +class LogitBoosts: + """Additive logit offsets for one channel's pad/BOS/EOS tokens. + + ``None`` and ``0.0`` both mean "leave this token alone"; config treats a + falsy value as unset. + """ + + pad: float | None = None + bos: float | None = None + eos: float | None = None + + def __bool__(self) -> bool: + return bool(self.pad or self.bos or self.eos) + + def as_dict(self) -> dict[str, float | None]: + return {"pad": self.pad, "bos": self.bos, "eos": self.eos} + + @classmethod + def from_dict(cls, values: dict[str, Any] | None) -> "LogitBoosts": + if not values: + return cls() + return cls( + pad=values.get("pad"), + bos=values.get("bos"), + eos=values.get("eos"), + ) + + @staticmethod + def _cfg_get(cfg: Any, key: str) -> Any: + """Read *key* from an OmegaConf/dict config or an HF ``PretrainedConfig``.""" + if hasattr(cfg, "get"): + try: + return cfg.get(key, None) + except TypeError: + pass + return getattr(cfg, key, None) + + @classmethod + def _from_cfg(cls, cfg: Any, keys: tuple[str, str, str]) -> "LogitBoosts": + if cfg is None: + return cls() + pad, bos, eos = (cls._cfg_get(cfg, key) for key in keys) + return cls( + pad=float(pad) if pad else None, + bos=float(bos) if bos else None, + eos=float(eos) if eos else None, + ) + + @classmethod + def agent_from_cfg(cls, cfg: Any) -> "LogitBoosts": + """Boosts for the agent text channel (``inference_*_boost``).""" + return cls._from_cfg(cfg, AGENT_BOOST_KEYS) + + @classmethod + def user_from_cfg(cls, cfg: Any) -> "LogitBoosts": + """Boosts for the user/ASR channel (``inference_user_*_boost``).""" + return cls._from_cfg(cfg, USER_BOOST_KEYS) + + +def apply_logit_boosts( + logits: torch.Tensor, + boosts: LogitBoosts, + *, + pad_id: int, + bos_id: int, + eos_id: int, +) -> torch.Tensor: + """Add *boosts* to the matching vocabulary entries of *logits*, in place. + + Indexing on the last dimension, so this accepts both the ``(B, T, V)`` + tensors the PyTorch heads produce and the ``(V,)`` slice a vLLM + logits processor receives. + """ + if not boosts: + return logits + for token_id, value in ((pad_id, boosts.pad), (bos_id, boosts.bos), (eos_id, boosts.eos)): + if value: + logits[..., token_id] += value + return logits diff --git a/nemo/collections/speechlm2/parts/precision.py b/nemo/collections/speechlm2/parts/precision.py index 99be94e2b5be..8c1fcc612256 100644 --- a/nemo/collections/speechlm2/parts/precision.py +++ b/nemo/collections/speechlm2/parts/precision.py @@ -11,11 +11,111 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import os from contextlib import contextmanager import torch +def inference_precision_in_effect( + *, + allow_tf32: bool, + matmul_precision: str, + deterministic: bool, +) -> bool: + """Whether torch is *currently* configured the way these arguments ask. + + Read straight back off torch rather than tracked in a module-level flag: + the question that matters to a caller is "are my settings applied?", not + "did someone call my context manager?". Reading the real state answers it + without adding mutable global state of our own, stays correct if the + switches were set some other way, and also catches a scope entered with + the *wrong* config. + + The flash/mem-efficient SDP kernels that :func:`inference_precision` + toggles are deliberately not checked: they follow from ``deterministic``, + and can be legitimately disabled for unrelated reasons. + """ + return ( + bool(torch.backends.cudnn.allow_tf32) == allow_tf32 + and bool(torch.backends.cuda.matmul.allow_tf32) == allow_tf32 + and torch.get_float32_matmul_precision() == matmul_precision + and torch.are_deterministic_algorithms_enabled() == deterministic + ) + + +@contextmanager +def inference_precision( + *, + allow_tf32: bool = True, + matmul_precision: str = "medium", + deterministic: bool = False, +): + """Apply the process-wide precision and determinism switches, then restore. + + These are torch-level globals rather than per-model state, so they have to + be applied before any weights load and stay in effect for the whole run. + Scoped rather than set-and-forget because leaving them on would silently + change every later computation in the process: a deterministic run inside + a test session would otherwise seed the RNGs and disable the fast + attention kernels for everything that follows it. + + ``deterministic`` guarantees identical text outputs across runs for the + same input even when sampling is enabled, by seeding the global RNGs and + forcing deterministic CUDA kernels. It costs speed, and vLLM engines + cannot honour it -- callers that offer a choice of engine must reject the + combination themselves. + """ + saved = ( + torch.backends.cudnn.allow_tf32, + torch.backends.cuda.matmul.allow_tf32, + torch.get_float32_matmul_precision(), + torch.backends.cuda.flash_sdp_enabled(), + torch.backends.cuda.mem_efficient_sdp_enabled(), + torch.are_deterministic_algorithms_enabled(), + torch.is_deterministic_algorithms_warn_only_enabled(), + ) + saved_cpu_rng = torch.get_rng_state() if deterministic else None + saved_cuda_rng = torch.cuda.get_rng_state_all() if deterministic and torch.cuda.is_available() else None + + torch.backends.cudnn.allow_tf32 = allow_tf32 + torch.backends.cuda.matmul.allow_tf32 = allow_tf32 + torch.set_float32_matmul_precision(matmul_precision) + + if deterministic: + # CuBLAS reads this once, at the first CUDA matmul in the process, so + # restoring it on exit would not undo anything. Left set: the only + # cost is a 32 KB workspace reservation. + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + torch.manual_seed(0) + torch.cuda.manual_seed_all(0) + + torch.backends.cuda.enable_flash_sdp(not deterministic) + torch.backends.cuda.enable_mem_efficient_sdp(not deterministic) + torch.use_deterministic_algorithms(deterministic, warn_only=False) + + try: + yield + finally: + ( + torch.backends.cudnn.allow_tf32, + torch.backends.cuda.matmul.allow_tf32, + matmul, + flash, + mem_efficient, + was_deterministic, + warn_only, + ) = saved + torch.set_float32_matmul_precision(matmul) + torch.backends.cuda.enable_flash_sdp(flash) + torch.backends.cuda.enable_mem_efficient_sdp(mem_efficient) + torch.use_deterministic_algorithms(was_deterministic, warn_only=warn_only) + if saved_cpu_rng is not None: + torch.set_rng_state(saved_cpu_rng) + if saved_cuda_rng is not None: + torch.cuda.set_rng_state_all(saved_cuda_rng) + + @contextmanager def fp32_precision(): """ diff --git a/nemo/collections/speechlm2/parts/pretrained.py b/nemo/collections/speechlm2/parts/pretrained.py index 6631673d3e3f..edb19cdcc13b 100644 --- a/nemo/collections/speechlm2/parts/pretrained.py +++ b/nemo/collections/speechlm2/parts/pretrained.py @@ -11,7 +11,8 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from contextlib import contextmanager +import json +from contextlib import contextmanager, nullcontext from pathlib import Path from typing import Dict @@ -65,7 +66,11 @@ def load_pretrained_nemo_config(cls, model_path_or_name: str): def load_pretrained_hf( - model_path_or_name: str, pretrained_weights: bool = True, dtype=torch.float32, trust_remote_code: bool = False + model_path_or_name: str, + pretrained_weights: bool = True, + dtype=torch.float32, + trust_remote_code: bool = False, + use_meta_device: bool = False, ): """ Load pretrained HuggingFace AutoModelForCausalLM. @@ -78,6 +83,8 @@ def load_pretrained_hf( pretrained_weights: Whether to load pretrained weights (True) or random init (False) dtype: Data type for the model trust_remote_code: Whether to trust remote code when loading model (needed for some models like Nemotron) + use_meta_device: If True, create the model on the meta device (no memory allocation). + The caller must handle materializing meta tensors from a checkpoint. """ if pretrained_weights: return AutoModelForCausalLM.from_pretrained( @@ -85,7 +92,43 @@ def load_pretrained_hf( ) else: config = AutoConfig.from_pretrained(model_path_or_name, trust_remote_code=trust_remote_code) - return AutoModelForCausalLM.from_config(config, torch_dtype=dtype, trust_remote_code=trust_remote_code) + with torch.device('meta') if use_meta_device else nullcontext(): + return AutoModelForCausalLM.from_config(config, torch_dtype=dtype, trust_remote_code=trust_remote_code) + + +def resolve_pretrained_config(cfg): + """Resolve pretrained config when pretrained_s2s_model points to an HF checkpoint. + + When the HF checkpoint contains a config.json with perception config, this function: + - Sets pretrained_weights to False (weights will be loaded from pretrained_s2s_model) + - Loads the perception config from the HF checkpoint into cfg + - Resolves the tokenizer path to local llm_artifacts if available + + Args: + cfg: DictConfig with model configuration (modified in-place for perception config). + + Returns: + Tuple of (pretrained_weights, tokenizer_path). + """ + tokenizer_path = cfg.pretrained_llm + pretrained_weights = cfg.pretrained_weights + pretrained_s2s = cfg.get("pretrained_s2s_model", None) + if pretrained_s2s is not None: + hf_config_path = Path(pretrained_s2s) / "config.json" + if hf_config_path.exists(): + with open(hf_config_path) as f: + hf_config = json.load(f) + if "perception" in hf_config: + pretrained_weights = False + with open_dict(cfg): + cfg.perception = hf_config["perception"] + logging.info(f"Loaded perception config from {hf_config_path}, skipping pretrained downloads") + # Use local tokenizer if available + llm_artifacts_dir = Path(pretrained_s2s) / "llm_artifacts" + if llm_artifacts_dir.is_dir(): + tokenizer_path = str(llm_artifacts_dir) + logging.info(f"Using local tokenizer from {llm_artifacts_dir}") + return pretrained_weights, tokenizer_path def load_pretrained_automodel_llm( diff --git a/nemo/collections/speechlm2/parts/text_utils.py b/nemo/collections/speechlm2/parts/text_utils.py index 8c2a36facf9b..c8ba7b1ac230 100644 --- a/nemo/collections/speechlm2/parts/text_utils.py +++ b/nemo/collections/speechlm2/parts/text_utils.py @@ -11,10 +11,214 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import re + import torch +from whisper_normalizer.english import EnglishTextNormalizer from nemo.collections.common.tokenizers import AutoTokenizer +_whisper_normalizer = EnglishTextNormalizer() + +# --------------------------------------------------------------------------- +# Timestamp constants +# +# The speechlm2 pipeline uses two distinct timestamp conventions: +# +# 1. **Training-data timestamps** — integer frame indices written into +# transcripts by the force-aligner, e.g. ``<|0|> Hey <|3|>``. +# Pattern: ``<|INT|>`` +# +# 2. **Inference-output timestamps** — floating-point seconds inserted by +# ``tokens_to_str(eval_text_turn_taking=True)`` to annotate where the +# model's BOS/EOS boundaries fall: +# * ``<|0.8|>`` — turn start (BOS position) +# * ``<$0.72$>`` — turn end (EOS position) +# These are NOT vocabulary tokens; they are annotations added after +# decoding. +# --------------------------------------------------------------------------- + +SECONDS_PER_FRAME = 0.08 + +# Training-data format: integer frame indices (with capture group for parsing) +TRAINING_TIMESTAMP_RE = re.compile(r"<\|(\d+)\|>") + +# Inference-output format: float seconds marking turn boundaries +TIMESTAMP_BOS_RE = re.compile(r"<\|[\d.]+\|>") # turn start +TIMESTAMP_EOS_RE = re.compile(r"<\$[\d.]+\$>") # turn end + +_MULTI_SPACE_RE = re.compile(r"\s+") + + +def format_bos_timestamp(frame_pos: int) -> str: + """Format a BOS (turn-start) timestamp annotation from a frame index.""" + return f"<|{round(frame_pos * SECONDS_PER_FRAME, 3)}|>" + + +def format_eos_timestamp(frame_pos: int) -> str: + """Format an EOS (turn-end) timestamp annotation from a frame index.""" + return f"<${round(frame_pos * SECONDS_PER_FRAME, 3)}$>" + + +def format_eot_timestamp(frame_pos: int) -> str: + """Format an EOT (end-of-text) timestamp annotation from a frame index.""" + return f"<{round(frame_pos * SECONDS_PER_FRAME, 3)}>" + + +def strip_timestamps(text: str) -> str: + """Strip all timestamp tokens (both training-data and inference-output formats). + + Handles: + - Training-data timestamps: ``<|0|>``, ``<|10|>`` (integer frames) + - Inference BOS timestamps: ``<|0.8|>`` (float seconds) + - Inference EOS timestamps: ``<$0.72$>`` (float seconds) + """ + text = TRAINING_TIMESTAMP_RE.sub("", text) + text = TIMESTAMP_EOS_RE.sub("", text) + # TIMESTAMP_BOS_RE is a superset of TRAINING_TIMESTAMP_RE, but we keep + # both calls so the intent is clear and either order is safe. + text = TIMESTAMP_BOS_RE.sub("", text) + return _MULTI_SPACE_RE.sub(" ", text).strip() + + +def get_special_token_strings(tokenizer, pad_id: int, model_cfg=None) -> set[str]: + """Collect all special token strings that should be stripped from decoded text. + + Derives tokens from the tokenizer and model config at runtime. + + Args: + tokenizer: Tokenizer with ``ids_to_tokens``, ``bos_token``, and + ``eos_token`` attributes. + pad_id: Pad token ID (typically from ``DuplexSTTModel.text_pad_id``). + model_cfg: Optional model config (OmegaConf or dict). When provided, + ``user_bos_token`` and ``user_eos_token`` are included in the set + (e.g. ``'^'`` and ``'$'`` for some checkpoints). + """ + pad_str = tokenizer.ids_to_tokens([pad_id])[0] + tokens = {pad_str} + if getattr(tokenizer, 'bos_token', None): + tokens.add(tokenizer.bos_token) + if getattr(tokenizer, 'eos_token', None): + tokens.add(tokenizer.eos_token) + if model_cfg is not None: + for key in ('user_bos_token', 'user_eos_token'): + tok = model_cfg.get(key, None) if hasattr(model_cfg, 'get') else None + if tok: + tokens.add(tok) + return tokens + + +def get_special_token_ids(tokenizer, pad_id: int, model_cfg=None) -> set[int]: + """Collect special token IDs that should bypass sampling (greedy-only). + + These tokens (pad, BOS, EOS, and optionally user turn markers) must not + be subject to top-p / temperature / repetition-penalty sampling, otherwise + EOS may be randomly sampled and generation may not stop properly. + + Args: + tokenizer: Tokenizer with ``bos_id``, ``eos_id`` attributes and a + ``text_to_ids`` method. + pad_id: Pad token ID (typically from ``DuplexSTTModel.text_pad_id``). + model_cfg: Optional model config (OmegaConf or dict). When provided, + ``user_bos_token`` and ``user_eos_token`` are resolved to IDs and + included in the set. + """ + ids = {pad_id} + if getattr(tokenizer, 'bos_id', None) is not None: + ids.add(tokenizer.bos_id) + if getattr(tokenizer, 'eos_id', None) is not None: + ids.add(tokenizer.eos_id) + if model_cfg is not None: + for key in ('user_bos_token', 'user_eos_token'): + tok = model_cfg.get(key, None) if hasattr(model_cfg, 'get') else None + if tok and hasattr(tokenizer, 'text_to_ids'): + tok_ids = tokenizer.text_to_ids(tok) + if tok_ids: + ids.add(tok_ids[0]) + return ids + + +def clean_pred_text( + text: str, + special_token_strings: set[str] | None = None, +) -> str: + """Clean prediction text for fair WER comparison. + + Strips special tokens (pad, BOS, EOS, user turn markers) and timestamp + annotations, then applies ``EnglishTextNormalizer`` -- the same normalizer + used by the offline eval metrics in ``speechlm2.parts.metrics.wer``. + + Args: + text: Raw decoded text that may contain special tokens and timestamps. + special_token_strings: Set of vocabulary token strings to remove + (pad, BOS, EOS, user turn markers such as ``'^'``). Obtain via + :func:`get_special_token_strings` at pipeline init time. + When ``None``, only timestamp annotations are stripped. + """ + if not text: + return "" + if special_token_strings: + for tok in special_token_strings: + text = text.replace(tok, '') + text = strip_timestamps(text) + return _whisper_normalizer(text) + + +def _decode_tokens_with_specials( + token_strings: list[str], + tokenizer, + pad_token_str: str, + keep_pad: bool = False, +) -> str: + """Decode token strings with proper byte-level BPE handling. + + Groups consecutive non-special tokens and decodes each group via + ``tokenizer.tokens_to_text()`` (HF ``convert_tokens_to_string``), which + properly reverses byte-level BPE encoding (e.g. ``âĢĻ`` -> ``'``). + Special tokens (BOS, EOS, PAD) are never passed to + ``convert_tokens_to_string``. BOS/EOS are always kept as literal + strings so that turn boundaries are visible. PAD tokens are kept + only when *keep_pad* is True. + + Args: + token_strings: Raw token strings from ``tokenizer.ids_to_tokens()``. + tokenizer: Tokenizer with ``tokens_to_text``, ``bos_token``, and + ``eos_token`` attributes (NeMo ``AutoTokenizer`` or similar). + pad_token_str: String representation of the pad token. + keep_pad: If True, preserve all special tokens as literal strings + in the output. If False, strip them. + """ + bos = getattr(tokenizer, 'bos_token', None) + eos = getattr(tokenizer, 'eos_token', None) + + # All tokens that must not go through convert_tokens_to_string. + special_tokens = {pad_token_str} + if bos: + special_tokens.add(bos) + if eos: + special_tokens.add(eos) + + result_parts: list[str] = [] + segment: list[str] = [] + + for tok in token_strings: + if tok in special_tokens: + if segment: + result_parts.append(tokenizer.tokens_to_text(segment)) + segment = [] + if tok == pad_token_str: + if keep_pad: + result_parts.append(tok) + else: + result_parts.append(tok) + else: + segment.append(tok) + + if segment: + result_parts.append(tokenizer.tokens_to_text(segment)) + + return ''.join(result_parts) + def tokens_to_str( tokens: torch.Tensor, @@ -23,9 +227,22 @@ def tokens_to_str( pad_id: int, eval_text_turn_taking: bool = False, show_eot_timestamps: bool = False, + keep_pad: bool = False, ) -> list[str]: """ - Convert token IDs to text strings, filtering out special tokens. + Convert token IDs to text strings with proper byte-level BPE decoding. + + When ``eval_text_turn_taking`` is True, BOS/EOS/EOT token positions are + replaced by timestamp annotations (these are **not** vocabulary tokens; + they are human-readable annotations added here): + + * ``<|t|>`` -- turn start (BOS position, seconds) + * ``<$t$>`` -- turn end (EOS position, seconds) + * ```` -- end-of-text (first pad after BOS, seconds) + + Note: training-data timestamps use integer frame indices (``<|10|>``), + while these inference-output timestamps use float seconds (``<|0.8|>``). + Both ``<|...|>`` formats are stripped by :func:`strip_timestamps`. Args: tokens: Token IDs tensor (B, T) @@ -34,14 +251,18 @@ def tokens_to_str( pad_id: Pad token ID to filter out eval_text_turn_taking: If True, insert timestamps at bos/eos positions show_eot_timestamps: If True, also insert timestamps at end-of-text (first pad after BOS) + keep_pad: If True, preserve all special tokens (including pad) as literal + strings in the output. Useful for "raw" output that shows the full + token stream. If False (default), special tokens are stripped. Returns: List of decoded text strings """ + pad_token_str = tokenizer.ids_to_tokens([pad_id])[0] ans = [] - # Helper function to filter special tokens from token IDs - # This filtering is applied regardless of eval_text_turn_taking mode + # Helper function to filter special tokens from token IDs. + # This filtering is applied regardless of eval_text_turn_taking mode. def filter_special_tokens(token_ids): # Filter out pad token_ids = token_ids[token_ids != pad_id] @@ -89,21 +310,19 @@ def filter_special_tokens(token_ids): # Filter out special tokens before converting to text text_ids = filter_special_tokens(text_ids) start_idx = pos - timestamp = round(float(pos) * 0.08, 3) out_str.append(tokenizer.ids_to_text(text_ids)) if pos_type == 'bos': - out_str.append(f"<|{timestamp}|>") + out_str.append(format_bos_timestamp(pos)) elif pos_type == 'eos': - out_str.append(f"<${timestamp}$>") + out_str.append(format_eos_timestamp(pos)) else: # eot - out_str.append(f"<{timestamp}>") + out_str.append(format_eot_timestamp(pos)) # Filter the remaining tokens after the last position remaining_ids = filter_special_tokens(hyp_ids[start_idx:]) out_str.append(tokenizer.ids_to_text(remaining_ids)) ans.append(" ".join(out_str)) else: - # For non-turn-taking mode: filter out ALL special tokens, return only pure text hyp_ids = hyp_ids[:hyp_len] - hyp_ids = filter_special_tokens(hyp_ids) - ans.append(tokenizer.ids_to_text(hyp_ids)) + toks = tokenizer.ids_to_tokens(hyp_ids.tolist()) + ans.append(_decode_tokens_with_specials(toks, tokenizer, pad_token_str=pad_token_str, keep_pad=keep_pad)) return ans diff --git a/nemo/collections/speechlm2/streaming/duplex_stt_inference.py b/nemo/collections/speechlm2/streaming/duplex_stt_inference.py index 9c717a9a5d49..0c486606a9b1 100644 --- a/nemo/collections/speechlm2/streaming/duplex_stt_inference.py +++ b/nemo/collections/speechlm2/streaming/duplex_stt_inference.py @@ -103,9 +103,6 @@ def _init_inference( else: T = T_local - input_embeds = source_encoded.clone() - input_embeds *= self.model.cfg.get("duplex_user_channel_weight", 1.0) - use_cache = True if 'Nemotron' in self.model.cfg.pretrained_llm: cache = None @@ -120,6 +117,9 @@ def _init_inference( gen_asr = torch.empty(B, T, device=self.model.device, dtype=torch.long) else: gen_asr = None + gen_function = None + if self.model.use_function_head: + gen_function = torch.full((B, T), self.model.text_pad_id, device=self.model.device, dtype=torch.long) if prompt_tokens is not None and prompt_token_lens is not None: for i, prompt_len in enumerate(prompt_token_lens): @@ -129,11 +129,16 @@ def _init_inference( if self.model.predict_user_text: gen_asr[i, :prompt_len] = self.model.text_pad_id - input_embeds[:, 0] += self.model._get_bos_embedding() * self.model.cfg.get("duplex_text_channel_weight", 1.0) - if self.model.predict_user_text: - input_embeds[:, 0] += self.model._get_asr_bos_embedding() * self.model.cfg.get( - "duplex_asr_text_weight", 1.0 - ) + has_prompt = prompt_token_lens is not None and prompt_token_lens.max().item() > 0 + input_embeds = source_encoded.clone() + input_embeds[:, 0:1] = self.model.build_input_embedding( + source_encoded[:, 0:1], + 0, + gen_text, + gen_asr, + gen_function, + has_prompt=has_prompt, + ) start_gen_pos = 0 if prompt_token_lens is not None: @@ -154,12 +159,15 @@ def _init_inference( "B": B, "T": T, "T_local": T_local, + "source_encoded": source_encoded, "input_embeds": input_embeds, "cache": cache, "use_cache": use_cache, "gen_text": gen_text, "gen_asr": gen_asr, + "gen_function": gen_function, "start_gen_pos": start_gen_pos, + "has_prompt": has_prompt, "is_prompt_position_mask": is_prompt_position_mask, } @@ -174,14 +182,71 @@ def _step_zero(self, inference_state): inference_state["gen_text"][:, 0] = ans["text_logits"][:, -1].argmax(dim=-1) if self.model.predict_user_text: inference_state["gen_asr"][:, 0] = ans["asr_logits"][:, -1].argmax(dim=-1) + if self.model.use_function_head: + inference_state["gen_function"][:, 0] = ans["function_logits"][:, -1].argmax(dim=-1) return ans, inference_state - def _maybe_apply_forced_turn_taking(self, t, inference_state, is_prompt_position): - """Apply forced turn-taking rules based on ASR channel tokens.""" + def _maybe_apply_forced_turn_taking( + self, + t, + gen_text, + gen_asr, + is_prompt_position=None, + ): + """Apply forced turn-taking rules based on ASR channel tokens. + + This mutates gen_text in place when a rule fires. + + This supports two types of turn-taking corrections: + + 1. User stopped talking -> model should start talking. + (a) If the ASR stream shows enough user silence after user text tokens, or + (b) user EOS, + then force agent BOS in the text stream. + + 2. User started talking -> model should stop talking. + (a) If the ASR stream shows user BOS, + then force agent EOS in the text stream. + + There are two parameters that control the behavior, which will be read from + self.model.cfg: + - force_turn_taking_threshold: how far back to look in the agent text stream + before allowing to force another BOS/EOS. + - force_turn_taking_pad_window: how many recent ASR tokens must all be PAD + before treating the user as silent (i.e. only relevant to type 1a described above). + The token before the pad window must be user text, not PAD or BOS; this + prevents startup silence after user BOS from making the agent start. + + Note: Runtime wrappers should sync their force_turn_taking* overrides into + self.model.cfg before calling this shared helper. + + Examples below of each forcing case. They assume the same time of forced turn-taking + did not happen in force_turn_taking_threshold, so forcing is allowed. + + 1a. User silence after text -> force agent BOS, with pad_window=3: + time step: t-4 t-3 t-2 t-1 t + ASR: non-special token PAD PAD PAD PAD + text: PAD PAD PAD PAD BOS + + 1b. User EOS -> force agent BOS: + time step: ... t + ASR: ... EOS + text: ... BOS + + 2a. User BOS -> force agent EOS: + time step: ... t + ASR: ... BOS + text: ... EOS + + """ if not self.model.cfg.get("force_turn_taking", False): return + B = gen_text.size(0) + if is_prompt_position is None: + is_prompt_position = torch.zeros(B, dtype=torch.bool, device=gen_text.device) + threshold = self.model.cfg.get("force_turn_taking_threshold", 40) pad_window_steps = self.model.cfg.get("force_turn_taking_pad_window", 25) @@ -200,64 +265,71 @@ def _maybe_apply_forced_turn_taking(self, t, inference_state, is_prompt_position else: legacy_user_eos_id = None - for batch_idx in range(inference_state["B"]): + for batch_idx in range(B): if is_prompt_position[batch_idx]: continue lookback_start = max(0, t - threshold) - agent_text_window = inference_state["gen_text"][batch_idx, lookback_start:t] - current_asr_token = inference_state["gen_asr"][batch_idx, t] + agent_text_window = gen_text[batch_idx, lookback_start:t] + current_asr_token = gen_asr[batch_idx, t] - # ASR EOS or ~1 sec of pad tokens → insert agent BOS if not present in window - # Skip if we don't have enough tokens at the beginning + # If we have not accumulated a full pad window yet, silence is not meaningful. if t < pad_window_steps: continue pad_lookback_start = t - pad_window_steps - asr_recent_tokens = inference_state["gen_asr"][batch_idx, pad_lookback_start:t] + asr_recent_tokens = gen_asr[batch_idx, pad_lookback_start:t] has_pad_window = ( (asr_recent_tokens == self.model.text_pad_id).all() if len(asr_recent_tokens) > 0 else False ) - # Require that the pad window starts after a non-pad token + # 1a. Silence-based user stopped trigger: the pad window must follow user text. if has_pad_window and pad_lookback_start > 0: - token_before_window = inference_state["gen_asr"][batch_idx, pad_lookback_start - 1] - has_pad_window = token_before_window != self.model.text_pad_id + token_before_window = gen_asr[batch_idx, pad_lookback_start - 1] + has_pad_window = (token_before_window != self.model.text_pad_id) and ( + token_before_window != self.model.text_bos_id + ) elif has_pad_window and pad_lookback_start == 0: # If the pad window starts at position 0, it doesn't meet the requirement has_pad_window = False - # Check for user EOS: either tokenizer.eos (new) or legacy user_eos (old models) + # 1b. Explicit user stopped trigger: ASR emits user EOS. is_user_eos = current_asr_token == self.model.tokenizer.eos if legacy_user_eos_id is not None: is_user_eos = is_user_eos or (current_asr_token == legacy_user_eos_id) - # Check for user BOS: either text_bos_id (new) or legacy user_bos (old models) + # 2. Explicit user started trigger: ASR emits user BOS. is_user_bos = current_asr_token == self.model.text_bos_id if legacy_user_bos_id is not None: is_user_bos = is_user_bos or (current_asr_token == legacy_user_bos_id) if is_user_eos or has_pad_window: - # User has finished talking or remains silent for a while + # 1a/1b. User stopped speaking: force agent BOS unless already started recently. if not (agent_text_window == self.model.text_bos_id).any(): - inference_state["gen_text"][batch_idx, t] = self.model.text_bos_id + gen_text[batch_idx, t] = self.model.text_bos_id + reason = "user EOS" if is_user_eos else "ASR pad window" + logging.debug( + f"Forced turn-taking at frame {t}, batch {batch_idx}: " + f"inserted agent BOS (reason: {reason})" + ) elif is_user_bos: - # User has started talking but agent has not stopped yet + # 2. User started speaking: force agent EOS unless already stopped recently. if not (agent_text_window == self.model.text_eos_id).any(): - inference_state["gen_text"][batch_idx, t] = self.model.text_eos_id + gen_text[batch_idx, t] = self.model.text_eos_id + logging.debug( + f"Forced turn-taking at frame {t}, batch {batch_idx}: " "inserted agent EOS (reason: user BOS)" + ) def _step_inference(self, t, inference_state, ans): """Perform inference for one step t in the autoregressive loop.""" - last_emb = self.model.embed_tokens(inference_state["gen_text"][:, t - 1]) * self.model.cfg.get( - "duplex_text_channel_weight", 1.0 + inference_state["input_embeds"][:, t : t + 1] = self.model.build_input_embedding( + inference_state["source_encoded"][:, t : t + 1], + t, + inference_state["gen_text"], + inference_state["gen_asr"], + inference_state["gen_function"], + has_prompt=inference_state["has_prompt"], ) - if self.model.predict_user_text: - last_asr_emb = self.model.embed_asr_tokens(inference_state["gen_asr"][:, t - 1]) * self.model.cfg.get( - "duplex_asr_text_weight", 1.0 - ) - last_emb += last_asr_emb - - inference_state["input_embeds"][:, t] += last_emb is_prompt_position = inference_state["is_prompt_position_mask"][:, t] @@ -288,7 +360,15 @@ def _step_inference(self, t, inference_state, ans): inference_state["gen_asr"][:, t] = torch.where( is_prompt_position, inference_state["gen_asr"][:, t], generated_asr ) - self._maybe_apply_forced_turn_taking(t, inference_state, is_prompt_position) + self._maybe_apply_forced_turn_taking( + t, inference_state["gen_text"], inference_state["gen_asr"], is_prompt_position + ) + + if self.model.use_function_head and not is_prompt_position.all(): + generated_function = ans["function_logits"][:, -1].argmax(dim=-1) + inference_state["gen_function"][:, t] = torch.where( + is_prompt_position, inference_state["gen_function"][:, t], generated_function + ) return ans diff --git a/pyproject.toml b/pyproject.toml index ab6680b9fe0d..b5c2c010ebfa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -350,8 +350,25 @@ speechlm2 = [ [tool.setuptools] py-modules = ["nemo"] +# vllm-omni spawns its stage children with ``multiprocessing`` start method +# ``spawn``, and spawned children inherit no Python state, so registering the +# NemotronDuplexH + EarTTS models and the pipeline in the parent is not enough. +# ``nemo_voicechat`` is declared in vLLM's group, not vllm-omni's, because vLLM +# loads its plugins earlier: vllm-omni resolves the pipeline for ``model_type`` +# while constructing AsyncOmniEngine, before it loads its own group. +# Registration is idempotent and no-ops when vllm-omni is absent, so an ordinary +# vLLM process pays only the import. +# Entry points are only discoverable from an installed distribution, so NeMo has +# to be pip-installed (``pip install -e .`` is enough); PYTHONPATH will not do. [project.entry-points."vllm.general_plugins"] nemo_speechlm = "nemo.collections.speechlm2.vllm.salm:register" +nemo_voicechat = "nemo.collections.speechlm2.inference.vllm_omni.register:register_nemo_voicechat" + +# Also declared in vllm-omni's own group, which is loaded in the stage children +# and worker processes. Harmless duplication: whichever loader runs first wins +# and the second call returns immediately. +[project.entry-points."vllm_omni.general_plugins"] +nemo_voicechat = "nemo.collections.speechlm2.inference.vllm_omni.register:register_nemo_voicechat" [project.urls] Download = "https://github.com/NVIDIA-NeMo/Speech/releases" diff --git a/tests/collections/speechlm2/nemo_inference_pipelines/conftest.py b/tests/collections/speechlm2/nemo_inference_pipelines/conftest.py new file mode 100644 index 000000000000..7cbbd14a3677 --- /dev/null +++ b/tests/collections/speechlm2/nemo_inference_pipelines/conftest.py @@ -0,0 +1,461 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared fixtures for streaming VoiceChat inference tests. + +Most GPU tests build a tiny random-weight checkpoint (fast enough for CI). +A few integration tests load ``nvidia/NVIDIA-NemotronLabs-VoiceChat-11B``, +downloading it into the Hugging Face cache if needed. + +Toy-weight training/offline tests live in +``tests/collections/speechlm2/test_voicechat.py`` and do not use this +conftest. +""" + +from __future__ import annotations + +import gc +import json +import logging +import os + +# nemotron_voicechat_pipeline_{parity,nocrash} tests set +# torch.use_deterministic_algorithms(True), which requires CuBLAS to have a +# deterministic workspace. CuBLAS reads this env var only once — at +# initialization (first CUDA matmul in the process) — so it must be set here, +# before any fixture or test triggers CUDA work. The setting is harmless for +# non-deterministic tests: it only reserves 32 KB of extra GPU workspace and +# has no effect unless deterministic mode is active. +os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + +from contextlib import ExitStack + +import numpy as np +import pytest +import soundfile as sf +import torch +from omegaconf import OmegaConf + +from nemo.collections.audio.parts.utils.transforms import resample +from nemo.collections.speechlm2.inference.factory.s2s_pipeline_builder import S2SPipelineBuilder +from nemo.collections.speechlm2.inference.model_wrappers.engine_selection import inference_precision_from_cfg +from nemo.collections.speechlm2.inference.pipelines.streaming_s2s_pipeline import StreamingS2SPipeline +from nemo.collections.speechlm2.models import NemotronVoiceChat +from nemo.collections.speechlm2.models.duplex_ear_tts import load_audio_librosa + +_pretrained_llm = "TinyLlama/TinyLlama_v1.1" +if os.path.exists("/home/TestData/speechlm/pretrained_models"): + _pretrained_llm = "/home/TestData/speechlm/pretrained_models/TinyLlama--TinyLlama_v1.1" + +# The config the example launcher ships, so tests exercise what users run. +CONF_YAML = os.path.abspath( + os.path.join( + os.path.dirname(__file__), + "../../../../examples/speechlm2/nemo_inference_pipelines/conf/s2s_streaming.yaml", + ) +) +_FORCE_ALIGN_AUDIO = os.path.abspath( + os.path.join(os.path.dirname(__file__), "..", "test_data", "force_align_test.mp3") +) +HF_VOICECHAT_11B = "nvidia/NVIDIA-NemotronLabs-VoiceChat-11B" +# Speaker name registered in the public 11B. The tiny checkpoint registers a +# random latent under the same name so every test passes the same +# ``s2s.speaker_name``. +DEFAULT_SPEAKER_NAME = "Aria" + + +def _merge_pipeline_cfg(model_path: str, audio_path: str, output_dir: str, *overrides: dict): + """Shipped config with *overrides* merged on top, in order. + + OmegaConf merges nested dicts recursively, so ``{"s2s": {"top_p": 0.9}}`` + overrides only that key. + """ + cfg = OmegaConf.merge( + OmegaConf.load(CONF_YAML), + {"audio_file": audio_path, "output_dir": output_dir, "s2s": {"model_path": model_path}}, + ) + for override in overrides: + if override: + cfg = OmegaConf.merge(cfg, override) + return cfg + + +_EMPTY_CACHE_AFTER_GIB = 16 + + +def _reclaim_gpu_after_large_load() -> None: + """Return GPU memory to the driver after a large native load. + + ``pipeline.shutdown`` only tears down the vLLM runtime. Native weights stay + on the wrapper until the pipeline is unreachable; PyTorch then keeps the + blocks in this process. vLLM engine cores are child processes and treat + that as used memory, so a native 11B test (~24 GiB) followed by vLLM/vLLM + OOMs on an 80 GiB card. + + Tiny-model tests are a few GiB and rebuild from the same size, so the + caching allocator is left warm. Both ``gc.collect`` and ``empty_cache`` + run only when reserved memory is still large — the 11B case, not the + nocrash sweep. Call after the pipeline has gone out of scope; collecting + while it is still a live local cannot free the weights. + """ + if not torch.cuda.is_available(): + return + reserved_gib = torch.cuda.memory_reserved() / 1024**3 + if reserved_gib >= _EMPTY_CACHE_AFTER_GIB: + gc.collect() + torch.cuda.empty_cache() + logging.info( + "GPU reclaim: %.1f GiB reserved -> %.1f GiB", + reserved_gib, + torch.cuda.memory_reserved() / 1024**3, + ) + + +@pytest.fixture +def build_pipeline(): + """Factory fixture that builds a pipeline scoped to the test. + + Holds both scopes the production callers hold, on an ``ExitStack`` so they + last to the end of the test: the precision globals (without which one + ``deterministic=true`` test would leave the whole session in deterministic + mode with seeded RNGs and the fast attention kernels off) and + ``pipeline.shutdown``, which releases any vLLM runtime. After the stack + unwinds the pipeline is unreachable, so a large leftover CUDA reservation + can be returned to the driver (see ``_reclaim_gpu_after_large_load``). + + A fixture rather than an import because pytest loads these test modules as + top-level modules with no parent package, so they cannot import from + ``conftest`` directly. + """ + with ExitStack() as stack: + + def build(model_path: str, audio_path: str, output_dir: str, *overrides: dict) -> StreamingS2SPipeline: + cfg = _merge_pipeline_cfg(model_path, audio_path, output_dir, *overrides) + stack.enter_context(inference_precision_from_cfg(cfg.s2s)) + pipeline = S2SPipelineBuilder.build_pipeline(cfg) + stack.callback(pipeline.shutdown) + return pipeline + + yield build + _reclaim_gpu_after_large_load() + + +def _tiny_voicechat_config( + *, + log_dir: str, + predict_user_text: bool = True, + streaming_encoder: bool = False, + use_function_head: bool = False, +) -> dict: + """Return a minimal NemotronVoiceChat config with random weights. + + Args: + log_dir: Base directory for the exp_manager and validation outputs. + Pass a per-test temporary directory so runs cannot collide. + predict_user_text: Enable ASR head for user text prediction. + streaming_encoder: When True, configure the conformer encoder for + cache-aware streaming (causal convolutions, chunked_limited + attention) matching the real checkpoint. When False, use + default (non-causal) settings suitable for offline tests. + """ + duplex_stt_log_dir = os.path.join(log_dir, "duplex_stt") + parity_log_dir = os.path.join(log_dir, "parity") + encoder_cfg: dict = { + "_target_": "nemo.collections.asr.modules.ConformerEncoder", + "feat_in": 80, + "d_model": 512, + "n_heads": 8, + "n_layers": 1, + "subsampling_factor": 8, + } + if streaming_encoder: + encoder_cfg.update( + { + "subsampling": "dw_striding", + "causal_downsampling": True, + "att_context_size": [70, 0], + "att_context_style": "chunked_limited", + "conv_kernel_size": 9, + "conv_context_size": "causal", + } + ) + + return { + "model": { + "scoring_asr": "stt_en_fastconformer_transducer_large", + "stt": { + "model": { + "pretrained_llm": _pretrained_llm, + "pretrained_weights": False, + "predict_user_text": predict_user_text, + "use_function_head": use_function_head, + "audio_loss_weight": 1, + "text_loss_weight": 3, + "duplex_function_channel_weight": 2.0, + "source_sample_rate": 16000, + "validation_save_path": duplex_stt_log_dir, + "perception": { + "_target_": "nemo.collections.speechlm2.modules.perception.AudioPerceptionModule", + "preprocessor": { + "_target_": "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor", + "features": 80, + }, + "encoder": encoder_cfg, + "modality_adapter": { + "_target_": "nemo.collections.speechlm2.modules.perception.IdentityConnector", + "d_model": 512, + }, + "output_dim": 2048, + }, + "optimizer": {"_target_": "torch.optim.AdamW"}, + }, + "data": {"source_sample_rate": 16000}, + "exp_manager": {"explicit_log_dir": duplex_stt_log_dir}, + }, + "speech_generation": { + "model": { + "pretrained_lm_name": _pretrained_llm, + "pretrained_ae_dir": None, + "pretrained_tts_model": None, + "scoring_asr": "stt_en_fastconformer_transducer_large", + "freeze_params": [r"^audio_codec\..+$", r"^embed_tokens\..+$"], + "bos_token": "", + "eos_token": "", + "pad_token": "", + "audio_codec_run_dtype": "float32", + "prevent_freeze_params": [], + "audio_save_path": "", + "inference_guidance_scale": 0.5, + "inference_noise_scale": 0.8, + "inference_top_p_or_k": 0.8, + "inference_guidance_enabled": False, + "subword_mask_exactly_as_eartts": False, + "context_hidden_mask_exactly_as_eartts": False, + "optimizer": { + "_target_": "torch.optim.AdamW", + "lr": 4e-5, + "betas": [0.9, 0.98], + "weight_decay": 0, + "foreach": True, + }, + "lr_scheduler": { + "_target_": "nemo.core.optim.lr_scheduler.InverseSquareRootAnnealing", + "warmup_steps": 2500, + "min_lr": 1e-6, + "max_steps": 100_000_000, + }, + "codec_config": { + "latent_size": 512, + "n_fft": 16, + "hop_length": 4, + "base_hidden_size": 384, + "channel_mult": [1, 2, 4], + "rates": [7, 7, 9], + "num_blocks": 3, + "kernel_size": 7, + "groups": 1, + "codebook_size": 1024, + "num_quantizers": 31, + "wav_to_token_ratio": 1764, + }, + "tts_config": { + # Required to construct audio_prompt_projection_W and + # register the fixture's speaker latent. + "use_audio_prompt_frozen_projection": True, + "use_gated_fusion_for_text_audio": True, + "disable_eos_prediction": True, + "use_bos_eos_emb": True, + "use_subword_flag_emb": True, + "num_delay_speech_tokens": 2, + "backbone_type": "gemma3_text", + "backbone_model_class": None, + "backbone_config_class": None, + "backbone_config": { + "hidden_size": 1152, + "intermediate_size": 4608, + "num_hidden_layers": 1, + "num_attention_heads": 16, + "num_key_value_heads": 16, + "head_dim": 72, + "attention_dropout": 0.1, + "use_cache": False, + }, + "latent_size": 512, + "codebook_size": 1024, + "num_quantizers": 31, + "context_hidden_size": None, + "cas_config": { + "backbone_type": "t5gemma", + "backbone_model_class": None, + "backbone_config_class": None, + "backbone_config": { + "is_encoder_decoder": False, + "encoder": { + "hidden_size": 1152, + "intermediate_size": 4608, + "num_hidden_layers": 1, + "num_attention_heads": 16, + "num_key_value_heads": 16, + "head_dim": 72, + "use_cache": False, + "attention_dropout": 0.1, + }, + }, + }, + "mog_head_config": { + "intermediate_size": 4608, + "num_layers": 3, + "low_rank": 64, + "num_predictions": 1024, + "min_log_std": -4.0, + "eps": 1e-6, + }, + "p_uncond": 0.1, + "label_smoothing": 0.01, + "max_training_rate": 0.8, + "quantizer_dropout": 0.5, + "random_target_masking": False, + "exponent": 3.0, + }, + }, + "data": { + "add_text_bos_and_eos_in_each_turn": True, + "add_audio_prompt": True, + "audio_prompt_duration": 3.0, + "frame_length": 0.08, + "source_sample_rate": 16000, + "target_sample_rate": 22050, + }, + "exp_manager": {"explicit_log_dir": duplex_stt_log_dir}, + }, + }, + "data": { + "frame_length": 0.08, + "source_sample_rate": 16000, + "target_sample_rate": 22050, + "input_roles": ["user", "User"], + "output_roles": ["agent", "Assistant", "assistant", "Agent"], + }, + "exp_manager": {"explicit_log_dir": parity_log_dir}, + } + + +def _build_tiny_model_artifacts(base, *, predict_user_text: bool, use_function_head: bool): + if not torch.cuda.is_available(): + pytest.skip("building the tiny checkpoint requires a GPU") + + audio_path = str(base / "test_audio.wav") + sf.write(audio_path, np.random.RandomState(42).randn(3 * 16000).astype(np.float32), 16000) + + speaker_ref_path = str(base / "speaker_ref.wav") + sf.write(speaker_ref_path, np.random.RandomState(99).randn(22050).astype(np.float32), 22050) + + cfg = _tiny_voicechat_config( + log_dir=str(base / "logs"), + predict_user_text=predict_user_text, + streaming_encoder=True, + use_function_head=use_function_head, + ) + model = NemotronVoiceChat(cfg) + model.to("cuda") + model.eval() + + speaker_audio, sr = load_audio_librosa(speaker_ref_path) + speaker_audio = resample(speaker_audio, sr, model.tts_model.target_sample_rate).to(model.device) + speaker_audio_lens = torch.tensor([speaker_audio.size(1)]).long().repeat(speaker_audio.size(0)).to(model.device) + with torch.no_grad(): + model.tts_model.set_audio_prompt_lantent( + speaker_audio, + speaker_audio_lens, + system_prompt=None, + batch_size=1, + name=DEFAULT_SPEAKER_NAME, + ) + + model_dir = str(base / "model") + model.save_pretrained(model_dir) + + # save_pretrained writes the tokenizer to llm_artifacts/, but config.json + # still references the HF hub name (e.g. "TinyLlama/TinyLlama_v1.1"). + # Save the LLM model config alongside the tokenizer so llm_artifacts/ + # is a complete local model reference, then rewrite config.json to point + # at it. This avoids HuggingFace network requests on every from_pretrained. + llm_artifacts = os.path.join(model_dir, "llm_artifacts") + model.stt_model.llm.config.save_pretrained(llm_artifacts) + cfg["model"]["stt"]["model"]["pretrained_llm"] = llm_artifacts + cfg["model"]["speech_generation"]["model"]["pretrained_lm_name"] = llm_artifacts + with open(os.path.join(model_dir, "config.json"), "w") as f: + json.dump(cfg, f) + + del model + torch.cuda.empty_cache() + + return model_dir, audio_path, DEFAULT_SPEAKER_NAME + + +@pytest.fixture(scope="session") +def tiny_model_artifacts(tmp_path_factory): + """Build the existing ASR-head tiny checkpoint used by pipeline tests.""" + return _build_tiny_model_artifacts( + tmp_path_factory.mktemp("tiny_model"), + predict_user_text=True, + use_function_head=False, + ) + + +@pytest.fixture(scope="session") +def tiny_function_model_artifacts(tmp_path_factory): + """Build a public-VoiceChat-style checkpoint: function head, no ASR head.""" + return _build_tiny_model_artifacts( + tmp_path_factory.mktemp("tiny_function_model"), + predict_user_text=False, + use_function_head=True, + ) + + +@pytest.fixture(scope="session") +def hf_voicechat_11b(): + """``nvidia/NVIDIA-NemotronLabs-VoiceChat-11B``, downloaded into ``HF_HOME`` if needed.""" + from huggingface_hub import snapshot_download + + return snapshot_download(HF_VOICECHAT_11B) + + +@pytest.fixture(scope="session") +def voicechat_audio_path(): + return _FORCE_ALIGN_AUDIO + + +@pytest.fixture(scope="session") +def voicechat_speaker_name(): + return DEFAULT_SPEAKER_NAME + + +@pytest.fixture(scope="session") +def real_vllm_omni_wrapper(tmp_path_factory, hf_voicechat_11b): + """Convert the public 11B snapshot for vLLM-Omni. + + ``NEMO_VLLM_WRAPPER_DIR`` is an optional prebuilt wrapper directory. + ``build_wrapper_checkpoint`` reuses it when complete. Otherwise one is + built under tmp. + """ + pytest.importorskip("vllm_omni") + if not torch.cuda.is_available(): + pytest.skip("converting the vLLM-Omni wrapper requires a GPU") + + from nemo.collections.speechlm2.inference.vllm_omni.checkpoint import build_wrapper_checkpoint + + wrapper_dir = os.environ.get("NEMO_VLLM_WRAPPER_DIR") or str(tmp_path_factory.mktemp("vllm_omni_wrapper")) + build_wrapper_checkpoint(hf_voicechat_11b, wrapper_dir) + return hf_voicechat_11b, wrapper_dir diff --git a/tests/collections/speechlm2/nemo_inference_pipelines/test_config_overrides.py b/tests/collections/speechlm2/nemo_inference_pipelines/test_config_overrides.py new file mode 100644 index 000000000000..91bc59a1e6d5 --- /dev/null +++ b/tests/collections/speechlm2/nemo_inference_pipelines/test_config_overrides.py @@ -0,0 +1,193 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Model config overrides and the logit boosts both backends share.""" + +from types import SimpleNamespace + +import pytest +import torch +from omegaconf import OmegaConf + +from nemo.collections.speechlm2.inference.model_wrappers.config_overrides import ( + COMPONENT_OF, + LLM, + TTS, + VLLM_FORCES_TRUE, + VLLM_IGNORES, + apply_model_cfg_overrides, +) +from nemo.collections.speechlm2.parts.logit_boosts import LogitBoosts, apply_logit_boosts + + +@pytest.fixture +def model(): + """VoiceChat stand-in exposing the two config objects overrides land in. + + Real ``DictConfig`` instances, because bridging goes through + ``OmegaConf.update``, which only accepts an OmegaConf container. + """ + return SimpleNamespace( + stt_model=SimpleNamespace(cfg=OmegaConf.create({})), + tts_model=SimpleNamespace(cfg=OmegaConf.create({})), + ) + + +@pytest.fixture +def warnings(monkeypatch): + from nemo.collections.speechlm2.inference.model_wrappers import config_overrides + + recorded: list[str] = [] + monkeypatch.setattr(config_overrides.logging, "warning", recorded.append) + return recorded + + +def test_overrides_land_in_the_config_their_consumer_reads(model, warnings): + """Each key reaches the submodel that reads it, and absent keys are left + alone so whatever the checkpoint carries stays in effect.""" + model.stt_model.cfg["inference_pad_boost"] = 1.5 + + effective = apply_model_cfg_overrides( + model, + { + "inference_user_pad_boost": 0.8, + "force_turn_taking": True, + "inference_top_p_or_k": 0.7, + }, + llm_engine_type="native", + tts_engine_type="native", + ) + + # DuplexSTTModel reads its own cfg; DuplexEARTTS reads its own. + assert model.stt_model.cfg["inference_user_pad_boost"] == 0.8 + assert model.stt_model.cfg["force_turn_taking"] is True + assert model.tts_model.cfg["inference_top_p_or_k"] == 0.7 + # Untouched by this call, and still reported as the effective value. + assert model.stt_model.cfg["inference_pad_boost"] == 1.5 + assert effective["inference_pad_boost"] == 1.5 + assert warnings == [] + + +@pytest.mark.parametrize( + ("overrides", "tts_engine_type", "expected"), + [ + # Boosts and turn-taking work on both backends, so they stay quiet. + ({"inference_user_pad_boost": 0.8, "inference_pad_boost": 0.3, "force_turn_taking": True}, "vllm_omni", None), + # vLLM EarTTS takes sampling from the converted checkpoint instead. + ({"inference_noise_scale": 0.9}, "vllm_omni", "inference_noise_scale"), + ({"inference_noise_scale": 0.0}, "vllm_omni", "inference_noise_scale"), + ({"inference_top_p_or_k": 0.0}, "vllm_omni", "inference_top_p_or_k"), + ({"inference_guidance_scale": 0.0}, "vllm_omni", "inference_guidance_scale"), + ({"inference_noise_scale": 0.9}, "native", None), + # It forces codec silence on EOS unconditionally: True is honoured, + # False cannot be. + ({"inference_force_speech_silence_on_eos": False}, "vllm_omni", "force_speech_silence"), + ({"inference_force_speech_silence_on_eos": True}, "vllm_omni", None), + ], +) +def test_a_backend_reports_exactly_the_keys_it_ignores(model, warnings, overrides, tts_engine_type, expected): + """No silent no-ops, and no noise about keys that do work.""" + apply_model_cfg_overrides( + model, + overrides, + llm_engine_type="native", + tts_engine_type=tts_engine_type, + ) + + if expected is None: + assert warnings == [] + else: + assert any(expected in warning for warning in warnings) + assert any(f"tts_engine_type={tts_engine_type}" in warning for warning in warnings) + + +@pytest.mark.parametrize( + ("overrides", "engines", "expected"), + [ + # Wrapper-consumed knobs are reported by the same table as the + # model-consumed ones, so there is one place to look and one format. + ({"use_llm_cache": True}, ("vllm_omni", "native"), "use_llm_cache"), + ({"use_llm_cache": True}, ("native", "native"), None), + ({"use_tts_torch_compile": True}, ("native", "vllm_omni"), "use_tts_torch_compile"), + ({"use_tts_subword_cache": True}, ("native", "vllm_omni"), "use_tts_subword_cache"), + # A falsy value is already a no-op, so it needs no warning. + ({"use_tts_torch_compile": False}, ("native", "vllm_omni"), None), + ], +) +def test_wrapper_knobs_report_through_the_same_table(model, warnings, overrides, engines, expected): + """Ignored wrapper knobs are reported through the same table as bridged keys.""" + llm_engine_type, tts_engine_type = engines + apply_model_cfg_overrides( + model, + overrides, + llm_engine_type=llm_engine_type, + tts_engine_type=tts_engine_type, + ) + + if expected is None: + assert warnings == [] + else: + assert any(expected in warning for warning in warnings) + + +def test_every_support_entry_is_well_formed_and_claimed_once(): + """Guards the table itself, which is now the single source of truth. + + A key in two tables would get two different verdicts, and a key bridged + into a model config must agree with that config's owning component -- both + are silent inconsistencies rather than crashes. + """ + tables = (VLLM_IGNORES, VLLM_FORCES_TRUE) + + seen: set[str] = set() + for table in tables: + for key, entry in table.items(): + component, why = entry + assert component in (LLM, TTS), f"{key} names an unknown component {component!r}" + assert why and not why.endswith("."), f"{key} reason is interpolated mid-sentence" + assert key not in seen, f"{key} appears in more than one support table" + seen.add(key) + # A bridged key must be attributed to the component whose config it + # is written into, or the warning names the wrong engine type. + if key in COMPONENT_OF: + assert COMPONENT_OF[key] == component + + +@pytest.mark.parametrize("shape", [(1, 2, 5), (5,)]) +def test_logit_boosts_are_read_and_applied_the_same_way_by_both_runtimes(shape): + """(B, T, V) for the PyTorch heads, (V,) for the vLLM logits processor.""" + from_mapping = LogitBoosts.agent_from_cfg({"inference_pad_boost": 0.8, "inference_eos_boost": None}) + assert (from_mapping.pad, from_mapping.bos, from_mapping.eos) == (0.8, None, None) + + # The converted Nemotron reads an HF PretrainedConfig, which has no .get(). + from_attrs = LogitBoosts.user_from_cfg( + SimpleNamespace( + inference_user_pad_boost=0.5, + inference_user_bos_boost=None, + inference_user_eos_boost=1.5, + ) + ) + assert (from_attrs.pad, from_attrs.bos, from_attrs.eos) == (0.5, None, 1.5) + + # Matches the truthiness gate the model heads have always used. + assert not LogitBoosts.agent_from_cfg({"inference_pad_boost": 0.0}) + + logits = torch.zeros(*shape) + apply_logit_boosts(logits, LogitBoosts(pad=0.5, eos=1.5), pad_id=0, bos_id=1, eos_id=2) + assert logits.reshape(-1, 5)[0].tolist() == [0.5, 0.0, 1.5, 0.0, 0.0] + + # Empty boosts touch nothing and need no token ids. + untouched = torch.zeros(*shape) + apply_logit_boosts(untouched, LogitBoosts(), pad_id=None, bos_id=None, eos_id=None) + assert untouched.reshape(-1, 5)[0].tolist() == [0.0] * 5 diff --git a/tests/collections/speechlm2/nemo_inference_pipelines/test_engine_selection.py b/tests/collections/speechlm2/nemo_inference_pipelines/test_engine_selection.py new file mode 100644 index 000000000000..f253e4c4470f --- /dev/null +++ b/tests/collections/speechlm2/nemo_inference_pipelines/test_engine_selection.py @@ -0,0 +1,77 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""What the s2s config resolves to, and which native weights that skips.""" + +import pytest + +from nemo.collections.speechlm2.inference.model_wrappers.engine_selection import ( + native_weight_skip_prefixes, + resolve_engine_types, +) + +NATIVE, VLLM = "native", "vllm_omni" + + +@pytest.mark.parametrize( + ("cfg", "engines", "skipped", "not_skipped"), + [ + ({}, (NATIVE, NATIVE), set(), {"stt_model.llm.", "tts_model.tts_model."}), + ( + {"llm_engine_type": VLLM, "tts_engine_type": VLLM}, + (VLLM, VLLM), + {"stt_model.llm.", "tts_model.tts_model."}, + set(), + ), + ( + {"llm_engine_type": NATIVE, "tts_engine_type": VLLM}, + (NATIVE, VLLM), + {"tts_model.tts_model."}, + {"stt_model.llm."}, + ), + ( + {"tts_engine_type": VLLM}, + (NATIVE, VLLM), + {"tts_model.tts_model."}, + {"stt_model.llm."}, + ), + ( + {"llm_engine_type": VLLM, "tts_engine_type": None}, + (VLLM, NATIVE), + {"stt_model.llm."}, + {"tts_model.tts_model."}, + ), + ], +) +def test_config_resolves_to_backends_and_the_weights_they_skip(cfg, engines, skipped, not_skipped): + """Each component key is independent; omitted keys default to native. + + A component on vLLM must also skip loading its native weights, so the two + decisions are checked together -- disagreeing would waste a full weight + load or, worse, leave a component with no weights at all. + """ + assert resolve_engine_types(cfg) == engines + + prefixes = native_weight_skip_prefixes(*engines) + assert skipped <= prefixes + assert prefixes.isdisjoint(not_skipped) + # The auxiliary RNN-T decoder is never needed by the streaming path. + assert {"stt_model.rnnt_decoder.", "stt_model.rnnt_joint."} <= prefixes + + +def test_unusable_engine_selection_is_named(): + with pytest.raises(ValueError, match="llm_engine_type='other'"): + resolve_engine_types({"llm_engine_type": "other"}) + with pytest.raises(ValueError, match="not a config key"): + resolve_engine_types({"engine_type": VLLM}) diff --git a/tests/collections/speechlm2/nemo_inference_pipelines/test_nemotron_voicechat_pipeline_nocrash.py b/tests/collections/speechlm2/nemo_inference_pipelines/test_nemotron_voicechat_pipeline_nocrash.py new file mode 100644 index 000000000000..71d554c02ef4 --- /dev/null +++ b/tests/collections/speechlm2/nemo_inference_pipelines/test_nemotron_voicechat_pipeline_nocrash.py @@ -0,0 +1,231 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""No-crash pipeline tests for NemotronVoiceChat streaming inference. + +The config sweep runs on a tiny random-weight model (CI). One extra case +loads ``nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`` (downloaded into the +Hugging Face cache if needed). + +Each test verifies only that the pipeline completes without raising — no +output quality checks. + +Run from the NeMo repo root:: + + CUDA_VISIBLE_DEVICES=0 pytest tests/collections/speechlm2/nemo_inference_pipelines/test_nemotron_voicechat_pipeline_nocrash.py -v -s +""" + +from __future__ import annotations + +import tempfile + +import pytest +import torch +from omegaconf import OmegaConf + +from nemo.collections.speechlm2.inference.model_wrappers.nemotron_voicechat_inference_wrapper import ( + NemotronVoicechatInferenceWrapper, +) +from nemo.collections.speechlm2.inference.utils.stepprogressbar import StepProgressBar + +MOCK_SYSTEM_PROMPT = "This is a mock prompt for the test" + +_TEST_DEFAULTS = { + "s2s": { + "llm_engine_type": "native", + "tts_engine_type": "native", + "compute_dtype": "float32", + "deterministic": False, + "decode_audio": False, + "use_perception_cache": False, + "use_perception_cudagraph": False, + "system_prompt": None, + "top_p": 1.0, + "repetition_penalty": 1.0, + "temperature": 1.0, + }, + "streaming": { + "chunk_size_in_secs": 0.08, + "buffer_size_in_secs": 71 * 0.08, + }, +} + +# --------------------------------------------------------------------------- +# Parametrized configs — each entry is a single overrides dict +# --------------------------------------------------------------------------- + +# Text-only configs (decode_audio=False): minimal STT-path smoke checks. +_TEXT_CONFIGS = [ + pytest.param({}, id="baseline"), + pytest.param( + {"s2s": {"use_perception_cache": True}}, + id="perception_cache", + ), + pytest.param({"pad_audio_by_sec": 2}, id="pad_by_sec"), +] + +# Audio configs (decode_audio=True): exercises the full STT + TTS pipeline. +_AUDIO_CONFIGS = [ + pytest.param({}, id="baseline"), + pytest.param( + { + "s2s": {"use_perception_cache": True, "system_prompt": MOCK_SYSTEM_PROMPT}, + "streaming": {"chunk_size_in_secs": 0.24}, + "pad_audio_to_sec": 5, + }, + id="perception_cache_prompt_multiframe_pad_to_sec", + ), + pytest.param( + { + "s2s": {"top_p": 0.9, "temperature": 0.7, "repetition_penalty": 1.1}, + "pad_silence_ratio": 0.5, + }, + id="sampling_pad_silence_ratio", + ), + pytest.param( + { + "s2s": {"use_tts_subword_cache": True, "use_tts_torch_compile": True}, + "pad_audio_by_sec": 2, + }, + id="tts_optimizations_pad_by_sec", + ), + pytest.param( + {"s2s": {"deterministic": True, "temperature": 0.0}}, + id="deterministic", + ), + pytest.param( + {"s2s": {"profile_timing": True}}, + id="profile_timing", + ), +] + + +def _run(pipeline, audio_path): + progress_bar = StepProgressBar.from_audio_filepaths( + [audio_path], + chunk_size_in_secs=pipeline.chunk_size_in_secs, + pad_audio_to_sec=pipeline.pad_audio_to_sec, + pad_silence_ratio=pipeline.pad_silence_ratio, + pad_audio_by_sec=pipeline.pad_audio_by_sec, + ) + result = pipeline.run([audio_path], progress_bar=progress_bar) + assert result is not None + return result + + +def test_speaker_reference_is_rejected(): + """Cloning from a wav is not a supported inference path.""" + cfg = OmegaConf.create( + { + "model_path": "unused", + "decode_audio": True, + "speaker_name": "Aria", + "speaker_reference": "/path/to/speaker.wav", + "llm_engine_type": "native", + "tts_engine_type": "native", + } + ) + with pytest.raises(ValueError, match="speaker_reference is not supported"): + NemotronVoicechatInferenceWrapper(cfg) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") +@pytest.mark.parametrize("overrides", _TEXT_CONFIGS) +def test_pipeline_no_crash(build_pipeline, tiny_model_artifacts, overrides): + """Run the streaming pipeline with various configs and verify it doesn't crash.""" + model_dir, audio_path, _ = tiny_model_artifacts + pipeline = build_pipeline( + model_dir, audio_path, tempfile.mkdtemp(prefix="no-crash-text-"), _TEST_DEFAULTS, overrides + ) + _run(pipeline, audio_path) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") +@pytest.mark.parametrize("overrides", _AUDIO_CONFIGS) +def test_pipeline_no_crash_decode_audio(build_pipeline, tiny_model_artifacts, overrides): + """Run the streaming pipeline with decode_audio=True and verify it doesn't crash.""" + model_dir, audio_path, speaker_name = tiny_model_artifacts + pipeline = build_pipeline( + model_dir, + audio_path, + tempfile.mkdtemp(prefix="no-crash-audio-"), + _TEST_DEFAULTS, + {"s2s": {"decode_audio": True, "speaker_name": speaker_name}}, + overrides, + ) + _run(pipeline, audio_path) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") +def test_pipeline_function_channel_without_asr(build_pipeline, tiny_function_model_artifacts): + """The public checkpoint's function channel is distinct from absent ASR.""" + model_dir, audio_path, _ = tiny_function_model_artifacts + pipeline = build_pipeline( + model_dir, + audio_path, + tempfile.mkdtemp(prefix="no-crash-function-no-asr-"), + _TEST_DEFAULTS, + {"s2s": {"decode_audio": False, "force_turn_taking": True}}, + ) + assert pipeline.s2s_model.model.stt_model.use_function_head + assert not pipeline.s2s_model.model.stt_model.predict_user_text + assert not pipeline.s2s_model.model.stt_model.cfg.force_turn_taking + + result = pipeline.run([audio_path]) + assert result is not None + assert result[0].token_asr_text is None + assert result[0].raw_asr_text is None + assert result[0].token_function is not None + assert result[0].raw_function_text is not None + assert result[0].capabilities.has_function_head + assert not result[0].capabilities.has_asr_head + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") +def test_pipeline_no_crash_hf_11b( + build_pipeline, hf_voicechat_11b, voicechat_audio_path, voicechat_speaker_name +): + """One native ``pipeline.run()`` on the public 11B: function channel, audio. + + The config sweep above stays on the tiny model so CI does not pay an 11B + load per case. This is the real-weight smoke check. + """ + pipeline = build_pipeline( + hf_voicechat_11b, + voicechat_audio_path, + tempfile.mkdtemp(prefix="no-crash-11b-"), + _TEST_DEFAULTS, + { + "s2s": { + "decode_audio": True, + "speaker_name": voicechat_speaker_name, + "system_prompt": MOCK_SYSTEM_PROMPT, + "force_turn_taking": True, + } + }, + ) + assert pipeline.s2s_model.model.stt_model.use_function_head + assert not pipeline.s2s_model.model.stt_model.predict_user_text + assert not pipeline.s2s_model.model.stt_model.cfg.force_turn_taking + + result = _run(pipeline, voicechat_audio_path) + output = result[0] + assert output.token_asr_text is None + assert output.raw_asr_text is None + assert output.token_function is not None + assert output.raw_function_text is not None + assert output.capabilities.has_function_head + assert not output.capabilities.has_asr_head + audio = output.audio_buffer + assert audio is not None and audio.numel() > 0 diff --git a/tests/collections/speechlm2/nemo_inference_pipelines/test_nemotron_voicechat_pipeline_parity.py b/tests/collections/speechlm2/nemo_inference_pipelines/test_nemotron_voicechat_pipeline_parity.py new file mode 100644 index 000000000000..991c53ecae2b --- /dev/null +++ b/tests/collections/speechlm2/nemo_inference_pipelines/test_nemotron_voicechat_pipeline_parity.py @@ -0,0 +1,323 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Offline vs. incremental inference parity tests for NemotronVoiceChat. + +``test_parity_tiny_model`` (cache × prompt) and +``test_parity_tiny_function_model_without_asr`` run on random-weight models. +``test_parity`` does one pass on ``nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`` +(downloaded into the Hugging Face cache if needed). + +Run from the NeMo repo root (use ``-s`` to see live progress):: + + CUDA_VISIBLE_DEVICES=0 pytest tests/collections/speechlm2/nemo_inference_pipelines/test_nemotron_voicechat_pipeline_parity.py -v -s +""" + +from __future__ import annotations + +import math +import tempfile +import time +from typing import Any + +import pytest +import torch +from omegaconf import OmegaConf + +from nemo.collections.speechlm2.inference.model_wrappers.nemotron_voicechat_inference_wrapper import ( + FRAME_SIZE_SAMPLES, + SAMPLE_RATE, +) +from nemo.collections.speechlm2.inference.pipelines.streaming_s2s_pipeline import StreamingS2SPipeline +from nemo.collections.speechlm2.inference.streaming.framing.s2s_request_options import S2SRequestOptions +from nemo.utils import logging + +MOCK_SYSTEM_PROMPT = "This is a mock prompt for the test" + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _compare_tensors( + a: torch.Tensor | None, + b: torch.Tensor | None, +) -> dict[str, Any]: + """Prefix-aware comparison of two tensors (tokens or logits, any shape).""" + if a is None or b is None: + return {"match": None, "note": "one or both tensors missing"} + a, b = a.detach().cpu().float(), b.detach().cpu().float() + T = min(a.shape[1], b.shape[1]) + if T == 0: + return {"prefix_len": 0, "match": True} + ap, bp = a[:, :T], b[:, :T] + diff = (ap - bp).abs() + match = bool(diff.max() == 0) + result: dict[str, Any] = {"prefix_len": T, "match": match, "max_abs_diff": float(diff.max())} + if not match: + reduce = tuple(i for i in range(diff.dim()) if i != 1) + per_step = diff.amax(dim=reduce) if reduce else diff.squeeze() + nonzero = (per_step > 0).nonzero(as_tuple=False) + if nonzero.numel(): + result["first_diff_step"] = int(nonzero[0].item()) + return result + + +def _merge_incremental_debug_steps(steps: list[dict[str, Any]]) -> dict[str, Any]: + """Merge per-step debug dicts from the pipeline into a single dict.""" + if not steps: + return {} + all_text_logits = [s["text_logits"] for s in steps if s.get("text_logits") is not None] + all_asr_logits = [s["asr_logits"] for s in steps if s.get("asr_logits") is not None] + return { + "text_logits": torch.cat(all_text_logits, dim=1) if all_text_logits else None, + "asr_logits": torch.cat(all_asr_logits, dim=1) if all_asr_logits else None, + } + + +def _load_and_pad_audio( + audio_path: str, + device: torch.device, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + """Load audio, zero-pad to a whole number of 80 ms frames, return ``(audio, lens)``.""" + import librosa + + audio_np, _ = librosa.load(audio_path, sr=SAMPLE_RATE) + padded_len = math.ceil(len(audio_np) / FRAME_SIZE_SAMPLES) * FRAME_SIZE_SAMPLES + audio = torch.nn.functional.pad( + torch.tensor(audio_np, device=device, dtype=dtype).unsqueeze(0), + (0, max(0, padded_len - len(audio_np))), + ) + return audio, torch.tensor([audio.shape[1]], device=device, dtype=torch.long) + + +def run_parity_check( + pipeline: StreamingS2SPipeline, + audio_path: str, + *, + system_prompt: str | None = None, +) -> dict[str, Any]: + """Run offline and incremental inference on the same audio, return comparison. + + Only STT-level tokens and logits are compared; TTS is irrelevant for + the core parity invariant. + """ + wrapper = pipeline.s2s_model + audio, audio_lens = _load_and_pad_audio(audio_path, wrapper.device, wrapper.dtype) + + prompt_tokens = prompt_token_lens = None + if system_prompt: + tok = wrapper.tokenizer + ids = [tok.bos_id] + tok.text_to_ids(system_prompt) + [tok.eos_id] + prompt_tokens = torch.tensor(ids, device=wrapper.device, dtype=torch.long).unsqueeze(0) + prompt_token_lens = torch.tensor([len(ids)], device=wrapper.device, dtype=torch.long) + + if wrapper.speaker_name is not None: + OmegaConf.update(wrapper.model.cfg, "inference_speaker_name", wrapper.speaker_name, force_add=True) + speaker_kw: dict[str, Any] = {} + if not wrapper.model.cfg.get("inference_speaker_name"): + speaker_kw["speaker_audio"] = torch.randn(1, 22050, device=wrapper.device) + speaker_kw["speaker_audio_lens"] = torch.tensor([22050], device=wrapper.device, dtype=torch.long) + + # -- Offline -- + logging.info("Running offline_inference ...") + t0 = time.time() + offline = wrapper.model.offline_inference( + input_signal=audio, + input_signal_lens=audio_lens, + prompt_tokens=prompt_tokens, + prompt_token_lens=prompt_token_lens, + decode_audio=False, + return_logits=True, + **speaker_kw, + ) + logging.info(f" offline done in {time.time() - t0:.2f}s") + + # -- Incremental -- + logging.info("Running incremental inference (pipeline.run) ...") + t0 = time.time() + pipeline.collect_debug = True + outputs = pipeline.run( + [audio_path], + options=[S2SRequestOptions(system_prompt=system_prompt)], + ) + logging.info(f" incremental done in {time.time() - t0:.2f}s") + + output = outputs[0] + inc_tokens = output.token_text + inc_asr_tokens = output.token_asr_text + assert output.debug_data, "collect_debug=True but no debug data was recorded" + inc_debug = _merge_incremental_debug_steps(output.debug_data) + + # -- Compare -- + # offline_inference returns logits for ALL positions (including prompt), + # while the incremental path only produces logits for audio positions. + # Trim the prompt prefix from offline logits so the two are aligned. + prompt_len = prompt_tokens.shape[1] if prompt_tokens is not None else 0 + + report: dict[str, Any] = { + "token_comparison": _compare_tensors(offline.get("tokens_text"), inc_tokens), + "asr_token_comparison": _compare_tensors(offline.get("tokens_text_src"), inc_asr_tokens), + } + for key, off_key, inc_key in [ + ("text_logit_comparison", "text_logits", "text_logits"), + ("asr_logit_comparison", "asr_logits", "asr_logits"), + ]: + off_t, inc_t = offline.get(off_key), inc_debug.get(inc_key) + if off_t is not None and prompt_len > 0: + off_t = off_t[:, prompt_len:] + if off_t is not None and inc_t is not None: + report[key] = _compare_tensors(off_t, inc_t) + + return report + + +def assert_parity( + report: dict[str, Any], + *, + strict: bool = True, + atol: float = 0.0, +) -> None: + """Raise ``AssertionError`` if parity checks in *report* fail.""" + failures: list[str] = [] + for key in ("token_comparison", "asr_token_comparison"): + c = report.get(key, {}) + if c.get("match") is False: + failures.append(f"{key}: diverge at step {c.get('first_diff_step')}") + if strict: + for key in ("text_logit_comparison", "asr_logit_comparison"): + c = report.get(key, {}) + if c.get("match") is False and c.get("max_abs_diff", 0) > atol: + failures.append(f"{key}: max_abs_diff={c['max_abs_diff']:.2e} > atol={atol:.2e}") + assert not failures, "Parity failed:\n " + "\n ".join(failures) + + +# Parity requires deterministic, float32, and greedy decoding. +_PARITY_DEFAULTS = { + "s2s": { + "llm_engine_type": "native", + "tts_engine_type": "native", + "compute_dtype": "float32", + "deterministic": True, + "decode_audio": False, + "use_perception_cache": False, + "use_perception_cudagraph": False, + "top_p": 1.0, + "repetition_penalty": 1.0, + "temperature": 1.0, + }, +} + + +def _build_parity_pipeline( + build_pipeline, + model_path: str, + audio_path: str, + output_dir: str, + *overrides: dict[str, Any], +) -> StreamingS2SPipeline: + """Build a pipeline configured for strict parity testing. + + The chunk size is set to cover the full audio in one step so that offline + and incremental paths see identical input. + """ + import librosa + + audio_np, _ = librosa.load(audio_path, sr=SAMPLE_RATE) + total_frames = math.ceil(len(audio_np) / FRAME_SIZE_SAMPLES) + chunk_secs = total_frames * FRAME_SIZE_SAMPLES / SAMPLE_RATE + + return build_pipeline( + model_path, + audio_path, + output_dir, + _PARITY_DEFAULTS, + { + "streaming": { + "chunk_size_in_secs": chunk_secs, + "buffer_size_in_secs": max(71 * 0.08, chunk_secs), + } + }, + *overrides, + ) + + +# --------------------------------------------------------------------------- +# Tiny random-weight models +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") +@pytest.mark.parametrize("use_llm_cache", [False, True], ids=["no_cache", "llm_cache"]) +@pytest.mark.parametrize("system_prompt", [None, MOCK_SYSTEM_PROMPT], ids=["no_prompt", "prompt"]) +def test_parity_tiny_model(build_pipeline, tiny_model_artifacts, use_llm_cache, system_prompt): + """Offline/incremental parity with a tiny random-weight model. + + Running both cache settings against the same offline reference is what + pins the invariant that the native KV cache is a speed path and not a + different model: if either matched offline and the other did not, one of + these cases would fail. + """ + model_dir, audio_path, _ = tiny_model_artifacts + pipeline = _build_parity_pipeline( + build_pipeline, + model_dir, + audio_path, + tempfile.mkdtemp(prefix="parity-tiny-"), + {"s2s": {"system_prompt": system_prompt, "use_llm_cache": use_llm_cache}}, + ) + report = run_parity_check(pipeline, audio_path, system_prompt=system_prompt) + assert_parity(report, strict=True, atol=0.0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") +def test_parity_tiny_function_model_without_asr(build_pipeline, tiny_function_model_artifacts): + """Function-token feedback matches offline inference without an ASR channel.""" + model_dir, audio_path, _ = tiny_function_model_artifacts + pipeline = _build_parity_pipeline( + build_pipeline, + model_dir, + audio_path, + tempfile.mkdtemp(prefix="parity-tiny-function-"), + {"s2s": {"system_prompt": MOCK_SYSTEM_PROMPT, "force_turn_taking": True}}, + ) + report = run_parity_check(pipeline, audio_path, system_prompt=MOCK_SYSTEM_PROMPT) + assert report["asr_token_comparison"]["match"] is None + assert_parity(report, strict=True, atol=0.0) + + +# --------------------------------------------------------------------------- +# Public 11B — one load, not the cache×prompt matrix +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") +def test_parity(build_pipeline, hf_voicechat_11b, voicechat_audio_path, voicechat_speaker_name): + """Offline/incremental parity on ``nvidia/NVIDIA-NemotronLabs-VoiceChat-11B``.""" + pipeline = _build_parity_pipeline( + build_pipeline, + hf_voicechat_11b, + voicechat_audio_path, + tempfile.mkdtemp(prefix="parity-11b-"), + { + "s2s": { + "system_prompt": MOCK_SYSTEM_PROMPT, + "speaker_name": voicechat_speaker_name, + } + }, + ) + report = run_parity_check(pipeline, voicechat_audio_path, system_prompt=MOCK_SYSTEM_PROMPT) + assert report["asr_token_comparison"]["match"] is None + assert_parity(report, strict=True, atol=0.0) diff --git a/tests/collections/speechlm2/nemo_inference_pipelines/test_nemotron_voicechat_pipeline_vllm.py b/tests/collections/speechlm2/nemo_inference_pipelines/test_nemotron_voicechat_pipeline_vllm.py new file mode 100644 index 000000000000..e5d188e27b4e --- /dev/null +++ b/tests/collections/speechlm2/nemo_inference_pipelines/test_nemotron_voicechat_pipeline_vllm.py @@ -0,0 +1,87 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""One live vLLM-Omni pipeline run on nvidia/NVIDIA-NemotronLabs-VoiceChat-11B. + +Skipped when ``vllm_omni`` is not installed. The 11B snapshot is downloaded +into the Hugging Face cache if needed. +""" + +from __future__ import annotations + +import tempfile + +import pytest +import torch + +pytest.importorskip("vllm_omni") + +from nemo.collections.speechlm2.inference.utils.stepprogressbar import StepProgressBar + +MOCK_SYSTEM_PROMPT = "This is a mock prompt for the test" +_VLLM = "vllm_omni" + +_VLLM_DEFAULTS = { + "s2s": { + "llm_engine_type": _VLLM, + "tts_engine_type": _VLLM, + "deterministic": False, + "decode_audio": True, + "system_prompt": MOCK_SYSTEM_PROMPT, + }, + "streaming": { + "chunk_size_in_secs": 0.08, + "buffer_size_in_secs": 71 * 0.08, + }, +} + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") +def test_pipeline_no_crash_vllm( + build_pipeline, hf_voicechat_11b, voicechat_audio_path, voicechat_speaker_name, real_vllm_omni_wrapper +): + """vLLM/vLLM ``pipeline.run()`` on the public 11B: text and audio exist.""" + _, wrapper_dir = real_vllm_omni_wrapper + output_dir = tempfile.mkdtemp(prefix="no-crash-vllm-") + + pipeline = build_pipeline( + hf_voicechat_11b, + voicechat_audio_path, + output_dir, + _VLLM_DEFAULTS, + { + "s2s": { + "speaker_name": voicechat_speaker_name, + "vllm_omni_config": {"wrapper_dir": wrapper_dir}, + } + }, + ) + wrapper = pipeline.s2s_model + assert wrapper.llm_engine_type == _VLLM + assert wrapper.tts_engine_type == _VLLM + + progress_bar = StepProgressBar.from_audio_filepaths( + [voicechat_audio_path], + chunk_size_in_secs=pipeline.chunk_size_in_secs, + pad_audio_to_sec=pipeline.pad_audio_to_sec, + pad_silence_ratio=pipeline.pad_silence_ratio, + pad_audio_by_sec=pipeline.pad_audio_by_sec, + ) + result = pipeline.run([voicechat_audio_path], progress_bar=progress_bar) + assert result is not None + assert len(result) == 1 + output = result[0] + assert output.token_text is not None and output.token_text.numel() > 0 + audio = output.audio_buffer + assert audio is not None and audio.numel() > 0 diff --git a/tests/collections/speechlm2/nemo_inference_pipelines/test_text_sampling.py b/tests/collections/speechlm2/nemo_inference_pipelines/test_text_sampling.py new file mode 100644 index 000000000000..bf9aebc32982 --- /dev/null +++ b/tests/collections/speechlm2/nemo_inference_pipelines/test_text_sampling.py @@ -0,0 +1,146 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from nemo.collections.speechlm2.inference.model_wrappers.text_sampling import sample_text_token + + +def test_special_tokens_bypass_sampling_and_never_penalise_history(monkeypatch): + """Special tokens are chosen greedily and excluded from the repetition + penalty, so a pad-heavy history cannot bias the next real token.""" + # All-ones parameters reduce to greedy. + assert sample_text_token( + torch.tensor([[0.1, 0.8, 0.3]]), + torch.empty((1, 0), dtype=torch.long), + 0, + top_p=1.0, + repetition_penalty=1.0, + temperature=1.0, + special_token_ids=set(), + ).tolist() == [1] + + penalty_kwargs = { + "top_p": 1.0, + "repetition_penalty": 1.5, + "temperature": 0.8, + "special_token_ids": {0}, + "special_ids_tensor": torch.tensor([0]), + } + logits = torch.tensor([[0.2, 1.1, 0.9, 0.4]]) + torch.manual_seed(9) + with_special_history = sample_text_token(logits, torch.tensor([[0]], dtype=torch.long), 1, **penalty_kwargs) + torch.manual_seed(9) + without_history = sample_text_token(logits, torch.empty((1, 0), dtype=torch.long), 0, **penalty_kwargs) + assert torch.equal(with_special_history, without_history) + + # A special token wins outright rather than going through multinomial. + def fail_multinomial(*_args, **_kwargs): + raise AssertionError("special-token bypass called torch.multinomial") + + monkeypatch.setattr(torch, "multinomial", fail_multinomial) + assert sample_text_token( + torch.tensor([[0.1, 1.5, 0.7]]), + torch.tensor([[2]], dtype=torch.long), + 1, + top_p=0.8, + repetition_penalty=1.2, + temperature=0.7, + special_token_ids={1}, + special_ids_tensor=torch.tensor([1]), + ).tolist() == [1] + + +def test_vllm_processor_reuses_shared_sampler_and_skips_prefill_history(): + pytest.importorskip("vllm") + from nemo.collections.speechlm2.inference.vllm_omni.nemotron_duplex_h.sampling import ( + SharedTextRequestSampler, + SharedTextSamplingState, + ) + + logits = torch.tensor([0.2, 1.3, 0.9, 0.5]) + params = { + "top_p": 0.9, + "repetition_penalty": 1.2, + "temperature": 0.75, + "special_token_ids": {0}, + } + + torch.manual_seed(17) + expected = sample_text_token( + logits.unsqueeze(0), + torch.tensor([[2, 2]], dtype=torch.long), + 2, + special_ids_tensor=torch.tensor([0]), + **params, + ) + + processor = SharedTextRequestSampler( + history_skip=1, + state=SharedTextSamplingState( + sample_count=3, + tokens=[2, 2], + ), + **params, + ) + torch.manual_seed(17) + forced_logits = processor([], logits.clone()) + + assert int(forced_logits.argmax().item()) == int(expected[0].item()) + assert torch.isfinite(forced_logits).sum().item() == 1 + + +def test_vllm_processor_history_survives_segment_readmission(): + """A vLLM request is re-admitted per streaming segment; the sampling + history has to outlive that or the repetition penalty resets mid-stream.""" + pytest.importorskip("vllm") + from types import SimpleNamespace + + from vllm import SamplingParams + + from nemo.collections.speechlm2.inference.vllm_omni.nemotron_duplex_h.sampling import ( + SHARED_TEXT_SAMPLING_ARG, + SharedTextSamplingLogitsProcessor, + ) + + adapter = SharedTextSamplingLogitsProcessor( + SimpleNamespace(scheduler_config=SimpleNamespace(max_num_seqs=1)), + torch.device("cpu"), + False, + ) + params = SamplingParams( + temperature=0.0, + extra_args={ + SHARED_TEXT_SAMPLING_ARG: { + "top_p": 1.0, + "temperature": 1.0, + "repetition_penalty": 1.0, + "special_token_ids": [0], + "history_skip": 1, + "history_key": "stream-1", + } + }, + ) + + first_segment = adapter.new_req_logits_processor(params) + assert first_segment is not None + first_segment([], torch.tensor([0.0, 2.0, 1.0])) + first_segment([], torch.tensor([0.0, 1.0, 2.0])) + + resumed_segment = adapter.new_req_logits_processor(params) + assert resumed_segment is not None + assert resumed_segment.state is first_segment.state + assert resumed_segment.state.sample_count == 2 + assert resumed_segment.state.tokens == [2] diff --git a/tests/collections/speechlm2/nemo_inference_pipelines/test_vllm_omni_checkpoint.py b/tests/collections/speechlm2/nemo_inference_pipelines/test_vllm_omni_checkpoint.py new file mode 100644 index 000000000000..e5d6dc9f8d4f --- /dev/null +++ b/tests/collections/speechlm2/nemo_inference_pipelines/test_vllm_omni_checkpoint.py @@ -0,0 +1,112 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Wrapper-checkpoint assembly on the filesystem. + +A real vLLM/vLLM ``pipeline.run()`` lives in +``test_nemotron_voicechat_pipeline_vllm.py``. +""" + +from types import SimpleNamespace + +import pytest + +from nemo.collections.speechlm2.inference.vllm_omni.checkpoint import build_wrapper_checkpoint + + +def _stub_source_and_partial_wrapper(tmp_path, ready_component): + """A stub source checkpoint plus a wrapper with one component already built.""" + source = tmp_path / "source" + source.mkdir() + (source / "config.json").write_text("{}") + (source / "model.safetensors").write_bytes(b"source") + + wrapper = tmp_path / "wrapper" + (wrapper / ready_component).mkdir(parents=True) + (wrapper / ready_component / "config.json").write_text("{}") + (wrapper / ready_component / "model.safetensors").write_bytes(b"ready") + (wrapper / "config.json").write_text('{"model_type": "nemotron_voicechat"}') + return source, wrapper + + +@pytest.mark.parametrize("component", ["nemotron", "eartts"]) +def test_wrapper_checkpoint_converts_only_the_requested_component(tmp_path, component): + """Reusing a ready component must not drag the other one along.""" + source, wrapper = _stub_source_and_partial_wrapper(tmp_path, component) + + result = build_wrapper_checkpoint( + str(source), + str(wrapper), + include_nemotron=component == "nemotron", + include_eartts=component == "eartts", + ) + + assert result == str(wrapper) + assert not (wrapper / ("eartts" if component == "nemotron" else "nemotron")).exists() + # Records which source it came from, which is what makes it verifiable. + assert (wrapper / ".nemo_source.json").is_file() + + +def test_wrapper_checkpoint_refuses_to_extend_an_unverified_partial_wrapper(tmp_path): + """A wrapper with no ``.nemo_source.json`` could have come from any source, + so adding a second component to it might silently mix two checkpoints. + """ + source, wrapper = _stub_source_and_partial_wrapper(tmp_path, "nemotron") + + with pytest.raises(ValueError, match="Cannot safely add a component"): + build_wrapper_checkpoint(str(source), str(wrapper), include_nemotron=False, include_eartts=True) + + +def test_converter_applies_voicechat_special_token_overrides(): + """The converted config must carry VoiceChat's BOS/EOS/PAD, not the LLM + backbone's; otherwise conversion succeeds with incorrect system-prompt + prefill token IDs. + """ + from nemo.collections.speechlm2.inference.vllm_omni.scripts.convert_duplex_stt_checkpoint import ( + _apply_source_special_tokens, + ) + + class FakeTokenizer: + def __init__(self): + self.vocab = {"": 0, "": 1, "": 2, "": 12} + + def get_vocab(self): + return self.vocab + + def add_special_tokens(self, values): + for name, token in values.items(): + setattr(self, name, token) + return 0 + + def convert_tokens_to_ids(self, token): + return self.vocab[token] + + config = SimpleNamespace(bos_token_id=1, eos_token_id=12, pad_token_id=0) + source = { + "model": { + "stt": { + "model": { + "override_tokens": { + "bos_token": "", + "eos_token": "", + "pad_token": "", + } + } + } + } + } + + _apply_source_special_tokens(config, FakeTokenizer(), source) + + assert (config.bos_token_id, config.eos_token_id, config.pad_token_id) == (1, 2, 12) diff --git a/tests/collections/speechlm2/nemo_inference_pipelines/test_vllm_omni_eartts_cfg.py b/tests/collections/speechlm2/nemo_inference_pipelines/test_vllm_omni_eartts_cfg.py new file mode 100644 index 000000000000..69c1a969b71e --- /dev/null +++ b/tests/collections/speechlm2/nemo_inference_pipelines/test_vllm_omni_eartts_cfg.py @@ -0,0 +1,289 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU-only tests for vLLM-Omni EarTTS classifier-free guidance. + +Guidance arithmetic and MaskGIT trajectory sharing are numerical properties +that an end-to-end run cannot localise, so they are checked directly against +the real model definition. Requires vllm-omni to be installed; skipped +otherwise rather than run against a stubbed vLLM, which would only test the +stubs. +""" + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + + +@pytest.fixture(scope="module") +def eartts(): + pytest.importorskip("vllm_omni") + from nemo.collections.speechlm2.inference.vllm_omni.eartts import eartts as eartts_module + + return eartts_module + + +def test_unconditional_embedding_replaces_only_text_branch(eartts): + config = SimpleNamespace( + hidden_size=2, + emb_vocab_size=4, + codebook_size=3, + latent_size=2, + num_quantizers=2, + use_gated_fusion_for_text_audio=False, + use_audio_prompt_frozen_projection=False, + ) + embedding = eartts.EarTTSInputEmbedding(config) + with torch.no_grad(): + for rvq in embedding.rvq_embs: + rvq.weight.zero_() + embedding.rvq_embs[0].weight[0] = torch.tensor([1.0, 0.0]) + embedding.rvq_embs[1].weight[0] = torch.tensor([0.0, 2.0]) + embedding.embed_code.weight.copy_(torch.eye(2)) + embedding.embed_subword.embed_subwords.weight.zero_() + embedding.embed_subword.embed_subwords.weight[1] = torch.tensor([3.0, 4.0]) + embedding.bos_emb.zero_() + embedding.null_emb.copy_(torch.tensor([9.0, 10.0])) + + kwargs = dict( + acoustic_tokens=torch.zeros(2, 2, dtype=torch.long), + text_tokens=torch.ones(2, dtype=torch.long), + text_mask=torch.ones(2, dtype=torch.long), + bos_mask=torch.zeros(2, dtype=torch.long), + speaker_latent=torch.zeros(2, 2), + ) + original = embedding(**kwargs) + explicit_no_cfg = embedding( + **kwargs, + cfg_is_uncond=torch.zeros(2, dtype=torch.bool), + ) + guided_input = embedding( + **kwargs, + cfg_is_uncond=torch.tensor([False, True]), + ) + + torch.testing.assert_close(original, explicit_no_cfg, rtol=0, atol=0) + torch.testing.assert_close(guided_input[0], torch.tensor([4.0, 6.0])) + torch.testing.assert_close(guided_input[1], torch.tensor([10.0, 12.0])) + torch.testing.assert_close( + guided_input[1] - embedding.null_emb, + torch.tensor([1.0, 2.0]), + ) + assert dict(embedding.named_parameters())["null_emb"] is embedding.null_emb + + +def test_cfg_rows_are_ordered_and_guided_after_mlp(eartts): + hidden = torch.tensor([[5.0], [10.0], [3.0], [20.0]]) + enabled = torch.ones(4, dtype=torch.bool) + is_uncond = torch.tensor([True, False, True, False]) + pair_id = torch.tensor([20, 10, 10, 20]) + scale = torch.tensor([1.5, 2.0, 2.0, 1.5]) + valid = torch.ones(4, dtype=torch.bool) + + ( + ordered, + ordered_is_uncond, + ordered_scale, + active, + partner, + conditional_rep, + inverse, + ) = eartts._prepare_cfg_sampling_batch( + hidden, + enabled, + is_uncond, + pair_id, + scale, + valid, + ) + torch.testing.assert_close( + ordered, + torch.tensor([[10.0], [20.0], [3.0], [5.0]]), + ) + assert ordered_is_uncond.tolist() == [False, False, True, True] + assert active.tolist() == [True, True, True, True] + assert partner.tolist() == [2, 3, 0, 1] + assert conditional_rep.tolist() == [0, 1, 0, 1] + + guided = eartts._apply_cfg_after_mlp( + ordered, + ordered_is_uncond, + ordered_scale, + active, + partner, + ) + torch.testing.assert_close( + guided, + torch.tensor([[24.0], [42.5], [24.0], [42.5]]), + ) + assert torch.equal(ordered[inverse], hidden) + + # An incomplete or disabled batch must fall straight through: no + # reordering, nothing active, and guidance a no-op. + plain = torch.randn(3, 4) + ordered, roles, scales, active, partner, _, inverse = eartts._prepare_cfg_sampling_batch( + plain, + cfg_enabled=torch.tensor([True, True, False]), + cfg_is_uncond=torch.tensor([False, True, False]), + cfg_pair_id=torch.tensor([7, 7, -1]), + cfg_scale=torch.ones(3), + valid=torch.ones(3, dtype=torch.bool), + ) + assert torch.equal(ordered, plain) + assert not active.any() + assert torch.equal(inverse, torch.arange(3)) + assert torch.equal(eartts._apply_cfg_after_mlp(ordered, roles, scales, active, partner), plain) + + +def test_maskgit_shares_one_code_trajectory_per_pair(eartts): + config = SimpleNamespace( + num_quantizers=2, + codebook_size=8, + noise_scale=0.7, + num_iter=2, + exponent=3.0, + latent_size=4, + hidden_size=4, + intermediate_size=8, + mog_num_layers=0, + mog_num_predictions=4, + mog_low_rank=None, + top_p_or_k=None, + mog_min_log_std=-4.0, + mog_eps=1e-6, + ) + sampler = eartts.MaskGITSampler(config) + torch.manual_seed(3) + with torch.no_grad(): + for parameter in sampler.parameters(): + parameter.normal_(mean=0.0, std=0.2) + + hidden = torch.randn(4, 4) + enabled = torch.ones(4, dtype=torch.bool) + roles = torch.tensor([True, False, True, False]) + pairs = torch.tensor([2, 1, 1, 2]) + scales = torch.full((4,), 1.25) + valid = torch.ones(4, dtype=torch.bool) + torch.manual_seed(11) + codes = sampler(hidden, enabled, roles, pairs, scales, valid) + + assert torch.equal(codes[0], codes[3]) + assert torch.equal(codes[1], codes[2]) + + no_cfg_hidden = torch.randn(3, 4) + torch.manual_seed(17) + implicit_no_cfg = sampler(no_cfg_hidden) + torch.manual_seed(17) + explicit_no_cfg = sampler( + no_cfg_hidden, + cfg_enabled=torch.zeros(3, dtype=torch.bool), + cfg_is_uncond=torch.zeros(3, dtype=torch.bool), + cfg_pair_id=torch.full((3,), -1, dtype=torch.long), + cfg_scale=torch.zeros(3), + valid=torch.ones(3, dtype=torch.bool), + ) + assert torch.equal(implicit_no_cfg, explicit_no_cfg) + + +def test_client_facing_stage_emits_the_drainable_audio_key(eartts): + """A final AR audio stage must publish codes under ``model_outputs``. + + vLLM-Omni remaps that key onto the drainable ``audio`` modality, so DELTA + streaming empties it every step. Any other key is retained across steps and + merged with ``CONCAT_LAST``, which widens a ``T x num_quantizers`` frame + instead of appending frames to it. + """ + hidden = torch.zeros(1, 4) + codes = torch.tensor([[3, 5]], dtype=torch.long) + + for single_stage_audio, expected_key in ((True, "model_outputs"), (False, "audio_codes")): + model = object.__new__(eartts.EarTTSForCausalLM) + nn.Module.__init__(model) + model._single_stage_audio = single_stage_audio + model._out_codes = codes.clone() + + output = model.make_omni_output(hidden) + + assert list(output.multimodal_outputs) == [expected_key] + torch.testing.assert_close(output.multimodal_outputs[expected_key], codes) + + stashed = model.postprocess(hidden, output.multimodal_outputs) + torch.testing.assert_close(stashed["last_acoustic_codes"], codes) + + # Per-request CFG metadata lands in model-owned buffers whose addresses + # must stay stable, because CUDA graphs capture them. + model = object.__new__(eartts.EarTTSForCausalLM) + nn.Module.__init__(model) + model.config = SimpleNamespace(guidance_scale=0.5) + model._cfg_enabled = torch.zeros(8, dtype=torch.bool) + model._cfg_is_uncond = torch.zeros(8, dtype=torch.bool) + model._cfg_pair_id = torch.full((8,), -1, dtype=torch.long) + model._cfg_scale = torch.zeros(8) + addresses = tuple( + value.data_ptr() + for value in ( + model._cfg_enabled, + model._cfg_is_uncond, + model._cfg_pair_id, + model._cfg_scale, + ) + ) + + model._write_cfg_state( + start=1, + span_len=2, + info_dict={ + "cfg_enabled": True, + "cfg_role": "cond", + "cfg_pair_id": "request-7", + "cfg_scale": 1.75, + }, + ) + model._write_cfg_state( + start=3, + span_len=1, + info_dict={ + "cfg_enabled": torch.tensor(True), + "cfg_role": ["uncond"], + "cfg_pair_id": "request-7", + "cfg_scale": torch.tensor(1.75), + }, + ) + + assert model._cfg_enabled[1:4].all() + assert model._cfg_is_uncond[1:4].tolist() == [False, False, True] + assert model._cfg_pair_id[1] == model._cfg_pair_id[3] + torch.testing.assert_close(model._cfg_scale[1:4], torch.full((3,), 1.75)) + assert addresses == tuple( + value.data_ptr() + for value in ( + model._cfg_enabled, + model._cfg_is_uncond, + model._cfg_pair_id, + model._cfg_scale, + ) + ) + + with pytest.raises(AssertionError, match="cfg_role"): + model._write_cfg_state( + start=4, + span_len=1, + info_dict={ + "cfg_enabled": True, + "cfg_role": "conditional", + "cfg_pair_id": 4, + }, + ) diff --git a/tests/collections/speechlm2/test_nemotron_voicechat.py b/tests/collections/speechlm2/test_voicechat.py similarity index 53% rename from tests/collections/speechlm2/test_nemotron_voicechat.py rename to tests/collections/speechlm2/test_voicechat.py index 4f2b431e4f25..97894bea7724 100644 --- a/tests/collections/speechlm2/test_nemotron_voicechat.py +++ b/tests/collections/speechlm2/test_voicechat.py @@ -12,7 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""NemotronVoiceChat model tests (training/offline), matching NeMo main. + +Streaming inference tests live under ``nemo_inference_pipelines/`` and load +the public 11B checkpoint instead of this toy-weight model. +""" + import os +import tempfile import pytest import torch @@ -22,6 +29,11 @@ from nemo.collections.common.data.utils import move_data_to_device from nemo.collections.speechlm2 import DuplexSTTDataset from nemo.collections.speechlm2.models import NemotronVoiceChat +from nemo.collections.speechlm2.models.nemotron_voicechat import ( + _apply_nemotron_labs_voicechat_release_config_shim, + _is_nemotron_labs_voicechat_release, +) +from nemo.collections.speechlm2.streaming.duplex_stt_inference import DuplexSTTStreamingInference if torch.cuda.is_available(): torch.set_default_device('cuda') @@ -37,6 +49,93 @@ target_sample_rate = 22050 +_RELEASE_BY_CONTENT = { + "_rnnt_merge_info": {}, + "model": { + "stt": { + "model": { + "pretrained_llm": "nvidia/NVIDIA-Nemotron-Nano-9B-v2", + "use_function_head": True, + "predict_user_text": False, + } + } + }, +} + + +@pytest.mark.parametrize( + ("model_id", "cfg", "is_release"), + [ + ("nvidia/NVIDIA-NemotronLabs-VoiceChat-11B", None, True), + ("/checkpoints/NVIDIA-NemotronLabs-VoiceChat-11B", None, True), + ("/checkpoints/NVIDIA-NemotronLabs-VoiceChat-11B/", None, True), + # Recognised by config content once downloaded under another name. + ("/checkpoints/arbitrary-name", _RELEASE_BY_CONTENT, True), + ("/checkpoints/some-other-voicechat-export", None, False), + ], +) +def test_identifies_the_nemotron_labs_voicechat_release(model_id, cfg, is_release): + """The shim must apply to the public 11B and to nothing else.""" + assert _is_nemotron_labs_voicechat_release(model_id, cfg) is is_release + + +def test_release_config_shim_flattens_nested_keys_without_clobbering_current_ones(): + """Values already in the flat position win; only absent ones are filled.""" + stale = { + "data": {"source_sample_rate": 8000}, + "exp_manager": {"explicit_log_dir": "/global"}, + "model": { + "stt": { + "model": { + "pretrained_llm": "dummy", + "base_model_name": "backbone", + "embed_tokens_name": "embeddings", + }, + "data": {"source_sample_rate": 16000}, + "exp_manager": {"explicit_log_dir": "/stt"}, + }, + "speech_generation": { + "model": {"tts_config": {"cas_config": {"pretrained_tokenizer_name": "legacy-tokenizer"}}} + }, + }, + } + + shimmed = _apply_nemotron_labs_voicechat_release_config_shim(stale) + normalized = shimmed["model"]["stt"]["model"] + + assert normalized["source_sample_rate"] == 16000 + assert normalized["validation_save_path"] == "/stt" + assert normalized["llm_attr_name"] == "model" + assert normalized["embed_tokens_attr_name"] == "embeddings" + assert "pretrained_tokenizer_name" not in ( + shimmed["model"]["speech_generation"]["model"]["tts_config"]["cas_config"] + ) + + already_flat = { + "model": { + "stt": { + "model": { + "source_sample_rate": 22050, + "validation_save_path": "/already-flat", + "llm_attr_name": "already-flat", + "embed_tokens_attr_name": "already-flat", + "base_model_name": "backbone", + "embed_tokens_name": "embeddings", + }, + "data": {"source_sample_rate": 16000}, + "exp_manager": {"explicit_log_dir": "/nested"}, + } + } + } + + preserved = _apply_nemotron_labs_voicechat_release_config_shim(already_flat)["model"]["stt"]["model"] + + assert preserved["source_sample_rate"] == 22050 + assert preserved["validation_save_path"] == "/already-flat" + assert preserved["llm_attr_name"] == "already-flat" + assert preserved["embed_tokens_attr_name"] == "already-flat" + + def create_model( predict_user_text=False, force_use_noise_augmentation=False, @@ -45,6 +144,8 @@ def create_model( old_noise_max_snr=0.0, ): """Helper function to create a model with configurable settings.""" + log_dir = tempfile.mkdtemp(prefix="test_nemotron_voicechat_logs-") + stt_log_dir = os.path.join(log_dir, "duplex_stt") test_stt_cfg = { "model": { "pretrained_llm": pretrained_llm, @@ -52,7 +153,7 @@ def create_model( "audio_loss_weight": 1, "text_loss_weight": 3, "source_sample_rate": source_sample_rate, - "validation_save_path": "/tmp/test_duplex_stt_logs", + "validation_save_path": stt_log_dir, "perception": { "_target_": "nemo.collections.speechlm2.modules.perception.AudioPerceptionModule", "preprocessor": { @@ -84,7 +185,7 @@ def create_model( "source_sample_rate": 16000, }, "exp_manager": { - "explicit_log_dir": "/tmp/test_duplex_stt_logs", + "explicit_log_dir": stt_log_dir, }, } @@ -203,7 +304,7 @@ def create_model( "target_sample_rate": target_sample_rate, }, "exp_manager": { - "explicit_log_dir": "/tmp/test_duplex_stt_logs", + "explicit_log_dir": stt_log_dir, }, } @@ -221,7 +322,7 @@ def create_model( "output_roles": ["agent", "Assistant", "assistant", "Agent"], }, "exp_manager": { - "explicit_log_dir": "/tmp/test_nemotron_voicechat_logs", + "explicit_log_dir": log_dir, }, } model = NemotronVoiceChat(test_config) @@ -287,18 +388,148 @@ def training_cutset_batch(): return CutSet([cut]) -def test_e2e_validation_step(model, dataset, training_cutset_batch): +def test_forced_turn_taking(): + pad_id = 0 + bos_id = 1 + eos_id = 2 + non_special_user_token_id = 42 + + threshold = 10 + pad_window = 3 + + t = 4 + + class DuplexSTTStreamingInferenceForTurnTakingTest(DuplexSTTStreamingInference): + def __init__(self): + class Tokenizer: + @staticmethod + def text_to_ids(token): + raise AssertionError(f"Unexpected legacy token lookup: {token}") + + Tokenizer.eos = eos_id + + self.text_pad_id = pad_id + self.text_bos_id = bos_id + self.text_eos_id = eos_id + self.tokenizer = Tokenizer() + self.cfg = { + "force_turn_taking": True, + "force_turn_taking_threshold": threshold, + "force_turn_taking_pad_window": pad_window, + } + super().__init__(model=self) + + streaming_inference = DuplexSTTStreamingInferenceForTurnTakingTest() + + def new_tokens(): + return ( + torch.full((1, 6), pad_id, dtype=torch.long), + torch.full((1, 6), pad_id, dtype=torch.long), + ) + + # The 1a/1b/2a labels match the cases documented in _maybe_apply_forced_turn_taking. + # 1a. Silence-based "user stopped" trigger. + # Negative cases: the rule should not fire when the token before the pad window is user BOS or PAD. + # A pad window immediately after BOS is startup silence, not user speech. + # time step: 0 1 2 3 4 + # <- pad_window=3 -> + # t-4 t-3 t-2 t-1 t + # ASR: BOS PAD PAD PAD PAD + # expected text: PAD PAD PAD PAD PAD + gen_text, gen_asr = new_tokens() + gen_asr[0, 0] = bos_id + streaming_inference._maybe_apply_forced_turn_taking(t, gen_text, gen_asr) + assert gen_text[0, t] == pad_id + + # Negative case: The rule also should not fire if user BOS is inside the pad window, + # since the window is not all PAD. + # time step: 0 1 2 3 4 + # <- pad_window=3 -> + # t-4 t-3 t-2 t-1 t + # ASR: PAD BOS PAD PAD PAD + # expected text: PAD PAD PAD PAD PAD + gen_text, gen_asr = new_tokens() + gen_asr[0, 1] = bos_id + streaming_inference._maybe_apply_forced_turn_taking(t, gen_text, gen_asr) + assert gen_text[0, t] == pad_id + + # 1a. Silence-based "user stopped" trigger. + # Now check the positive case: the rule should fire after user speech. + # A pad window after a real user token means the user stopped speaking. + # time step: 0 1 2 3 4 + # <- pad_window=3 -> + # t-4 t-3 t-2 t-1 t + # ASR: non-special token PAD PAD PAD PAD + # expected text: PAD PAD PAD PAD BOS + gen_text, gen_asr = new_tokens() + gen_asr[0, 0] = non_special_user_token_id + streaming_inference._maybe_apply_forced_turn_taking(t, gen_text, gen_asr) + assert gen_text[0, t] == bos_id + + # 1b. Explicit user EOS means the user stopped speaking, so force agent BOS. + # time step: 0 1 2 3 4 + # ASR: PAD PAD PAD PAD EOS + # expected text: PAD PAD PAD PAD BOS + gen_text, gen_asr = new_tokens() + gen_asr[0, t] = eos_id + streaming_inference._maybe_apply_forced_turn_taking(t, gen_text, gen_asr) + assert gen_text[0, t] == bos_id + + # Threshold negative case: do not force another agent BOS if one already + # appears in the recent text-channel lookback. + # time step: 0 1 2 3 4 + # ASR: PAD PAD PAD PAD EOS + # text before: BOS PAD PAD PAD PAD + # expected text: BOS PAD PAD PAD PAD + gen_text, gen_asr = new_tokens() + gen_text[0, 0] = bos_id + gen_asr[0, t] = eos_id + streaming_inference._maybe_apply_forced_turn_taking(t, gen_text, gen_asr) + assert gen_text[0, t] == pad_id + + # 2a. Explicit user BOS means the user started speaking, so force agent EOS. + # time step: 0 1 2 3 4 + # ASR: PAD PAD PAD PAD BOS + # expected text: PAD PAD PAD PAD EOS + gen_text, gen_asr = new_tokens() + gen_asr[0, t] = bos_id + streaming_inference._maybe_apply_forced_turn_taking(t, gen_text, gen_asr) + assert gen_text[0, t] == eos_id + + # Threshold negative case: do not force another agent EOS if one already + # appears in the recent text-channel lookback. + # time step: 0 1 2 3 4 + # ASR: PAD PAD PAD PAD BOS + # text before: EOS PAD PAD PAD PAD + # expected text: EOS PAD PAD PAD PAD + gen_text, gen_asr = new_tokens() + gen_text[0, 0] = eos_id + gen_asr[0, t] = bos_id + streaming_inference._maybe_apply_forced_turn_taking(t, gen_text, gen_asr) + assert gen_text[0, t] == pad_id + + +def test_e2e_validation_step_feeds_its_metrics(model, dataset, training_cutset_batch): + """One validation step must reach the metric and result buffers. + + ``validation_step`` returns nothing, so asserting on its return value + cannot fail. What it is actually for is the chain behind it: offline + inference -> ASR transcription -> BLEU/ASR-BLEU update -> result logging. + Asserting that the buffers filled pins that chain, and fails if inference + silently produces nothing. + """ model.eval() model.on_validation_epoch_start() batch = dataset[training_cutset_batch] batch = move_data_to_device(batch, device=model.device) - results = model.validation_step( + model.validation_step( {"dummy_val_set": batch}, batch_idx=0, speaker_audio=torch.randn(1, 22050, device=model.device), speaker_audio_lens=torch.tensor([22050], device=model.device), ) - assert results is None # no return value + + assert model.results_logger.cached_results, "validation_step logged no results" def test_e2s_offline_generation(model): From 1c06a8f2f34e86d085b83bc2bc431934685ce566 Mon Sep 17 00:00:00 2001 From: Elena Rastorgueva Date: Tue, 1 Sep 2026 17:46:56 +0000 Subject: [PATCH 2/6] chore(speechlm2): stub vLLM-Omni VoiceChat backends for the native PR Keep DuplexLLM/DuplexTTS and the Vllm* classes so the native frame loop already matches the combined form. Selecting vllm_omni raises NotImplementedError; the runtime is the parent commit on duplex-vllm-omni-on-main. Signed-off-by: Elena Rastorgueva --- docs/source/speechlm2/streaming_inference.rst | 24 +- .../conf/s2s_streaming.yaml | 8 +- .../model_wrappers/backend/vllm/__init__.py | 26 +- .../model_wrappers/backend/vllm/eartts.py | 42 +- .../model_wrappers/backend/vllm/llm.py | 53 +- .../model_wrappers/engine_selection.py | 14 + .../nemotron_voicechat_inference_wrapper.py | 152 +- .../speechlm2/inference/vllm_omni/__init__.py | 59 - .../inference/vllm_omni/checkpoint.py | 311 --- .../inference/vllm_omni/deploy/eartts.yaml | 45 - .../vllm_omni/deploy/nemotron_voicechat.yaml | 62 - .../inference/vllm_omni/eartts/__init__.py | 17 - .../vllm_omni/eartts/configuration_eartts.py | 217 -- .../inference/vllm_omni/eartts/eartts.py | 1945 ----------------- .../inference/vllm_omni/eartts/pipeline.py | 63 - .../inference/vllm_omni/eartts/scheduler.py | 451 ---- .../vllm_omni/nemotron_duplex_h/__init__.py | 19 - .../nemotron_duplex_h/nemotron_duplex_h.py | 771 ------- .../vllm_omni/nemotron_duplex_h/sampling.py | 247 --- .../vllm_omni/nemotron_voicechat/__init__.py | 27 - .../vllm_omni/nemotron_voicechat/pipeline.py | 80 - .../vllm_omni/nemotron_voicechat/scheduler.py | 68 - .../speechlm2/inference/vllm_omni/outputs.py | 128 -- .../speechlm2/inference/vllm_omni/register.py | 154 -- .../speechlm2/inference/vllm_omni/runtime.py | 322 --- .../inference/vllm_omni/scripts/__init__.py | 13 - .../convert_duplex_eartts_checkpoint.py | 357 --- .../scripts/convert_duplex_stt_checkpoint.py | 343 --- .../speechlm2/inference/vllm_omni/session.py | 592 ----- pyproject.toml | 17 - .../nemo_inference_pipelines/conftest.py | 28 +- .../test_engine_selection.py | 16 + .../test_nemotron_voicechat_pipeline_vllm.py | 87 - .../test_text_sampling.py | 84 - .../test_vllm_omni_checkpoint.py | 112 - .../test_vllm_omni_eartts_cfg.py | 289 --- 36 files changed, 104 insertions(+), 7139 deletions(-) delete mode 100644 nemo/collections/speechlm2/inference/vllm_omni/__init__.py delete mode 100644 nemo/collections/speechlm2/inference/vllm_omni/checkpoint.py delete mode 100644 nemo/collections/speechlm2/inference/vllm_omni/deploy/eartts.yaml delete mode 100644 nemo/collections/speechlm2/inference/vllm_omni/deploy/nemotron_voicechat.yaml delete mode 100644 nemo/collections/speechlm2/inference/vllm_omni/eartts/__init__.py delete mode 100644 nemo/collections/speechlm2/inference/vllm_omni/eartts/configuration_eartts.py delete mode 100644 nemo/collections/speechlm2/inference/vllm_omni/eartts/eartts.py delete mode 100644 nemo/collections/speechlm2/inference/vllm_omni/eartts/pipeline.py delete mode 100644 nemo/collections/speechlm2/inference/vllm_omni/eartts/scheduler.py delete mode 100644 nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/__init__.py delete mode 100644 nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/nemotron_duplex_h.py delete mode 100644 nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/sampling.py delete mode 100644 nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/__init__.py delete mode 100644 nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/pipeline.py delete mode 100644 nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/scheduler.py delete mode 100644 nemo/collections/speechlm2/inference/vllm_omni/outputs.py delete mode 100644 nemo/collections/speechlm2/inference/vllm_omni/register.py delete mode 100644 nemo/collections/speechlm2/inference/vllm_omni/runtime.py delete mode 100644 nemo/collections/speechlm2/inference/vllm_omni/scripts/__init__.py delete mode 100644 nemo/collections/speechlm2/inference/vllm_omni/scripts/convert_duplex_eartts_checkpoint.py delete mode 100644 nemo/collections/speechlm2/inference/vllm_omni/scripts/convert_duplex_stt_checkpoint.py delete mode 100644 nemo/collections/speechlm2/inference/vllm_omni/session.py delete mode 100644 tests/collections/speechlm2/nemo_inference_pipelines/test_nemotron_voicechat_pipeline_vllm.py delete mode 100644 tests/collections/speechlm2/nemo_inference_pipelines/test_vllm_omni_checkpoint.py delete mode 100644 tests/collections/speechlm2/nemo_inference_pipelines/test_vllm_omni_eartts_cfg.py diff --git a/docs/source/speechlm2/streaming_inference.rst b/docs/source/speechlm2/streaming_inference.rst index 5a049b3e7955..853cbc5e5f84 100644 --- a/docs/source/speechlm2/streaming_inference.rst +++ b/docs/source/speechlm2/streaming_inference.rst @@ -577,8 +577,15 @@ a vLLM-Omni implementation: llm.py # PyTorchLLM (wraps the DuplexSTT forward pass) eartts.py # PyTorchEarTTS (wraps DuplexEARTTS.infer_codes_one_step) vllm/ - llm.py # VllmLLM (drives OmniStreamingSession.step_llm) - eartts.py # VllmEarTTS (drives OmniStreamingSession.step_tts) + llm.py # VllmLLM (NotImplementedError in this PR) + eartts.py # VllmEarTTS (NotImplementedError in this PR) + +This PR ships the native engines. ``VllmLLM`` and ``VllmEarTTS`` exist so the +wrapper's frame loop already matches the combined form; selecting +``llm_engine_type`` / ``tts_engine_type`` of ``vllm_omni`` raises +``NotImplementedError`` at construction. The runtime (``inference/vllm_omni/``, +``OmniRuntime``, wrapper-checkpoint conversion) is the parent commit on +``duplex-vllm-omni-on-main``. ``NemotronVoicechatInferenceWrapper`` selects one implementation per component at construction and stores them as ``llm_backend`` and ``tts_backend``. Its @@ -654,9 +661,12 @@ reported at load time rather than silently doing nothing. vLLM-Omni Integration """"""""""""""""""""" -Each component selected as ``vllm_omni`` gets the ``Vllm*`` implementation of -its contract, and the corresponding native class is not created. The wrapper -starts only the selected one-stage ``AsyncOmni`` engine or engines: +Selecting ``vllm_omni`` is rejected at wrapper construction in this PR. The +intended shape is: each component selected as ``vllm_omni`` gets the ``Vllm*`` +implementation of its contract, and the corresponding native class is not +created. The wrapper starts only the selected one-stage ``AsyncOmni`` engine +or engines. That runtime lives on the parent commit +(``duplex-vllm-omni-on-main``): - **Nemotron** -- ``NemotronDuplexHForCausalLM``, which consumes the per-step acoustic embedding and samples a text token plus the checkpoint's optional @@ -667,8 +677,8 @@ starts only the selected one-stage ``AsyncOmni`` engine or engines: The split keeps the component boundary in NeMo, so either component can be replaced without changing the other engine. Nemotron settings come from ``inference/vllm_omni/deploy/nemotron_voicechat.yaml`` and EarTTS settings from -``inference/vllm_omni/deploy/eartts.yaml``. Override them independently with -``vllm_omni_config.stage_overrides`` and +``inference/vllm_omni/deploy/eartts.yaml`` on that parent commit. Override them +independently with ``vllm_omni_config.stage_overrides`` and ``vllm_omni_config.eartts_stage_overrides``. Nemotron text sampling uses vLLM's custom logits-processor hook to call the diff --git a/examples/speechlm2/nemo_inference_pipelines/conf/s2s_streaming.yaml b/examples/speechlm2/nemo_inference_pipelines/conf/s2s_streaming.yaml index a305a074bff3..d49efe7782ca 100644 --- a/examples/speechlm2/nemo_inference_pipelines/conf/s2s_streaming.yaml +++ b/examples/speechlm2/nemo_inference_pipelines/conf/s2s_streaming.yaml @@ -33,10 +33,12 @@ s2s: model_path: ??? decode_audio: true # Whether to conduct the TTS portion of inference, or just the STT portion speaker_name: null # Required when decode_audio is true. Must match a speaker registered in the checkpoint. - llm_engine_type: native # 'native' or 'vllm_omni' - tts_engine_type: native # 'native' or 'vllm_omni' + llm_engine_type: native # 'native' or 'vllm_omni' (vllm_omni raises NotImplementedError in this PR) + tts_engine_type: native # 'native' or 'vllm_omni' (vllm_omni raises NotImplementedError in this PR) - # vllm_omni engine configuration (ignored unless a component is vllm_omni). + # vllm_omni engine configuration. Selecting vllm_omni is rejected at wrapper + # construction in this PR; the keys document the combined-form API. The + # implementation is the parent commit on duplex-vllm-omni-on-main. # # On the first run a wrapper directory is built lazily under # ``$TMPDIR/_vllm_omni_wrapper`` containing diff --git a/nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/__init__.py b/nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/__init__.py index bc50bcde6b4a..774f3079090a 100644 --- a/nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/__init__.py +++ b/nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/__init__.py @@ -14,26 +14,12 @@ """vLLM-Omni implementations of the two component contracts. -Both are thin: the engines are process-scoped and owned by ``OmniRuntime``, -while the request state lives in the per-stream ``OmniStreamingSession`` these -classes read off the decode state. Importing this package does not import -vLLM. +Construction raises ``NotImplementedError`` in this PR. The classes exist so +the native frame loop already matches the combined-form loop; the parent +commit on ``duplex-vllm-omni-on-main`` has the runtime. """ -from typing import Any +from nemo.collections.speechlm2.inference.model_wrappers.backend.vllm.eartts import VllmEarTTS +from nemo.collections.speechlm2.inference.model_wrappers.backend.vllm.llm import VllmLLM - -def require_session(state: Any): - """Return the stream's ``OmniStreamingSession``, or say why there isn't one. - - ``omni_session`` is a declared field on ``StreamingDecodeState``, so this - reads it directly: a missing attribute is a programming error worth an - AttributeError, while ``None`` is the real case worth explaining. - """ - session = state.omni_session - if session is None: - raise RuntimeError( - "A vllm_omni component requires a per-stream OmniStreamingSession; " - "make sure begin_stream(...) ran for this stream before the first frame." - ) - return session +__all__ = ["VllmLLM", "VllmEarTTS"] diff --git a/nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/eartts.py b/nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/eartts.py index fba82d28a3b5..795ddfa59ba8 100644 --- a/nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/eartts.py +++ b/nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/eartts.py @@ -14,9 +14,12 @@ """vLLM-Omni backend for the TTS (EarTTS) component of NemotronVoiceChat. -Implements :class:`~nemo.collections.speechlm2.inference.model_wrappers.backend.eartts.DuplexTTS` -against the per-stream ``OmniStreamingSession``; the PyTorch sibling lives in -``backend/pytorch/eartts.py``. +Implements :class:`~nemo.collections.speechlm2.inference.model_wrappers.backend.eartts.DuplexTTS`. +The PyTorch sibling lives in ``backend/pytorch/eartts.py``. + +This PR stubs the class: construction raises so a ``vllm_omni`` engine +selection cannot silently fall through to native. The implementation is the +parent commit on ``duplex-vllm-omni-on-main``. """ from typing import Any @@ -24,17 +27,14 @@ import torch from nemo.collections.speechlm2.inference.model_wrappers.backend.eartts import DuplexTTS -from nemo.collections.speechlm2.inference.model_wrappers.backend.vllm import require_session +from nemo.collections.speechlm2.inference.model_wrappers.engine_selection import VLLM_OMNI, reject_unimplemented_vllm class VllmEarTTS(DuplexTTS): """Runs EarTTS in a vLLM-Omni engine, one text token per step. - Stateless itself, like its LLM counterpart: the engine belongs to - ``OmniRuntime`` and the request belongs to the session on the decode state. - With classifier-free guidance enabled, the session's conditional and - unconditional requests are kept in lockstep by the custom scheduler, so one - submission still yields one acoustic frame. + Not implemented in this PR. Same contract as :class:`PyTorchEarTTS`, so the + wrapper's frame loop does not branch on engine type. """ def __init__(self, device: torch.device): @@ -43,24 +43,10 @@ def __init__(self, device: torch.device): device: Device the native audio codec decodes on, so the codes this backend returns land where the codec expects them. """ - self.device = device + del device + reject_unimplemented_vllm("native", VLLM_OMNI) def step(self, state: Any, current_frame_idx: int, request_id: str) -> torch.Tensor: - """One EarTTS step -- see ``DuplexTTS.step``. - - ``inference_force_speech_silence_on_eos`` is not applied here: the - converted EarTTS substitutes codec silence itself when the incoming - text token is EOS, matching what DuplexEARTTS does natively. It has no - flag for it, so it cannot honour a ``False`` setting; the wrapper - reports that at load time. - """ - del request_id # The session already owns this stream's request ids. - - session = require_session(state) - text_token = int(state.gen_text[:, current_frame_idx].item()) - session.step_tts(text_token) - audio_chunks = session.drain_audio_codes() - if not audio_chunks: - raise RuntimeError("vLLM EarTTS produced no audio codes for the submitted text token") - # The native codec helpers consume [B, T, num_quantizers]. - return torch.cat(audio_chunks, dim=0).to(self.device, dtype=torch.long).unsqueeze(0) + """One EarTTS step -- see ``DuplexTTS.step``.""" + del state, current_frame_idx, request_id + reject_unimplemented_vllm("native", VLLM_OMNI) diff --git a/nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/llm.py b/nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/llm.py index 98679f9644b1..38f7a440077b 100644 --- a/nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/llm.py +++ b/nemo/collections/speechlm2/inference/model_wrappers/backend/vllm/llm.py @@ -14,9 +14,12 @@ """vLLM-Omni backend for the LLM component of NemotronVoiceChat. -Implements :class:`~nemo.collections.speechlm2.inference.model_wrappers.backend.llm.DuplexLLM` -against the per-stream ``OmniStreamingSession``; the PyTorch sibling lives in -``backend/pytorch/llm.py``. +Implements :class:`~nemo.collections.speechlm2.inference.model_wrappers.backend.llm.DuplexLLM`. +The PyTorch sibling lives in ``backend/pytorch/llm.py``. + +This PR stubs the class: construction raises so a ``vllm_omni`` engine +selection cannot silently fall through to native. The implementation is the +parent commit on ``duplex-vllm-omni-on-main``. """ from typing import Any @@ -24,17 +27,20 @@ import torch from nemo.collections.speechlm2.inference.model_wrappers.backend.llm import DuplexLLM, LlmStepResult -from nemo.collections.speechlm2.inference.model_wrappers.backend.vllm import require_session +from nemo.collections.speechlm2.inference.model_wrappers.engine_selection import VLLM_OMNI, reject_unimplemented_vllm class VllmLLM(DuplexLLM): """Runs Nemotron in a vLLM-Omni engine, one acoustic frame per step. - Stateless itself: the engine is process-scoped and owned by - ``OmniRuntime``, and everything request-scoped lives in the session that - the pipeline attached to the decode state at prefill. + Not implemented in this PR. The native frame loop already calls + :meth:`DuplexLLM.step` without inspecting the engine type, so landing the + runtime later does not reshape the per-frame path. """ + def __init__(self) -> None: + reject_unimplemented_vllm(VLLM_OMNI, "native") + def step( self, frame_embedding: torch.Tensor, @@ -47,32 +53,7 @@ def step( sampling_params: dict[str, float] | None = None, debug_logger: Any = None, ) -> LlmStepResult: - """One Nemotron step -- see ``DuplexLLM.step``. - - Nemotron builds its own duplex input embedding from the acoustic frame, - so there is no ``build_input_embedding`` and no history replay here; - ``frame_offset`` and ``has_prompt`` do not apply. Per-stream sampling - was fixed when the session was created, and logits stay inside the - engine, so ``return_debug`` cannot be honoured either -- the result's - logit fields stay None. - - The previous frame's committed text token is fed back explicitly. That - is what carries a forced-turn-taking rewrite into Nemotron's own - history, the role ``gen_text`` plays for the PyTorch backend. - """ - del frame_offset, has_prompt, sampling_params, return_debug - - session = require_session(state) - if debug_logger is not None: - debug_logger.log_input_embeds(frame_embedding) - - prev_text_token = None - if current_frame_idx > 0: - prev_text_token = int(state.gen_text[0, current_frame_idx - 1].item()) - - tokens = session.step_llm(frame_embedding.reshape(-1), prev_text_token=prev_text_token) - return LlmStepResult( - predicted_token=tokens.text, - asr_predicted_token=tokens.asr, - function_predicted_token=tokens.function, - ) + """One Nemotron step -- see ``DuplexLLM.step``.""" + del frame_embedding, state, frame_offset, current_frame_idx, has_prompt + del return_debug, sampling_params, debug_logger + reject_unimplemented_vllm(VLLM_OMNI, "native") diff --git a/nemo/collections/speechlm2/inference/model_wrappers/engine_selection.py b/nemo/collections/speechlm2/inference/model_wrappers/engine_selection.py index eae8b78fc2c3..3252fe814c15 100644 --- a/nemo/collections/speechlm2/inference/model_wrappers/engine_selection.py +++ b/nemo/collections/speechlm2/inference/model_wrappers/engine_selection.py @@ -29,6 +29,20 @@ SUPPORTED_S2S_ENGINE_TYPES = frozenset({NATIVE, VLLM_OMNI}) +# vllm_omni remains a legal config value so the native loop already matches the +# combined-form loop. This PR does not ship the runtime; selecting it raises. +VLLM_NOT_IMPLEMENTED = ( + "vLLM-Omni VoiceChat backends are not implemented in this PR. " + "Set llm_engine_type='native' and tts_engine_type='native'. " + "The implementation is the parent commit on duplex-vllm-omni-on-main." +) + + +def reject_unimplemented_vllm(llm_engine_type: str, tts_engine_type: str) -> None: + """Raise if a vLLM component was selected. This PR only ships native engines.""" + if llm_engine_type == VLLM_OMNI or tts_engine_type == VLLM_OMNI: + raise NotImplementedError(VLLM_NOT_IMPLEMENTED) + def _component_engine(model_cfg: Mapping, key: str) -> str: value = model_cfg.get(key, NATIVE) diff --git a/nemo/collections/speechlm2/inference/model_wrappers/nemotron_voicechat_inference_wrapper.py b/nemo/collections/speechlm2/inference/model_wrappers/nemotron_voicechat_inference_wrapper.py index f35b7cd208de..03d82bc8444a 100644 --- a/nemo/collections/speechlm2/inference/model_wrappers/nemotron_voicechat_inference_wrapper.py +++ b/nemo/collections/speechlm2/inference/model_wrappers/nemotron_voicechat_inference_wrapper.py @@ -13,8 +13,6 @@ # limitations under the License. import copy -import json -import os import time import torch @@ -42,6 +40,7 @@ VLLM_OMNI, native_weight_skip_prefixes, precision_matches_cfg, + reject_unimplemented_vllm, reject_unsupported_determinism, resolve_engine_types, ) @@ -51,7 +50,6 @@ PerceptionCacheState, ) from nemo.collections.speechlm2.models.nemotron_voicechat import NemotronVoiceChat -from nemo.collections.speechlm2.parts.logit_boosts import LogitBoosts from nemo.collections.speechlm2.parts.text_utils import ( _decode_tokens_with_specials, get_special_token_ids, @@ -94,6 +92,7 @@ def __init__(self, model_cfg: DictConfig): self.llm_engine_type, self.tts_engine_type = resolve_engine_types(model_cfg) self._deterministic = bool(model_cfg.get("deterministic", False)) reject_unsupported_determinism(self.llm_engine_type, self.tts_engine_type, self._deterministic) + reject_unimplemented_vllm(self.llm_engine_type, self.tts_engine_type) if not precision_matches_cfg(model_cfg): # Direct construction warns: only S2SPipelineBuilder requires the # precision scope. These torch globals are not applied here. @@ -357,86 +356,14 @@ def _initialize_model(self): # ------------------------------------------------------------------ def _initialize_vllm_omni_backend(self): - """Build the wrapper checkpoint, start the AsyncOmni runtime, and - pre-load the speaker latent. + """Start the AsyncOmni runtime for selected vLLM components. - The wrapper checkpoint (``config.json`` + ``nemotron/`` + ``eartts/``) - is converted lazily on first use under - ``$TMPDIR/_vllm_omni_wrapper`` and reused afterwards; set - ``vllm_omni_config.wrapper_dir`` to keep it somewhere persistent. + Not implemented in this PR. The parent commit on + ``duplex-vllm-omni-on-main`` has the wrapper-checkpoint conversion, + ``OmniRuntime``, and session wiring. Kept as the construction hook so + the native ``_initialize_model`` path already matches the combined form. """ - # Deferred because these reach vLLM-Omni, an optional dependency. - from nemo.collections.speechlm2.inference.vllm_omni.checkpoint import ( - EARTTS_SUBDIR, - build_wrapper_checkpoint, - load_speaker_latent, - write_nemotron_inference_overrides, - ) - from nemo.collections.speechlm2.inference.vllm_omni.runtime import OmniRuntime - - cfg = self.vllm_omni_config or {} - - wrapper_dir = build_wrapper_checkpoint( - self.model_path, - wrapper_dir=cfg.get("wrapper_dir", None), - nemotron_dtype=cfg.get("nemotron_dtype", "float32"), - eartts_precompute_batch_size=int(cfg.get("eartts_precompute_batch_size", 256)), - include_nemotron=self.use_vllm_llm, - include_eartts=self.use_vllm_tts, - ) - self.omni_wrapper_dir = wrapper_dir - - if self.use_vllm_llm: - # Must happen before the stage child loads the checkpoint. - user_boosts = LogitBoosts.user_from_cfg(self.model.stt_model.cfg) - write_nemotron_inference_overrides( - wrapper_dir, - { - "inference_user_pad_boost": user_boosts.pad, - "inference_user_bos_boost": user_boosts.bos, - "inference_user_eos_boost": user_boosts.eos, - }, - ) - - self.omni_runtime = OmniRuntime( - wrapper_dir, - stage_configs_path=cfg.get("stage_configs_path", None), - eartts_stage_configs_path=cfg.get("eartts_stage_configs_path", None), - stage_overrides=cfg.get("stage_overrides", None), - eartts_stage_overrides=cfg.get("eartts_stage_overrides", None), - log_stats=bool(cfg.get("log_stats", False)), - stage_init_timeout=int(cfg.get("stage_init_timeout", 600)), - enable_llm=self.use_vllm_llm, - enable_tts=self.use_vllm_tts, - ) - - if self.use_vllm_tts: - eartts_dir = os.path.join(wrapper_dir, EARTTS_SUBDIR) - with open(os.path.join(eartts_dir, "config.json"), encoding="utf-8") as fh: - eartts_config = json.load(fh) - guidance_enabled = cfg.get("guidance_enabled") - if guidance_enabled is None: - guidance_enabled = eartts_config.get("enable_guidance", True) - guidance_scale = cfg.get("guidance_scale") - if guidance_scale is None: - guidance_scale = eartts_config.get("guidance_scale", 0.5) - self.omni_guidance_enabled = bool(guidance_enabled) - self.omni_guidance_scale = float(guidance_scale) - speaker_name = self.speaker_name or cfg.get("speaker_name") - if speaker_name is None: - raise ValueError( - "tts_engine_type='vllm_omni' requires a speaker_name (set " - "s2s.speaker_name or s2s.vllm_omni_config.speaker_name); the " - "speaker latent is read from the converted EarTTS checkpoint." - ) - self.omni_speaker_latent = load_speaker_latent(eartts_dir, speaker_name) - logging.info( - "vllm_omni speaker_latent: name='%s', shape=%s; CFG=%s scale=%s", - speaker_name, - tuple(self.omni_speaker_latent.shape), - self.omni_guidance_enabled, - self.omni_guidance_scale, - ) + reject_unimplemented_vllm(self.llm_engine_type, self.tts_engine_type) def start_vllm_omni_session( self, @@ -446,66 +373,9 @@ def start_vllm_omni_session( request_id: str, sampling_params: dict[str, float] | None = None, ) -> None: - """Create and attach a per-stream :class:`OmniStreamingSession`. - - Called by the streaming pipeline once it has the system prompt for - the new stream; this replaces native-engine prefill. The session only - enqueues the prefill chunk, so Nemotron does not actually run until - the first :meth:`infer_one_step` call. Per-stream sampling parameters - are fixed when this long-lived vLLM request is created. - - One session class covers all three vLLM combinations: it reads which - components exist off the runtime, which is the same decision that built - the runtime in the first place. - """ - if not self.use_vllm_omni: - return - if self.omni_runtime is None: - raise RuntimeError("vllm_omni backend was not initialized; call _initialize_model first.") - from nemo.collections.speechlm2.inference.vllm_omni.checkpoint import ( - NEMOTRON_SUBDIR, - compute_prefill_len, - ) - from nemo.collections.speechlm2.inference.vllm_omni.session import OmniStreamingSession - - t_prefill = 0 - if self.use_vllm_llm: - nemotron_dir = os.path.join(self.omni_wrapper_dir, NEMOTRON_SUBDIR) - t_prefill = compute_prefill_len(nemotron_dir, system_prompt or "") - - effective_sampling_params = { - "temperature": float(self.temperature), - "top_p": float(self.top_p), - "repetition_penalty": float(self.repetition_penalty), - } - if sampling_params: - effective_sampling_params.update( - {key: float(value) for key, value in sampling_params.items() if key in effective_sampling_params} - ) - - stt = self.model.stt_model - state.omni_session = OmniStreamingSession( - self.omni_runtime, - request_id=request_id, - system_prompt=system_prompt or "", - speaker_latent=self.omni_speaker_latent, - t_prefill=t_prefill, - sampling_params=effective_sampling_params, - special_token_ids=self.special_token_ids, - guidance_enabled=self.omni_guidance_enabled, - guidance_scale=self.omni_guidance_scale, - step_timeout=float((self.vllm_omni_config or {}).get("step_timeout", 60.0)), - profile=self._profile_timing, - # The agent-channel boosts ride along with sampling; the ASR-channel - # ones are applied by the converted model, which reads them from its - # own config at load time. - agent_logit_boosts=LogitBoosts.agent_from_cfg(stt.cfg), - text_token_ids={ - "pad_id": int(stt.text_pad_id), - "bos_id": int(stt.text_bos_id), - "eos_id": int(stt.text_eos_id), - }, - ) + """Attach a per-stream vLLM session (not implemented in this PR).""" + del state, system_prompt, request_id, sampling_params + reject_unimplemented_vllm(self.llm_engine_type, self.tts_engine_type) # ------------------------------------------------------------------ # Per-stream lifecycle diff --git a/nemo/collections/speechlm2/inference/vllm_omni/__init__.py b/nemo/collections/speechlm2/inference/vllm_omni/__init__.py deleted file mode 100644 index 472e8728604d..000000000000 --- a/nemo/collections/speechlm2/inference/vllm_omni/__init__.py +++ /dev/null @@ -1,59 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""NemotronDuplexH + EarTTS model/pipeline package for vLLM-Omni. - -This package lives inside NeMo and plugs into ``vllm-omni`` at runtime -through :func:`register_nemo_voicechat`. The function is auto-invoked in -every vllm-omni subprocess via the ``vllm_omni.general_plugins`` entry -point declared in NeMo's ``pyproject.toml`` (vllm-omni uses ``spawn`` for -stage children, so the entry-point hook is required — PYTHONPATH alone -is not enough). - -The plugin registers three things: - -* HF config ``"eartts"`` → :class:`EarTTSConfig` -* Model arch ``"NemotronDuplexHForCausalLM"`` → - :class:`nemo.collections.speechlm2.inference.vllm_omni.nemotron_duplex_h.nemotron_duplex_h.NemotronDuplexHForCausalLM` -* Model arch ``"EarTTSForCausalLM"`` → - :class:`nemo.collections.speechlm2.inference.vllm_omni.eartts.eartts.EarTTSForCausalLM` -* One-stage pipelines ``model_type = "nemotron_voicechat"`` and ``"eartts"``. - -Bundled deploy YAMLs for the independent engines live under ``deploy/``. -""" - -from __future__ import annotations - -from pathlib import Path - - -def default_deploy_yaml() -> Path: - """Return the absolute path to the bundled ``nemotron_voicechat.yaml``.""" - return Path(__file__).resolve().parent / "deploy" / "nemotron_voicechat.yaml" - - -def default_eartts_deploy_yaml() -> Path: - """Return the absolute path to the bundled single-stage ``eartts.yaml``.""" - return Path(__file__).resolve().parent / "deploy" / "eartts.yaml" - - -from nemo.collections.speechlm2.inference.vllm_omni.register import ( - register_nemo_voicechat, -) - -__all__ = [ - "default_deploy_yaml", - "default_eartts_deploy_yaml", - "register_nemo_voicechat", -] diff --git a/nemo/collections/speechlm2/inference/vllm_omni/checkpoint.py b/nemo/collections/speechlm2/inference/vllm_omni/checkpoint.py deleted file mode 100644 index 00264d78afae..000000000000 --- a/nemo/collections/speechlm2/inference/vllm_omni/checkpoint.py +++ /dev/null @@ -1,311 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Building and reading the vLLM-Omni wrapper checkpoint. - -``AsyncOmni(model=...)`` wants a directory laid out as:: - - / - config.json # {"model_type": "nemotron_voicechat"} - nemotron/ # converted NemotronDuplexH checkpoint - eartts/ # converted EarTTS checkpoint + speaker_latents/ - -Everything that reads or writes that layout lives here: conversion, the source -fingerprint that makes incremental builds safe, the small config patches -applied before an engine starts, and the two loaders that read values back out -of a built wrapper. Nothing here imports vLLM, so it can be exercised without -an engine. -""" - -import hashlib -import json -import os -import shutil -import tempfile -from pathlib import Path -from typing import Any - -import torch - -from nemo.utils import logging - -NEMOTRON_SUBDIR = "nemotron" -EARTTS_SUBDIR = "eartts" -_WRAPPER_CONFIG = {"model_type": "nemotron_voicechat"} -_SOURCE_MANIFEST = ".nemo_source.json" - - -def _checkpoint_fingerprint(model_path: str) -> dict[str, Any]: - """Cheap content identity for safe incremental wrapper construction.""" - root = Path(model_path) - config_path = root / "config.json" - if not config_path.is_file(): - raise FileNotFoundError(f"Checkpoint config not found: {config_path}") - weights = sorted(root.glob("*.safetensors")) - if not weights: - raise FileNotFoundError(f"No safetensors weights found in checkpoint: {root}") - return { - "config_sha256": hashlib.sha256(config_path.read_bytes()).hexdigest(), - "weights": [{"name": path.name, "size": path.stat().st_size} for path in weights], - } - - -def _read_source_manifest(path: str) -> dict[str, Any] | None: - try: - with open(path, encoding="utf-8") as fh: - value = json.load(fh) - return value if isinstance(value, dict) else None - except (FileNotFoundError, json.JSONDecodeError, OSError): - return None - - -def build_wrapper_checkpoint( - model_path: str, - wrapper_dir: str | None = None, - *, - nemotron_dtype: str = "float32", - eartts_precompute_batch_size: int = 256, - include_nemotron: bool = True, - include_eartts: bool = True, -) -> str: - """Build a wrapper checkpoint directory consumed by ``AsyncOmni(model=...)``. - - Layout:: - - / - config.json # {"model_type": "nemotron_voicechat"} - nemotron/ # converted NemotronDuplexH checkpoint - eartts/ # converted EarTTS checkpoint + speaker_latents/ - - Args: - model_path: Path to the source NemotronVoiceChat HF-format checkpoint - directory (``config.json`` + ``model.safetensors``). - wrapper_dir: Where to put the wrapper directory. Defaults to - ``/_vllm_omni_wrapper``, where ```` is - :func:`tempfile.gettempdir` and so honours ``$TMPDIR``. - nemotron_dtype: dtype for the converted Nemotron checkpoint. - eartts_precompute_batch_size: batch size used when baking out the - EarTTS subword-encoder lookup table. - - Returns: - Absolute path to the wrapper directory. If the wrapper directory - already exists and looks complete, the existing one is returned and - nothing is re-converted. - """ - src = os.path.normpath(model_path) - if wrapper_dir is None: - wrapper_dir = os.path.join(tempfile.gettempdir(), os.path.basename(src) + "_vllm_omni_wrapper") - wrapper_dir = os.path.abspath(wrapper_dir) - - nemotron_dir = os.path.join(wrapper_dir, NEMOTRON_SUBDIR) - eartts_dir = os.path.join(wrapper_dir, EARTTS_SUBDIR) - config_path = os.path.join(wrapper_dir, "config.json") - manifest_path = os.path.join(wrapper_dir, _SOURCE_MANIFEST) - source_fingerprint = _checkpoint_fingerprint(src) - - nemotron_ready = ( - os.path.isdir(nemotron_dir) - and os.path.isfile(os.path.join(nemotron_dir, "config.json")) - and os.path.isfile(os.path.join(nemotron_dir, "model.safetensors")) - ) - eartts_ready = ( - os.path.isdir(eartts_dir) - and os.path.isfile(os.path.join(eartts_dir, "config.json")) - and os.path.isfile(os.path.join(eartts_dir, "model.safetensors")) - ) - config_ready = os.path.isfile(config_path) - - if not include_nemotron and not include_eartts: - raise ValueError("At least one vLLM-Omni component must be requested") - - manifest = _read_source_manifest(manifest_path) - if manifest is None: - adding_to_unverified_partial_wrapper = (include_nemotron and not nemotron_ready and eartts_ready) or ( - include_eartts and not eartts_ready and nemotron_ready - ) - if adding_to_unverified_partial_wrapper: - raise ValueError( - f"Cannot safely add a component to wrapper {wrapper_dir}: " - f"{_SOURCE_MANIFEST} is missing, so the existing component's " - "source checkpoint cannot be verified. Use a fresh wrapper_dir." - ) - elif manifest.get("source") != source_fingerprint: - logging.warning( - "Wrapper source checkpoint changed; rebuilding converted components in %s", - wrapper_dir, - ) - for component_dir in (nemotron_dir, eartts_dir): - if os.path.isdir(component_dir): - shutil.rmtree(component_dir) - nemotron_ready = False - eartts_ready = False - manifest = None - else: - if include_nemotron and nemotron_ready and manifest.get("nemotron", {}).get("dtype") != nemotron_dtype: - shutil.rmtree(nemotron_dir) - nemotron_ready = False - if ( - include_eartts - and eartts_ready - and manifest.get("eartts", {}).get("precompute_batch_size") != eartts_precompute_batch_size - ): - shutil.rmtree(eartts_dir) - eartts_ready = False - - if (not include_nemotron or nemotron_ready) and (not include_eartts or eartts_ready) and config_ready: - if manifest is None: - # Wrapper is complete but carries no source manifest. Stamp one - # now: no component is being added, so this cannot mix checkpoints. - logging.warning( - "Adopting vLLM-Omni wrapper without source manifest: %s", - wrapper_dir, - ) - adopted_manifest: dict[str, Any] = {"source": source_fingerprint} - if nemotron_ready: - adopted_manifest["nemotron"] = {"dtype": nemotron_dtype} - if eartts_ready: - adopted_manifest["eartts"] = {"precompute_batch_size": eartts_precompute_batch_size} - with open(manifest_path, "w", encoding="utf-8") as fh: - json.dump(adopted_manifest, fh, indent=2, sort_keys=True) - logging.info(f"Reusing existing vllm-omni wrapper checkpoint at {wrapper_dir}") - return wrapper_dir - - os.makedirs(wrapper_dir, exist_ok=True) - - if include_nemotron and not nemotron_ready: - # Convert the Nemotron LLM with the existing DuplexSTT converter. - # That converter's output (HF NemotronH config + filtered weights) - # is consumed directly by NemotronDuplexHForCausalLM's WeightsMapper. - if os.path.isdir(nemotron_dir): - shutil.rmtree(nemotron_dir) - logging.info(f"Converting Nemotron LLM into {nemotron_dir} ...") - from nemo.collections.speechlm2.inference.vllm_omni.scripts.convert_duplex_stt_checkpoint import ( - convert_to_vllm_format as convert_nemotron, - ) - - convert_nemotron( - checkpoint_path=src, - output_dir=nemotron_dir, - dtype=nemotron_dtype, - ) - - if include_eartts and not eartts_ready: - if os.path.isdir(eartts_dir): - shutil.rmtree(eartts_dir) - logging.info(f"Converting EarTTS into {eartts_dir} ...") - from nemo.collections.speechlm2.inference.vllm_omni.scripts.convert_duplex_eartts_checkpoint import ( - convert_to_vllm_format as convert_eartts, - ) - - convert_eartts( - outdir=eartts_dir, - config=os.path.join(src, "config.json"), - model_path=os.path.join(src, "model.safetensors"), - precompute_batch_size=eartts_precompute_batch_size, - ) - - if not config_ready: - with open(config_path, "w", encoding="utf-8") as fh: - json.dump(_WRAPPER_CONFIG, fh, indent=2) - - completed_manifest: dict[str, Any] = {"source": source_fingerprint} - if os.path.isfile(os.path.join(nemotron_dir, "model.safetensors")): - completed_manifest["nemotron"] = {"dtype": nemotron_dtype} - if os.path.isfile(os.path.join(eartts_dir, "model.safetensors")): - completed_manifest["eartts"] = {"precompute_batch_size": eartts_precompute_batch_size} - with open(manifest_path, "w", encoding="utf-8") as fh: - json.dump(completed_manifest, fh, indent=2, sort_keys=True) - - return wrapper_dir - - -def write_nemotron_inference_overrides(wrapper_dir: str, overrides: dict[str, Any]) -> None: - """Update inference settings in the converted Nemotron ``config.json``. - - ``NemotronDuplexHForCausalLM`` reads some settings off its HF config at - load time -- the user-channel logit boosts, which cannot be delivered per - request because the ASR head's logits never reach vLLM's sampler. Those are - still chosen per run in the inference yaml, so the small JSON is rewritten - here before the stage child starts, rather than re-converting weights. - - Keys whose value is ``None`` are removed, so clearing a boost in the config - clears it in the engine too. - """ - config_path = os.path.join(wrapper_dir, "nemotron", "config.json") - if not os.path.isfile(config_path): - return - with open(config_path, encoding="utf-8") as fh: - config = json.load(fh) - - changed = False - for key, value in overrides.items(): - if value is None: - if config.pop(key, None) is not None: - changed = True - elif config.get(key) != value: - config[key] = value - changed = True - if not changed: - return - - with open(config_path, "w", encoding="utf-8") as fh: - json.dump(config, fh, indent=2) - logging.info(f"Updated Nemotron inference overrides in {config_path}: {overrides}") - - -def load_speaker_latent(eartts_dir: str, speaker_name: str) -> torch.Tensor: - """Load ``/speaker_latents/.pt`` (saved by the - EarTTS converter) and return a contiguous CPU tensor of shape - ``(Tref, hidden_size)``. - """ - latents_dir = os.path.join(eartts_dir, "speaker_latents") - latent_path = os.path.join(latents_dir, f"{speaker_name}.pt") - if not os.path.isfile(latent_path): - available = [] - if os.path.isdir(latents_dir): - available = sorted( - os.path.splitext(name)[0] for name in os.listdir(latents_dir) if name.endswith(".pt") - ) - raise FileNotFoundError( - f"Speaker latent for '{speaker_name}' not found at {latent_path}. " - f"Registered speakers: {available or '(none)'}. " - "Pick a speaker_name present in the EarTTS checkpoint, or re-run the " - "EarTTS converter on a checkpoint that contains the requested " - "audio_prompt_latents." - ) - latent = torch.load(latent_path, weights_only=False) - if isinstance(latent, torch.Tensor) and latent.dim() == 3: - latent = latent[0] - if not isinstance(latent, torch.Tensor) or latent.dim() != 2: - raise ValueError( - f"Expected speaker latent at {latent_path} to be a 2-D tensor [Tref, hidden], " - f"got {type(latent).__name__} with shape " - f"{tuple(latent.shape) if isinstance(latent, torch.Tensor) else 'n/a'}" - ) - return latent.detach().to(torch.float32).cpu().contiguous() - - -def compute_prefill_len(model_dir: str, system_prompt: str) -> int: - """Length of the prefill chunk fed to NemotronDuplexH for a - given system prompt. Mirrors the in-model tokenization: - ``[BOS] + tokenizer.encode(prompt) + [EOS]``. - """ - from transformers import AutoTokenizer - - from nemo.collections.speechlm2.inference.vllm_omni.nemotron_duplex_h.nemotron_duplex_h import ( - NemotronDuplexHForCausalLM, - ) - - tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True) - return NemotronDuplexHForCausalLM.compute_prefix_len(tokenizer, system_prompt) diff --git a/nemo/collections/speechlm2/inference/vllm_omni/deploy/eartts.yaml b/nemo/collections/speechlm2/inference/vllm_omni/deploy/eartts.yaml deleted file mode 100644 index 34750c82eb1d..000000000000 --- a/nemo/collections/speechlm2/inference/vllm_omni/deploy/eartts.yaml +++ /dev/null @@ -1,45 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# Single-stage EarTTS runtime. The model path passed to AsyncOmni is the -# converted ``eartts/`` directory itself. -async_chunk: false -trust_remote_code: true -enable_prefix_caching: false -enable_chunked_prefill: false -distributed_executor_backend: uni - -stages: - - stage_id: 0 - # One conditional + one unconditional request per VoiceChat stream. - # The no-CFG path uses only one slot. - max_num_seqs: 2 - max_num_batched_tokens: 2048 - max_model_len: 2048 - gpu_memory_utilization: 0.30 - enforce_eager: false - # Keep the synchronous scheduler: paired streaming requests must observe - # the same completed step before their next updates are admitted. - async_scheduling: false - skip_tokenizer_init: true - dtype: float32 - devices: "0" - compilation_config: - cudagraph_mode: PIECEWISE - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 2048 - detokenize: false diff --git a/nemo/collections/speechlm2/inference/vllm_omni/deploy/nemotron_voicechat.yaml b/nemo/collections/speechlm2/inference/vllm_omni/deploy/nemotron_voicechat.yaml deleted file mode 100644 index 536231160e94..000000000000 --- a/nemo/collections/speechlm2/inference/vllm_omni/deploy/nemotron_voicechat.yaml +++ /dev/null @@ -1,62 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# Deploy config for the single-stage ``nemotron_voicechat`` NemotronDuplexH -# streaming pipeline. EarTTS is served by ``eartts.yaml`` in a second -# AsyncOmni engine and is coordinated by NeMo. -# -# The model directory passed to ``AsyncOmni(model=...)`` is a small -# user-managed wrapper directory with this layout:: -# -# / -# config.json # {"model_type": "nemotron_voicechat"} -# nemotron/ # directory or symlink → Nemotron ckpt -# eartts/ # loaded separately by the EarTTS runtime -# -# ``model_type = nemotron_voicechat`` dispatches to the one-stage pipeline. -# ``model_subdir`` / ``tokenizer_subdir`` tell the engine to load the nested -# Nemotron checkpoint; see -# ``vllm_omni/engine/stage_init_utils.py:_resolve_model_tokenizer_paths``. -async_chunk: false -trust_remote_code: true -enable_prefix_caching: false -enable_chunked_prefill: false -distributed_executor_backend: uni - -stages: - # Nemotron-Duplex-H (autoregressive text + optional ASR/function channel). - # PIECEWISE compilation only; see the model docstring for why FULL - # cudagraph mode is unsafe with this streaming setup. - - stage_id: 0 - max_num_seqs: 1 - # Prompt tokens plus one token per 80 ms frame. - max_num_batched_tokens: 2048 - max_model_len: 2048 - gpu_memory_utilization: 0.45 - enforce_eager: false - async_scheduling: true - devices: "0" - model_subdir: nemotron - tokenizer_subdir: nemotron - engine_extras: - logits_processors: - - nemo.collections.speechlm2.inference.vllm_omni.nemotron_duplex_h.sampling:SharedTextSamplingLogitsProcessor - compilation_config: - cudagraph_mode: PIECEWISE - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 2048 - detokenize: false diff --git a/nemo/collections/speechlm2/inference/vllm_omni/eartts/__init__.py b/nemo/collections/speechlm2/inference/vllm_omni/eartts/__init__.py deleted file mode 100644 index b8f12a085699..000000000000 --- a/nemo/collections/speechlm2/inference/vllm_omni/eartts/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from nemo.collections.speechlm2.inference.vllm_omni.eartts.configuration_eartts import EarTTSConfig - -__all__ = ["EarTTSConfig"] diff --git a/nemo/collections/speechlm2/inference/vllm_omni/eartts/configuration_eartts.py b/nemo/collections/speechlm2/inference/vllm_omni/eartts/configuration_eartts.py deleted file mode 100644 index be85e27cde3b..000000000000 --- a/nemo/collections/speechlm2/inference/vllm_omni/eartts/configuration_eartts.py +++ /dev/null @@ -1,217 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""HuggingFace-style configuration for the EarTTS model. - -The configuration mirrors the fields that -:class:`nemo.collections.speechlm2.inference.vllm_omni.eartts.eartts.EarTTSForCausalLM` -reads from ``vllm_config.model_config.hf_config``: - -* Gemma3 backbone fields consumed by ``Gemma3Model`` (``hidden_size``, - ``intermediate_size``, ``num_hidden_layers``, ``num_attention_heads``, - ``num_key_value_heads``, ``head_dim``, ``vocab_size``, - ``max_position_embeddings``, ``query_pre_attn_scalar``, - ``attention_bias``, ``rms_norm_eps``, ``layer_types``, - ``sliding_window``, ``rope_local_base_freq``, ``rope_theta``, - ``rope_scaling``, ``hidden_activation``, ``tie_word_embeddings``, - ``final_logit_softcapping``, ``attn_logits_soft_cap``, - ``use_bidirectional_attention``, ``is_causal``). - -* MaskGIT sampler fields (``num_quantizers``, ``codebook_size``, - ``num_iter``, ``top_p_or_k``, ``noise_scale``, ``exponent``, - ``latent_size``, ``mog_low_rank``, ``mog_num_layers``, - ``mog_num_predictions``, ``mog_min_log_std``, ``mog_eps``). - -* Subword embedding / fusion fields consumed by - :class:`EarTTSInputEmbedding` (``emb_vocab_size``, - ``use_gated_fusion_for_text_audio``, - ``use_audio_prompt_frozen_projection``). The original NeMo model used - a character-aware subword encoder + subword-flag + BOS/EOS additive - embeddings; all of those operations are deterministic per token id - and are baked out at checkpoint-conversion time into a single - ``nn.Embedding`` of size ``(emb_vocab_size, hidden_size)``. - -vLLM's ``patch_rope_parameters`` (transformers-v4 path) auto-populates -``config.rope_parameters`` from ``rope_scaling`` + ``rope_theta`` during -config loading, so the Gemma3 backbone (which expects -``config.rope_parameters``) works without any extra plumbing here. -""" - -from typing import Optional - -from transformers import AutoConfig, PretrainedConfig - - -class EarTTSConfig(PretrainedConfig): - model_type = "eartts" - - def __init__( - self, - # Gemma 3 backbone - hidden_size: int = 1152, - context_hidden_size: int = 1536, - intermediate_size: int = 4608, - num_hidden_layers: int = 28, - num_attention_heads: int = 16, - num_key_value_heads: int = 16, - head_dim: int = 72, - # ``vocab_size`` controls the width of the logits tensor returned - # by ``EarTTSForCausalLM.compute_logits`` — vLLM's sampler and - # ``LogitsProcessor`` size their working buffers from - # ``config.vocab_size``, so it must match the dummy logits the - # model produces. The model emits a 2-class placeholder - # (``[0, -inf]``) so the sampler's argmax always picks index 0; - # the real audio output is the codes tensor exposed via - # ``make_omni_output``. ``2`` is the minimum that keeps vLLM's - # sampler happy. - vocab_size: int = 2, - max_position_embeddings: int = 131072, - # MaskGIT / MoG sampling - num_quantizers: int = 31, - codebook_size: int = 1024, - num_iter: int = 8, - top_p_or_k: float = 0.8, - noise_scale: float = 0.8, - exponent: float = 3.0, - latent_size: int = 512, - mog_low_rank: int = 64, - mog_num_layers: int = 3, - mog_num_predictions: int = 1024, - mog_min_log_std: float = -4.0, - mog_eps: float = 1e-6, - # Classifier-free guidance. The converter exports both fields and - # ``null_emb``; the runtime reads them unless the request overrides - # ``cfg_scale``. - enable_guidance: bool = False, - guidance_scale: float = 0.5, - # Gemma3-specific attributes required by Gemma3Model - query_pre_attn_scalar: float = 256.0, - attention_bias: bool = False, - rms_norm_eps: float = 1e-6, - layer_types: Optional[list] = None, - sliding_window: Optional[int] = 4096, - rope_local_base_freq: float = 10000.0, - # NeMo / EarTTS uses 1M for the global-attention RoPE base. - rope_theta: float = 1000000.0, - rope_scaling: Optional[dict] = None, - hidden_activation: str = "gelu_pytorch_tanh", - tie_word_embeddings: bool = True, - final_logit_softcapping: Optional[float] = None, - attn_logits_soft_cap: Optional[float] = None, - use_bidirectional_attention: bool = False, - is_causal: bool = True, - # Subword encoding. The character-aware subword encoder / - # subword-flag / BOS-EOS embedding tables that NeMo applied at - # runtime are precomputed at checkpoint conversion time into a - # single ``(emb_vocab_size, hidden_size)`` lookup, so only the - # vocab size and the audio-side fusion / projection toggles - # remain as runtime config. - emb_vocab_size: int = 151936, - use_gated_fusion_for_text_audio: bool = True, - use_audio_prompt_frozen_projection: bool = False, - # HF-canonical model dtype (replaces the deprecated - # ``torch_dtype``). Forwarded to ``PretrainedConfig`` so it is - # converted into a real ``torch.dtype`` and exposed as - # ``config.dtype``. - dtype: str = "float32", - # Text-channel specials, copied from the source VoiceChat tokenizer - # at conversion time. Used to pad prefill text and to force codec - # silence when the incoming text token is EOS. - pad_token_id: Optional[int] = None, - eos_token_id: Optional[int] = None, - **kwargs, - ): - # Gemma3 backbone - self.hidden_size = hidden_size - self.context_hidden_size = context_hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.num_key_value_heads = num_key_value_heads - self.head_dim = head_dim - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - - # MaskGIT / MoG sampling - self.num_quantizers = num_quantizers - self.codebook_size = codebook_size - self.num_iter = num_iter - self.top_p_or_k = top_p_or_k - self.noise_scale = noise_scale - self.exponent = exponent - self.latent_size = latent_size - self.mog_low_rank = mog_low_rank - self.mog_num_layers = mog_num_layers - self.mog_num_predictions = mog_num_predictions - self.mog_min_log_std = mog_min_log_std - self.mog_eps = mog_eps - self.enable_guidance = enable_guidance - self.guidance_scale = guidance_scale - - # Gemma3-specific attributes - self.query_pre_attn_scalar = query_pre_attn_scalar - self.attention_bias = attention_bias - self.rms_norm_eps = rms_norm_eps - # Default all layers to global attention if not specified. - self.layer_types = ( - layer_types if layer_types is not None else ["full_attention"] * num_hidden_layers - ) - self.sliding_window = sliding_window - self.rope_local_base_freq = rope_local_base_freq - self.rope_theta = rope_theta - self.rope_scaling = rope_scaling - self.hidden_activation = hidden_activation - self.final_logit_softcapping = final_logit_softcapping - self.attn_logits_soft_cap = attn_logits_soft_cap - self.use_bidirectional_attention = use_bidirectional_attention - self.is_causal = is_causal - - # Subword encoding (precomputed lookup; see class docstring). - self.emb_vocab_size = emb_vocab_size - self.use_gated_fusion_for_text_audio = use_gated_fusion_for_text_audio - self.use_audio_prompt_frozen_projection = use_audio_prompt_frozen_projection - - # Forward HF-owned fields (``tie_word_embeddings`` and ``dtype``) - # to ``PretrainedConfig`` so they round-trip through - # save/load_pretrained and are visible as ``config.dtype`` / - # ``config.tie_word_embeddings``. Without this, user-supplied - # values silently fall back to PretrainedConfig's defaults. - super().__init__( - tie_word_embeddings=tie_word_embeddings, - dtype=dtype, - pad_token_id=pad_token_id, - eos_token_id=eos_token_id, - **kwargs, - ) - - -def register_eartts_config() -> None: - """Make ``model_type: "eartts"`` resolvable by ``AutoConfig``. - - Idempotent, because it has two callers by necessity and either may run - first: :func:`register_nemo_voicechat` (the plugin entry point) and the - import below. The import-time call is what covers subprocesses that reach - this module without the plugin -- ``StageEngineCoreProc`` unpickles a - config and calls ``AutoConfig.from_pretrained`` -- and mirrors the pattern - used by other vllm-omni custom configs (voxcpm, fish_speech, - mammoth_moda2, ...). - """ - try: - AutoConfig.register(EarTTSConfig.model_type, EarTTSConfig) - except ValueError: - # transformers raises when the model_type is already registered, which - # is the expected outcome for whichever caller runs second. - pass - - -register_eartts_config() diff --git a/nemo/collections/speechlm2/inference/vllm_omni/eartts/eartts.py b/nemo/collections/speechlm2/inference/vllm_omni/eartts/eartts.py deleted file mode 100644 index d88d9e066313..000000000000 --- a/nemo/collections/speechlm2/inference/vllm_omni/eartts/eartts.py +++ /dev/null @@ -1,1945 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Inference-only EarTTS model definition for vLLM-Omni. - -The model architecture (RMSNorm, MLP, MLPLayer, GatedProjectedSumRMSNorm, -PrecomputedSubwordEmbedding, EarTTSInputEmbedding, MoGHead, -MaskGITSampler, EarTTSModel) matches the PyTorch EarTTS modules. -Classifier-free guidance (CFG) is driven by per-request metadata and keeps -its role, pair, and scale contract in model-owned stable-address buffers. - -The original NeMo model used a character-aware subword encoder (a small -transformer over per-character embeddings) followed by additive -subword-continuation and BOS/EOS flag embeddings to embed text tokens. -Those operations are deterministic per token id, so the checkpoint -converter runs them once over the full vocabulary and stores the result -as a single ``nn.Embedding`` (see :class:`PrecomputedSubwordEmbedding`). - -The outer :class:`EarTTSForCausalLM` exposes the minimal vLLM-Omni -preprocess/postprocess hooks. Per-request inputs are passed via -``additional_information``. - -Inputs (only one mode — streaming text token ids), named after the -categories of :class:`~vllm_omni.data_entry_keys.OmniPayload` so that the -same two keys work whether they arrive on the request or over an -inter-stage connector, which accepts nothing outside that schema: - -* ``embed.voice`` (prefill only): Tensor of shape - ``(Tref, hidden_size)``. The user-supplied speaker latent that - replaces ``embed_code(rvq_sum(acoustic_tokens))`` on every pre-BOS - prefill position. ``Tref`` is also the prefill placeholder length - (the user passes ``prompt_token_ids = [0] * Tref``). -* ``ids.output`` (every decode step): Python ``list[int]`` of the text - tokens the producer has sent most recently. :meth:`preprocess` takes - its **last** entry, so decode step ``k`` consumes ``t_k`` whether the - producer sends one token per step or a growing history. - -There is no whole-utterance text path: callers must always provide token -ids per step via the streaming-text contract above. - -Per-step flow: - -1. ``preprocess`` writes the per-token tensors consumed by - :class:`EarTTSInputEmbedding` — ``acoustic_tokens (BTx31)``, - ``text_tokens (BT)``, ``text_mask (BT)``, ``bos_mask (BT)``, - ``speaker_latent (BT x hidden_size)`` — into the model-owned - static-address buffers at the request's flat-batch offset. Returns - placeholder ``input_ids`` and the ``inputs_embeds`` slice it - received from the runner unchanged (the actual embedding is - computed inside the compiled ``forward``; the buffer's contents - are ignored). - - Prefill is fully derived from ``speaker_latent``: - - * ``acoustic_tokens`` = ``model.sil_tokens`` broadcast to every - prefill position (only the BOS frame's audio embedding actually - contributes to the model output; the rest are replaced by the - speaker latent inside :class:`EarTTSInputEmbedding`). - * ``text_tokens`` = ``[PAD] * (Tref - 1) + [EOS]``. - * ``text_mask`` = ``[0] * (Tref - 2) + [1, 1]``. - * ``bos_mask`` = ``[0] * (Tref - 1) + [1]``. - * ``speaker_latent`` = the user-supplied ``embed.voice`` tensor. - - Decode each step (chooses ``acoustic_tokens`` in this order): - - * ``text_token == EOS`` (``2``) → ``model.sil_tokens``. - * First decode step (``ear_decode_offset == 0``) → the acoustic pad - id (``codebook_size``) broadcast across all quantizers. - * Otherwise → previous-step codes stashed by :meth:`postprocess` - as ``last_acoustic_codes``. - - ``text_tokens = ids.output[-1]``, - ``text_mask = 1``, ``bos_mask = 0``, - ``speaker_latent = 0`` (no replacement on decode). - -2. ``forward`` slices the buffers up to ``num_tokens`` and calls the - compiled :class:`EarTTSModel` (embedding + Gemma3 backbone). The - compiled :class:`EarTTSSamplerModel` (MaskGIT) is invoked - conditionally on decode positions. Generated codes are copied into - a stable ``_out_codes`` buffer for :meth:`make_omni_output`. - -3. ``compute_logits`` returns trivial logits so vLLM's standard - sampler always picks index ``0`` — the actual audio output is the - codes tensor exposed via :meth:`make_omni_output`. - -4. ``postprocess`` stashes the last frame's codes as - ``last_acoustic_codes`` for the next step's :meth:`preprocess`. -""" - -import bisect -import hashlib -import math -from collections.abc import Iterable -from typing import Any, Optional, Union - -import numpy as np -import torch -from torch import nn -from transformers.generation.logits_process import ( - TopKLogitsWarper, - TopPLogitsWarper, -) -from vllm.compilation.backends import set_model_tag -from vllm.compilation.decorators import ( - ignore_torch_compile, - support_torch_compile, -) -from vllm.config import CUDAGraphMode, VllmConfig -from vllm.forward_context import BatchDescriptor, get_forward_context -from vllm.model_executor.models.gemma3 import Gemma3Model -from vllm.model_executor.models.interfaces import SupportsPP -from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper -from vllm.sequence import IntermediateTensors - -from vllm_omni.model_executor.models.output_templates import OmniOutput - - - -def _prepare_cfg_sampling_batch( - hidden_states: torch.Tensor, - cfg_enabled: torch.Tensor, - cfg_is_uncond: torch.Tensor, - cfg_pair_id: torch.Tensor, - cfg_scale: torch.Tensor, - valid: torch.Tensor, -) -> tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, -]: - """Order a complete CFG batch and describe its logical pairs. - - All operations stay on-device: CUDA graph replay can therefore change the - role/pair/scale buffers without depending on Python-side state. A batch is - guided only when every valid row has exactly one enabled opposite-role - companion with the same pair id. Padded rows are ignored; the conditional - row's scale is authoritative for each pair. - - Returns ordered hidden states, ordered role/scale tensors, an active-row - mask, each row's partner and conditional representative indices, and the - inverse permutation used to restore the runner's original row order. - """ - batch_size = int(hidden_states.shape[0]) - identity = torch.arange(batch_size, device=hidden_states.device) - - pair_match = ( - valid[:, None] - & valid[None, :] - & cfg_enabled[:, None] - & cfg_enabled[None, :] - & (cfg_pair_id[:, None] == cfg_pair_id[None, :]) - & (cfg_is_uncond[:, None] != cfg_is_uncond[None, :]) - ) - partner_count = pair_match.sum(dim=1) - partner = pair_match.to(torch.long).argmax(dim=1) - complete = valid.any() & ((~valid) | (cfg_enabled & (partner_count == 1))).all() - - # Lexicographic order: valid rows first, then conditional before - # unconditional, with pair ids sorted identically inside both role blocks. - order = torch.argsort(cfg_pair_id, stable=True) - order = order[torch.argsort(cfg_is_uncond[order].to(torch.long), stable=True)] - order = order[torch.argsort((~valid[order]).to(torch.long), stable=True)] - order = torch.where(complete, order, identity) - inverse_order = torch.argsort(order) - - hidden_states = hidden_states[order] - cfg_is_uncond = cfg_is_uncond[order] - cfg_pair_id = cfg_pair_id[order] - cfg_scale = cfg_scale[order] - valid = valid[order] - active = complete & valid - - ordered_pair_match = ( - active[:, None] - & active[None, :] - & (cfg_pair_id[:, None] == cfg_pair_id[None, :]) - & (cfg_is_uncond[:, None] != cfg_is_uncond[None, :]) - ) - partner = ordered_pair_match.to(torch.long).argmax(dim=1) - conditional_rep = torch.where(cfg_is_uncond, partner, identity) - conditional_scale = torch.where( - cfg_is_uncond, - cfg_scale[partner], - cfg_scale, - ) - return ( - hidden_states, - cfg_is_uncond, - conditional_scale, - active, - partner, - conditional_rep, - inverse_order, - ) - - -def _apply_cfg_after_mlp( - x: torch.Tensor, - cfg_is_uncond: torch.Tensor, - cfg_scale: torch.Tensor, - cfg_active: torch.Tensor, - cfg_partner: torch.Tensor, -) -> torch.Tensor: - """Apply EarTTS CFG to MoG MLP outputs before all projections.""" - partner_x = x[cfg_partner] - conditional_x = torch.where(cfg_is_uncond[:, None], partner_x, x) - unconditional_x = torch.where(cfg_is_uncond[:, None], x, partner_x) - guided_x = conditional_x + cfg_scale[:, None].to(x.dtype) * ( - conditional_x - unconditional_x - ) - return torch.where(cfg_active[:, None], guided_x, x) - - -# --------------------------------------------------------------------------- -# Shared EarTTS building blocks, matching the native DuplexEARTTS modules. -# --------------------------------------------------------------------------- - - -class RMSNorm(nn.Module): - def __init__(self, dim: int, eps: float = 1e-6): - super().__init__() - self.eps = eps - self.weight = nn.Parameter(torch.zeros(dim)) - - def _norm(self, x): - return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) - - def forward(self, x): - # Normalize in fp32 and cast back at the end, so low-precision - # activations do not lose the mean-square accumulation. - output = self._norm(x.float()) - # Llama does x.to(float16) * w whilst Gemma3 is (x * w).to(float16) - output = output * (1.0 + self.weight.float()) - return output.type_as(x) - - -class MLP(nn.Module): - def __init__( - self, - hidden_size: int, - intermediate_size: int, - ): - super().__init__() - self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) - self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) - self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) - self.act_fn = nn.GELU(approximate="tanh") - - def forward(self, x: torch.Tensor) -> torch.Tensor: - down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) - return down_proj - - -class MLPLayer(nn.Module): - def __init__( - self, - hidden_size: int, - intermediate_size: int, - eps: float = 1e-6, - ): - super().__init__() - self.pre_norm = RMSNorm(hidden_size, eps=eps) - self.mlp = MLP(hidden_size, intermediate_size) - self.post_norm = RMSNorm(hidden_size, eps=eps) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - y = self.pre_norm(x) - y = self.mlp(y) - y = self.post_norm(y) - x = x + y - return x - - -class GatedProjectedSumRMSNorm(nn.Module): - def __init__( - self, - audio_dim, - text_dim, - hidden_dim, - final_norm=True, - num_codebooks=31, - init_residual_scale=0.5, - ): - super().__init__() - self.num_codebooks = num_codebooks - - self.audio_proj = nn.Linear(audio_dim, hidden_dim) - self.text_proj = nn.Linear(text_dim, hidden_dim) - - nn.init.normal_(self.audio_proj.weight, mean=0.0, std=0.015) - nn.init.zeros_(self.audio_proj.bias) - nn.init.normal_(self.text_proj.weight, mean=0.0, std=0.015) - nn.init.zeros_(self.text_proj.bias) - - # FP32 gate params - self.gate = nn.Parameter( - torch.zeros(hidden_dim, dtype=torch.float32), requires_grad=False - ) - self.residual_scale = nn.Parameter( - torch.tensor(init_residual_scale, dtype=torch.float32), - requires_grad=False, - ) - - self.final_norm = RMSNorm(hidden_dim) if final_norm else nn.Identity() - - def forward(self, audio_emb, text_emb): - 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) - - dtype = audio_h.dtype - - gate = torch.sigmoid(self.gate) # FP32 - res = torch.sigmoid(self.residual_scale) # FP32 - - h = gate.to(dtype) * audio_h + (1 - gate).to(dtype) * text_h - h = res.to(dtype) * h - h = self.final_norm(h.float()).to(dtype) - - return h - - -class PrecomputedSubwordEmbedding(nn.Module): - """Per-token text embedding lookup baked out at checkpoint-conversion time. - - The original NeMo model embeds text with a character-aware subword - encoder (a small transformer over per-character embeddings) followed - by additive subword-continuation and BOS/EOS flag embeddings. All of - those operations are deterministic per token id, so the converter - runs them once over the full vocabulary and stores the result here - as a single ``nn.Embedding``. - """ - - def __init__(self, vocab_size: int, hidden_size: int): - super().__init__() - self.embed_subwords = nn.Embedding(vocab_size, hidden_size) - - def forward(self, subword_ids: torch.Tensor) -> torch.Tensor: - return self.embed_subwords(subword_ids) - - -class EarTTSInputEmbedding(nn.Module): - """Module that takes text tokens, audio tokens and prepares input - embedding for EarTTS model. - """ - - def __init__(self, config): - super().__init__() - - hidden_size = config.hidden_size - vocab_size = config.emb_vocab_size - - # allows to embed acoustic tokens into a single embeddings - self.rvq_embs = nn.ModuleList( - [ - nn.Embedding(config.codebook_size + 1, config.latent_size) - for _ in range(config.num_quantizers) - ] - ) - self.embed_code = nn.Linear(config.latent_size, hidden_size, bias=False) - # Pre-computed per-token text embedding lookup. Replaces the - # original char-aware subword encoder + subword-flag + BOS/EOS - # additive embeddings; all of those are deterministic per token - # id and are baked into this single table by the checkpoint - # converter. - self.embed_subword = PrecomputedSubwordEmbedding(vocab_size, hidden_size) - self.bos_emb = nn.Parameter(torch.empty(hidden_size)) - # Learned classifier-free text-conditioning embedding. The audio and - # speaker branches remain unchanged for unconditional rows. - self.null_emb = nn.Parameter(torch.empty(hidden_size)) - - self.use_gated_fusion_for_text_audio = config.use_gated_fusion_for_text_audio - if self.use_gated_fusion_for_text_audio: - self.gated_fusion_audio_text = GatedProjectedSumRMSNorm( - hidden_size, hidden_size, hidden_size, config.num_quantizers - ) - - self.use_audio_prompt_frozen_projection = ( - config.use_audio_prompt_frozen_projection - ) - if self.use_audio_prompt_frozen_projection: - self.audio_prompt_projection_W = nn.Parameter( - torch.empty(hidden_size, hidden_size), - requires_grad=False, - ) - - def forward( - self, - acoustic_tokens: torch.Tensor, - text_tokens: torch.Tensor, - text_mask: torch.Tensor, - bos_mask: torch.Tensor, - speaker_latent: Optional[torch.Tensor] = None, - cfg_is_uncond: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """ - Works for context and generation phases to prepare total input - embeddings for EarTTS model. - - Inputs: - acoustic_tokens: (BT x 31) - audio tokens - text_tokens: (BT) - text token to embed - text_mask: (BT) - masks text embeddings for prefill - bos_mask: (BT) - specifies where BOS is applied (first frame of prefill) - speaker_latent: (BT x hidden_size) - external speaker latent. - Non-zero rows replace ``embed_code(rvq_sum(...))`` at - pre-BOS prefill positions; zero rows (decode steps and - the BOS frame) leave ``audio_emb`` untouched. Pass an - all-zero tensor on decode. - - Returns: - embedding of shape (BT x dim) - """ - - # prepare bos emb that is applied to audio embedding - bos_emb = bos_mask.unsqueeze(1) * self.bos_emb # BT x dim - - acoustic_tokens = acoustic_tokens.transpose(0, 1) # 31 x BT - audio_emb = sum( - emb(acoustic_tokens[i]) for i, emb in enumerate(self.rvq_embs) - ) # BT x latent_size - audio_emb = self.embed_code(audio_emb) # BT x hidden_size - - if self.use_audio_prompt_frozen_projection: - if speaker_latent is None: - # No external latent -> derive one from the acoustic - # tokens, matching DuplexEARTTS when no speaker prompt - # is supplied. vLLM-Omni callers always pass a real or - # zero latent, so they take the other branch. - latent_provided = torch.zeros_like(bos_mask).unsqueeze(-1) - latent = torch.nn.functional.linear( - audio_emb, self.audio_prompt_projection_W.T - ) - else: - # ``latent_provided`` is non-zero exactly on the rows - # the user populated with a real speaker latent - # (prefill pre-BOS positions). Decode rows are filled - # with zeros by ``preprocess``, so they read as "not - # provided" here. - latent_provided = ( - speaker_latent.abs().sum(-1, keepdim=True) > 0 - ) # (BT, 1) - latent = speaker_latent - - # Replace only at pre-BOS positions of prefill -- i.e. - # ``bos_mask == 0 AND latent was actually provided``. This - # excludes: - # * the BOS frame (``bos_mask == 1``), where - # ``embed_code(acoustic_tokens)`` of ``sil_tokens`` - # survives (this is the audio_emb the backbone sees on - # the BOS frame). - # * AR decode steps (``latent_provided == False`` because - # ``speaker_latent`` is all zeros). - replace_mask = (bos_mask.unsqueeze(-1) == 0) & latent_provided - audio_emb = torch.where(replace_mask, latent, audio_emb) - - audio_emb = audio_emb + bos_emb - - # Embed text tokens via the pre-computed lookup (subword-flag and - # BOS/EOS additions are baked into the table at conversion time). - # Apply the mask that zeroes this embedding on prefill positions. - text_emb = self.embed_subword(text_tokens) * text_mask.unsqueeze(1) # BT x dim - if cfg_is_uncond is not None: - text_emb = torch.where( - cfg_is_uncond.unsqueeze(1), - self.null_emb.to(text_emb.dtype), - text_emb, - ) - - # prepare total embedding by combining audio and text branches - if self.use_gated_fusion_for_text_audio: - # Gated fusion needs ``audio_emb`` and ``text_emb`` as - # separate inputs (it learns a per-feature gate to mix - # them), which is why neither branch can be folded into a - # single precomputed ``inputs_embeds`` tensor outside the - # compiled forward. - total_emb = self.gated_fusion_audio_text(audio_emb, text_emb) - else: - total_emb = audio_emb + text_emb # BT x dim - return total_emb - - -def gumbel_like(tensor: torch.Tensor, eps: float = 1e-8) -> torch.Tensor: - """ - Generates a tensor of Gumbel noise with the same shape as the input - tensor. Used for the Gumbel-Max trick. - """ - u = torch.rand_like(tensor) - return -torch.log(-torch.log(u + eps) + eps) - - -def batch_matmul(x: torch.Tensor, w: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - """Performs a batched matrix multiplication using PyTorch's native functions. - In NeMo this is implemented as a custom kernel using triton. - - Args: - x: ``[batch_size, d_in]`` - w: ``[num_weights, d_out, d_in]`` - y: ``[batch_size]`` - - Returns: - Tensor of shape ``[batch_size, d_out]``. - """ - return torch.bmm(w[y], x.unsqueeze(2)).squeeze(2) - - -class MoGHead(nn.Module): - """A Mixture of Gaussians (MoG) prediction head. - - This module takes a hidden state and predicts the parameters for a - mixture of Gaussian distributions. It's suitable for modeling - continuous, multi-modal data. - """ - - def __init__( - self, - hidden_size: int, - intermediate_size: int, - out_size: int, - num_layers: int, - num_predictions: int, - low_rank: Optional[int] = 64, - top_p_or_k: Optional[Union[float, int]] = 1.0, - min_log_std: float = -4.0, - eps: float = 1e-6, - ): - super().__init__() - self.out_size = out_size - self.low_rank = low_rank - self.num_predictions = num_predictions - self.min_log_std = min_log_std - self.top_p_or_k = top_p_or_k - - self.logits_processor = ( - TopPLogitsWarper(self.top_p_or_k) - if isinstance(self.top_p_or_k, float) - else ( - TopKLogitsWarper(self.top_p_or_k) - if isinstance(self.top_p_or_k, int) - else None - ) - ) - - self.mlp_stack = nn.Sequential( - *[ - MLPLayer(hidden_size, intermediate_size, eps=eps) - for _ in range(num_layers) - ], - RMSNorm(hidden_size, eps=eps), - ) - - if low_rank is None: - self.proj_logits = nn.Linear(hidden_size, num_predictions, bias=False) - self.proj_mus = nn.Linear( - hidden_size, num_predictions * out_size, bias=False - ) - self.proj_logs = nn.Linear(hidden_size, 1, bias=False) - else: - assert low_rank < out_size - self.proj_logits = nn.Linear(hidden_size, num_predictions, bias=False) - self.proj_mus = nn.Linear( - hidden_size, num_predictions * low_rank, bias=False - ) - self.proj_logs = nn.Linear(hidden_size, 1, bias=False) - self.proj_else = nn.Linear(hidden_size, out_size, bias=False) - self.low_mat = nn.Parameter( - torch.empty(num_predictions, out_size, low_rank) - ) - - def forward( - self, - x: torch.Tensor, - cfg_is_uncond: Optional[torch.Tensor] = None, - cfg_scale: Optional[torch.Tensor] = None, - cfg_active: Optional[torch.Tensor] = None, - cfg_partner: Optional[torch.Tensor] = None, - cfg_conditional_rep: Optional[torch.Tensor] = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - bt = x.size(0) - n, d = self.num_predictions, self.low_rank or self.out_size - - x = self.mlp_stack(x) - if ( - cfg_is_uncond is not None - and cfg_scale is not None - and cfg_active is not None - and cfg_partner is not None - ): - # Native EarTTS guidance is applied after the MoG MLP stack and - # before proj_logits/proj_mus/proj_logs/proj_else. - x = _apply_cfg_after_mlp( - x, - cfg_is_uncond=cfg_is_uncond, - cfg_scale=cfg_scale, - cfg_active=cfg_active, - cfg_partner=cfg_partner, - ) - - logits = self.proj_logits(x) - - # Apply top-p or top-k filtering to the mixture logits - if self.logits_processor is not None: - logits = self.logits_processor(None, logits.view(-1, n)).view_as(logits) - - # Sample a mixture component using the Gumbel-Max trick - gumbel = gumbel_like(logits) - if cfg_active is not None and cfg_conditional_rep is not None: - gumbel = torch.where( - cfg_active[:, None], - gumbel[cfg_conditional_rep], - gumbel, - ) - mixture_indices = (nn.functional.log_softmax(logits, dim=-1) + gumbel).argmax( - -1 - ) - - # Select the mean corresponding to the sampled component - mu = batch_matmul( - x.view(bt, -1), - self.proj_mus.weight.detach().view(n, d, -1), - mixture_indices.view(bt), - ).view(bt, d) - if self.proj_mus.bias is not None: - mu += self.proj_mus.bias.detach().view(n, d)[mixture_indices] - - if self.low_rank: - mu = batch_matmul( - mu.view(bt, -1), - self.low_mat.detach().view(n, self.out_size, -1), - mixture_indices.view(bt), - ).view(bt, self.out_size) - mu_res = self.proj_else(x) - else: - mu_res = torch.zeros((bt, d), device=x.device) - - logs = self.proj_logs(x).clamp_min(self.min_log_std) - return mu * torch.exp(logs) + mu_res, logs - - -class MaskGITSampler(nn.Module): - def __init__(self, config): - super().__init__() - self.config = config - self.num_quantizers = self.config.num_quantizers - self.codebook_size = self.config.codebook_size - self.noise_scale = self.config.noise_scale - self.debug_cfg_contract = bool( - getattr(self.config, "debug_cfg_contract", False) - ) - - # pre-compute how many tokens are unmasked at each iteration - rates = np.linspace(0.0, 1.0, self.config.num_iter + 1)[:-1].reshape(-1, 1) - masking_rates = np.power( - 1 - np.power(rates, self.config.exponent), 1 / self.config.exponent - ) - num_maskings = np.ceil(masking_rates * self.num_quantizers).astype(int) - num_maskings_shifted = np.pad( - num_maskings[1:], ((0, 1), (0, 0)), constant_values=0 - ) - sampling_per_step = num_maskings - num_maskings_shifted - sampling_per_step_flat = sampling_per_step.flatten() - # Drop any values at the beginning that are 0 - first_nonzero = np.argmax(sampling_per_step_flat != 0) - self.num_to_sample = sampling_per_step_flat[first_nonzero:].tolist() - - # create layers used for acoustic tokens embedding - self.rvq_embs = nn.Parameter( - torch.empty( - self.config.num_quantizers, - self.config.codebook_size, - self.config.latent_size, - ) - ) - self.embed_code = nn.Linear( - self.config.latent_size, self.config.hidden_size, bias=False - ) - # MoG head for generation (uncompiled part) - self.mog_head = MoGHead( - hidden_size=self.config.hidden_size, - intermediate_size=self.config.intermediate_size, - out_size=self.config.latent_size, - num_layers=self.config.mog_num_layers, - num_predictions=self.config.mog_num_predictions, - low_rank=self.config.mog_low_rank, - top_p_or_k=self.config.top_p_or_k, - min_log_std=self.config.mog_min_log_std, - eps=self.config.mog_eps, - ) - - def _depthsum_embedding(self, code: torch.Tensor) -> torch.Tensor: - """Embeds all codes into a single embedding.""" - embs = nn.functional.pad( - self.rvq_embs, [0, 0, 0, 1] - ) # num_quantizers x (codebook_size + 1) x latent_size - res = nn.functional.embedding(code[0], embs[0]) - for i in range(1, len(embs)): - res = res + nn.functional.embedding(code[i], embs[i]) - return res - - def _depthsum_encoding_step_reshaped( - self, - r: torch.Tensor, # [B*T, hidden_size] - code: torch.Tensor, # [num_quantizers, B*T] - depth_str: int, - k: int, - ) -> torch.Tensor: - """RVQ encoding with reshaped code tensor.""" - for i in range(depth_str, depth_str + k): - # Compute distances: ||emb||² - 2⟨r, emb⟩ - idx_sel = ( - self.rvq_embs[i].pow(2).sum(-1) # [vocab_size] - - 2 * (r @ self.rvq_embs[i].T) # [B*T, vocab_size] - ).argmin(-1) # [B*T] - - # Update residual - emb_i = nn.functional.embedding( - idx_sel, - self.rvq_embs[i], - ) # [B*T, latent_size] - r = r - emb_i - - # Store selected indices - code[i] = idx_sel - - return code - - def forward( - self, - hidden_states: torch.Tensor, - cfg_enabled: Optional[torch.Tensor] = None, - cfg_is_uncond: Optional[torch.Tensor] = None, - cfg_pair_id: Optional[torch.Tensor] = None, - cfg_scale: Optional[torch.Tensor] = None, - valid: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """Performs the iterative unmasking process for a single - generation step. - """ - - device = hidden_states.device - batch_size = int(hidden_states.shape[0]) - if cfg_enabled is None: - cfg_enabled = torch.zeros(batch_size, dtype=torch.bool, device=device) - if cfg_is_uncond is None: - cfg_is_uncond = torch.zeros(batch_size, dtype=torch.bool, device=device) - if cfg_pair_id is None: - cfg_pair_id = torch.full( - (batch_size,), - -1, - dtype=torch.long, - device=device, - ) - if cfg_scale is None: - cfg_scale = torch.zeros(batch_size, dtype=torch.float32, device=device) - if valid is None: - valid = torch.ones(batch_size, dtype=torch.bool, device=device) - - ( - hidden_states, - cfg_is_uncond, - cfg_scale, - cfg_active, - cfg_partner, - cfg_conditional_rep, - inverse_order, - ) = _prepare_cfg_sampling_batch( - hidden_states, - cfg_enabled=cfg_enabled, - cfg_is_uncond=cfg_is_uncond, - cfg_pair_id=cfg_pair_id, - cfg_scale=cfg_scale, - valid=valid, - ) - complete_contract = ((~cfg_enabled) | cfg_active).all() - if self.debug_cfg_contract and not bool( - complete_contract.item() - ): - raise RuntimeError( - "Incomplete EarTTS CFG model batch: " - f"enabled={cfg_enabled.tolist()} " - f"is_uncond={cfg_is_uncond.tolist()} " - f"pair_id={cfg_pair_id.tolist()} " - f"valid={valid.tolist()} " - f"active={cfg_active.tolist()}" - ) - - # Initialize the full code tensor - code = ( - torch.zeros( - (self.num_quantizers, hidden_states.shape[0]), - dtype=torch.long, - device=device, - ) - + self.codebook_size - ) - # Iteratively unmask the continuous part of the code - cnt = 0 - for k in self.num_to_sample: - # Prepare input for the MoG head - mog_input_embeds = self.embed_code( - self._depthsum_embedding(code) - ) # (BT x hidden_size) - mog_input_embeds += hidden_states - - mog_mu, mog_logs = self.mog_head( - mog_input_embeds, - cfg_is_uncond=cfg_is_uncond, - cfg_scale=cfg_scale, - cfg_active=cfg_active, - cfg_partner=cfg_partner, - cfg_conditional_rep=cfg_conditional_rep, - ) - normal = torch.randn_like(mog_mu) - normal = torch.where( - cfg_active[:, None], - normal[cfg_conditional_rep], - normal, - ) - z = mog_mu + torch.exp(mog_logs) * normal * self.noise_scale - code = self._depthsum_encoding_step_reshaped(z, code, cnt, k) - # Match PyTorch EarTTS CFG: every MaskGIT iteration feeds the - # conditional code trajectory back into both KV streams. - code = torch.where( - cfg_active.unsqueeze(0), - code[:, cfg_conditional_rep], - code, - ) - - cnt += k - return code.transpose(0, 1)[inverse_order] # BT x num_quantizers - - -@support_torch_compile -class EarTTSModel(nn.Module): - """Embedding preparation + Gemma3 backbone (compiled together). - - MaskGIT sampling lives in :class:`EarTTSSamplerModel` so the iterative - sampler can be skipped on prefill positions while still being CUDA-graph - captured for decode-only batches. See :meth:`EarTTSForCausalLM.forward`. - """ - - def __init__( - self, - *, - vllm_config: VllmConfig, - prefix: str = "", - ): - super().__init__() - config = vllm_config.model_config.hf_config - self.total_emb = EarTTSInputEmbedding(config) - self.backbone = Gemma3Model(vllm_config=vllm_config, prefix=prefix) - - # Per-codebook silence acoustic tokens. Registered as a - # persistent int32 buffer (loaded from the checkpoint under - # ``model.sil_tokens``) rather than nn.Parameter so that vLLM's - # automatic float dtype casting (e.g. ``model.to(bfloat16)``) - # leaves it untouched. - self.register_buffer( - "sil_tokens", - # Zero is a safe dummy-loader default; production checkpoints - # overwrite this persistent buffer in ``load_weights``. - torch.zeros(int(config.num_quantizers), dtype=torch.int32), - persistent=True, - ) - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - intermediate_tensors: Optional[IntermediateTensors], - acoustic_tokens: torch.Tensor, - text_tokens: torch.Tensor, - text_mask: torch.Tensor, - bos_mask: torch.Tensor, - speaker_latent: torch.Tensor, - cfg_is_uncond: torch.Tensor, - ) -> torch.Tensor: - """Forward pass through embeddings and backbone transformer. - Returns the backbone's ``hidden_states``. - """ - total_emb = self.total_emb( - acoustic_tokens=acoustic_tokens, - text_tokens=text_tokens, - text_mask=text_mask, - bos_mask=bos_mask, - speaker_latent=speaker_latent, - cfg_is_uncond=cfg_is_uncond, - ) - hidden_states = self.backbone( - input_ids, positions, intermediate_tensors, inputs_embeds=total_emb - ) - return hidden_states - - -@support_torch_compile -class EarTTSSamplerModel(nn.Module): - """MaskGIT sampler in its own compile group. - - Hosting the sampler in a separate ``@support_torch_compile`` module - is what makes it possible for :meth:`EarTTSForCausalLM.forward` to: - - * Capture and replay a CUDA-graph for decode-only batches (where - every position needs sampling). - * Skip the sampler entirely on prefill positions, where the audio - output isn't actually needed. - * Run the sampler on a sliced subset of positions in mixed - prefill+decode batches, with a ``BatchDescriptor`` override so the - sampler's CUDA-graph cache is hit at the padded decode-batch size. - - The :meth:`forward` operates on a stable-address scratch buffer - (:attr:`_sampler_input`) so callers can pass a transient slice - (e.g. ``hidden_states[decode_idx]``) without breaking CUDA-graph - replay. The non-compiled :meth:`sample` wrapper does that copy and - then invokes the compiled :meth:`forward`. - """ - - def __init__( - self, - *, - vllm_config: VllmConfig, - prefix: str = "", - ): - super().__init__() - config = vllm_config.model_config.hf_config - self.sampler = MaskGITSampler(config) - - # Stable-address scratch buffer for the sampler's input. Every - # CUDA-graph replay must read from the same ``data_ptr()``; the - # caller may pass either the full backbone output or a fresh - # ``hidden_states[decode_idx]`` slice, so we copy into this - # buffer (in :meth:`sample`) before invoking :meth:`forward`. - max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens - hidden_size = config.hidden_size - dtype = vllm_config.model_config.dtype - self._sampler_input = torch.zeros( - max_num_tokens, hidden_size, dtype=dtype - ) - self._sampler_cfg_enabled = torch.zeros(max_num_tokens, dtype=torch.bool) - self._sampler_cfg_is_uncond = torch.zeros(max_num_tokens, dtype=torch.bool) - self._sampler_cfg_pair_id = torch.full((max_num_tokens,), -1, dtype=torch.long) - self._sampler_cfg_scale = torch.zeros(max_num_tokens, dtype=torch.float32) - self._sampler_valid = torch.zeros(max_num_tokens, dtype=torch.bool) - - def forward( - self, - hidden_states: torch.Tensor, - cfg_enabled: torch.Tensor, - cfg_is_uncond: torch.Tensor, - cfg_pair_id: torch.Tensor, - cfg_scale: torch.Tensor, - valid: torch.Tensor, - ) -> torch.Tensor: - """Compiled — runs MaskGIT on a (stable-address) hidden buffer.""" - return self.sampler( - hidden_states, - cfg_enabled=cfg_enabled, - cfg_is_uncond=cfg_is_uncond, - cfg_pair_id=cfg_pair_id, - cfg_scale=cfg_scale, - valid=valid, - ) - - def sample( - self, - hidden_states: torch.Tensor, - *, - cfg_enabled: torch.Tensor, - cfg_is_uncond: torch.Tensor, - cfg_pair_id: torch.Tensor, - cfg_scale: torch.Tensor, - valid: torch.Tensor, - ) -> torch.Tensor: - """Non-compiled wrapper — copies into the stable buffer first. - - Mirrors the qwen3-tts code-predictor pattern: transient inputs - are first written into a model-owned static-address buffer so - the captured CUDA-graph for the compiled :meth:`forward` always - reads from the recorded ``data_ptr()``. - """ - seq_len = int(hidden_states.shape[0]) - buf = self._sampler_input[:seq_len] - buf.copy_(hidden_states) - enabled_buf = self._sampler_cfg_enabled[:seq_len] - role_buf = self._sampler_cfg_is_uncond[:seq_len] - pair_buf = self._sampler_cfg_pair_id[:seq_len] - scale_buf = self._sampler_cfg_scale[:seq_len] - valid_buf = self._sampler_valid[:seq_len] - enabled_buf.copy_(cfg_enabled) - role_buf.copy_(cfg_is_uncond) - pair_buf.copy_(cfg_pair_id) - scale_buf.copy_(cfg_scale) - valid_buf.copy_(valid) - return self( - buf, - enabled_buf, - role_buf, - pair_buf, - scale_buf, - valid_buf, - ) - - -# --------------------------------------------------------------------------- -# Outer model — the vLLM-Omni preprocess/postprocess entry point. -# --------------------------------------------------------------------------- - - -# Placeholder token id used to fill the per-step ``input_ids`` returned -# by :meth:`preprocess`. Must be a valid id in ``[0, config.vocab_size)`` -# but is otherwise unused — the actual decode-vs-prefill behaviour is -# driven by the per-token buffers populated in :meth:`preprocess`. -# -# The width of the dummy logits tensor returned by -# :meth:`compute_logits` is taken from ``config.vocab_size`` (see -# :class:`EarTTSConfig`) so vLLM's sampler / ``LogitsProcessor`` and the -# model agree on the logits shape. ``compute_logits`` returns -# ``[0, -inf, ..., -inf]`` so the sampler's argmax always picks index 0 -# regardless of how wide ``vocab_size`` is — the real audio output is -# the codes tensor exposed via :meth:`make_omni_output`. -_DUMMY_TOKEN_ID = 0 - - -@ignore_torch_compile -@support_torch_compile -class EarTTSForCausalLM(nn.Module, SupportsPP): - """EarTTS for vLLM-Omni. - - Inputs (passed via ``additional_information``): - - * ``embed.voice`` (prefill chunk 0 only) — Tensor of shape - ``(Tref, hidden_size)`` carrying the user-supplied speaker - latent. The user must also pass ``prompt_token_ids = [0] * - Tref`` so the prefill placeholder length matches. - * ``ids.output`` — Python ``list[int]`` of the most recently sent - text tokens; preprocess consumes the last entry, so decode step - ``k`` consumes ``t_k``. - * CFG metadata (on every prefill/decode chunk): ``cfg_enabled``, - ``cfg_role`` (``"cond"`` or ``"uncond"``), ``cfg_pair_id``, and - ``cfg_scale``. Unconditional rows replace only text conditioning - with ``model.total_emb.null_emb``. Complete decode pairs are sampled - with native EarTTS guidance and receive identical acoustic codes. - - Per-step flow (see module docstring for details): - - ``preprocess`` populates five model-owned buffers - (:attr:`_acoustic_tokens`, :attr:`_text_tokens`, :attr:`_text_mask`, - :attr:`_bos_mask`, :attr:`_speaker_latent`) at each request's - flat-batch offset. ``forward`` slices them up to ``num_tokens`` and - runs the compiled :class:`EarTTSModel` (embedding + Gemma3 - backbone) for every position, then conditionally invokes the - compiled :class:`EarTTSSamplerModel` (MaskGIT) to produce codes. - The sampler is skipped on prefill positions — see :meth:`forward` - for details. The generated codes (BTx31) are written to - :attr:`_out_codes` and exposed as a multimodal output by - :meth:`make_omni_output` (see there for the key it uses). - ``postprocess`` stashes the final-frame codes under - ``last_acoustic_codes`` for the next decode step's - :meth:`preprocess`. - - Sampler skipping mirrors the qwen3-tts code-predictor pattern: - - * **Profile / dummy run** (``attn_metadata is None``) and - **decode-only batches** (``max_query_len == 1``) run the sampler - on every token so the captured CUDA graph covers all of - ``cudagraph_capture_sizes``. - * **Mixed prefill+decode batches**: only decode-token positions go - through the sampler. The sampler's ``BatchDescriptor`` is - overridden to the padded decode-batch size so the right captured - graph is replayed. - * **Prefill-only batches**: the sampler is skipped entirely. - * Prefill rows of :attr:`_out_codes` are intentionally not - written. ``last_acoustic_codes`` returned by :meth:`postprocess` - after prefill is therefore undefined — :meth:`_preprocess_decode` - seeds the first decode step's acoustic input with the acoustic - pad id (``codebook_size``) so this never matters. - """ - - # ``model.sampler.*`` lands on :attr:`sampler_module` (the MaskGIT compile - # group). Other prefixes (``model.total_emb.``, ``model.backbone.``) match - # the module layout 1:1. - hf_to_vllm_mapper = WeightsMapper( - orig_to_new_prefix={ - "model.sampler.": "sampler_module.sampler.", - } - ) - - # Omni preprocess/postprocess hooks (consumed by the gpu model runner). - has_preprocess = True - has_postprocess = True - have_multimodal_outputs = True - - # No ``gpu_resident_buffer_keys``: vLLM-Omni's opt-out from the - # ``model_intermediate_buffer`` D2H round-trip is keyed by - # ``(type_key, qualifier)`` pairs and is only consulted for *nested* - # payload entries, so a model whose payloads are flat -- as both stages - # here are -- cannot express its keys in that form. Declaring flat names - # is not merely inert, it breaks: the runner unpacks every declared key - # as a pair as soon as any nested entry arrives, and 0.24 onwards always - # sends one (``meta``). - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - self.config = vllm_config.model_config.hf_config - self.vllm_config = vllm_config - - # Embedding + Gemma3 backbone — runs on every position. Built - # under the default ``"backbone"`` model tag (vLLM's compile - # cache key for the main model). We don't wrap this in a - # ``set_model_tag`` block because :func:`set_model_tag` asserts - # the new tag differs from the current one and the default is - # already ``"backbone"``. - self.model = EarTTSModel( - vllm_config=vllm_config, - prefix=prefix, - ) - - # MaskGIT sampler in its own compile group, so it can be invoked - # conditionally (decode positions only, or skipped entirely on - # prefill-only batches) while still being CUDA-graph captured - # for decode-only batches over ``cudagraph_capture_sizes``. The - # ``"sampler"`` tag keys the sampler's compile cache separately - # from the backbone's. - with set_model_tag("sampler"): - self.sampler_module = EarTTSSamplerModel( - vllm_config=vllm_config, - prefix=prefix, - ) - - # Pad ids used by buffers / preprocess. Match the conventions of - # the original EarTTSInputEmbedding: an acoustic token id of - # ``codebook_size`` is the trailing "no audio" pad row in - # ``rvq_embs`` (which has ``codebook_size + 1`` entries). - # How the sampled codes are surfaced from ``make_omni_output``, driven - # by the stage's pipeline-config ``engine_output_type``: - # - # * "audio" (this stage is the final, client-facing one, which is the - # split VoiceChat layout): emit under the ``model_outputs`` key. - # vLLM-Omni's output processor remaps ``model_outputs`` to the - # drainable ``audio`` modality key, so DELTA streaming drains it - # after every step and the client receives one frame per step. - # Any other key is retained across steps *and* concatenated along - # the last dimension (``get_accumulation_strategy`` maps the audio - # modality to ``CONCAT_LAST``), which for a ``T x num_quantizers`` - # code tensor silently widens the per-step frame instead of - # appending to it. - # * otherwise: emit ``audio_codes`` for a downstream stage to consume. - engine_output_type = getattr( - vllm_config.model_config, "engine_output_type", None - ) - self._single_stage_audio = str(engine_output_type or "").lower() == "audio" - - self._num_quantizers: int = int(self.config.num_quantizers) - self._hidden_size: int = int(self.config.hidden_size) - self._acoustic_pad_id: int = int(self.config.codebook_size) - text_pad_id = getattr(self.config, "pad_token_id", None) - eos_token_id = getattr(self.config, "eos_token_id", None) - if text_pad_id is None or eos_token_id is None: - raise ValueError( - "EarTTS config.json must set pad_token_id and eos_token_id from the " - "source VoiceChat tokenizer. Re-run convert_duplex_eartts_checkpoint.py." - ) - self._text_pad_id: int = int(text_pad_id) - self._eos_token_id: int = int(eos_token_id) - - # ── Persistent stable-address buffers ──────────────────────── - # Plain tensor attributes (not nn.Parameter / not register_buffer): - # * AutoWeightsLoader only walks named_parameters() and persistent - # registered buffers, so plain attributes are invisible to it - # (no spurious "missing weight" errors during load_weights). - # * vLLM constructs models inside - # ``with torch.device(device_config.device):`` so a bare - # ``torch.zeros(...)`` here is allocated directly on the GPU. - # * Addresses stay stable across CUDA graph replays as long as - # we never re-assign these names (only do in-place writes via - # copy_/fill_/indexed assignment), which is what the rest of - # this class does. The piecewise CUDAGraphWrapper records - # data_ptr() at capture time and expects the same pointer at - # replay time — that holds with plain tensors. - max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens - model_dtype = vllm_config.model_config.dtype - - self._acoustic_tokens = torch.full( - (max_num_tokens, self._num_quantizers), - self._acoustic_pad_id, - dtype=torch.long, - ) - self._text_tokens = torch.full( - (max_num_tokens,), self._text_pad_id, dtype=torch.long - ) - self._text_mask = torch.zeros(max_num_tokens, dtype=torch.long) - self._bos_mask = torch.zeros(max_num_tokens, dtype=torch.long) - # Speaker latent buffer — model dtype, hidden_size wide. - # Decode rows stay all-zero (which the embedding module reads - # as "latent not provided" so ``audio_emb`` is preserved). - # Prefill rows are populated from the user-supplied tensor. - self._speaker_latent = torch.zeros( - max_num_tokens, self._hidden_size, dtype=model_dtype - ) - # Per-token CFG contract. These plain tensors follow the same - # stable-address rules as the model input buffers above and are copied - # into the sampler's own CUDA-graph scratch buffers before sampling. - self._cfg_enabled = torch.zeros(max_num_tokens, dtype=torch.bool) - self._cfg_is_uncond = torch.zeros(max_num_tokens, dtype=torch.bool) - self._cfg_pair_id = torch.full((max_num_tokens,), -1, dtype=torch.long) - self._cfg_scale = torch.zeros(max_num_tokens, dtype=torch.float32) - # vLLM-Omni 0.26 computes per-request flat-batch slices but does not - # pass ``start``/``end`` into preprocess. The CFG scheduler guarantees - # cond then uncond order, so this cursor reconstructs those slices. - self._preprocess_cursor = 0 - self._out_codes = torch.zeros( - max_num_tokens, self._num_quantizers, dtype=torch.long - ) - - # ------------------------------------------------------------------ - # Helpers - # ------------------------------------------------------------------ - - def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: - """Compatibility shim — not actually consumed at runtime since - every forward goes through ``inputs_embeds`` assembled inside - :meth:`forward`. - """ - return self.model.backbone.embed_input_ids(input_ids) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.get_input_embeddings(input_ids) - - @staticmethod - def _unwrap_singleton(value: Any) -> Any: - """Unwrap a possibly list-wrapped scalar (e.g. ``[tensor]``).""" - if isinstance(value, list): - return value[0] if value else None - return value - - @staticmethod - def _payload_get(info_dict: dict[str, Any], category: str, qualifier: str) -> Any: - """Read one ``OmniPayload`` field, e.g. ``("embed", "voice")``. - - Categories arrive as sub-dicts, which is also what lets the runner - keep a prefill-only field alive while a per-step field in another - category is replaced. - """ - sub = info_dict.get(category) - return sub.get(qualifier) if isinstance(sub, dict) else None - - @classmethod - def _cfg_scalar(cls, value: Any) -> Any: - value = cls._unwrap_singleton(value) - if isinstance(value, torch.Tensor): - assert value.numel() == 1, ( - "EarTTS CFG metadata tensors must contain one scalar; " - f"got shape={tuple(value.shape)}." - ) - return value.item() - return value - - @classmethod - def _stable_cfg_pair_id(cls, value: Any) -> int: - """Map request-provided pair ids to a deterministic signed int64.""" - value = cls._cfg_scalar(value) - assert value is not None and not isinstance( - value, bool - ), "EarTTS CFG requires a non-empty ``cfg_pair_id``." - if isinstance(value, int): - assert ( - -(1 << 63) <= value < (1 << 63) - ), f"EarTTS cfg_pair_id={value} does not fit in int64." - return value - if isinstance(value, float): - assert ( - value.is_integer() - ), f"EarTTS cfg_pair_id must be integral or a string; got {value}." - return cls._stable_cfg_pair_id(int(value)) - encoded = str(value).encode("utf-8") - assert encoded, "EarTTS CFG requires a non-empty ``cfg_pair_id``." - return int.from_bytes( - hashlib.blake2b(encoded, digest_size=8).digest(), - "little", - ) & ((1 << 63) - 1) - - def _write_cfg_state( - self, - *, - start: int, - span_len: int, - info_dict: dict[str, Any], - ) -> None: - """Validate one request/chunk's CFG metadata and fill static rows.""" - enabled_value = self._cfg_scalar(info_dict.get("cfg_enabled", False)) - if isinstance(enabled_value, str): - normalized = enabled_value.strip().lower() - assert normalized in { - "true", - "false", - "1", - "0", - }, f"EarTTS cfg_enabled must be boolean; got {enabled_value!r}." - cfg_enabled = normalized in {"true", "1"} - else: - cfg_enabled = bool(enabled_value) - - cfg_is_uncond = False - cfg_pair_id = -1 - scale_value = self._cfg_scalar(info_dict.get("cfg_scale")) - if scale_value is None: - scale_value = getattr(self.config, "guidance_scale", 0.5) - cfg_scale = float(scale_value) - assert math.isfinite( - cfg_scale - ), f"EarTTS cfg_scale must be finite; got {cfg_scale}." - - if cfg_enabled: - role = str(self._cfg_scalar(info_dict.get("cfg_role")) or "").lower() - assert role in {"cond", "uncond"}, ( - "EarTTS CFG requires cfg_role='cond' or 'uncond'; " f"got {role!r}." - ) - cfg_is_uncond = role == "uncond" - cfg_pair_id = self._stable_cfg_pair_id(info_dict.get("cfg_pair_id")) - - end = start + span_len - self._cfg_enabled[start:end].fill_(cfg_enabled) - self._cfg_is_uncond[start:end].fill_(cfg_is_uncond) - self._cfg_pair_id[start:end].fill_(cfg_pair_id) - self._cfg_scale[start:end].fill_(cfg_scale) - - def _validate_speaker_latent(self, value: Any) -> torch.Tensor: - """Assert ``speaker_latent`` has shape ``(Tref, hidden_size)``.""" - x = self._unwrap_singleton(value) - assert isinstance(x, torch.Tensor), ( - f"speaker_latent must be a torch.Tensor; got {type(x).__name__}." - ) - assert x.ndim == 2 and x.shape[1] == self._hidden_size, ( - "speaker_latent must have shape (Tref, hidden_size=" - f"{self._hidden_size}); got {tuple(x.shape)}." - ) - return x.to(dtype=self._speaker_latent.dtype).contiguous() - - def _build_prefill_tensors( - self, - speaker_latent: torch.Tensor, - device: torch.device, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Build the cached prefill ``(text_tokens, text_mask, bos_mask, - speaker_latent)`` of length ``prefill_len = speaker_latent.shape[0]``. - - Layout: ``text_tokens = [PAD] * (n - 1) + [EOS]``, - ``text_mask = [0] * (n - 2) + [1, 1]``, - ``bos_mask = [0] * (n - 1) + [1]``. Acoustic tokens are not - cached — :meth:`preprocess` broadcasts ``model.sil_tokens`` at - every prefill position. - """ - prefill_len = int(speaker_latent.shape[0]) - assert prefill_len > 0, ( - "speaker_latent must have at least one frame " - f"(got shape={tuple(speaker_latent.shape)})." - ) - - text_tokens = torch.full( - (prefill_len,), self._text_pad_id, dtype=torch.long, device=device - ) - text_tokens[-1] = self._eos_token_id - - text_mask = torch.zeros(prefill_len, dtype=torch.long, device=device) - text_mask[max(0, prefill_len - 2):] = 1 - - bos_mask = torch.zeros(prefill_len, dtype=torch.long, device=device) - bos_mask[-1] = 1 - - speaker_latent = speaker_latent.to( - device=device, dtype=self._speaker_latent.dtype, non_blocking=True - ).contiguous() - - return text_tokens, text_mask, bos_mask, speaker_latent - - # ------------------------------------------------------------------ - # preprocess - # ------------------------------------------------------------------ - - def preprocess( - self, - input_ids: torch.Tensor, - input_embeds: Optional[torch.Tensor], - *, - start: int = 0, - end: int = 0, - **info_dict: Any, - ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: - """Build per-request ``(input_ids, inputs_embeds)`` for this step. - - Prefill (``span_len > 1``): - On the first prefill chunk, constructs the per-position - prefill tensors of length - ``prefill_len = speaker_latent.shape[0]``: - - * ``text_tokens`` = ``[PAD] * (prefill_len - 1) + [EOS]`` - * ``text_mask`` = ``[0] * (prefill_len - 2) + [1, 1]`` - * ``bos_mask`` = ``[0] * (prefill_len - 1) + [1]`` - * ``acoustic_tokens`` = ``model.sil_tokens`` broadcast at - every position (only the BOS frame's ``audio_emb`` is - actually consumed; the others are replaced by - ``speaker_latent`` inside the embedding module). - * ``speaker_latent`` = the user-supplied tensor. - - ``embed.voice`` is the only required - ``additional_information`` field. Multi-chunk prefill is - tracked by ``ear_prefill_offset``; the cached - ``ear_prefill_speaker_latent`` is sliced into each chunk. - - Decode (``span_len == 1``): - Takes the newest text token from ``ids.output``, which the - producer (a user-driven :class:`StreamingInput` or an - upstream stage in an ``async_chunk`` pipeline such as - ``nemotron_voicechat``) refreshes on every step. - - Acoustic input rules (in order): - * ``text_token == EOS`` → ``model.sil_tokens``. - * First decode (``ear_decode_offset == 0``) → - broadcast acoustic pad id (``codebook_size``). - * Otherwise → ``last_acoustic_codes`` (stashed by - :meth:`postprocess` after the previous step). - - ``text_mask = 1``, ``bos_mask = 0``, - ``speaker_latent = 0``. - """ - # Normalize: some runner paths still pass per-request state - # nested under ``additional_information`` instead of flattened. - nested = info_dict.get("additional_information") - if isinstance(nested, dict): - merged = { - k: v for k, v in info_dict.items() if k != "additional_information" - } - for k, v in nested.items(): - merged.setdefault(k, v) - info_dict = merged - - device = input_ids.device - span_len = int(input_ids.shape[0]) - if span_len <= 0: - base = ( - input_embeds - if input_embeds is not None - else self.embed_input_ids(input_ids) - ) - return input_ids, base, {} - - explicit_start = int(start) - explicit_end = int(end) - if explicit_end > explicit_start: - flat_start = explicit_start - else: - role = str( - self._cfg_scalar(info_dict.get("cfg_role")) - or "" - ).lower() - cfg_enabled = bool( - self._cfg_scalar( - info_dict.get("cfg_enabled", False) - ) - ) - if not cfg_enabled or role == "cond": - self._preprocess_cursor = 0 - flat_start = self._preprocess_cursor - self._preprocess_cursor += span_len - - self._write_cfg_state( - start=flat_start, - span_len=span_len, - info_dict=info_dict, - ) - - if span_len > 1: - return self._preprocess_prefill( - input_ids=input_ids, - input_embeds=input_embeds, - start=flat_start, - span_len=span_len, - device=device, - info_dict=info_dict, - ) - return self._preprocess_decode( - input_ids=input_ids, - input_embeds=input_embeds, - start=flat_start, - device=device, - info_dict=info_dict, - ) - - def _preprocess_prefill( - self, - *, - input_ids: torch.Tensor, - input_embeds: Optional[torch.Tensor], - start: int, - span_len: int, - device: torch.device, - info_dict: dict[str, Any], - ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: - """Prefill branch of :meth:`preprocess`. Writes one chunk-slice of - the cached prefill tensors into the static buffers.""" - cached_speaker_latent = info_dict.get("ear_prefill_speaker_latent") - - info_update: dict[str, Any] = {} - if not isinstance(cached_speaker_latent, torch.Tensor): - # First chunk: build & cache prefill tensors from the - # user-supplied speaker latent. - speaker_latent = self._validate_speaker_latent( - self._payload_get(info_dict, "embed", "voice") - ) - ( - cached_text_tokens, - cached_text_mask, - cached_bos_mask, - cached_speaker_latent, - ) = self._build_prefill_tensors(speaker_latent, device=device) - - info_update["ear_prefill_text_tokens"] = cached_text_tokens - info_update["ear_prefill_text_mask"] = cached_text_mask - info_update["ear_prefill_bos_mask"] = cached_bos_mask - info_update["ear_prefill_speaker_latent"] = cached_speaker_latent - info_update["ear_prefill_offset"] = 0 - info_update["ear_decode_offset"] = 0 - else: - cached_text_tokens = info_dict["ear_prefill_text_tokens"] - cached_text_mask = info_dict["ear_prefill_text_mask"] - cached_bos_mask = info_dict["ear_prefill_bos_mask"] - - offset = int(info_dict.get("ear_prefill_offset", 0) or 0) - full_len = int(cached_speaker_latent.shape[0]) - s, e = offset, offset + span_len - assert 0 <= s and e <= full_len, ( - "prefill chunk overshoots cached prefill: offset=" - f"{offset}, span_len={span_len}, prefill_len={full_len}. " - "User must pass prompt_token_ids of length " - "speaker_latent.shape[0]." - ) - - buf_s = start - buf_e = buf_s + span_len - self._text_tokens[buf_s:buf_e].copy_(cached_text_tokens[s:e]) - self._text_mask[buf_s:buf_e].copy_(cached_text_mask[s:e]) - self._bos_mask[buf_s:buf_e].copy_(cached_bos_mask[s:e]) - self._speaker_latent[buf_s:buf_e].copy_(cached_speaker_latent[s:e]) - # Acoustic input is sil_tokens broadcast — only the BOS-frame's - # audio_emb is consumed (the rest get replaced by speaker_latent - # inside EarTTSInputEmbedding). - self._acoustic_tokens[buf_s:buf_e] = self.model.sil_tokens.to( - self._acoustic_tokens.dtype - ) - - info_update["ear_prefill_offset"] = offset + span_len - - # Placeholder input_ids; compiled forward reads the buffers, not these. - input_ids_out = torch.full_like(input_ids, _DUMMY_TOKEN_ID) - return input_ids_out, input_embeds, info_update - - def _preprocess_decode( - self, - *, - input_ids: torch.Tensor, - input_embeds: Optional[torch.Tensor], - start: int, - device: torch.device, - info_dict: dict[str, Any], - ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: - """Decode branch of :meth:`preprocess`. Writes one row at ``start``.""" - chunk_text_tokens = self._payload_get(info_dict, "ids", "output") - assert isinstance(chunk_text_tokens, list) and chunk_text_tokens, ( - "EarTTS decode requires a non-empty ``ids.output`` list in " - f"additional_information; got {type(chunk_text_tokens).__name__} " - f"with available keys {sorted(info_dict)}." - ) - - decode_offset = int(info_dict.get("ear_decode_offset", 0) or 0) - # The newest token is the one this step consumes: the producer sends - # exactly one per step, and a producer that resends a history still - # has the current token at the end. - text_token_id = int(chunk_text_tokens[-1]) - - buf_s = start - - # Acoustic input selection: - # * EOS subword → force sil_tokens (return to silence). - # * First decode after prefill → seed with the acoustic pad id - # (codebook_size) broadcast across all quantizers. - # * Otherwise → previous-step predicted codes. - if text_token_id == self._eos_token_id: - self._acoustic_tokens[buf_s].copy_( - self.model.sil_tokens.to(self._acoustic_tokens.dtype) - ) - elif decode_offset == 0: - self._acoustic_tokens[buf_s].fill_(self._acoustic_pad_id) - else: - last_codes = info_dict.get("last_acoustic_codes") - assert isinstance(last_codes, torch.Tensor) and last_codes.numel() > 0, ( - "EarTTS decode (offset > 0) requires " - "``last_acoustic_codes`` from the previous step's " - "postprocess." - ) - ac = ( - last_codes.to(device=device, dtype=torch.long) - .reshape(-1)[: self._num_quantizers] - ) - self._acoustic_tokens[buf_s, : ac.shape[0]].copy_(ac) - - self._text_tokens[buf_s] = text_token_id - self._text_mask[buf_s] = 1 - self._bos_mask[buf_s] = 0 - # Decode never replaces audio_emb with a latent. - self._speaker_latent[buf_s].zero_() - - info_update: dict[str, Any] = {"ear_decode_offset": decode_offset + 1} - return input_ids, input_embeds, info_update - - # ------------------------------------------------------------------ - # forward — runs the compiled embedding + backbone, then the sampler - # only on decode positions (skipping the expensive MaskGIT loop on - # prefill positions). - # ------------------------------------------------------------------ - - def _get_decode_idxs(self): - """Return ``(decode_token_indices, num_requests)`` for sampler dispatch. - - Mirrors the qwen3-tts code-predictor pattern: - - * ``(None, 0)`` → run sampler on every token. Used during - profile / dummy runs (no ``attn_metadata``) and decode-only - batches (``max_query_len == 1``), so the captured CUDA graph - covers all of ``cudagraph_capture_sizes``. - * ``(decode_token_indices, num_requests)`` → run sampler only on - the listed positions. ``decode_token_indices`` is padded up to - the next captured CUDA-graph size (so the sampler's graph - cache is hit) and ``num_requests`` is the unpadded count of - real decode tokens (used to scatter codes back into the right - rows of :attr:`_out_codes`). - """ - ctx = get_forward_context() - attn_metadata = ctx.attn_metadata - if attn_metadata is None: - # Profile / dummy run. Apply sampler everywhere so capture - # covers every cudagraph_capture_sizes value. - return None, 0 - - if isinstance(attn_metadata, dict): - any_layer_meta = next(iter(attn_metadata.values())) - else: - any_layer_meta = attn_metadata - - if any_layer_meta.max_query_len == 1: - # Decode-only batch: every position is a decode position, - # so just run the sampler over the whole flat batch. - return None, 0 - - start_loc = any_layer_meta.query_start_loc - tokens_per_req = start_loc[1:] - start_loc[:-1] - is_decode = (tokens_per_req == 1) - decode_token_indices = start_loc[:-1][is_decode] - - num_requests = decode_token_indices.shape[0] - padded_num_requests = num_requests - if ( - self.vllm_config.compilation_config.cudagraph_mode - != CUDAGraphMode.NONE - ): - sizes = self.vllm_config.compilation_config.cudagraph_capture_sizes - idx = bisect.bisect_left(sizes, num_requests) - if idx < len(sizes): - padded_num_requests = sizes[idx] - if padded_num_requests != num_requests: - decode_token_indices = torch.nn.functional.pad( - decode_token_indices, - (0, padded_num_requests - num_requests), - ) - return decode_token_indices, num_requests - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - intermediate_tensors: Optional[IntermediateTensors] = None, - inputs_embeds: Optional[torch.Tensor] = None, - **_: Any, - ) -> torch.Tensor: - """Run the compiled embedding + backbone over every position, - then run the compiled MaskGIT sampler only on decode positions - (the sampler is skipped on prefill-only batches and on prefill - rows of mixed batches). ``inputs_embeds`` is ignored — the - actual embedding is assembled inside the compiled - :class:`EarTTSInputEmbedding` from the per-token buffers - populated by :meth:`preprocess`. - """ - num_tokens = int(input_ids.shape[0]) - - acoustic_tokens = self._acoustic_tokens[:num_tokens] - text_tokens = self._text_tokens[:num_tokens] - text_mask = self._text_mask[:num_tokens] - bos_mask = self._bos_mask[:num_tokens] - speaker_latent = self._speaker_latent[:num_tokens] - cfg_is_uncond = self._cfg_is_uncond[:num_tokens] - - hidden_states = self.model( - input_ids=input_ids, - positions=positions, - intermediate_tensors=intermediate_tensors, - acoustic_tokens=acoustic_tokens, - text_tokens=text_tokens, - text_mask=text_mask, - bos_mask=bos_mask, - speaker_latent=speaker_latent, - cfg_is_uncond=cfg_is_uncond, - ) - - decode_idx, num_req = self._get_decode_idxs() - if decode_idx is None: - # Dummy/profile run or decode-only batch: sample everywhere. - codes = self.sampler_module.sample( - hidden_states, - cfg_enabled=self._cfg_enabled[:num_tokens], - cfg_is_uncond=self._cfg_is_uncond[:num_tokens], - cfg_pair_id=self._cfg_pair_id[:num_tokens], - cfg_scale=self._cfg_scale[:num_tokens], - valid=torch.ones( - num_tokens, - dtype=torch.bool, - device=hidden_states.device, - ), - ) - self._out_codes[:num_tokens].copy_(codes.to(dtype=torch.long)) - elif num_req > 0: - # Mixed batch: gather decode positions, override the - # BatchDescriptor so the sampler's CUDA-graph cache is hit - # at the padded decode-batch size. - ctx = get_forward_context() - orig_batch_descriptor = ctx.batch_descriptor - ctx.batch_descriptor = BatchDescriptor( - num_tokens=decode_idx.shape[0], - ) - decode_hidden = hidden_states[decode_idx] - sampler_valid = ( - torch.arange( - decode_idx.shape[0], - device=decode_idx.device, - ) - < num_req - ) - codes = self.sampler_module.sample( - decode_hidden, - cfg_enabled=self._cfg_enabled[decode_idx], - cfg_is_uncond=self._cfg_is_uncond[decode_idx], - cfg_pair_id=self._cfg_pair_id[decode_idx], - cfg_scale=self._cfg_scale[decode_idx], - valid=sampler_valid, - ) - ctx.batch_descriptor = orig_batch_descriptor - - valid_dec_idx = decode_idx[:num_req] - self._out_codes[valid_dec_idx] = codes[:num_req].to( - dtype=torch.long - ) - # Prefill-only batch: sampler skipped. ``_out_codes`` rows for - # those positions are not written here on purpose — callers - # must not rely on them; the public per-decode-step contract - # is driven by ``last_acoustic_codes`` from postprocess and - # the seed rules in :meth:`_preprocess_decode`. - - return hidden_states - - # ------------------------------------------------------------------ - # compute_logits — sampler bypass (the real output is ``codes``) - # ------------------------------------------------------------------ - - def compute_logits( - self, - hidden_states: Union[torch.Tensor, OmniOutput], - sampling_metadata: Any = None, - ) -> Optional[torch.Tensor]: - """Return zero logits of width ``config.vocab_size``. - - ``config.vocab_size`` is what vLLM's sampler / ``LogitsProcessor`` - use to size their working buffers, so deriving the width from - the same field guarantees the two agree. The sampled token id - is irrelevant: ``input_ids`` are never consumed by the model - (the per-step decode behaviour is driven by the buffers - populated in :meth:`preprocess`), and the real audio output is - the codes tensor exposed via :meth:`make_omni_output`. - """ - if isinstance(hidden_states, OmniOutput): - hidden_states = hidden_states.text_hidden_states - if hidden_states is None: - return None - batch_size = hidden_states.shape[0] - return hidden_states.new_zeros(batch_size, int(self.config.vocab_size)) - - # ------------------------------------------------------------------ - # multimodal output plumbing - # ------------------------------------------------------------------ - - def make_omni_output( - self, - model_outputs: Union[torch.Tensor, OmniOutput], - **_: Any, - ) -> OmniOutput: - """Wrap backbone hidden states with the codes generated by the - sampler (BTx31). - - The key depends on whether this stage is the client-facing one; see - ``self._single_stage_audio`` in :meth:`__init__`. ``postprocess`` - accepts either. - """ - if isinstance(model_outputs, OmniOutput): - return model_outputs - - hidden = model_outputs - num_tokens = int(hidden.shape[0]) - audio_codes = self._out_codes[:num_tokens].clone() - key = "model_outputs" if self._single_stage_audio else "audio_codes" - return OmniOutput( - text_hidden_states=hidden, - multimodal_outputs={key: audio_codes}, - ) - - # ------------------------------------------------------------------ - # postprocess — stash last-frame codes for the next decode step - # ------------------------------------------------------------------ - - def postprocess( - self, - hidden_states: torch.Tensor, - multimodal_outputs: Optional[dict[str, Any]] = None, - **_: Any, - ) -> dict[str, Any]: - """Pull the last-frame codes out of the multimodal output (or, - as a fallback, out of :attr:`_out_codes` using the slice's - storage offset) and stash them under ``last_acoustic_codes`` so - the next step's :meth:`preprocess` can use them as the decode - input. - """ - if hidden_states.numel() == 0: - return {} - - mm = multimodal_outputs or {} - audio_codes = mm.get("audio_codes") - if audio_codes is None: - audio_codes = mm.get("model_outputs") - if isinstance(audio_codes, torch.Tensor) and audio_codes.numel() > 0: - # ``hidden_states`` is a slice of the flat batch. Recover - # the request's last position via storage_offset and pick - # the corresponding row from ``audio_codes``. - stride0 = hidden_states.stride(0) or 1 - req_start = hidden_states.storage_offset() // stride0 - last = req_start + hidden_states.shape[0] - 1 - last_codes = audio_codes[last : last + 1].detach() - return {"last_acoustic_codes": last_codes} - - return {} - - # ------------------------------------------------------------------ - # weight loading - # ------------------------------------------------------------------ - - def load_weights( - self, weights: Iterable[tuple[str, torch.Tensor]] - ) -> set[str]: - skip_prefixes: list[str] = [] - if self.config.tie_word_embeddings: - skip_prefixes.append("lm_head.") - - # The Gemma3 backbone keeps a vestigial ``embed_tokens`` layer, - # but this model never consumes ``input_ids`` (every forward - # goes through ``inputs_embeds`` assembled from the audio / text - # buffers in :meth:`preprocess`). We expose a 2-class placeholder - # ``vocab_size`` purely so the vLLM sampler's working buffers - # match the dummy logits returned by :meth:`compute_logits`. - # The checkpoint, however, ships ``embed_tokens.weight`` at the - # original tokenizer vocab size — which trips - # ``VocabParallelEmbedding``'s - # ``loaded_weight.shape[output_dim] == self.org_vocab_size`` - # assertion. Truncate (or pad) the loaded weight to - # ``(config.vocab_size, hidden_size)`` so the assertion passes; - # the surviving rows are never consumed at runtime. - target_vocab = int(self.config.vocab_size) - embed_weight_name = "model.backbone.embed_tokens.weight" - - def _adjusted_weights() -> Iterable[tuple[str, torch.Tensor]]: - for name, w in weights: - if name == embed_weight_name and w.dim() >= 1 and w.shape[0] != target_vocab: - if w.shape[0] >= target_vocab: - yield name, w[:target_vocab].contiguous() - else: - pad = torch.zeros( - target_vocab - w.shape[0], - *w.shape[1:], - dtype=w.dtype, - device=w.device, - ) - yield name, torch.cat([w, pad], dim=0).contiguous() - else: - yield name, w - - # ``AutoWeightsLoader`` only dispatches into child modules and - # ``nn.Parameter``s, so any registered buffer (e.g. - # ``model.sil_tokens``) needs to be routed manually. Resolve the - # buffer names *after* applying ``hf_to_vllm_mapper`` so the - # checkpoint key matches the in-model attribute path. - buffers_dict = dict(self.named_buffers()) - loaded_buffer_names: set[str] = set() - - def _route_buffers( - stream: Iterable[tuple[str, torch.Tensor]], - ) -> Iterable[tuple[str, torch.Tensor]]: - for name, weight in stream: - if name in buffers_dict: - buf = buffers_dict[name] - with torch.no_grad(): - buf.copy_(weight.to(buf.dtype)) - loaded_buffer_names.add(name) - continue - yield name, weight - - # ``hf_to_vllm_mapper`` rewrites ``model.sampler.*`` to - # ``sampler_module.sampler.*`` so the upstream EarTTS checkpoint - # (which still places the MaskGIT sampler under ``model.``) lands - # on the dedicated :attr:`sampler_module` compile group. - loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) - loaded = loader.load_weights( - _route_buffers(_adjusted_weights()), mapper=self.hf_to_vllm_mapper - ) - loaded.update(loaded_buffer_names) - return loaded diff --git a/nemo/collections/speechlm2/inference/vllm_omni/eartts/pipeline.py b/nemo/collections/speechlm2/inference/vllm_omni/eartts/pipeline.py deleted file mode 100644 index 59dfa2adf490..000000000000 --- a/nemo/collections/speechlm2/inference/vllm_omni/eartts/pipeline.py +++ /dev/null @@ -1,63 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Single-stage streaming EarTTS pipeline. - -NeMo submits text tokens directly to this engine after NemotronDuplexH has -produced them. CFG is represented by two explicit requests in the same -engine, giving the conditional and unconditional streams independent vLLM KV -caches while allowing :class:`EarTTSCFGScheduler` to keep them in lockstep. -""" - -from vllm_omni.config.stage_config import ( - PipelineConfig, - StageExecutionType, - StagePipelineConfig, -) - -_CFG_SCHEDULER = ( - "nemo.collections.speechlm2.inference.vllm_omni." - "eartts.scheduler.EarTTSCFGScheduler" -) - - -EARTTS_PIPELINE = PipelineConfig( - model_type="eartts", - model_arch="EarTTSForCausalLM", - hf_architectures=("EarTTSForCausalLM",), - stages=( - StagePipelineConfig( - stage_id=0, - model_stage="eartts", - execution_type=StageExecutionType.LLM_AR, - input_sources=(), - final_output=True, - final_output_type="audio", - # vLLM-Omni derives the entry-stage ``generate`` task from - # ``owns_tokenizer`` (runtime name: ``is_comprehension``). EarTTS - # consumes token-id placeholders and deploy keeps - # ``skip_tokenizer_init: true``, but this flag must still be true - # for a direct SamplingParams request to pass task validation. - owns_tokenizer=True, - model_arch="EarTTSForCausalLM", - engine_output_type="audio", - retains_state_across_chunks=True, - scheduler_cls=_CFG_SCHEDULER, - sampling_constraints={"detokenize": False}, - ), - ), -) - - -__all__ = ["EARTTS_PIPELINE"] diff --git a/nemo/collections/speechlm2/inference/vllm_omni/eartts/scheduler.py b/nemo/collections/speechlm2/inference/vllm_omni/eartts/scheduler.py deleted file mode 100644 index 3f559dbe3ebd..000000000000 --- a/nemo/collections/speechlm2/inference/vllm_omni/eartts/scheduler.py +++ /dev/null @@ -1,451 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Pair-aware synchronous vLLM-Omni scheduler for EarTTS CFG. - -EarTTS classifier-free guidance is represented by two ordinary top-level -requests. Their ``SamplingParams.extra_args`` must contain:: - - { - "cfg_enabled": True, - "cfg_role": "cond" | "uncond", - "cfg_pair_id": "", - "cfg_scale": 0.5, - } - -The model runner is responsible for blending the pair's model outputs. This -scheduler supplies the ordering and lock-step contract required by that -operation. It deliberately subclasses the synchronous Omni AR scheduler: -async placeholder scheduling can put the two members on different token -positions before either result reaches the scheduler. -""" - -from __future__ import annotations - -import math -from typing import Any - -from vllm.logger import init_logger -from vllm.v1.core.sched.output import SchedulerOutput -from vllm.v1.request import Request, RequestStatus, StreamingUpdate -from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler - -logger = init_logger(__name__) - -_COND = "cond" -_UNCOND = "uncond" -_ROLES = (_COND, _UNCOND) - - -def _extra_args(request: Any) -> dict[str, Any]: - sampling_params = getattr(request, "sampling_params", None) - extra_args = getattr(sampling_params, "extra_args", None) - return extra_args if isinstance(extra_args, dict) else {} - - -def _normalized_sampled_tokens(value: Any) -> tuple[int, ...]: - if value is None: - return () - if isinstance(value, int): - return (value,) - if hasattr(value, "tolist"): - value = value.tolist() - if isinstance(value, int): - return (value,) - return tuple(int(token_id) for token_id in value) - - -class EarTTSCFGScheduler(OmniARScheduler): - """Synchronous Omni AR scheduler that keeps EarTTS CFG pairs lock-step.""" - - def __init__(self, *args: Any, **kwargs: Any) -> None: - super().__init__(*args, **kwargs) - self._cfg_pairs: dict[str, dict[str, str]] = {} - self._cfg_req_to_pair: dict[str, str] = {} - self._cfg_pair_scales: dict[str, float] = {} - - max_num_seqs = int(getattr(self.scheduler_config, "max_num_seqs", 0) or 0) - if max_num_seqs and max_num_seqs < 2: - raise ValueError("EarTTSCFGScheduler requires max_num_seqs >= 2") - - @staticmethod - def _cfg_metadata(request: Request) -> tuple[str, str, float] | None: - extra_args = _extra_args(request) - if not bool(extra_args.get("cfg_enabled", False)): - return None - - role = extra_args.get("cfg_role") - if role not in _ROLES: - raise ValueError(f"CFG request {request.request_id!r}: cfg_role must be 'cond' or 'uncond', got {role!r}") - - raw_pair_id = extra_args.get("cfg_pair_id") - if raw_pair_id is None or not str(raw_pair_id): - raise ValueError(f"CFG request {request.request_id!r}: cfg_pair_id must be non-empty") - pair_id = str(raw_pair_id) - - raw_scale = extra_args.get("cfg_scale") - if isinstance(raw_scale, bool) or not isinstance(raw_scale, int | float): - raise ValueError(f"CFG request {request.request_id!r}: cfg_scale must be a finite non-negative number") - scale = float(raw_scale) - if not math.isfinite(scale) or scale < 0.0: - raise ValueError(f"CFG request {request.request_id!r}: cfg_scale must be a finite non-negative number") - return pair_id, role, scale - - def add_request(self, request: Request) -> None: - metadata = self._cfg_metadata(request) - if metadata is not None: - pair_id, role, scale = metadata - roles = self._cfg_pairs.get(pair_id, {}) - existing = roles.get(role) - if existing is not None and existing != request.request_id: - raise ValueError(f"CFG pair {pair_id!r} already has {role} request {existing!r}") - existing_scale = self._cfg_pair_scales.get(pair_id) - if existing_scale is not None and existing_scale != scale: - raise ValueError( - f"CFG pair {pair_id!r} has inconsistent cfg_scale values: {existing_scale} and {scale}" - ) - - super().add_request(request) - - if metadata is not None: - pair_id, role, scale = metadata - self._cfg_pairs.setdefault(pair_id, {})[role] = request.request_id - self._cfg_req_to_pair[request.request_id] = pair_id - self._cfg_pair_scales[pair_id] = scale - # Conditional and null prompts need deterministic, identical - # progress; asymmetric prefix-cache hits violate that contract. - if hasattr(request, "skip_reading_prefix_cache"): - request.skip_reading_prefix_cache = True - - def _pair_requests(self, pair_id: str) -> tuple[Request | None, Request | None]: - roles = self._cfg_pairs.get(pair_id, {}) - return self.requests.get(roles.get(_COND, "")), self.requests.get(roles.get(_UNCOND, "")) - - def _drop_pair(self, pair_id: str) -> None: - for request_id in self._cfg_pairs.pop(pair_id, {}).values(): - self._cfg_req_to_pair.pop(request_id, None) - self._cfg_pair_scales.pop(pair_id, None) - - @staticmethod - def _remove_from_queue(queue: Any, requests: list[Request]) -> None: - if not requests: - return - if hasattr(queue, "remove_requests"): - queue.remove_requests(requests) - return - for request in requests: - queue.remove(request) - - @staticmethod - def _prepend_to_queue(queue: Any, requests: list[Request]) -> None: - if not requests: - return - if hasattr(queue, "prepend_requests"): - queue.prepend_requests(requests) - return - if hasattr(queue, "prepend_request"): - for request in reversed(requests): - queue.prepend_request(request) - return - for request in reversed(requests): - queue.insert(0, request) - - @classmethod - def _replace_queue(cls, queue: Any, requests: list[Request]) -> None: - current = list(queue) - cls._remove_from_queue(queue, current) - for request in requests: - queue.add_request(request) if hasattr(queue, "add_request") else queue.append(request) - - def _available_sequence_slots(self) -> int: - capacity = getattr(self, "max_num_running_reqs", None) - if capacity is None: - capacity = getattr(self.scheduler_config, "max_num_seqs", 0) - return max(0, int(capacity or 0) - len(self.running)) - - def _prepare_cfg_waiting(self) -> list[tuple[Any, list[Request]]]: - """Hide unsafe pairs and put admissible pairs first and adjacent.""" - skipped_queue = getattr(self, "skipped_waiting", None) - promoted_held: list[Request] = [] - - # Streaming updates are applied by vLLM while traversing - # ``skipped_waiting``. The first traversal promotes each pair member - # from WAITING_FOR_STREAMING_REQ to WAITING but deliberately skips - # scheduling it (see _try_promote_blocked_waiting_request below). - # Once both members are promoted, move them to the ordinary waiting - # queue together so the admission logic below sees one atomic pair. - if skipped_queue is not None: - skipped_by_id = { - request.request_id: request - for request in list(skipped_queue) - } - for pair_id, roles in self._cfg_pairs.items(): - if set(roles) != set(_ROLES): - continue - cond = skipped_by_id.get(roles[_COND]) - uncond = skipped_by_id.get(roles[_UNCOND]) - if cond is None or uncond is None: - continue - ready = [ - request - for request in (cond, uncond) - if request.status == RequestStatus.WAITING - ] - if len(ready) == 2: - self._remove_from_queue( - skipped_queue, [cond, uncond] - ) - self.waiting.add_request(cond) - self.waiting.add_request(uncond) - elif len(ready) == 1: - self._remove_from_queue(skipped_queue, ready) - promoted_held.extend(ready) - - waiting_items = list(self.waiting) - skipped_items = list(skipped_queue) if skipped_queue is not None else [] - waiting_ids = {request.request_id for request in waiting_items} - skipped_ids = {request.request_id for request in skipped_items} - running_ids = {request.request_id for request in self.running} - - held: list[tuple[Any, list[Request]]] = [] - if promoted_held: - held.append((skipped_queue, promoted_held)) - hold_waiting: set[str] = set() - complete_waiting_pairs: list[str] = [] - - for pair_id, roles in self._cfg_pairs.items(): - pair_ids = {request_id for request_id in roles.values()} - is_complete = set(roles) == set(_ROLES) and all(request_id in self.requests for request_id in pair_ids) - in_waiting = pair_ids & waiting_ids - in_skipped = pair_ids & skipped_ids - in_running = pair_ids & running_ids - - if in_running and (in_waiting or in_skipped): - raise RuntimeError(f"EarTTS CFG pair {pair_id!r} was split across running and waiting queues") - if len(in_running) == 1: - raise RuntimeError(f"EarTTS CFG pair {pair_id!r} has only one running member") - - if not is_complete or len(in_waiting) == 1: - hold_waiting.update(in_waiting) - elif len(in_waiting) == 2: - complete_waiting_pairs.append(pair_id) - - # A pair consumes two sequence slots. Expose only whole pairs to the - # upstream scheduler and put them before ordinary requests so another - # admission cannot consume the second slot between pair members. - admitted_pair_count = self._available_sequence_slots() // 2 - allowed_pairs = set(complete_waiting_pairs[:admitted_pair_count]) - for pair_id in complete_waiting_pairs[admitted_pair_count:]: - hold_waiting.update(self._cfg_pairs[pair_id].values()) - - held_waiting = [request for request in waiting_items if request.request_id in hold_waiting] - self._remove_from_queue(self.waiting, held_waiting) - if held_waiting: - held.append((self.waiting, held_waiting)) - - remaining = list(self.waiting) - ordinary = [request for request in remaining if request.request_id not in self._cfg_req_to_pair] - ordered: list[Request] = [] - for pair_id in complete_waiting_pairs: - if pair_id not in allowed_pairs: - continue - cond, uncond = self._pair_requests(pair_id) - if cond is not None and uncond is not None: - ordered.extend((cond, uncond)) - ordered.extend(ordinary) - if ordered != remaining: - self._replace_queue(self.waiting, ordered) - - actual_ids = [request.request_id for request in self.waiting] - for pair_id in allowed_pairs: - roles = self._cfg_pairs[pair_id] - cond_id, uncond_id = roles[_COND], roles[_UNCOND] - try: - cond_index = actual_ids.index(cond_id) - except ValueError: - continue - if cond_index + 1 >= len(actual_ids) or actual_ids[cond_index + 1] != uncond_id: - raise RuntimeError( - "EarTTSCFGScheduler requires an FCFS-compatible waiting " - f"queue; CFG pair {pair_id!r} could not be made adjacent" - ) - - return held - - def _try_promote_blocked_waiting_request( - self, request: Request - ) -> bool: - promoted = super()._try_promote_blocked_waiting_request( - request - ) - if promoted and request.request_id in self._cfg_req_to_pair: - # The base scheduler would immediately schedule this first - # promoted member. Return False once so it is parked back in - # skipped_waiting; after its peer is promoted, - # _prepare_cfg_waiting moves both to waiting atomically. - return False - return promoted - - def _should_defer_waiting_admission(self) -> bool: - """Install the pair guard after Omni has processed pending inputs. - - ``OmniARScheduler.schedule`` invokes this hook immediately before the - stock vLLM scheduler. Preparing here is important: doing it at the - start of this class's :meth:`schedule` would hide requests from Omni's - chunk/input processing and could leave a streaming pair parked in - ``skipped_waiting`` forever. - """ - self._cfg_decode_ready_before = { - request.request_id - for request in self.running - if self._get_confirmed_num_computed_tokens(request) >= request.num_prompt_tokens - } - self._cfg_held_for_schedule = self._prepare_cfg_waiting() - self._cfg_waiting_before_schedule = {request.request_id for request in self.waiting} - return super()._should_defer_waiting_admission() - - def _restore_held(self, held: list[tuple[Any, list[Request]]]) -> None: - for queue, requests in reversed(held): - self._prepend_to_queue(queue, requests) - - def _equalize_pair_progress(self, scheduler_output: SchedulerOutput) -> None: - scheduled = scheduler_output.num_scheduled_tokens - for pair_id in self._cfg_pairs: - cond, uncond = self._pair_requests(pair_id) - if cond is None or uncond is None: - continue - cond_scheduled = int(scheduled.get(cond.request_id, 0) or 0) - uncond_scheduled = int(scheduled.get(uncond.request_id, 0) or 0) - if not cond_scheduled or not uncond_scheduled: - continue - - target = min(cond.num_computed_tokens, uncond.num_computed_tokens) - feasible = all( - request.num_computed_tokens - target < count - for request, count in ((cond, cond_scheduled), (uncond, uncond_scheduled)) - ) - if not feasible: - continue - - for request, count in ((cond, cond_scheduled), (uncond, uncond_scheduled)): - difference = request.num_computed_tokens - target - if difference <= 0: - continue - request.num_computed_tokens = target - if hasattr(request, "num_in_flight_tokens"): - request.num_in_flight_tokens = max(0, request.num_in_flight_tokens - difference) - scheduler_output.num_scheduled_tokens[request.request_id] = count - difference - scheduler_output.total_num_scheduled_tokens -= difference - - def _assert_atomic_admission(self, scheduler_output: SchedulerOutput, waiting_before: set[str]) -> None: - scheduled = scheduler_output.num_scheduled_tokens - for pair_id, roles in self._cfg_pairs.items(): - pair_ids = {roles.get(_COND), roles.get(_UNCOND)} - pair_ids.discard(None) - if len(pair_ids & waiting_before) != 2: - continue - admitted = {request_id for request_id in pair_ids if scheduled.get(request_id, 0)} - if admitted and admitted != pair_ids: - raise RuntimeError( - f"EarTTS CFG pair {pair_id!r} was not admitted atomically: scheduled={sorted(admitted)}" - ) - - def _assert_complete_decode_pairs(self, scheduler_output: SchedulerOutput, decode_ready_before: set[str]) -> None: - scheduled = scheduler_output.num_scheduled_tokens - for pair_id, roles in self._cfg_pairs.items(): - cond_id = roles.get(_COND) - uncond_id = roles.get(_UNCOND) - if cond_id is None or uncond_id is None: - continue - cond_count = int(scheduled.get(cond_id, 0) or 0) - uncond_count = int(scheduled.get(uncond_id, 0) or 0) - if not cond_count and not uncond_count: - continue - if cond_count != uncond_count: - raise RuntimeError( - f"EarTTS scheduler split CFG decode pair {pair_id!r}: " - f"{cond_id}={cond_count}, {uncond_id}={uncond_count}" - ) - - def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: - self._cfg_held_for_schedule: list[tuple[Any, list[Request]]] = [] - self._cfg_waiting_before_schedule: set[str] = set() - self._cfg_decode_ready_before: set[str] = set() - - scheduler_config = getattr(self, "scheduler_config", None) - original_threshold = getattr(scheduler_config, "long_prefill_token_threshold", None) - if self._cfg_pairs and original_threshold is not None: - budget = int(getattr(self, "max_num_scheduled_tokens", 0) or 0) - if budget: - scheduler_config.long_prefill_token_threshold = max(1, budget // 2) - try: - scheduler_output = super().schedule(throttle_prefills) - finally: - if original_threshold is not None: - scheduler_config.long_prefill_token_threshold = original_threshold - self._restore_held(self._cfg_held_for_schedule) - - self._equalize_pair_progress(scheduler_output) - self._assert_atomic_admission(scheduler_output, self._cfg_waiting_before_schedule) - self._assert_complete_decode_pairs(scheduler_output, self._cfg_decode_ready_before) - return scheduler_output - - def _assert_matching_sampled_tokens(self, scheduler_output: SchedulerOutput, model_runner_output: Any) -> None: - sampled = getattr(model_runner_output, "sampled_token_ids", None) - req_id_to_index = getattr(model_runner_output, "req_id_to_index", None) - if sampled is None or not isinstance(req_id_to_index, dict): - return - - scheduled = scheduler_output.num_scheduled_tokens - for pair_id, roles in self._cfg_pairs.items(): - cond_id = roles.get(_COND) - uncond_id = roles.get(_UNCOND) - if ( - cond_id not in scheduled - or uncond_id not in scheduled - or cond_id not in req_id_to_index - or uncond_id not in req_id_to_index - ): - continue - cond_tokens = _normalized_sampled_tokens(sampled[req_id_to_index[cond_id]]) - uncond_tokens = _normalized_sampled_tokens(sampled[req_id_to_index[uncond_id]]) - if cond_tokens != uncond_tokens: - raise RuntimeError( - f"EarTTS CFG pair {pair_id!r} sampled-token mismatch: " - f"cond={cond_tokens}, uncond={uncond_tokens}" - ) - - def update_from_output(self, scheduler_output: SchedulerOutput, model_runner_output: Any) -> Any: - self._assert_matching_sampled_tokens(scheduler_output, model_runner_output) - # A length stop on a resumable StreamingInput segment is not terminal: - # Omni parks the request until its next chunk. Do not mirror the - # segment's internal free/park operation onto its CFG peer here. - # Explicit abort/final teardown still flows through finish_requests, - # which expands either member to the complete pair below. - return super().update_from_output( - scheduler_output, model_runner_output - ) - - def _update_request_as_session(self, session: Request, update: StreamingUpdate) -> None: - super()._update_request_as_session(session, update) - # Stage 0 receives direct StreamingInput updates. Upstream extends the - # prompt but does not copy this per-chunk payload. Downstream stages - # are intentionally untouched; their connector owns payload delivery. - if self.vllm_config.model_config.stage_id == 0: - additional_information = getattr(update, "additional_information", None) - if additional_information is not None: - session.additional_information = additional_information - - -__all__ = ["EarTTSCFGScheduler"] diff --git a/nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/__init__.py b/nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/__init__.py deleted file mode 100644 index 33bb11b24f6d..000000000000 --- a/nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/__init__.py +++ /dev/null @@ -1,19 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from nemo.collections.speechlm2.inference.vllm_omni.nemotron_duplex_h.nemotron_duplex_h import ( - NemotronDuplexHForCausalLM, -) - -__all__ = ["NemotronDuplexHForCausalLM"] diff --git a/nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/nemotron_duplex_h.py b/nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/nemotron_duplex_h.py deleted file mode 100644 index ee7df96cbaac..000000000000 --- a/nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/nemotron_duplex_h.py +++ /dev/null @@ -1,771 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Inference-only NemotronDuplexH model for vLLM-Omni. - -A minimal extension of the upstream :class:`NemotronHForCausalLM` that: - -1. Accepts pre-computed acoustic encoder embeddings per step via - ``acoustic_embedding`` in the per-request payload (one row per - scheduled token). The prefill step receives the *system prompt as - raw text* via ``system_prompt`` in the same payload; the - model's :meth:`preprocess` tokenizes it in-process (using a - HuggingFace tokenizer loaded once in ``__init__`` from the - checkpoint dir) and constructs the prefill combined embedding - itself as - - prompt_embed = embed_tokens([BOS] + text_ids + [EOS]) - + embed_tokens(pad_id) - + embed_asr_tokens(pad_id) - - It then *clears* the buffer entry by returning - ``{"system_prompt": None}`` as its update dict so that subsequent - decode steps fall through to the decode branch. The producer - should still send ``system_prompt=None`` on every decode chunk to - make the intent explicit, but the actual clearing happens - consumer-side because the orchestrator's serialization filters - ``None`` values. -2. Embeds up to two additional per-step token id streams, each fed - **autoregressively from the model itself** via a per-request buffer - that ``postprocess`` keeps populated after every step: - - - ``input_asr_ids`` – the ASR channel (``predict_user_text`` - checkpoints), embedded with its own ``embed_asr_tokens`` table. - - ``input_function_ids`` – the function channel - (``use_function_head`` checkpoints), embedded with the *text* - ``embed_tokens`` table and scaled by - ``duplex_function_channel_weight``, mirroring - ``DuplexSTTModel.build_input_embedding``. - - Which channels exist is checkpoint-dependent. ASR and function can - both be enabled. The converter records whichever heads it found. - -3. Combines the enabled signals into the input embedding fed to the - NemotronH backbone: - - hidden_in = embed_tokens(input_ids) - [+ embed_asr_tokens(input_asr_ids)] - [+ embed_tokens(input_function_ids) * weight] - + acoustic_embedding - -4. Adds a parallel head per enabled channel (``asr_head`` / - ``function_head``) that produces one token at every decoding step. - The head matmul and the ``argmax`` run in :meth:`make_omni_output` - (which the runner invokes *outside* the CUDA-graph wrapper) on the - full-batch ``hidden_states`` returned by :meth:`forward`. The tokens - are exposed under ``OmniOutput.multimodal_outputs["asr_tokens"]`` and - ``["function_tokens"]``, and :meth:`postprocess` stashes the - request's last id of each back into the corresponding buffer so the - next step's :meth:`preprocess` can read it as that channel's - autoregressive input. - - Returning a dict-with-tensor directly from :meth:`forward` is - unsafe under FULL CUDA graphs: ``weak_ref_tensors`` cannot weak-ref - tensors nested inside dicts, and the wrapper coerces ``NamedTuple`` - to a plain ``tuple`` on replay. Routing the multimodal output - through :meth:`make_omni_output` keeps every cudagraph-replayed - value a plain ``Tensor``. - -Text token sampling uses a custom vLLM logits processor which calls the same -PyTorch sampler as the native backend, then forces vLLM's greedy sampler to -the selected token. The auxiliary channels are always greedy. -""" - -from collections.abc import Iterable -from typing import Any - -import torch -from transformers import AutoTokenizer, PreTrainedTokenizerBase -from vllm.config import VllmConfig -from vllm.model_executor.layers.vocab_parallel_embedding import ( - DEFAULT_VOCAB_PADDING_SIZE, - ParallelLMHead, - VocabParallelEmbedding, -) -from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM -from vllm.model_executor.models.utils import ( - AutoWeightsLoader, - WeightsMapper, - maybe_prefix, -) -from vllm.sequence import IntermediateTensors - -from vllm_omni.model_executor.models.output_templates import OmniOutput - -from nemo.collections.speechlm2.parts.logit_boosts import ( - LogitBoosts, - apply_logit_boosts, -) -from nemo.utils import logging as logger - - -def _is_system_prompt_prefill( - system_prompt: Any, - runner_is_prefill: bool, - request_id: str, - prompt_token_cache: dict[str, list[int]], -) -> bool: - """Distinguish initial prompt slices from streaming prompt extensions.""" - has_system_prompt = isinstance(system_prompt, str) and bool( - system_prompt.strip() - ) - return has_system_prompt or ( - runner_is_prefill and request_id in prompt_token_cache - ) - - -def _is_internal_prefill_token( - is_prompt_prefill: bool, - runner_is_prefill: bool, - has_acoustic_embedding: bool, -) -> bool: - """True for vLLM's 1-token generate after the system prompt is in KV. - - The runner labels every streaming extension as prefill because the - prompt grows before that token is computed. After a prompt longer - than the 64-token long-prefill threshold, the last prompt slice pops - the token cache; the engine then schedules one more token (internal - ``t_0``) with ``_omni_is_prefill=True`` and no ``acoustic_embedding``. - A prefill-only ``generate_step`` (empty ``is_first`` frame) does not - queue audio before that happens, so the decode path would assert. - Real client steps always carry an acoustic frame. - """ - return ( - not is_prompt_prefill - and runner_is_prefill - and not has_acoustic_embedding - ) - - -class NemotronDuplexHForCausalLM(NemotronHForCausalLM): - """NemotronH + optional per-step ASR and function token channels.""" - - have_multimodal_outputs = True - has_preprocess = True - has_postprocess = True - - # No ``gpu_resident_buffer_keys``; see the note in ``eartts.py``. The keys - # this stage passes between ``postprocess`` and the next ``preprocess`` are - # flat, which that mechanism cannot express. - - hf_to_vllm_mapper = WeightsMapper( - orig_to_new_prefix={ - # NemotronH backbone weights live under - # `stt_model.llm.backbone.*` in the duplex checkpoint and need to - # land under our `model.*`. - "stt_model.llm.backbone": "model", - "stt_model.llm": "model", - "stt_model.embed_tokens": "model.embed_tokens", - "stt_model.embed_asr_tokens": "embed_asr_tokens", - "stt_model.lm_head": "lm_head", - "stt_model.asr_head": "asr_head", - "stt_model.function_head": "function_head", - # Bare-NemotronH naming, kept as a fallback. - "backbone": "model", - }, - orig_to_new_substr={"A_log": "A", "embeddings": "embed_tokens"}, - # Fusing q/k/v into ``qkv_proj`` is done here. This class replaces - # ``NemotronHForCausalLM.hf_to_vllm_mapper`` wholesale, so the stacked - # mapping has to be restated or attention projections fail to load. - orig_to_new_stacked={ - ".q_proj": (".qkv_proj", "q"), - ".k_proj": (".qkv_proj", "k"), - ".v_proj": (".qkv_proj", "v"), - }, - ) - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__(vllm_config=vllm_config, prefix=prefix) - - config = vllm_config.model_config.hf_config - - # Missing flags default to ASR on / function off. Fresh conversions - # always write both flags from the checkpoint weights. - self.use_asr_head = bool(getattr(config, "use_asr_head", True)) - self.use_function_head = bool(getattr(config, "use_function_head", False)) - self.function_channel_weight = float(getattr(config, "duplex_function_channel_weight", 1.0)) - - if self.use_asr_head: - self.embed_asr_tokens = VocabParallelEmbedding( - config.vocab_size, - config.hidden_size, - org_num_embeddings=config.vocab_size, - ) - - self.asr_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - org_num_embeddings=config.vocab_size, - padding_size=DEFAULT_VOCAB_PADDING_SIZE, - prefix=maybe_prefix(prefix, "asr_head"), - ) - - # The function channel has no embedding table of its own: its feedback - # token is embedded with the text ``embed_tokens``, exactly as - # ``DuplexSTTModel`` does. - if self.use_function_head: - self.function_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - org_num_embeddings=config.vocab_size, - padding_size=DEFAULT_VOCAB_PADDING_SIZE, - prefix=maybe_prefix(prefix, "function_head"), - ) - - # Tokenizer is used in ``preprocess`` to convert the - # ``additional_information["system_prompt"]`` text into token - # IDs on the prefill chunk. Loaded from the checkpoint dir so - # the vocabulary aligns with ``embed_tokens``. Cached once on - # init to keep the per-step preprocess fast. - model_path = vllm_config.model_config.model - self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) - # The runner can split a prompt at its 64-token long-prefill threshold. - # Cache the full tokenization in this model process until every slice - # for a request has been consumed. - self._prompt_token_cache: dict[str, list[int]] = {} - self._seeded_channel_requests: set[str] = set() - self._channel_state: dict[str, dict[str, torch.Tensor]] = {} - # This pipeline is configured with max_num_seqs=1. Keep a fallback - # for vLLM-Omni paths that rewrite the internal request id between - # streaming segments. - self._last_channel_state: dict[str, torch.Tensor] = {} - self._current_is_prompt_prefill = False - - # Special token IDs used to construct the prefill prompt: - # ``[BOS] + text_ids + [EOS]`` and the pad embedding added to - # every prefill position (mirrors the reference STT recipe - # where the BOS / pad embeddings are both ``embed_tokens(pad_id)``). - self.pad_token_id = int(config.pad_token_id) - self.bos_token_id = int(config.bos_token_id) - self.eos_token_id = int(config.eos_token_id) - - # User (ASR) channel boosts, read from the converted config so this - # model applies them exactly as DuplexSTTModel does. The agent-channel - # boosts arrive per request through the shared text sampling hook. - self.user_logit_boosts = LogitBoosts.user_from_cfg(config) - if self.user_logit_boosts: - logger.info( - "NemotronDuplexH user logit boosts: " - f"{self.user_logit_boosts.as_dict()}" - ) - self._last_asr_token = torch.full( - (1,), self.pad_token_id, dtype=torch.long - ) - self._last_function_token = torch.full( - (1,), self.pad_token_id, dtype=torch.long - ) - - # Per-position pad embedding added on every prefill step: the pad id - # embedded once per *enabled* auxiliary channel, shape - # ``(hidden_size,)``. Materialized at the end of - # :meth:`load_weights` because the embedding tables are not - # populated yet in ``__init__``. Registered as a *non-persistent* - # buffer so it follows ``.to(device)`` / dtype casts with the - # rest of the module but is **not** saved in the state_dict - # (it is fully derived from ``embed_tokens`` / - # ``embed_asr_tokens`` which are already saved — duplicating it - # in the checkpoint would just be a footgun). - self.register_buffer("_pad_combined_emb", None, persistent=False) - - # ------------------------------------------------------------------ # - # producer-side helper # - # ------------------------------------------------------------------ # - - @staticmethod - def compute_prefix_len(tokenizer: PreTrainedTokenizerBase, system_prompt: str) -> int: - """Length of the prefill chunk for a given system prompt. - - Mirrors the in-model tokenization done by :meth:`preprocess`: - - [BOS] + tokenizer.encode(system_prompt, add_special_tokens=False) + [EOS] - - The streaming producer needs this number to size the - placeholder ``prompt_token_ids`` it hands vLLM on the prefill - chunk — vLLM schedules off that list's length, while the - actual embedding is constructed inside :meth:`preprocess` - from the ``system_prompt`` string. - - Exposed as a ``@staticmethod`` so callers can compute the - length without instantiating the model (which would download - the full checkpoint). They just need any tokenizer compatible - with the model's vocabulary — typically - ``AutoTokenizer.from_pretrained()``, the - same instance used to decode output tokens. - """ - text_ids = tokenizer.encode(system_prompt, add_special_tokens=False) - return len(text_ids) + 2 # +2 for BOS / EOS wrapped in preprocess - - # ------------------------------------------------------------------ # - # preprocess # - # ------------------------------------------------------------------ # - - def _materialize_pad_combined_emb(self) -> None: - embed_weight = self.model.embed_tokens.weight - device = embed_weight.device - dtype = embed_weight.dtype - pad_tokens = torch.full( - (1,), self.pad_token_id, device=device, dtype=torch.long - ) - pad_emb = self.model.embed_tokens(pad_tokens).to(dtype).squeeze(0) - combined_pad = pad_emb - if self.use_asr_head: - combined_pad = combined_pad + self.embed_asr_tokens( - pad_tokens - ).to(dtype).squeeze(0) - if self.use_function_head: - combined_pad = ( - combined_pad - + pad_emb * self.function_channel_weight - ) - self._pad_combined_emb = combined_pad.detach() - - def _embeds_without_acoustic( - self, input_ids: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: - """Text + pad-channel embeddings, used for prompt slices and internal ``t_0``.""" - if self._pad_combined_emb is None: - self._materialize_pad_combined_emb() - target_dtype = self.model.embed_tokens.weight.dtype - text_emb = self.model.embed_tokens(input_ids).to(target_dtype) - return input_ids, text_emb + self._pad_combined_emb, {"system_prompt": None} - - def preprocess( - self, - input_ids: torch.Tensor, - input_embeds: torch.Tensor | None, - **info_dict: Any, - ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: - """Combine text/asr/speech embeddings into a single per-token vector. - - Three paths: - - * **Prefill construction.** When - ``additional_information["system_prompt"]`` is a non-empty - string, this is the prefill chunk. We tokenize the prompt - in-process as ``[BOS] + tokenizer.encode(prompt) + [EOS]``, - embed it with ``model.embed_tokens``, and add the pad - embedding of each enabled auxiliary channel (which the - reference STT recipe folds in uniformly across every prefill - position): - - prefill_combined = embed_tokens(prompt_token_ids) - + _pad_combined_emb - - The producer-supplied ``input_ids`` for this chunk are - placeholders (their length must match the tokenized prompt - length so vLLM's scheduling sees the right prefill size); - they are returned unchanged so vLLM's bookkeeping is - consistent. We then *clear* the buffer entry by returning - ``{"system_prompt": None}`` in the update dict. Clearing has - to happen here because the orchestrator's serialization - (:func:`vllm_omni.data_entry_keys.serialize_payload`) - silently drops ``None`` values, so the producer cannot - overwrite the buffer with ``None`` via the streaming-input - merge. Decode chunks may carry ``system_prompt=None``, but it - has no effect; the state transition happens only here. - - * **Internal ``t_0`` after a chunked prompt.** vLLM may split the - prompt at its 64-token long-prefill threshold and then - schedule a one-token continuation labeled - ``_omni_is_prefill`` with no ``acoustic_embedding``. That - token is discarded by the session (``output_count <= 1``); - embed it like prefill (text + pad, no user acoustics) so a - prefill-only ``generate_step`` cannot kill the engine. - - * **Decode (single-token step).** Builds the combined embedding - per scheduled token from: - - - ``input_ids`` – per-step text token id (one per - scheduled token; standard vLLM - autoregressive feedback). - - ``input_asr_ids`` – per-step ASR token id, written - back by :meth:`postprocess` on - every step. ASR channel only. - - ``input_function_ids`` – per-step function token id, same - write-back path. Function - channel only. - - ``acoustic_embedding`` – per-step acoustic encoder - embedding, sourced from - ``additional_information``. - - ``input_embeds`` is the runner's pre-allocated scratch buffer - on this path and its contents are ignored. - """ - device = input_ids.device - n = int(input_ids.shape[0]) - - # Prefill vs decode is detected directly on the value of - # ``system_prompt``: a non-empty string means prefill, anything - # else (``None`` / missing / empty) means decode. The buffer - # flips from str → ``None`` inside this method itself (see the - # update dict returned below) because serialization drops - # ``None`` and the producer's "send None on each decode chunk" - # pattern alone is not enough to clear the slot. - system_prompt = info_dict.get("system_prompt") - is_prefill = bool(info_dict.get("_omni_is_prefill", False)) - request_id = str( - info_dict.get("global_request_id") - or info_dict.get("request_id") - or "" - ) - has_system_prompt = isinstance(system_prompt, str) and bool( - system_prompt.strip() - ) - is_prompt_prefill = _is_system_prompt_prefill( - system_prompt, - is_prefill, - request_id, - self._prompt_token_cache, - ) - has_acoustic = isinstance(info_dict.get("acoustic_embedding"), torch.Tensor) - is_internal_t0 = _is_internal_prefill_token( - is_prompt_prefill, is_prefill, has_acoustic - ) - # vLLM labels every one-token streaming extension as prefill because - # the prompt grows before that token is computed. Prompt slices and - # the engine-internal t_0 after them seed auxiliary feedback with PAD; - # client-visible decode steps (which always carry acoustic_embedding) - # do not. - self._current_is_prompt_prefill = is_prompt_prefill or is_internal_t0 - if is_prompt_prefill: - # [BOS] + encode(text, add_special_tokens=False) + [EOS]. - # ``add_special_tokens=False`` keeps full control of which - # specials get wrapped around the text (the underlying HF - # tokenizer would otherwise prepend its own BOS, which may - # or may not equal ``config.bos_token_id``). - if has_system_prompt: - text_ids = self.tokenizer.encode( - system_prompt, add_special_tokens=False - ) - prompt_token_ids = [ - self.bos_token_id, - *text_ids, - self.eos_token_id, - ] - self._prompt_token_cache[request_id] = prompt_token_ids - else: - prompt_token_ids = self._prompt_token_cache[request_id] - prompt_len = len(prompt_token_ids) - expected_prompt_len = info_dict.get("duplex_prompt_len") - if expected_prompt_len is not None: - assert prompt_len == int(expected_prompt_len), ( - f"system_prompt tokenizes to {prompt_len} ids but vLLM " - f"tracks a prompt of length {expected_prompt_len}" - ) - offset = int(info_dict.get("duplex_token_offset", 0) or 0) - end = offset + n - if not (0 <= offset < end <= prompt_len): - # Cache still populated but the runner scheduled past the - # prompt (internal t_0). Same pad-embedding path as below. - self._prompt_token_cache.pop(request_id, None) - return self._embeds_without_acoustic(input_ids) - - prompt_tokens = torch.tensor( - prompt_token_ids[offset:end], - device=device, - dtype=torch.long, - ) - _, prefill_combined, updates = self._embeds_without_acoustic( - prompt_tokens - ) - if end == prompt_len: - self._prompt_token_cache.pop(request_id, None) - return input_ids, prefill_combined, updates - - if is_internal_t0: - return self._embeds_without_acoustic(input_ids) - - combined = self.model.embed_tokens(input_ids) - - request_id = str( - info_dict.get("global_request_id") - or info_dict.get("request_id") - or "" - ) - cached_channel_state = self._channel_state.get( - request_id - ) or self._last_channel_state - for key, value in cached_channel_state.items(): - if not isinstance(info_dict.get(key), torch.Tensor): - info_dict[key] = value - if self.use_asr_head and not isinstance( - info_dict.get("input_asr_ids"), torch.Tensor - ): - info_dict["input_asr_ids"] = self._last_asr_token - if self.use_function_head and not isinstance( - info_dict.get("input_function_ids"), torch.Tensor - ): - info_dict[ - "input_function_ids" - ] = self._last_function_token - if request_id not in self._seeded_channel_requests: - # The initial prompt output is not guaranteed to run postprocess - # before the first direct StreamingInput update on a one-stage - # engine. Seed optional feedback channels exactly as native does; - # subsequent missing state remains an error. - if self.use_asr_head and not isinstance( - info_dict.get("input_asr_ids"), torch.Tensor - ): - info_dict["input_asr_ids"] = torch.full( - (n,), - self.pad_token_id, - device=device, - dtype=torch.long, - ) - if self.use_function_head and not isinstance( - info_dict.get("input_function_ids"), torch.Tensor - ): - info_dict["input_function_ids"] = torch.full( - (n,), - self.pad_token_id, - device=device, - dtype=torch.long, - ) - self._seeded_channel_requests.add(request_id) - - if self.use_asr_head: - asr_ids = self._channel_ids(info_dict, "input_asr_ids", n, device) - combined = combined + self.embed_asr_tokens(asr_ids) - - if self.use_function_head: - function_ids = self._channel_ids(info_dict, "input_function_ids", n, device) - combined = combined + self.model.embed_tokens(function_ids) * self.function_channel_weight - - # Per-step acoustic encoder embedding, sourced from - # ``additional_information["acoustic_embedding"]``. - acoustic = info_dict.get("acoustic_embedding") - assert isinstance(acoustic, torch.Tensor), ( - "acoustic_embedding is required in the per-step payload on every decode step; " - f"got {type(acoustic).__name__} with available keys {sorted(info_dict)}" - ) - acoustic = acoustic.to(device=device, dtype=combined.dtype) - assert acoustic.dim() == 2, f"acoustic_embedding must be 2D, got shape {tuple(acoustic.shape)}" - assert acoustic.shape[0] == n, ( - f"acoustic_embedding length {acoustic.shape[0]} does not match scheduled token count {n}" - ) - combined = combined + acoustic - - return input_ids, combined, {} - - @staticmethod - def _channel_ids(info_dict: dict[str, Any], key: str, n: int, device: torch.device) -> torch.Tensor: - """Read one auxiliary channel's per-step feedback ids from the payload.""" - ids = info_dict.get(key) - assert isinstance(ids, torch.Tensor), ( - f"{key} is required on every decode step but is " - f"{type(ids).__name__}; available keys {sorted(info_dict)}" - ) - ids = ids.to(device=device, dtype=torch.long).reshape(-1) - assert ids.numel() == n, f"{key} length {ids.numel()} does not match scheduled token count {n}" - return ids - - # ------------------------------------------------------------------ # - # postprocess - autoregressive feedback for the auxiliary channels # - # ------------------------------------------------------------------ # - - def postprocess( - self, - hidden_states: torch.Tensor, - multimodal_outputs: dict[str, Any] | None = None, - **info_dict: Any, - ) -> dict[str, Any]: - """Stash this request's last auxiliary tokens as the next step's input. - - ``hidden_states`` is a slice of the full-batch hidden_states tensor, - and each ``multimodal_outputs`` entry is the corresponding full-batch - token tensor produced by :meth:`make_omni_output`. We pick the token - aligned with the last position of this request's slice. - - On the prefill chunk the function channel is seeded with the pad id - instead of the prompt's own prediction, because native starts decoding - with ``gen_function`` still at its ``text_pad_id`` fill value and only - feeds back real function tokens from the second frame onwards. - """ - assert multimodal_outputs - start = hidden_states.storage_offset() // hidden_states.stride(0) - last_idx = start + hidden_states.shape[0] - 1 - # The initial prompt can be split into a 64-token slice plus a - # one-token continuation, so tensor length alone cannot identify it. - is_prefill = self._current_is_prompt_prefill - - def last_token(key: str) -> torch.Tensor: - tokens = multimodal_outputs.get(key) - assert isinstance(tokens, torch.Tensor), f"{key} missing from multimodal_outputs" - return tokens[last_idx : last_idx + 1].detach().to(torch.long) - - updates: dict[str, Any] = {} - if self.use_asr_head: - updates["input_asr_ids"] = last_token("asr_tokens") - if self.use_function_head: - if is_prefill: - updates["input_function_ids"] = torch.full( - (1,), self.pad_token_id, device=hidden_states.device, dtype=torch.long - ) - else: - updates["input_function_ids"] = last_token("function_tokens") - request_id = str( - info_dict.get("global_request_id") - or info_dict.get("request_id") - or "" - ) - if request_id: - channel_state = { - key: value.detach() - for key, value in updates.items() - if isinstance(value, torch.Tensor) - } - self._channel_state[request_id] = channel_state - self._last_channel_state = channel_state - return updates - - # ------------------------------------------------------------------ # - # forward # - # ------------------------------------------------------------------ # - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - **kwargs: Any, - ) -> torch.Tensor | IntermediateTensors: - """Run the backbone and return its hidden states. - - ASR tokens are produced by :meth:`make_omni_output`, which the - runner invokes *outside* the CUDA-graph wrapper. This keeps the - captured graph's output a plain ``Tensor`` (or - ``IntermediateTensors``) — both types that ``weak_ref_tensors`` - handles correctly. Returning a ``NamedTuple`` containing a - ``dict[str, Tensor]`` directly here would corrupt the dict's - tensors on FULL graph replay (the wrapper coerces - ``NamedTuple`` -> plain ``tuple`` and cannot weak-ref tensors - nested in dicts). - - IMPORTANT — cudagraph mode requirement - -------------------------------------- - This model must be run with ``cudagraph_mode="PIECEWISE"`` (or - ``enforce_eager=True``). The streaming-input pattern used here - keeps extending each request's prompt with every audio chunk, - so ``num_computed_tokens < num_prompt_tokens`` is permanently - true and Mamba's metadata builder always classifies the request - as a *prefill* (because - :func:`split_decodes_and_prefills` is called with - ``treat_short_extends_as_decodes=False`` in - ``Mamba2AttentionMetadataBuilder._compute_common_metadata``). - - With FULL cudagraph mode, the persistent - ``state_indices_tensor_d`` buffer is only updated when - ``num_prefills == 0``, so for streaming it stays at the - capture-time dummy value (0) while the FULL decode graph is - still dispatched (the dispatcher only checks ``query_len``). - The captured Mamba kernel then reads slot 0 of ``mamba_cache`` - instead of the real slot, producing garbage hidden states. - PIECEWISE side-steps this because the Mamba layer runs eagerly - and reads the freshly-computed metadata tensor, and the prefill - code path correctly *writes* the chunk into Mamba state on - every step (which is essential — there is no separate "prefill" - phase in this streaming setup). - """ - hidden_states = self.model(input_ids, positions, intermediate_tensors, inputs_embeds) - return hidden_states - - # ------------------------------------------------------------------ # - # make_omni_output - runs eagerly outside the CUDA graph wrapper # - # ------------------------------------------------------------------ # - - def make_omni_output( - self, - model_outputs: torch.Tensor | IntermediateTensors | OmniOutput, - **_: Any, - ) -> OmniOutput: - """Wrap backbone hidden states with the auxiliary channel tokens. - - Invoked by :class:`OmniGPUModelRunner._model_forward` after the - CUDA-graph wrapper has returned, so the auxiliary head matmuls + - ``argmax`` here run eagerly. They operate on the full-batch - ``hidden_states`` tensor in a single GEMM each, so the cost is - negligible relative to the backbone forward. - """ - if isinstance(model_outputs, OmniOutput): - return model_outputs - if isinstance(model_outputs, IntermediateTensors): - return OmniOutput( - text_hidden_states=model_outputs, - intermediate_tensors=model_outputs, - ) - - hidden = model_outputs - multimodal_outputs: dict[str, torch.Tensor] = {} - if self.use_asr_head: - asr_logits = self.logits_processor(self.asr_head, hidden) - # The ASR head's logits never reach vLLM's sampler, so the - # user-channel boosts are applied here rather than in the shared - # text logits processor. Same arithmetic as DuplexSTTModel. - apply_logit_boosts( - asr_logits, - self.user_logit_boosts, - pad_id=self.pad_token_id, - bos_id=self.bos_token_id, - eos_id=self.eos_token_id, - ) - multimodal_outputs["asr_tokens"] = torch.argmax(asr_logits, dim=-1).to(torch.long) - self._last_channel_state["input_asr_ids"] = ( - multimodal_outputs["asr_tokens"][-1:] - .detach() - ) - self._last_asr_token.copy_( - multimodal_outputs["asr_tokens"][-1:] - ) - if self.use_function_head: - function_logits = self.logits_processor(self.function_head, hidden) - multimodal_outputs["function_tokens"] = torch.argmax(function_logits, dim=-1).to(torch.long) - if self._current_is_prompt_prefill: - function_state = torch.full( - (1,), - self.pad_token_id, - device=hidden.device, - dtype=torch.long, - ) - else: - function_state = multimodal_outputs[ - "function_tokens" - ][-1:].detach() - self._last_channel_state[ - "input_function_ids" - ] = function_state - self._last_function_token.copy_(function_state) - - return OmniOutput( - text_hidden_states=hidden, - multimodal_outputs=multimodal_outputs, - ) - - # ------------------------------------------------------------------ # - # weight loading # - # ------------------------------------------------------------------ # - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self, skip_prefixes=["mtp"]) - loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) - - # Now that the embedding tables are populated, materialize the - # per-prefill pad embedding into the ``_pad_combined_emb`` buffer - # declared in ``__init__``. See the buffer's registration site - # for why this lives here rather than in ``__init__`` (embedding - # tables are empty there) and why it's non-persistent (derived - # from weights that are already saved). - self._materialize_pad_combined_emb() - - return loaded diff --git a/nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/sampling.py b/nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/sampling.py deleted file mode 100644 index 94d90b48c9ff..000000000000 --- a/nemo/collections/speechlm2/inference/vllm_omni/nemotron_duplex_h/sampling.py +++ /dev/null @@ -1,247 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""VoiceChat text sampling, shared with the PyTorch backend. - -Both backends decode the text head with -:func:`~nemo.collections.speechlm2.inference.model_wrappers.text_sampling.sample_text_token`. -The PyTorch backend calls it directly; vLLM reaches it through the -logits-processor hook implemented here, so the two cannot drift. -""" - -import math -from collections import OrderedDict -from dataclasses import dataclass, field -from typing import Any - -import torch -from vllm import SamplingParams -from vllm.v1.sample.logits_processor import AdapterLogitsProcessor - -from nemo.collections.speechlm2.inference.model_wrappers.text_sampling import sample_text_token -from nemo.collections.speechlm2.parts.logit_boosts import LogitBoosts, apply_logit_boosts - -SHARED_TEXT_SAMPLING_ARG = "nemo_shared_text_sampling" - - -def _sampling_config(params: SamplingParams) -> dict[str, Any] | None: - extra_args = params.extra_args - if not isinstance(extra_args, dict): - return None - value = extra_args.get(SHARED_TEXT_SAMPLING_ARG) - return value if isinstance(value, dict) else None - - -@dataclass -class SharedTextSamplingState: - """Sampling history that survives vLLM streaming segment re-admission.""" - - sample_count: int = 0 - tokens: list[int] = field(default_factory=list) - - -class SharedTextRequestSampler: - """Select with NeMo's sampler, then force vLLM greedy to that token.""" - - def __init__( - self, - *, - top_p: float, - repetition_penalty: float, - temperature: float, - special_token_ids: set[int], - history_skip: int, - state: SharedTextSamplingState | None = None, - boosts: LogitBoosts | None = None, - pad_id: int | None = None, - bos_id: int | None = None, - eos_id: int | None = None, - ) -> None: - self.top_p = top_p - self.repetition_penalty = repetition_penalty - self.temperature = temperature - self.special_token_ids = special_token_ids - self.history_skip = history_skip - self.state = state or SharedTextSamplingState() - self.boosts = boosts or LogitBoosts() - self.pad_id = pad_id - self.bos_id = bos_id - self.eos_id = eos_id - if self.boosts and None in (pad_id, bos_id, eos_id): - raise ValueError("Agent logit boosts require pad_id, bos_id and eos_id") - self._special_ids_tensor = ( - torch.tensor(sorted(special_token_ids), dtype=torch.long) if special_token_ids else None - ) - - def __call__( - self, - output_ids: list[int], - logits: torch.Tensor, - ) -> torch.Tensor: - del output_ids - history = self.state.tokens - generated_tokens = torch.tensor( - history, - device=logits.device, - dtype=torch.long, - ).unsqueeze(0) - if self._special_ids_tensor is not None and self._special_ids_tensor.device != logits.device: - self._special_ids_tensor = self._special_ids_tensor.to(logits.device) - - # Same order as DuplexSTTModel: boost the special tokens, then sample. - apply_logit_boosts( - logits, - self.boosts, - pad_id=self.pad_id, - bos_id=self.bos_id, - eos_id=self.eos_id, - ) - - sampled = sample_text_token( - logits.unsqueeze(0), - generated_tokens, - len(history), - top_p=self.top_p, - repetition_penalty=self.repetition_penalty, - temperature=self.temperature, - special_token_ids=self.special_token_ids, - special_ids_tensor=self._special_ids_tensor, - ) - selected_token = int(sampled[0].item()) - if self.state.sample_count >= self.history_skip: - self.state.tokens.append(selected_token) - self.state.sample_count += 1 - logits.fill_(float("-inf")) - logits[selected_token] = 0.0 - return logits - - -class SharedTextSamplingLogitsProcessor(AdapterLogitsProcessor): - """Batch adapter enabling shared VoiceChat text sampling per request.""" - - def __init__( - self, - vllm_config: Any, - device: torch.device, - is_pin_memory: bool, - ) -> None: - super().__init__(vllm_config, device, is_pin_memory) - max_num_seqs = int(getattr(vllm_config.scheduler_config, "max_num_seqs", 1) or 1) - self._max_history_states = max(1, max_num_seqs) - self._history_states: OrderedDict[str, SharedTextSamplingState] = OrderedDict() - - @classmethod - def validate_params(cls, sampling_params: SamplingParams) -> None: - extra_args = sampling_params.extra_args - if not isinstance(extra_args, dict): - return - raw_config = extra_args.get(SHARED_TEXT_SAMPLING_ARG) - if raw_config is None: - return - if not isinstance(raw_config, dict): - raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG} must be a mapping") - - top_p = raw_config.get("top_p") - temperature = raw_config.get("temperature") - repetition_penalty = raw_config.get("repetition_penalty") - for name, value in ( - ("top_p", top_p), - ("temperature", temperature), - ("repetition_penalty", repetition_penalty), - ): - if isinstance(value, bool) or not isinstance(value, int | float): - raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG}.{name} must be numeric") - if not math.isfinite(float(value)): - raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG}.{name} must be finite") - if not 0.0 < float(top_p) <= 1.0: - raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG}.top_p must be in (0, 1]") - if float(temperature) < 0.0: - raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG}.temperature must be >= 0") - if float(repetition_penalty) <= 0.0: - raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG}.repetition_penalty must be > 0") - - special_ids = raw_config.get("special_token_ids") - if not isinstance(special_ids, list) or any( - isinstance(token_id, bool) or not isinstance(token_id, int) for token_id in special_ids - ): - raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG}.special_token_ids " "must be a list of integers") - history_skip = raw_config.get("history_skip") - if isinstance(history_skip, bool) or not isinstance(history_skip, int) or history_skip < 0: - raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG}.history_skip " "must be a non-negative integer") - history_key = raw_config.get("history_key") - if not isinstance(history_key, str) or not history_key: - raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG}.history_key " "must be a non-empty string") - - boosts = raw_config.get("boosts") - if boosts is None: - return - if not isinstance(boosts, dict): - raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG}.boosts must be a mapping") - for name in ("pad", "bos", "eos"): - value = boosts.get(name) - if value is None: - continue - if isinstance(value, bool) or not isinstance(value, int | float): - raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG}.boosts.{name} must be numeric") - if not math.isfinite(float(value)): - raise ValueError(f"{SHARED_TEXT_SAMPLING_ARG}.boosts.{name} must be finite") - if not any(boosts.get(name) for name in ("pad", "bos", "eos")): - return - for name in ("pad_id", "bos_id", "eos_id"): - token_id = raw_config.get(name) - if isinstance(token_id, bool) or not isinstance(token_id, int) or token_id < 0: - raise ValueError( - f"{SHARED_TEXT_SAMPLING_ARG}.{name} must be a non-negative " "integer when boosts are set" - ) - - def is_argmax_invariant(self) -> bool: - return False - - def new_req_logits_processor( - self, - params: SamplingParams, - ) -> SharedTextRequestSampler | None: - config = _sampling_config(params) - if config is None: - return None - history_key = str(config["history_key"]) - state = self._history_states.get(history_key) - if state is None: - while len(self._history_states) >= self._max_history_states: - self._history_states.popitem(last=False) - state = SharedTextSamplingState() - self._history_states[history_key] = state - else: - self._history_states.move_to_end(history_key) - boosts = LogitBoosts.from_dict(config.get("boosts")) - return SharedTextRequestSampler( - top_p=float(config["top_p"]), - repetition_penalty=float(config["repetition_penalty"]), - temperature=float(config["temperature"]), - special_token_ids=set(config["special_token_ids"]), - history_skip=int(config["history_skip"]), - state=state, - boosts=boosts, - pad_id=config.get("pad_id"), - bos_id=config.get("bos_id"), - eos_id=config.get("eos_id"), - ) - - -__all__ = [ - "SHARED_TEXT_SAMPLING_ARG", - "SharedTextRequestSampler", - "SharedTextSamplingState", - "SharedTextSamplingLogitsProcessor", -] diff --git a/nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/__init__.py b/nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/__init__.py deleted file mode 100644 index 374ca75eab75..000000000000 --- a/nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/__init__.py +++ /dev/null @@ -1,27 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Single-stage ``nemotron_voicechat`` NemotronDuplexH Omni pipeline. - -EarTTS is registered as a separate one-stage pipeline; NeMo coordinates the -two engines. -""" - -from nemo.collections.speechlm2.inference.vllm_omni.nemotron_voicechat.pipeline import ( - NEMOTRON_VOICECHAT_PIPELINE, -) - -__all__ = [ - "NEMOTRON_VOICECHAT_PIPELINE", -] diff --git a/nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/pipeline.py b/nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/pipeline.py deleted file mode 100644 index 0be71e9110cb..000000000000 --- a/nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/pipeline.py +++ /dev/null @@ -1,80 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Single-stage ``nemotron_voicechat`` Omni pipeline. - -The VoiceChat wrapper intentionally runs NemotronDuplexH and EarTTS as -independent one-stage engines. This pipeline is the Nemotron half: it -consumes one acoustic encoder embedding per :class:`StreamingInput` update -and emits the text token plus whichever auxiliary channel (ASR or function) -the checkpoint contains. NeMo forwards the text token to the separate -``eartts`` pipeline. - -The pipeline is registered against ``model_type = "nemotron_voicechat"``, -which the component checkpoint does not report natively, so the converted -wrapper directory remains the model root:: - - / - config.json # {"model_type": "nemotron_voicechat"} - nemotron/ # directory or symlink → Nemotron ckpt - eartts/ # directory or symlink → EarTTS ckpt - -Only ``nemotron/`` is loaded by this pipeline. ``eartts/`` is passed -directly to a second :class:`AsyncOmni` instance. The bundled deploy YAML at -``nemo/collections/speechlm2/inference/vllm_omni/deploy/nemotron_voicechat.yaml`` -points this stage at its component via ``model_subdir`` / ``tokenizer_subdir``. -""" - -from __future__ import annotations - -from vllm_omni.config.stage_config import ( - PipelineConfig, - StageExecutionType, - StagePipelineConfig, -) - -_SCHED_ASYNC = ( - "nemo.collections.speechlm2.inference.vllm_omni." "nemotron_voicechat.scheduler.NemotronVoicechatARAsyncScheduler" -) - - -NEMOTRON_VOICECHAT_PIPELINE = PipelineConfig( - model_type="nemotron_voicechat", - model_arch="NemotronDuplexHForCausalLM", - stages=( - StagePipelineConfig( - stage_id=0, - model_stage="nemotron", - execution_type=StageExecutionType.LLM_AR, - input_sources=(), - final_output=True, - final_output_type="text", - owns_tokenizer=True, - model_arch="NemotronDuplexHForCausalLM", - # Stock vLLM-Omni 0.26 only includes single-stage AR requests in - # the client multimodal pooler payload for the "audio" engine - # output path. The final output remains text, so sampled text - # tokens stay on RequestOutput while OmniOutput.multimodal_outputs - # carries the optional ASR/function token beside it. - engine_output_type="audio", - scheduler_cls=_SCHED_ASYNC, - sampling_constraints={"detokenize": False}, - ), - ), -) - - -__all__ = [ - "NEMOTRON_VOICECHAT_PIPELINE", -] diff --git a/nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/scheduler.py b/nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/scheduler.py deleted file mode 100644 index 8d4dcbf21ced..000000000000 --- a/nemo/collections/speechlm2/inference/vllm_omni/nemotron_voicechat/scheduler.py +++ /dev/null @@ -1,68 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Streaming scheduler for the one-stage Nemotron VoiceChat engine. - -The only deviation from vLLM-Omni's async AR scheduler is forwarding each -direct ``StreamingInput`` chunk's ``additional_information`` payload onto the -session. This carries the current acoustic embedding into the model runner -without modifying the installed vLLM-Omni package. -""" - -from __future__ import annotations - -from vllm.v1.request import Request, StreamingUpdate -from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler - - -class NemotronVoicechatSchedulerMixin: - """Forward direct per-chunk payloads on the one-stage session.""" - - def _update_request_as_session(self, session: Request, update: StreamingUpdate) -> None: - super()._update_request_as_session(session, update) - - # Forward the chunk's payload onto the session, which is the courier - # that carries it to ``OmniNewRequestData`` and from there into the - # runner's ``model_intermediate_buffer``. Upstream propagates - # ``model_intermediate_buffer`` itself but not ``additional_information``, - # and this pipeline has to use the latter: the per-chunk payload is a - # tensor (``acoustic_embedding``), and ``model_intermediate_buffer`` is - # typed ``dict[str, Any]`` on the request, so vLLM's msgpack decoder has - # no declared type to rebuild a tensor from and would hand the model a - # ``[dtype, shape, bytes]`` list instead. ``additional_information`` is - # the transport with an explicit tensor encoding, which is why upstream's - # own duplex example keeps ``model_intermediate_buffer`` to plain lists. - # - # Replace rather than merge: this field is a per-chunk message, and - # accumulating whole payloads across chunks would keep stale - # prefill-only keys alive. ``None`` means "this chunk omitted the - # field" rather than "clear the session", so placeholder chunks do not - # drop the initial request's state. The runner does the actual merge - # into the cached buffer, one sub-key at a time. Only stage 0 does this: - # in a downstream stage the chunk transfer adapter is the sole writer of - # the payload, so upstream returns early there. - if self.vllm_config.model_config.stage_id == 0: - new_info = getattr(update, "additional_information", None) - if new_info is not None: - session.additional_information = new_info - - -class NemotronVoicechatARAsyncScheduler(NemotronVoicechatSchedulerMixin, OmniARAsyncScheduler): - """Default: matches upstream's ``async_scheduling=True`` for LLM_AR stages.""" - - -__all__ = [ - "NemotronVoicechatARAsyncScheduler", - "NemotronVoicechatSchedulerMixin", -] diff --git a/nemo/collections/speechlm2/inference/vllm_omni/outputs.py b/nemo/collections/speechlm2/inference/vllm_omni/outputs.py deleted file mode 100644 index df63105b108f..000000000000 --- a/nemo/collections/speechlm2/inference/vllm_omni/outputs.py +++ /dev/null @@ -1,128 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Reading vLLM-Omni stage outputs. - -The shape of an ``OmniStageOutput`` is not part of any contract we control: -the multimodal payload may sit on the stage output or nested on its -completion, and the key it arrives under depends on whether the stage is -client-facing. These readers are the one place that tolerates that, so the -version-sensitivity is contained rather than spread through the session. -""" - -from typing import Any, NamedTuple - -import torch - - -class StepTokens(NamedTuple): - """Per-frame tokens sampled by Nemotron for one acoustic frame. - - ``asr`` and ``function`` are *None* when the checkpoint has no such - channel. ASR and function channels are independently optional. - """ - - text: int - asr: int | None = None - function: int | None = None - - -def _multimodal_output(stage_output: Any, req_out: Any) -> Any: - """Return a stage's multimodal payload, whichever level carries it. - - vLLM-Omni attaches the payload to the ``MultimodalCompletionOutput`` in - ``request_output.outputs[0]`` and also lifts it onto the stage output - itself, so check both (mirroring - ``vllm_omni.metrics.utils.first_multimodal_output``). - - Stock vLLM-Omni 0.26 surfaces EarTTS audio codes here. The registered - Nemotron pipeline also uses the final multimodal engine-output route so - optional ASR/function tensors accompany its text ``RequestOutput``. - """ - mm = getattr(stage_output, "multimodal_output", None) - if mm: - return mm - outputs = getattr(req_out, "outputs", None) or () - for completion in outputs: - nested = getattr(completion, "multimodal_output", None) - if nested: - return nested - return {} - - -def _audio_codes(mm: Any) -> Any: - """Return this step's EarTTS acoustic codes from a multimodal payload. - - ``EarTTSForCausalLM.make_omni_output`` publishes them under - ``model_outputs``, which vLLM-Omni's output processor remaps to the - drainable ``audio`` modality key: in DELTA mode that key is emptied after - every step, so each payload carries only the frames computed this step. - Keys other than the modality's own are retained across steps and merged - with :class:`TensorAccumulationStrategy` ``CONCAT_LAST`` for audio, which - widens a ``T x num_quantizers`` frame instead of appending to it — so - ``audio_codes`` is read last, only for a stage that is not client-facing. - """ - if mm is None: - return None - for key in ("audio", "model_outputs", "audio_codes"): - value = mm.get(key) - if value is not None: - return value - return None - - -def _step_delta(value: Any, finished: bool, *, skip_finished: bool = True): - """Mirror of ``_step_delta`` in the vllm-omni example: pull the - new-this-step multimodal chunk from an :class:`OmniStageOutput`'s - ``multimodal_output`` value (which may be a tensor, a list of tensors, - ``None``, or absent). - - ``skip_finished`` drops a terminal duplicate. The split streaming - requests use one-token segments, so callers pass ``False`` and separately - skip each request's prefill output. - """ - if finished and skip_finished: - return None - if isinstance(value, torch.Tensor): - return value if value.numel() > 0 else None - if isinstance(value, list) and value: - last = value[-1] - return last if isinstance(last, torch.Tensor) and last.numel() > 0 else None - return None - - -def _step_tokens(stage_output: Any) -> StepTokens: - """Extract text and optional auxiliary tokens from one Nemotron output. - - The text token remains on the stock vLLM ``RequestOutput`` even when the - stage uses the multimodal engine-output path. Auxiliary tensors may be - lifted onto ``OmniStageOutput`` or remain nested on its completion. - """ - req_out = stage_output.request_output - mm = _multimodal_output(stage_output, req_out) - finished = bool(getattr(req_out, "finished", False)) - if req_out and req_out.outputs and req_out.outputs[0].token_ids: - text_tok = int(req_out.outputs[0].token_ids[-1]) - else: - text_tok = 0 - - def last_token(key: str) -> int | None: - delta = _step_delta(mm.get(key), finished, skip_finished=False) - return int(delta[-1].item()) if delta is not None else None - - return StepTokens( - text=text_tok, - asr=last_token("asr_tokens"), - function=last_token("function_tokens"), - ) diff --git a/nemo/collections/speechlm2/inference/vllm_omni/register.py b/nemo/collections/speechlm2/inference/vllm_omni/register.py deleted file mode 100644 index e48eaa644329..000000000000 --- a/nemo/collections/speechlm2/inference/vllm_omni/register.py +++ /dev/null @@ -1,154 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Runtime registration of the NemotronDuplexH + EarTTS models and the -``nemotron_voicechat`` pipeline with vLLM and vLLM-Omni. - -Call :func:`register_nemo_voicechat` once, before constructing -``vllm_omni.AsyncOmni`` / ``vllm_omni.Omni``. It is idempotent. - -NeMo's ``pyproject.toml`` exposes this function as an entry point in **both** -the ``vllm.general_plugins`` and ``vllm_omni.general_plugins`` groups, so a -plugin loader invokes it automatically in every process (orchestrator + each -spawned ``StageEngineCoreProc`` child + workers). vllm-omni uses -``multiprocessing`` with start method ``spawn`` for stage children, so spawned -processes do NOT inherit Python state from the parent — registering in the -parent alone is not enough, which is why a plugin entry point is the correct -hook. vLLM's group is the one that matters in the parent: vllm-omni resolves -the pipeline for ``model_type`` while constructing ``AsyncOmniEngine``, before -it loads its own plugin group. - -Being in vLLM's group means this also runs in processes that have no vllm-omni -installed at all (an ordinary ``vllm serve``), so the vllm-omni half is -optional and skipped when the import fails. - -Three things get registered: - -1. ``EarTTSConfig`` with ``transformers.AutoConfig`` (so ``AutoConfig.from_pretrained`` - resolves ``model_type = "eartts"``) and with vLLM's ``_CONFIG_REGISTRY``. -2. Model architectures (``NemotronDuplexHForCausalLM``, ``EarTTSForCausalLM``) - with both ``vllm.model_executor.models.ModelRegistry`` and - ``vllm_omni.model_executor.models.OmniModelRegistry``. The two registries - serve different lookups inside vLLM-Omni so both have to know about the - new arches. -3. The one-stage :data:`NEMOTRON_VOICECHAT_PIPELINE` and - :data:`EARTTS_PIPELINE` with ``vllm_omni.config.register_pipeline``. -""" - -from __future__ import annotations - -import logging - -logger = logging.getLogger(__name__) - - -_PKG = "nemo.collections.speechlm2.inference.vllm_omni" - - -_ARCH_MAP: dict[str, tuple[str, str]] = { - # arch_name -> (module_path, class_name) - "NemotronDuplexHForCausalLM": ( - f"{_PKG}.nemotron_duplex_h.nemotron_duplex_h", - "NemotronDuplexHForCausalLM", - ), - "EarTTSForCausalLM": ( - f"{_PKG}.eartts.eartts", - "EarTTSForCausalLM", - ), -} - - -_registered = False - - -def register_nemo_voicechat() -> None: - """Register the NemotronDuplexH + EarTTS models, ``EarTTSConfig``, - and the ``nemotron_voicechat`` pipeline with vLLM / vLLM-Omni. - - Safe to call multiple times. - """ - global _registered - if _registered: - return - - _register_hf_configs() - _register_model_archs() - omni = _register_omni() - - _registered = True - logger.info( - "nemo_voicechat: registered NemotronDuplexH + EarTTS%s.", - " + split nemotron_voicechat/eartts pipelines" - if omni - else " (vllm-omni absent, pipelines not registered)", - ) - - -def _register_hf_configs() -> None: - from nemo.collections.speechlm2.inference.vllm_omni.eartts.configuration_eartts import ( - EarTTSConfig, - register_eartts_config, - ) - - register_eartts_config() - - try: - from vllm.transformers_utils.config import _CONFIG_REGISTRY - except ImportError: - _CONFIG_REGISTRY = None - - if _CONFIG_REGISTRY is not None and EarTTSConfig.model_type not in _CONFIG_REGISTRY: - _CONFIG_REGISTRY[EarTTSConfig.model_type] = EarTTSConfig - - -def _register_model_archs() -> None: - # vLLM's public model registry — needed for ``ModelRegistry.is_*`` - # checks and for the arch → module resolution used outside of - # OmniModelConfig. - from vllm.model_executor.models import ModelRegistry - - supported_archs = ModelRegistry.get_supported_archs() - for arch, (module_path, class_name) in _ARCH_MAP.items(): - if arch not in supported_archs: - ModelRegistry.register_model(arch, f"{module_path}:{class_name}") - - -def _register_omni() -> bool: - """Register with vLLM-Omni, if it is installed. Returns whether it was.""" - try: - # OmniModelRegistry is the mirror registry ``OmniModelConfig.registry`` - # returns, used to load the model class for each pipeline stage. - from vllm_omni.config import register_pipeline - from vllm_omni.model_executor.models import OmniModelRegistry - except ImportError: - return False - - omni_supported = OmniModelRegistry.get_supported_archs() - for arch, (module_path, class_name) in _ARCH_MAP.items(): - if arch not in omni_supported: - OmniModelRegistry.register_model(arch, f"{module_path}:{class_name}") - - from nemo.collections.speechlm2.inference.vllm_omni.nemotron_voicechat.pipeline import ( - NEMOTRON_VOICECHAT_PIPELINE, - ) - from nemo.collections.speechlm2.inference.vllm_omni.eartts.pipeline import ( - EARTTS_PIPELINE, - ) - - register_pipeline(NEMOTRON_VOICECHAT_PIPELINE) - register_pipeline(EARTTS_PIPELINE) - return True - - -__all__ = ["register_nemo_voicechat"] diff --git a/nemo/collections/speechlm2/inference/vllm_omni/runtime.py b/nemo/collections/speechlm2/inference/vllm_omni/runtime.py deleted file mode 100644 index 6db1da775e97..000000000000 --- a/nemo/collections/speechlm2/inference/vllm_omni/runtime.py +++ /dev/null @@ -1,322 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Process-scoped vLLM-Omni engines on one background asyncio loop. - -``AsyncOmni`` is asynchronous and its engines are expensive, so exactly one -:class:`OmniRuntime` is built per process and shared by every stream. It owns -a daemon thread running a dedicated event loop, the independent Nemotron and -EarTTS engines, and the deploy-YAML overrides they were started with. - -Nemotron and EarTTS get separate one-stage engines so NeMo can hand tokens -between them and give EarTTS a classifier-free-guidance companion request -without duplicating the much larger Nemotron request. - -Which components exist is decided here, at construction, and read off the -runtime afterwards (``llm_engine``/``tts_engine`` being None) -- callers do not -pass backend flags around. -""" - -import asyncio -import logging as stdlib_logging -import os -import tempfile -import threading -from pathlib import Path -from typing import Any - -import yaml - -from nemo.collections.speechlm2.inference.vllm_omni import default_deploy_yaml, default_eartts_deploy_yaml -from nemo.collections.speechlm2.inference.vllm_omni.checkpoint import EARTTS_SUBDIR -from nemo.utils import logging - - -class _ExpectedJanusShutdownFilter(stdlib_logging.Filter): - """Hide only vLLM-Omni's expected output-queue shutdown traceback.""" - - def filter(self, record: stdlib_logging.LogRecord) -> bool: - if "[AsyncOmni] final_output_loop failed." not in record.getMessage(): - return True - exc = record.exc_info[1] if record.exc_info else None - return not ( - exc is not None - and exc.__class__.__module__.startswith("janus") - and exc.__class__.__name__ - in { - "ShutDown", - "SyncQueueShutDown", - "AsyncQueueShutDown", - } - ) - - -class OmniRuntime: - """Long-lived split AsyncOmni engines + one background asyncio loop. - - Constructed once by the inference wrapper and shared across streams. - Nemotron and EarTTS have independent one-stage engines so NeMo can hand - tokens between them and create an EarTTS CFG companion without duplicating - the much larger Nemotron request. - """ - - def __init__( - self, - wrapper_dir: str, - *, - stage_configs_path: str | None = None, - eartts_stage_configs_path: str | None = None, - stage_overrides: dict | None = None, - eartts_stage_overrides: dict | None = None, - log_stats: bool = False, - stage_init_timeout: int = 600, - enable_llm: bool = True, - enable_tts: bool = True, - ) -> None: - if not enable_llm and not enable_tts: - raise ValueError("OmniRuntime requires at least one enabled component") - self.enable_llm = bool(enable_llm) - self.enable_tts = bool(enable_tts) - - llm_yaml = Path(stage_configs_path) if stage_configs_path else default_deploy_yaml() - tts_yaml = Path(eartts_stage_configs_path) if eartts_stage_configs_path else default_eartts_deploy_yaml() - required_yamls = [] - if self.enable_llm: - required_yamls.append(llm_yaml) - if self.enable_tts: - required_yamls.append(tts_yaml) - for deploy_yaml in required_yamls: - if not deploy_yaml.is_file(): - raise FileNotFoundError(f"Deploy YAML not found: {deploy_yaml}") - - # Accept single-pipeline override keys as well: ``stage_0`` addresses - # the Nemotron engine and ``stage_1`` the EarTTS engine's stage 0. - llm_overrides, legacy_tts_overrides = self._split_stage_overrides(stage_overrides) - if eartts_stage_overrides is None: - eartts_stage_overrides = legacy_tts_overrides - self._llm_stage_yaml_path = ( - self._maybe_write_overridden_yaml(llm_yaml, llm_overrides, prefix="nemotron_") if self.enable_llm else None - ) - self._tts_stage_yaml_path = ( - self._maybe_write_overridden_yaml(tts_yaml, eartts_stage_overrides, prefix="eartts_") - if self.enable_tts - else None - ) - self._wrapper_dir = wrapper_dir - self._eartts_dir = os.path.join(wrapper_dir, EARTTS_SUBDIR) - self._shutdown = False - - # Start the background loop in a daemon thread first; ``AsyncOmni`` - # is constructed *on* that loop (its ``__init__`` allocates - # ``asyncio.Condition`` / ``asyncio.Queue`` and the orchestrator - # binds them to the current event loop, so the engine must be - # built from inside that loop's thread). - self._loop = asyncio.new_event_loop() - self._ready_evt = threading.Event() - self._thread = threading.Thread( - target=self._loop_runner, - name="OmniRuntimeLoop", - daemon=True, - ) - self._thread.start() - self._ready_evt.wait() - - # Register in this process before constructing the engine: - # ``AsyncOmniEngine.__init__`` resolves ``model_type`` before loading - # plugin groups. An unregistered model type selects the default diffusion - # pipeline, which expects ``model_index.json``. Entry points register the - # same pipeline in spawned stage processes. - from vllm_omni import AsyncOmni - - from nemo.collections.speechlm2.inference.vllm_omni.register import register_nemo_voicechat - - register_nemo_voicechat() - - logging.info( - "Creating split AsyncOmni engines from wrapper=%s (Nemotron) and %s (EarTTS) ...", - wrapper_dir, - self._eartts_dir, - ) - - async def _build_engines() -> tuple[Any | None, Any | None]: - llm_engine = None - tts_engine = None - try: - if self.enable_llm: - llm_engine = AsyncOmni( - model=wrapper_dir, - stage_configs_path=str(self._llm_stage_yaml_path), - log_stats=log_stats, - stage_init_timeout=stage_init_timeout, - ) - if self.enable_tts: - tts_engine = AsyncOmni( - model=self._eartts_dir, - stage_configs_path=str(self._tts_stage_yaml_path), - log_stats=log_stats, - stage_init_timeout=stage_init_timeout, - ) - return llm_engine, tts_engine - except BaseException: - if llm_engine is not None: - llm_engine.shutdown() - raise - - fut = asyncio.run_coroutine_threadsafe(_build_engines(), self._loop) - self.llm_engine, self.tts_engine = fut.result() - logging.info( - "Split AsyncOmni ready (Nemotron=%s, EarTTS=%s)", - f"{self.llm_engine.num_stages} stage" if self.llm_engine is not None else "native", - f"{self.tts_engine.num_stages} stage" if self.tts_engine is not None else "native", - ) - - # ------------------------------------------------------------------ # - # YAML override # - # ------------------------------------------------------------------ # - - @staticmethod - def _split_stage_overrides( - stage_overrides: dict | None, - ) -> tuple[dict | None, dict | None]: - if not stage_overrides: - return None, None - common = dict(stage_overrides.get("common", {}) or {}) - llm: dict[str, Any] = {} - tts: dict[str, Any] = {} - if common: - llm["common"] = common - tts["common"] = common - if stage_overrides.get("stage_0"): - llm["stage_0"] = dict(stage_overrides["stage_0"]) - if stage_overrides.get("stage_1"): - tts["stage_0"] = dict(stage_overrides["stage_1"]) - return llm or None, tts or None - - @staticmethod - def _maybe_write_overridden_yaml( - deploy_yaml: Path, - stage_overrides: dict | None, - *, - prefix: str, - ) -> Path: - """Apply per-stage overrides to the deploy YAML, write to a tmp file. - - ``stage_overrides`` shape:: - - { - "common": {}, - "stage_0": {}, - "stage_1": {}, - } - - Returns the path that ``AsyncOmni`` should load; the original YAML - is returned untouched when no overrides are supplied. - """ - if not stage_overrides: - return deploy_yaml - - with open(deploy_yaml, encoding="utf-8") as fh: - cfg = yaml.safe_load(fh) - - common = stage_overrides.get("common", {}) or {} - per_stage = {int(k.split("_", 1)[1]): v for k, v in stage_overrides.items() if k.startswith("stage_") and v} - - for stage in cfg.get("stages", []): - for key, value in common.items(): - stage[key] = value - sid = int(stage.get("stage_id", -1)) - for key, value in per_stage.get(sid, {}).items(): - stage[key] = value - - tmp = tempfile.NamedTemporaryFile( - mode="w", - suffix=".yaml", - prefix=prefix, - delete=False, - ) - yaml.dump(cfg, tmp, default_flow_style=False, sort_keys=False) - tmp.close() - logging.info(f"Wrote overridden stage config to {tmp.name}") - return Path(tmp.name) - - # ------------------------------------------------------------------ # - # Background loop # - # ------------------------------------------------------------------ # - - def _loop_runner(self) -> None: - asyncio.set_event_loop(self._loop) - self._ready_evt.set() - try: - self._loop.run_forever() - finally: - # ``run_forever`` returns when ``loop.stop()`` is called from - # ``shutdown``. Tear down any pending tasks before closing. - try: - pending = asyncio.all_tasks(self._loop) - for task in pending: - task.cancel() - except RuntimeError: - pass - try: - self._loop.close() - except Exception: - pass - - def submit(self, coro): - """Schedule a coroutine on the background loop, return the concurrent ``Future``.""" - return asyncio.run_coroutine_threadsafe(coro, self._loop) - - def shutdown(self) -> None: - """Stop both engines and the background loop.""" - if self._shutdown: - return - self._shutdown = True - - async def _shutdown_engines() -> None: - for name in ("tts_engine", "llm_engine"): - engine = getattr(self, name, None) - if engine is None: - continue - try: - final_output_task = getattr(engine, "final_output_task", None) - if final_output_task is not None and not final_output_task.done(): - final_output_task.cancel() - await asyncio.gather( - final_output_task, - return_exceptions=True, - ) - engine.final_output_task = None - engine.shutdown() - except Exception as exc: - logging.warning(f"{name}.shutdown() raised: {exc!r}") - - shutdown_filter = _ExpectedJanusShutdownFilter() - async_omni_logger = stdlib_logging.getLogger("vllm_omni.entrypoints.async_omni") - async_omni_logger.addFilter(shutdown_filter) - root_handlers = list(stdlib_logging.getLogger().handlers) - for handler in root_handlers: - handler.addFilter(shutdown_filter) - try: - self.submit(_shutdown_engines()).result(timeout=120) - except Exception as exc: - logging.warning(f"Split AsyncOmni shutdown raised: {exc!r}") - finally: - async_omni_logger.removeFilter(shutdown_filter) - for handler in root_handlers: - handler.removeFilter(shutdown_filter) - try: - self._loop.call_soon_threadsafe(self._loop.stop) - except Exception: - pass - self._thread.join(timeout=10) diff --git a/nemo/collections/speechlm2/inference/vllm_omni/scripts/__init__.py b/nemo/collections/speechlm2/inference/vllm_omni/scripts/__init__.py deleted file mode 100644 index 9e3fb699d9f6..000000000000 --- a/nemo/collections/speechlm2/inference/vllm_omni/scripts/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. diff --git a/nemo/collections/speechlm2/inference/vllm_omni/scripts/convert_duplex_eartts_checkpoint.py b/nemo/collections/speechlm2/inference/vllm_omni/scripts/convert_duplex_eartts_checkpoint.py deleted file mode 100644 index 6e7cce134690..000000000000 --- a/nemo/collections/speechlm2/inference/vllm_omni/scripts/convert_duplex_eartts_checkpoint.py +++ /dev/null @@ -1,357 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Convert the DuplexEARTTS component of a NemotronVoiceChat checkpoint to vLLM format. - -The converter expects the HuggingFace-format NemotronVoiceChat checkpoint layout: -``config.json`` contains ``model.speech_generation`` and ``model.stt`` entries, -and ``model.safetensors`` contains nested ``tts_model.tts_model.*`` weights. - -The character-aware subword encoder (``embed_subword``) is collapsed into a -single pre-computed lookup table mapping ``token_id -> hidden_size`` embedding. -The character/transformer weights of that encoder are dropped, since the lookup -fully captures their deterministic per-token output (including the additive -subword-flag and BOS/EOS contributions). -""" - -import argparse -import json -import os -import tqdm - -import torch -from omegaconf import DictConfig, OmegaConf -from safetensors.torch import load_file, save_file -from transformers import AutoConfig - -from nemo.collections.speechlm2.models.duplex_ear_tts import DuplexEARTTS -from nemo.utils import logging - - -def parse_args(): - parser = argparse.ArgumentParser() - parser.add_argument("--config", type=str, required=True) - parser.add_argument("--model", type=str, required=True) - parser.add_argument("--outdir", type=str, required=True) - parser.add_argument( - "--precompute-batch-size", - type=int, - default=256, - help="Batch size for pre-computing per-token embeddings.", - ) - return parser.parse_args() - - -def _precompute_subword_embeddings(model: DuplexEARTTS, batch_size: int) -> torch.Tensor: - """Run ``embed_subword`` over the entire vocabulary to bake out a lookup table. - - The character-aware subword encoder is fully deterministic per token id - (it takes ``subword_ids`` only and adds id-conditioned flag/BOS-EOS - embeddings). Running it once per id and storing the result lets vLLM - replace the whole encoder with a single ``nn.Embedding`` lookup. - - Returns: - Tensor of shape ``[vocab_size, hidden_size]`` matching the dtype of the - encoder's parameters. - """ - embed_subword = model.tts_model.embed_subword - embed_subword.eval() - - # Run the precomputation on GPU when available; the encoder is small but - # the vocabulary loop is long, so this is a meaningful speedup. Only the - # subword encoder is moved (not the full model) to keep peak memory low. - device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") - embed_subword.to(device) - logging.info(f"Precomputing subword embeddings on {device}") - - dtype = next(embed_subword.parameters()).dtype - - subword_ids_map = embed_subword.subword_id_to_char_ids - vocab_size = max(int(k) for k in subword_ids_map.keys()) + 1 - hidden_size = embed_subword.proj_embedding.out_features - - table = torch.zeros((vocab_size, hidden_size), dtype=dtype, device=device) - - with torch.no_grad(): - for start in tqdm.tqdm(range(0, vocab_size, batch_size), desc="Precomputing subword embeddings"): - end = min(start + batch_size, vocab_size) - ids = torch.arange(start, end, dtype=torch.long, device=device).unsqueeze(0) - mask = torch.ones_like(ids, dtype=torch.bool) - embeds = embed_subword(ids, mask) - table[start:end] = embeds.squeeze(0).to(dtype) - - return table.cpu() - - -def convert_to_vllm_format(outdir: str, config: str, model_path: str, precompute_batch_size: int = 256) -> None: - """Convert DuplexEARTTS weights from a NemotronVoiceChat HF checkpoint for vLLM. - - Args: - outdir: Directory where the vLLM-compatible checkpoint will be written. - config: Path to the NemotronVoiceChat ``config.json`` file. - model_path: Path to the NemotronVoiceChat ``model.safetensors`` file. - precompute_batch_size: Batch size used while running the subword encoder - once per token id to construct the lookup table. - """ - os.makedirs(outdir, exist_ok=True) - - with open(config, "r") as f: - full_config = json.load(f) - - # This converter builds a real DuplexEARTTS from the checkpoint's own - # config, so it needs the same config normalization that - # NemotronVoiceChat.from_pretrained applies to the published VoiceChat - # release. Both helpers are no-ops for any other checkpoint. - from nemo.collections.speechlm2.models.nemotron_voicechat import ( - _apply_nemotron_labs_voicechat_release_config_shim, - _is_nemotron_labs_voicechat_release, - ) - - if _is_nemotron_labs_voicechat_release(os.path.dirname(os.path.abspath(config)), full_config): - logging.info("Applying NemotronLabs VoiceChat release shim before building DuplexEARTTS") - _apply_nemotron_labs_voicechat_release_config_shim(full_config) - - config_dict = full_config["model"]["speech_generation"] - cfg = DictConfig(config_dict) - # Inference-only overrides, so the model built here matches the settings - # DuplexEARTTS.generate() uses when the embeddings are precomputed below. - cfg.model.tts_config.use_unshifthed_prompt = True - cfg.data.add_audio_prompt_after_description = True - cfg.model.tts_config.use_unshifthed_prompt = True - cfg.model.subword_mask_exactly_as_eartts = False - cfg.model.context_hidden_mask_exactly_as_eartts = False - cfg.model.tts_config.disable_eos_prediction = True - cfg.model.inference_force_speech_silence_on_eos = True - cfg.model.use_word_sep_tokenizer = False - cfg.model.num_delay_speech_tokens = 0 - cfg.data.source_sample_rate = 22050 - cfg.data.target_sample_rate = 22050 - cfg.model.pretrained_model = None - - model = DuplexEARTTS(OmegaConf.to_container(cfg, resolve=True)).eval() - hidden_size = cfg.model.tts_config.backbone_config.hidden_size - - # Load the HuggingFace-format NemotronVoiceChat safetensors checkpoint. - raw_weights = load_file(model_path) - # The checkpoint is wrapped by an outer module (NemotronVoiceChat) whose TTS - # attribute is also called ``tts_model``. Strip a single ``tts_model.`` prefix - # to land in the DuplexEARTTS state-dict namespace. - weights = {k[len("tts_model.") :]: v for k, v in raw_weights.items() if k.startswith("tts_model.")} - - # Load the real weights into the DuplexEARTTS model so that running - # ``embed_subword`` produces the trained per-token outputs (otherwise we - # would just bake out random init values). - missing, unexpected = model.load_state_dict(weights, strict=False) - # Some keys (e.g. the unused language model / audio codec heads) may be - # missing or unexpected; that is fine for the embedding sub-tree we care - # about. Surface the diagnostics anyway. - if missing: - logging.info(f"load_state_dict missing keys (expected for unused submodules): {len(missing)}") - if unexpected: - logging.info(f"load_state_dict unexpected keys: {len(unexpected)}") - - # Pre-compute the subword lookup table once per token id. This collapses - # the entire char-aware encoder (char embedding + transformer + projection - # + subword/BOS-EOS flag adds) into a single ``nn.Embedding`` lookup that - # vLLM can use directly. - precomputed_subword_emb = _precompute_subword_embeddings(model, precompute_batch_size) - vocab_size, _ = precomputed_subword_emb.shape - - # Codec silence tokens are produced once at training time by encoding a - # zero waveform with the audio codec and picking the most common frame. - # Bake the resulting per-codebook ids into the vLLM checkpoint so the - # runtime does not need to load / run the codec to know what "silence" - # looks like (used e.g. when forcing silence on EOS). - codec_silence_tokens = model.codec_silence_tokens.detach().clone().cpu().to(torch.int32) - - # Strip the ``tts_model.`` prefix so the renaming below operates on - # RVQEARTTSModel state-dict keys. - weights = {k[len("tts_model.") :]: v for k, v in weights.items() if k.startswith("tts_model.")} - - # duplicate weights for rvq embeddings and embed code - rvq_embs_weight = weights["rvq_embs"].clone() # 31 x codebook_size x latent_size - rvq_embs_weight_pad = torch.nn.functional.pad( - rvq_embs_weight, [0, 0, 0, 1] - ) # 31 x (codebook_size + 1) x latent_size - embed_code_weight = weights["embed_code.weight"].clone() # latent_size x hidden_size - - # ====================== - # embedding module weights - bos_emb = weights["bos_emb"] - null_emb = weights["null_emb"] - - embedding_module_weights = {} - embedding_module_weights["bos_emb"] = bos_emb - # CFG-only (unconditional text branch). Always exported so a wrapper - # works with guidance on or off without reconverting. - embedding_module_weights["null_emb"] = null_emb - - # Single pre-computed lookup replacing the entire char-aware encoder. - embedding_module_weights["embed_subword.embed_subwords.weight"] = precomputed_subword_emb - - # Keep gated fusion + audio prompt projection: these depend on runtime - # tensors, not on token id, so they cannot be pre-computed. - for key, weight in weights.items(): - if key.startswith("gated_fusion_audio_text."): - embedding_module_weights[key] = weight - if "audio_prompt_projection_W" in weights: - embedding_module_weights["audio_prompt_projection_W"] = weights["audio_prompt_projection_W"] - - for i in range(rvq_embs_weight_pad.shape[0]): - embedding_module_weights[f"rvq_embs.{i}.weight"] = rvq_embs_weight_pad[i] - embedding_module_weights["embed_code.weight"] = embed_code_weight - embedding_module_weights = {f"total_emb.{k}": v for k, v in embedding_module_weights.items()} - - # ====================== - # gemma backbone weights - backbone_module_weights = {k: v for k, v in weights.items() if k.startswith("backbone.")} - backbone_module_weights["backbone.embed_tokens.weight"] = torch.randn( - 1, hidden_size, dtype=bos_emb.dtype, device=bos_emb.device - ) - - # ====================== - # sampler weights - used_keys = ("rvq_embs", "embed_code", "mog_head") - sampler_weights = {"sampler." + k: v for k, v in weights.items() if k.startswith(used_keys)} - - # combine embedding module and backbone module weights - weights = {**embedding_module_weights, **backbone_module_weights, **sampler_weights} - weights = {"model." + k: v for k, v in weights.items()} - - # Top-level silence token buffer (int32 tensor of shape [num_quantizers]). - # Stored under ``model.sil_tokens`` so the vLLM model can register it as a - # plain buffer at the top of its module tree. - weights["model.sil_tokens"] = codec_silence_tokens - - # save weights - safetensors_path = os.path.join(outdir, "model.safetensors") - save_file(weights, safetensors_path) - logging.info("Saved weights for vllm model") - weight_map = {name: "model.safetensors" for name in weights.keys()} - index = { - "metadata": {"total_size": sum(w.numel() * w.element_size() for w in weights.values())}, - "weight_map": weight_map, - } - index_path = os.path.join(outdir, "model.safetensors.index.json") - with open(index_path, "w") as f: - json.dump(index, f, indent=2) - logging.info("Saved model index") - - # save config.json - flat_config = {"architectures": ["EarTTSForCausalLM"], "model_type": "eartts"} - # not using vocab size of the backbone model, but need 2 for dummy sampling to work - flat_config["vocab_size"] = 2 - - # Parse backbone config exactly as NeMo does to get all defaults from transformers - backbone_type = cfg.model.tts_config.get("backbone_type", None) - backbone_config_dict = ( - OmegaConf.to_container(cfg.model.tts_config.backbone_config, resolve=True) - if cfg.model.tts_config.get("backbone_config") - else {} - ) - - # Create AutoConfig the same way NeMo does - this fills in all defaults - parsed_backbone_config = AutoConfig.for_model(backbone_type, **backbone_config_dict) - - # Store the backbone type for vllm to use - flat_config["backbone_type"] = backbone_type - - # Forward all backbone configs from the parsed AutoConfig (includes defaults) - for key in [ - "hidden_size", - "intermediate_size", - "num_hidden_layers", - "num_attention_heads", - "num_key_value_heads", - "head_dim", - "max_position_embeddings", - "rope_theta", - "rope_local_base_freq", - "sliding_window", - "layer_types", - ]: - if hasattr(parsed_backbone_config, key): - value = getattr(parsed_backbone_config, key) - # convert to list if it's a tuple or other iterable (except str) - if hasattr(value, '__iter__') and not isinstance(value, (str, dict)): - value = list(value) - flat_config[key] = value - # forward overall configs - for key in ["latent_size", "codebook_size", "num_quantizers", "exponent"]: - flat_config[key] = cfg.model.tts_config[key] - # forward mog head configs - for key in ["num_layers", "low_rank", "num_predictions", "min_log_std", "eps"]: - flat_config[f"mog_{key}"] = cfg.model.tts_config.mog_head_config[key] - - # forward inference configs (with name mapping for vLLM model) - # MaskGIT unmasking iterations, fixed at 8 to match DuplexEARTTS inference. - flat_config["num_iter"] = 8 - flat_config["noise_scale"] = cfg.model.get("inference_noise_scale", 0.8) - flat_config["top_p_or_k"] = cfg.model.get("inference_top_p_or_k", 0.8) - - # Classifier-free guidance. ``guidance_scale`` is the checkpoint default; - # individual requests may override it through ``additional_information["cfg_scale"]``. - flat_config["guidance_scale"] = cfg.model.get("inference_guidance_scale", 0.5) - flat_config["enable_guidance"] = True - - # Text-channel specials from the source tokenizer, not the Gemma backbone. - flat_config["pad_token_id"] = int(model.text_pad_id) - flat_config["eos_token_id"] = int(model.text_eos_id) - - # Embedding module configuration. The char-aware encoder is gone; vLLM only - # needs to know the size of the pre-computed lookup table. - flat_config["emb_vocab_size"] = vocab_size - - flat_config["use_gated_fusion_for_text_audio"] = cfg.model.tts_config.use_gated_fusion_for_text_audio - flat_config["use_audio_prompt_frozen_projection"] = cfg.model.tts_config.use_audio_prompt_frozen_projection - - # configuring custom inputs/outputs - flat_config["custom_input_specs"] = [ - { - "name": "acoustic_tokens", - "dim": flat_config["num_quantizers"], - "dtype": "int32", - }, - {"name": "text_tokens", "dtype": "int32"}, - {"name": "text_mask"}, - {"name": "bos_mask"}, - {"name": "speaker_latent", "dim": flat_config["hidden_size"]}, - ] - flat_config["custom_outputs"] = ["acoustic_tokens"] - - with open(os.path.join(outdir, "config.json"), "w") as f: - json.dump(flat_config, f, indent=2) - logging.info("Saved vllm config") - - # Extract and save pre-computed speaker latents (audio_prompt_latents.*) - # from the NeMo checkpoint so they can be used at inference time. - speaker_latents_dir = os.path.join(outdir, "speaker_latents") - found_latents = False - for key, tensor in raw_weights.items(): - if "audio_prompt_latents." in key: - speaker_name = key.split("audio_prompt_latents.")[-1] - os.makedirs(speaker_latents_dir, exist_ok=True) - latent_path = os.path.join(speaker_latents_dir, f"{speaker_name}.pt") - torch.save(tensor, latent_path) - logging.info(f"Saved speaker latent '{speaker_name}' to {latent_path} (shape={tensor.shape})") - found_latents = True - if not found_latents: - logging.warning( - "No audio_prompt_latents found in checkpoint. " "speaker_name will not work unless latents are added." - ) - - -if __name__ == "__main__": - args = parse_args() - convert_to_vllm_format(args.outdir, args.config, args.model, args.precompute_batch_size) diff --git a/nemo/collections/speechlm2/inference/vllm_omni/scripts/convert_duplex_stt_checkpoint.py b/nemo/collections/speechlm2/inference/vllm_omni/scripts/convert_duplex_stt_checkpoint.py deleted file mode 100644 index 07e2bb2769ac..000000000000 --- a/nemo/collections/speechlm2/inference/vllm_omni/scripts/convert_duplex_stt_checkpoint.py +++ /dev/null @@ -1,343 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -Convert the DuplexSTT component of a NemotronVoiceChat checkpoint to vLLM format. - -This script extracts weights from a HuggingFace-format NemotronVoiceChat -checkpoint with tensors such as: -- stt_model.llm.layers.* -- stt_model.lm_head.* -- stt_model.asr_head.* -- stt_model.embed_asr_tokens.* -- stt_model.function_head.* -- stt_model.embed_tokens.* - -And converts them to a HuggingFace layout that can be loaded by vLLM with the -custom WeightsMapper defined in nemotron_duplex_h.py. - -Which auxiliary channels a checkpoint carries varies: - -- ``predict_user_text=True`` gives ``asr_head`` + ``embed_asr_tokens`` -- ``use_function_head=True`` gives ``function_head`` (reusing ``embed_tokens`` - for its feedback, so it has no embedding table of its own) - -The converter records whichever heads are present as -``use_asr_head`` / ``use_function_head``, because ``NemotronDuplexHForCausalLM`` -has to decide which modules to build *before* it sees any weights. -""" - -import argparse -import json -import os -from pathlib import Path -import torch -from safetensors.torch import load_file, save_file -from transformers import AutoConfig, AutoTokenizer -from nemo.utils import logging - - -def load_checkpoint(checkpoint_path: str) -> dict[str, torch.Tensor]: - """ - Load a NemotronVoiceChat checkpoint state dict. - - Args: - checkpoint_path: Path to a checkpoint directory, safetensors file, or PyTorch checkpoint file. - - Returns: - Dictionary of tensor names to tensors - """ - if os.path.isdir(checkpoint_path): - checkpoint_path = os.path.join(checkpoint_path, "model.safetensors") - - if checkpoint_path.endswith('.safetensors'): - logging.info(f"Loading safetensors from {checkpoint_path}") - return load_file(checkpoint_path) - else: - logging.info(f"Loading PyTorch checkpoint from {checkpoint_path}") - ckpt = torch.load(checkpoint_path, map_location='cpu') - # Handle different checkpoint formats - if 'state_dict' in ckpt: - return ckpt['state_dict'] - elif 'model' in ckpt: - return ckpt['model'] - else: - return ckpt - - -def filter_tensors(state_dict: dict[str, torch.Tensor], prefixes_to_keep: list[str]) -> dict[str, torch.Tensor]: - """ - Filter tensors to keep only those with specified prefixes. - - Args: - state_dict: Full state dictionary - prefixes_to_keep: List of prefixes to keep (e.g., ["stt_model.llm", "stt_model.asr_head"]) - - Returns: - Filtered state dictionary - """ - filtered_dict = {} - for name, tensor in state_dict.items(): - if any(name.startswith(prefix) for prefix in prefixes_to_keep): - filtered_dict[name] = tensor - logging.debug(f"Keeping: {name} with shape {tensor.shape}") - else: - logging.debug(f"Skipping: {name}") - - logging.info(f"Total tensors kept: {len(filtered_dict)}") - return filtered_dict - - -def _apply_source_special_tokens(base_config, tokenizer, source_config: dict | None) -> None: - """Match the converted tokenizer/config to the VoiceChat channel tokens. - - VoiceChat training overrides the LLM-backbone tokenizer specials - (typically ```` for EOS and ```` for PAD). Keeping the - backbone's original EOS/PAD ids corrupts system-prompt prefill even - though text tokenization itself appears valid. Copy whatever the source - VoiceChat config actually used. - """ - try: - model_config = source_config["model"]["stt"]["model"] - except (KeyError, TypeError): - return - - overrides = model_config.get("override_tokens", {}) or {} - special_tokens = { - name: overrides.get(name) or model_config.get(name) - for name in ("bos_token", "eos_token", "pad_token") - } - special_tokens = {name: token for name, token in special_tokens.items() if token} - if not special_tokens: - return - - vocabulary = tokenizer.get_vocab() - missing = [token for token in special_tokens.values() if token not in vocabulary] - if missing: - raise ValueError( - "VoiceChat special tokens must already exist in the backbone vocabulary; " - f"missing={missing}" - ) - added = tokenizer.add_special_tokens(special_tokens) - if added: - raise ValueError( - "VoiceChat special-token overrides unexpectedly expanded the vocabulary; " - f"added={added}" - ) - - for name, token in special_tokens.items(): - token_id = int(tokenizer.convert_tokens_to_ids(token)) - setattr(base_config, f"{name}_id", token_id) - logging.info( - "VoiceChat special-token IDs: bos=%s eos=%s pad=%s", - getattr(base_config, "bos_token_id", None), - getattr(base_config, "eos_token_id", None), - getattr(base_config, "pad_token_id", None), - ) - - -def convert_to_vllm_format( - checkpoint_path: str, - output_dir: str, - config_path: str | None = None, - pretrained_llm: str | None = None, - tensors_to_keep: list[str] | None = None, - dtype: str = "float32", -) -> None: - """ - Convert the DuplexSTT component to vLLM-compatible HuggingFace format. - - Args: - checkpoint_path: Path to the NeMo checkpoint (.safetensors or .pt) - output_dir: Directory to save the converted checkpoint - config_path: Path to config.json (if None, will look in same dir as checkpoint) - pretrained_llm: HuggingFace model name to get base config from - tensors_to_keep: List of tensor prefixes to keep (default: all stt_model.* tensors) - dtype: Data type for tensors ("float32", "float16", "bfloat16") - """ - # Default prefixes to keep. The auxiliary-channel entries are only present - # in some checkpoints; absent ones simply match nothing. - if tensors_to_keep is None: - tensors_to_keep = [ - "stt_model.llm", - "stt_model.lm_head", - "stt_model.asr_head", - "stt_model.embed_asr_tokens", - "stt_model.function_head", - "stt_model.embed_tokens", - ] - - # Load config to get pretrained_llm if not provided - if config_path is None: - ckpt_dir = checkpoint_path if os.path.isdir(checkpoint_path) else os.path.dirname(checkpoint_path) - config_path = os.path.join(ckpt_dir, "config.json") - - config = None - if os.path.exists(config_path): - logging.info(f"Loading config from {config_path}") - with open(config_path, "r") as f: - config = json.load(f) - - try: - pretrained_llm = config["model"]["stt"]["model"]["pretrained_llm"] - logging.info(f"Found pretrained_llm in config: {pretrained_llm}") - except KeyError: - if pretrained_llm is None: - raise ValueError("Could not find pretrained_llm in config and none provided via argument") - else: - if pretrained_llm is None: - raise ValueError(f"Config file not found at {config_path} and pretrained_llm not provided") - - # Create output directory - output_path = Path(output_dir) - output_path.mkdir(parents=True, exist_ok=True) - - # Load base config from pretrained model - logging.info(f"Loading base config from {pretrained_llm}") - base_config = AutoConfig.from_pretrained(pretrained_llm, trust_remote_code=True) - - # Load tokenizer from pretrained model - logging.info(f"Loading tokenizer from {pretrained_llm}") - tokenizer = AutoTokenizer.from_pretrained(pretrained_llm, trust_remote_code=True) - _apply_source_special_tokens(base_config, tokenizer, config) - - # Load checkpoint - logging.info(f"Loading checkpoint from {checkpoint_path}") - state_dict = load_checkpoint(checkpoint_path) - - # Filter tensors - logging.info(f"Filtering tensors to keep prefixes: {tensors_to_keep}") - filtered_state_dict = filter_tensors(state_dict, tensors_to_keep) - - if len(filtered_state_dict) == 0: - raise ValueError( - f"No tensors found with prefixes {tensors_to_keep}. " - f"Available prefixes: {set(k.split('.')[0] for k in state_dict.keys())}" - ) - - # Record which auxiliary channels this checkpoint actually carries, so the - # vLLM model builds exactly those modules. Detected from the weights that - # made it through the filter rather than from the source config, so a - # narrowed --tensors-to-keep stays consistent with what gets saved. - has_asr_head = any(name.startswith("stt_model.asr_head") for name in filtered_state_dict) - has_function_head = any(name.startswith("stt_model.function_head") for name in filtered_state_dict) - - if has_asr_head and not any(name.startswith("stt_model.embed_asr_tokens") for name in filtered_state_dict): - raise ValueError( - "Checkpoint has stt_model.asr_head but no stt_model.embed_asr_tokens; " - "the ASR channel needs both (the head to predict the token and the " - "embedding table to feed it back on the next step)." - ) - - # The function channel scales its feedback embedding by this weight, matching - # DuplexSTTModel.build_input_embedding. - function_channel_weight = 1.0 - if has_function_head and config is not None: - try: - function_channel_weight = float(config["model"]["stt"]["model"].get("duplex_function_channel_weight", 1.0)) - except (KeyError, TypeError): - logging.warning("Could not read duplex_function_channel_weight from source config; defaulting to 1.0") - - custom_outputs = ["text_logits"] - if has_asr_head: - custom_outputs += ["asr_tokens", "asr_logits"] - if has_function_head: - custom_outputs += ["function_tokens", "function_logits"] - - base_config.update( - { - "custom_input_specs": [{"name": "combined_embeds", "dtype": dtype, "dim": base_config.hidden_size}], - "custom_outputs": custom_outputs, - "use_asr_head": has_asr_head, - "use_function_head": has_function_head, - "duplex_function_channel_weight": function_channel_weight, - } - ) - logging.info( - f"Auxiliary channels: asr_head={has_asr_head}, function_head={has_function_head} " - f"(function_channel_weight={function_channel_weight})" - ) - - # Save tensors - output_model_path = output_path / "model.safetensors" - logging.info(f"Saving tensors to {output_model_path}") - save_file(filtered_state_dict, str(output_model_path)) - - # Save config - output_config_path = output_path / "config.json" - logging.info(f"Saving config to {output_config_path}") - base_config.save_pretrained(str(output_path)) - - # Save tokenizer - logging.info(f"Saving tokenizer to {output_path}") - tokenizer.save_pretrained(str(output_path)) - - logging.info(f"Conversion completed successfully! Output saved to: {output_path}") - - -def main(): - parser = argparse.ArgumentParser(description="Convert NeMo STT checkpoint to HuggingFace format for vLLM") - parser.add_argument( - "--checkpoint", - type=str, - required=True, - help="Path to NeMo checkpoint file (.safetensors or .pt/.pth)", - ) - parser.add_argument( - "--output-dir", - type=str, - required=True, - help="Directory to save converted checkpoint", - ) - parser.add_argument( - "--config", - type=str, - default=None, - help="Path to config.json (default: same directory as checkpoint)", - ) - parser.add_argument( - "--pretrained-llm", - type=str, - default=None, - help="HuggingFace model name to use as base (default: read from config)", - ) - parser.add_argument( - "--tensors-to-keep", - type=str, - nargs="+", - default=None, - help="Tensor prefixes to keep (default: all stt_model.* backbone llm related tensors)", - ) - parser.add_argument( - "--dtype", - type=str, - default="float32", - choices=["float32", "float16", "bfloat16", "fp32", "fp16", "bf16"], - help="Target dtype for tensors (default: float32)", - ) - - args = parser.parse_args() - - convert_to_vllm_format( - checkpoint_path=args.checkpoint, - output_dir=args.output_dir, - config_path=args.config, - pretrained_llm=args.pretrained_llm, - tensors_to_keep=args.tensors_to_keep, - dtype=args.dtype, - ) - - -if __name__ == "__main__": - main() diff --git a/nemo/collections/speechlm2/inference/vllm_omni/session.py b/nemo/collections/speechlm2/inference/vllm_omni/session.py deleted file mode 100644 index e30c3689a755..000000000000 --- a/nemo/collections/speechlm2/inference/vllm_omni/session.py +++ /dev/null @@ -1,592 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Synchronous per-stream bridge onto the vLLM-Omni engines. - -The S2S wrapper drives inference from a synchronous PyTorch loop (perception --> per-frame text + ASR -> audio codec decode), while ``AsyncOmni.generate`` -is an ``async for``. Spawning an event loop per chunk would re-pay the -request-init cost every 80 ms and break vllm-omni's session semantics, so a -:class:`OmniStreamingSession` runs consumer tasks on the runtime's shared loop -and hands frames back across two synchronous queues. - -Per-step protocol: - -1. **Prefill** -- Nemotron receives the system prompt, EarTTS the speaker - latent. With CFG enabled EarTTS receives a conditional and an unconditional - prefill, each with its own KV cache. Nemotron's first token ``t_0`` is fed - back internally rather than exposed. -2. **Decode step k** (``k >= 1``) -- ``prompt_token_ids = [t_{k-1}]``, - ``additional_information.acoustic_embedding = ac_emb[k-1]``, producing - ``t_k`` plus whichever auxiliary channel the checkpoint carries. - -The components are stepped independently: :meth:`OmniStreamingSession.step_llm` -returns Nemotron's tokens and :meth:`OmniStreamingSession.step_tts` submits a -text token, so the caller can rewrite that token (forced turn-taking) in -between and both TTS backends see the same value. - -The synchronous side is single-threaded: only one step may be in flight at a -time. :meth:`OmniStreamingSession.finish` closes the request cleanly and -:meth:`OmniStreamingSession.abort` drops it. -""" - -import asyncio -import threading -import time -import uuid -from queue import Queue -from typing import Any - -import torch - -from nemo.collections.speechlm2.inference.vllm_omni.outputs import ( - StepTokens, - _audio_codes, - _multimodal_output, - _step_delta, - _step_tokens, -) -from nemo.collections.speechlm2.inference.vllm_omni.runtime import OmniRuntime -from nemo.collections.speechlm2.parts.logit_boosts import LogitBoosts -from nemo.utils import logging - - -class _Sentinel: - """Marker placed on the sync output queues to signal end-of-stream or an error.""" - - __slots__ = ("exc",) - - def __init__(self, exc: BaseException | None = None): - self.exc = exc - - -_END_OF_STREAM = _Sentinel() - - -class OmniStreamingSession: - """One split Nemotron/EarTTS streaming request. - - The two components are driven independently: :meth:`step_llm` submits an - acoustic frame and returns Nemotron's tokens, :meth:`step_tts` submits a - text token and produces one acoustic frame. A session owning both exposes - both, which lets the caller act on the text token in between rather than - having EarTTS consume it inside the session. With CFG enabled, conditional - and unconditional EarTTS requests run in the same engine and the custom - scheduler keeps their independent KV-cache streams in lockstep. - """ - - def __init__( - self, - runtime: OmniRuntime, - request_id: str, - system_prompt: str = "", - speaker_latent: torch.Tensor | None = None, - t_prefill: int = 0, - *, - sampling_params: dict | None = None, - special_token_ids: set[int] | None = None, - guidance_enabled: bool = True, - guidance_scale: float = 0.5, - step_timeout: float = 60.0, - profile: bool = False, - agent_logit_boosts: "LogitBoosts | None" = None, - text_token_ids: dict[str, int] | None = None, - ) -> None: - self._has_llm = runtime.llm_engine is not None - self._has_tts = runtime.tts_engine is not None - if not self._has_llm and not self._has_tts: - raise ValueError("OmniStreamingSession requires at least one vLLM component") - if self._has_tts and (speaker_latent is None or speaker_latent.numel() == 0): - raise ValueError("speaker_latent is required for OmniStreamingSession") - if self._has_llm and t_prefill <= 0: - raise ValueError(f"t_prefill must be > 0 (got {t_prefill})") - - self._runtime = runtime - self.request_id = request_id - self._llm_request_id = f"{request_id}:nemotron" - self._sampling_history_key = f"{self._llm_request_id}:{uuid.uuid4().hex}" - self._cfg_pair_id = f"{request_id}:eartts" - self._tts_cond_request_id = f"{self._cfg_pair_id}:cond" - self._tts_uncond_request_id = f"{self._cfg_pair_id}:uncond" - self._step_timeout = step_timeout - self._system_prompt = system_prompt - self._speaker_latent = speaker_latent.detach().cpu().contiguous() if speaker_latent is not None else None - self._t_prefill = int(t_prefill) - self._sampling_overrides = dict(sampling_params or {}) - self._special_token_ids = tuple(sorted(special_token_ids or ())) - self._guidance_enabled = bool(guidance_enabled) - self._guidance_scale = float(guidance_scale) - self._agent_logit_boosts = agent_logit_boosts or LogitBoosts() - self._text_token_ids = dict(text_token_ids or {}) - if self._has_llm and self._agent_logit_boosts: - missing = {"pad_id", "bos_id", "eos_id"} - set(self._text_token_ids) - if missing: - raise ValueError(f"agent_logit_boosts requires text_token_ids {sorted(missing)}") - - # Separate completion queues per component: a session may have both, - # and step_llm()/step_tts() must not consume each other's items. - self._text_out_q: "Queue[StepTokens | _Sentinel]" = Queue() - self._tts_done_q: "Queue[StepTokens | _Sentinel]" = Queue() - self._audio_buf_lock = threading.Lock() - self._audio_buf: list[torch.Tensor] = [] - self._closed = False - self._error: BaseException | None = None - self._loop = runtime._loop - self._queues_ready = threading.Event() - - self._input_q: asyncio.Queue | None = None - self._llm_internal_q: asyncio.Queue | None = None - self._tts_cond_input_q: asyncio.Queue | None = None - self._tts_uncond_input_q: asyncio.Queue | None = None - self._pending_tokens_q: asyncio.Queue | None = None - self._uncond_audio_q: asyncio.Queue | None = None - - self._prof: dict[str, list[float]] | None = {} if profile else None - self._t_put = 0.0 - self._t_yield = 0.0 - self._t_llm = 0.0 - self._t_done = 0.0 - - self._consumer_future = runtime.submit(self._run_consumer()) - if not self._queues_ready.wait(timeout=30): - self._consumer_future.cancel() - raise TimeoutError("Timed out creating split vLLM-Omni session queues") - - def _rec(self, name: str, dt_s: float) -> None: - if self._prof is not None: - self._prof.setdefault(name, []).append(dt_s * 1000.0) - - def _rec_ms(self, name: str, dt_ms: float) -> None: - if self._prof is not None: - self._prof.setdefault(name, []).append(float(dt_ms)) - - def log_timing_summary(self) -> None: - if not self._prof: - return - parts = [] - for name, values in self._prof.items(): - mean = sum(values) / len(values) - parts.append( - f"{name}: mean={mean:.1f}ms min={min(values):.1f}ms " f"max={max(values):.1f}ms n={len(values)}" - ) - logging.info(f"OmniStreamingSession {self.request_id} per-frame timing:\n " + "\n ".join(parts)) - - def _cfg_payload(self, role: str) -> dict[str, Any]: - return { - "cfg_enabled": self._guidance_enabled, - "cfg_role": role, - "cfg_pair_id": self._cfg_pair_id, - "cfg_scale": self._guidance_scale, - } - - def _record_stage_metrics(self, prefix: str, stage_output: Any) -> None: - if self._prof is None: - return - for key, value in (getattr(stage_output, "stage_durations", None) or {}).items(): - self._rec_ms(f"{prefix}.{key}", value) - - async def _run_consumer(self) -> None: - tasks: list[asyncio.Task] = [] - try: - self._input_q = asyncio.Queue() - self._llm_internal_q = asyncio.Queue() - self._tts_cond_input_q = asyncio.Queue() - self._tts_uncond_input_q = asyncio.Queue() - self._pending_tokens_q = asyncio.Queue() - self._uncond_audio_q = asyncio.Queue() - self._queues_ready.set() - - if self._has_llm: - tasks.append(asyncio.create_task(self._consume_llm())) - if self._has_tts: - tasks.append(asyncio.create_task(self._consume_tts("cond", self._tts_cond_input_q))) - if self._has_tts and self._guidance_enabled: - tasks.append(asyncio.create_task(self._consume_tts("uncond", self._tts_uncond_input_q))) - await asyncio.gather(*tasks) - except asyncio.CancelledError: - raise - except BaseException as exc: - self._error = exc - # Both queues, so a caller blocked in either step never hangs. - self._text_out_q.put(_Sentinel(exc)) - self._tts_done_q.put(_Sentinel(exc)) - raise - finally: - for task in tasks: - if not task.done(): - task.cancel() - if tasks: - await asyncio.gather(*tasks, return_exceptions=True) - self._queues_ready.set() - self._text_out_q.put(_END_OF_STREAM) - self._tts_done_q.put(_END_OF_STREAM) - self.log_timing_summary() - - async def _consume_llm(self) -> None: - from vllm import SamplingParams - from vllm.engine.protocol import StreamingInput - from vllm.sampling_params import RequestOutputKind - - from nemo.collections.speechlm2.inference.vllm_omni.nemotron_duplex_h.sampling import SHARED_TEXT_SAMPLING_ARG - - shared_sampling = { - "temperature": float(self._sampling_overrides.get("temperature", 1.0)), - "top_p": float(self._sampling_overrides.get("top_p", 1.0)), - "repetition_penalty": float(self._sampling_overrides.get("repetition_penalty", 1.0)), - "special_token_ids": list(self._special_token_ids), - # The first vLLM output is the internal prefill token t_0. The - # repetition history starts with the first client-visible frame. - "history_skip": 1, - "history_key": self._sampling_history_key, - # Agent-channel boosts, applied before sampling exactly as - # DuplexSTTModel does. The user-channel ones cannot travel this way - # because the ASR head's logits never reach vLLM's sampler; the - # model applies those itself. - "boosts": self._agent_logit_boosts.as_dict(), - **self._text_token_ids, - } - params = SamplingParams( - temperature=0.0, - top_p=1.0, - repetition_penalty=1.0, - max_tokens=1, - detokenize=False, - ignore_eos=True, - output_kind=RequestOutputKind.DELTA, - extra_args={SHARED_TEXT_SAMPLING_ARG: shared_sampling}, - ) - - async def inputs(): - yield StreamingInput( - prompt={ - "prompt_token_ids": [0] * self._t_prefill, - "additional_information": { - "system_prompt": self._system_prompt, - }, - }, - sampling_params=params, - ) - while True: - submission = await self._input_q.get() - if submission is None: - return - acoustic, committed_text = submission - # Always drain the internal queue to stay one output per input, - # then let the caller's committed token win. That is how the - # caller's forced-turn-taking rewrite reaches Nemotron's own - # history, matching what native feedback does through gen_text. - prev_tokens = await self._llm_internal_q.get() - if committed_text is not None: - prev_tokens = prev_tokens._replace(text=int(committed_text)) - additional_information: dict[str, Any] = { - "system_prompt": None, - "acoustic_embedding": acoustic, - } - if prev_tokens.asr is not None: - additional_information["input_asr_ids"] = torch.tensor([prev_tokens.asr], dtype=torch.long) - if prev_tokens.function is not None: - additional_information["input_function_ids"] = torch.tensor( - [prev_tokens.function], dtype=torch.long - ) - self._t_yield = time.perf_counter() - yield StreamingInput( - prompt={ - "prompt_token_ids": [int(prev_tokens.text)], - "additional_information": additional_information, - }, - sampling_params=params, - ) - - output_count = 0 - try: - async for stage_output in self._runtime.llm_engine.generate( - inputs(), - sampling_params_list=[params], - request_id=self._llm_request_id, - ): - now = time.perf_counter() - self._record_stage_metrics("llm", stage_output) - current_tokens = _step_tokens(stage_output) - await self._llm_internal_q.put(current_tokens) - - output_count += 1 - if output_count <= 1: - continue - - self._rec("pull", self._t_yield - self._t_put) - self._rec("llm_engine", now - self._t_yield) - self._t_llm = now - # The token is returned to the caller, never forwarded to - # EarTTS from here. The caller owns what happens in between - # (forced turn-taking rewrites the text token) and submits it - # with step_tts, so both TTS backends see the same token. - self._t_done = now - self._text_out_q.put(current_tokens) - finally: - # Safety net rather than the normal path: finish() closes the TTS - # inputs itself. This covers Nemotron ending first, which would - # otherwise leave the EarTTS consumer waiting and stall the - # gather() in _run_consumer. - if self._has_tts: - await self._tts_cond_input_q.put(None) - if self._guidance_enabled: - await self._tts_uncond_input_q.put(None) - - async def _consume_tts(self, role: str, input_q: asyncio.Queue) -> None: - from vllm import SamplingParams - from vllm.engine.protocol import StreamingInput - from vllm.sampling_params import RequestOutputKind - - extra_args = self._cfg_payload(role) if self._guidance_enabled else {} - params = SamplingParams( - temperature=0.0, - top_p=1.0, - max_tokens=1, - detokenize=False, - ignore_eos=True, - output_kind=RequestOutputKind.DELTA, - extra_args=extra_args, - ) - request_id = self._tts_cond_request_id if role == "cond" else self._tts_uncond_request_id - - async def inputs(): - prefill_info = { - "embed": {"voice": self._speaker_latent.clone()}, - **self._cfg_payload(role), - } - if not self._guidance_enabled: - prefill_info["cfg_enabled"] = False - yield StreamingInput( - prompt={ - "prompt_token_ids": [0] * int(self._speaker_latent.shape[0]), - "additional_information": prefill_info, - }, - sampling_params=params, - ) - while True: - text_tok = await input_q.get() - if text_tok is None: - return - yield StreamingInput( - prompt={ - "prompt_token_ids": [0], - "additional_information": { - "ids": {"output": [int(text_tok)]}, - **self._cfg_payload(role), - }, - }, - sampling_params=params, - ) - - output_count = 0 - async for stage_output in self._runtime.tts_engine.generate( - inputs(), - sampling_params_list=[params], - request_id=request_id, - ): - now = time.perf_counter() - req_out = stage_output.request_output - mm = _multimodal_output(stage_output, req_out) - finished = bool(getattr(req_out, "finished", False)) - self._record_stage_metrics(f"tts_{role}", stage_output) - output_count += 1 - if output_count <= 1: - continue - audio = _step_delta(_audio_codes(mm), finished, skip_finished=False) - if audio is None or audio.ndim != 2 or audio.shape[0] < 1: - continue - # EarTTS emits exactly one acoustic frame per streaming update, so - # keep only the newest row: a prefill update covers the whole - # speaker latent, and a non-drained key would arrive cumulative. - audio = audio[-1:] - audio = audio.detach().cpu().to(torch.long) - if role == "uncond": - await self._uncond_audio_q.put(audio) - continue - - # A final empty update can race after all text tokens have been - # consumed. It has no corresponding wrapper step and must not - # synthesize another frame. - if self._pending_tokens_q.empty(): - continue - tokens = await self._pending_tokens_q.get() - if self._guidance_enabled: - uncond_audio = await self._uncond_audio_q.get() - if not torch.equal(audio, uncond_audio): - raise RuntimeError( - "EarTTS CFG pair produced divergent client-visible " - "codes: " - f"cond_shape={tuple(audio.shape)} " - f"uncond_shape={tuple(uncond_audio.shape)} " - f"cond={audio.tolist()} " - f"uncond={uncond_audio.tolist()}" - ) - with self._audio_buf_lock: - self._audio_buf.append(audio) - if self._has_llm: - self._rec("tts_after_llm", now - self._t_llm) - else: - self._rec("tts_engine", now - self._t_put) - self._t_done = now - self._tts_done_q.put(tokens) - - def step_llm( - self, - acoustic_embedding: torch.Tensor, - *, - prev_text_token: int | None = None, - ) -> StepTokens: - """Submit one acoustic frame to Nemotron and return its tokens. - - Returns as soon as Nemotron has produced the frame's tokens. EarTTS is - not driven from here even when this session owns both components: the - caller submits the (possibly rewritten) text token with - :meth:`step_tts`. - - Args: - acoustic_embedding: This frame's encoded audio. - prev_text_token: Text token to feed back as the previous step's - output, letting a caller that rewrote it (forced turn-taking) - keep Nemotron's history consistent with its own. ``None`` - keeps whatever Nemotron last produced, which is what the first - frame after prefill needs since its predecessor is the - engine-internal prefill token. - """ - if not self._has_llm: - raise RuntimeError("This vLLM-Omni session has no Nemotron component") - if self._closed: - raise RuntimeError(f"OmniStreamingSession {self.request_id} is closed") - ac_emb = acoustic_embedding.detach().cpu().contiguous() - if ac_emb.dim() == 1: - ac_emb = ac_emb.unsqueeze(0) - elif ac_emb.dim() == 3: - ac_emb = ac_emb.reshape(-1, ac_emb.shape[-1]) - if ac_emb.dim() != 2: - raise ValueError( - "acoustic_embedding must be shapeable to 2D [n, hidden], " f"got {tuple(acoustic_embedding.shape)}" - ) - ac_emb = ac_emb.to(torch.float32) - - self._t_put = time.perf_counter() - asyncio.run_coroutine_threadsafe( - self._input_q.put((ac_emb, prev_text_token)), - self._loop, - ).result() - self._rec("put", time.perf_counter() - self._t_put) - - item = self._text_out_q.get(timeout=self._step_timeout) - returned = time.perf_counter() - if not isinstance(item, _Sentinel): - self._rec("deliver", returned - self._t_done) - self._rec("step_total", returned - self._t_put) - if isinstance(item, _Sentinel): - if item.exc is not None: - raise RuntimeError(f"OmniStreamingSession {self.request_id} consumer raised") from item.exc - raise RuntimeError(f"OmniStreamingSession {self.request_id} ended before producing a token") - return item - - def step_tts(self, text_token: int) -> None: - """Submit one text token to EarTTS and wait for its acoustic frame. - - Valid whether or not this session also owns Nemotron; the resulting - codes are collected by :meth:`drain_audio_codes`. - """ - if not self._has_tts: - raise RuntimeError("This vLLM-Omni session has no EarTTS component") - if self._closed: - raise RuntimeError(f"OmniStreamingSession {self.request_id} is closed") - - tokens = StepTokens(int(text_token)) - self._t_put = time.perf_counter() - - async def _put_tts_input() -> None: - await self._pending_tokens_q.put(tokens) - await self._tts_cond_input_q.put(tokens.text) - if self._guidance_enabled: - await self._tts_uncond_input_q.put(tokens.text) - - asyncio.run_coroutine_threadsafe(_put_tts_input(), self._loop).result() - item = self._tts_done_q.get(timeout=self._step_timeout) - if isinstance(item, _Sentinel): - if item.exc is not None: - raise RuntimeError(f"OmniStreamingSession {self.request_id} consumer raised") from item.exc - raise RuntimeError(f"OmniStreamingSession {self.request_id} ended before producing audio") - - def drain_audio_codes(self) -> list[torch.Tensor]: - with self._audio_buf_lock: - out = self._audio_buf - self._audio_buf = [] - return out - - def _abort_engine_requests(self) -> None: - requests = ( - (self._runtime.llm_engine, self._llm_request_id), - (self._runtime.tts_engine, self._tts_cond_request_id), - (self._runtime.tts_engine, self._tts_uncond_request_id), - ) - for engine, request_id in requests: - if engine is None: - continue - if request_id == self._tts_uncond_request_id and not self._guidance_enabled: - continue - try: - abort_result = engine.abort(request_id) - if asyncio.iscoroutine(abort_result): - asyncio.run_coroutine_threadsafe(abort_result, self._loop).result(timeout=5) - except Exception as exc: - logging.debug(f"AsyncOmni.abort({request_id}) raised: {exc!r}") - - def finish(self, *, drain_remaining_audio_s: float = 0.0) -> None: - if self._closed: - return - self._closed = True - try: - - async def _close_inputs() -> None: - # Close every input this session owns. A session with both - # components drives them independently, so closing only one - # would leave the other consumer waiting forever. - if self._has_llm: - await self._input_q.put(None) - if self._has_tts: - await self._tts_cond_input_q.put(None) - if self._guidance_enabled: - await self._tts_uncond_input_q.put(None) - - asyncio.run_coroutine_threadsafe(_close_inputs(), self._loop).result(timeout=5) - except Exception: - pass - try: - self._consumer_future.result(timeout=max(drain_remaining_audio_s, 1.0)) - except Exception as exc: - logging.debug(f"OmniStreamingSession {self.request_id} consumer ended with: {exc!r}") - self._abort_engine_requests() - self._consumer_future.cancel() - try: - self._consumer_future.result(timeout=5) - except Exception: - pass - for queue in (self._text_out_q, self._tts_done_q): - try: - while True: - queue.get_nowait() - except Exception: - pass - - def abort(self) -> None: - if self._closed: - return - self._closed = True - self._abort_engine_requests() - self._consumer_future.cancel() diff --git a/pyproject.toml b/pyproject.toml index b5c2c010ebfa..ab6680b9fe0d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -350,25 +350,8 @@ speechlm2 = [ [tool.setuptools] py-modules = ["nemo"] -# vllm-omni spawns its stage children with ``multiprocessing`` start method -# ``spawn``, and spawned children inherit no Python state, so registering the -# NemotronDuplexH + EarTTS models and the pipeline in the parent is not enough. -# ``nemo_voicechat`` is declared in vLLM's group, not vllm-omni's, because vLLM -# loads its plugins earlier: vllm-omni resolves the pipeline for ``model_type`` -# while constructing AsyncOmniEngine, before it loads its own group. -# Registration is idempotent and no-ops when vllm-omni is absent, so an ordinary -# vLLM process pays only the import. -# Entry points are only discoverable from an installed distribution, so NeMo has -# to be pip-installed (``pip install -e .`` is enough); PYTHONPATH will not do. [project.entry-points."vllm.general_plugins"] nemo_speechlm = "nemo.collections.speechlm2.vllm.salm:register" -nemo_voicechat = "nemo.collections.speechlm2.inference.vllm_omni.register:register_nemo_voicechat" - -# Also declared in vllm-omni's own group, which is loaded in the stage children -# and worker processes. Harmless duplication: whichever loader runs first wins -# and the second call returns immediately. -[project.entry-points."vllm_omni.general_plugins"] -nemo_voicechat = "nemo.collections.speechlm2.inference.vllm_omni.register:register_nemo_voicechat" [project.urls] Download = "https://github.com/NVIDIA-NeMo/Speech/releases" diff --git a/tests/collections/speechlm2/nemo_inference_pipelines/conftest.py b/tests/collections/speechlm2/nemo_inference_pipelines/conftest.py index 7cbbd14a3677..040942582547 100644 --- a/tests/collections/speechlm2/nemo_inference_pipelines/conftest.py +++ b/tests/collections/speechlm2/nemo_inference_pipelines/conftest.py @@ -97,11 +97,10 @@ def _merge_pipeline_cfg(model_path: str, audio_path: str, output_dir: str, *over def _reclaim_gpu_after_large_load() -> None: """Return GPU memory to the driver after a large native load. - ``pipeline.shutdown`` only tears down the vLLM runtime. Native weights stay - on the wrapper until the pipeline is unreachable; PyTorch then keeps the - blocks in this process. vLLM engine cores are child processes and treat - that as used memory, so a native 11B test (~24 GiB) followed by vLLM/vLLM - OOMs on an 80 GiB card. + Native weights stay on the wrapper until the pipeline is unreachable; + PyTorch then keeps the blocks in this process. A native 11B test (~24 GiB) + followed by another large load OOMs on an 80 GiB card if the allocator is + not reclaimed. Tiny-model tests are a few GiB and rebuild from the same size, so the caching allocator is left warm. Both ``gc.collect`` and ``empty_cache`` @@ -440,22 +439,3 @@ def voicechat_audio_path(): @pytest.fixture(scope="session") def voicechat_speaker_name(): return DEFAULT_SPEAKER_NAME - - -@pytest.fixture(scope="session") -def real_vllm_omni_wrapper(tmp_path_factory, hf_voicechat_11b): - """Convert the public 11B snapshot for vLLM-Omni. - - ``NEMO_VLLM_WRAPPER_DIR`` is an optional prebuilt wrapper directory. - ``build_wrapper_checkpoint`` reuses it when complete. Otherwise one is - built under tmp. - """ - pytest.importorskip("vllm_omni") - if not torch.cuda.is_available(): - pytest.skip("converting the vLLM-Omni wrapper requires a GPU") - - from nemo.collections.speechlm2.inference.vllm_omni.checkpoint import build_wrapper_checkpoint - - wrapper_dir = os.environ.get("NEMO_VLLM_WRAPPER_DIR") or str(tmp_path_factory.mktemp("vllm_omni_wrapper")) - build_wrapper_checkpoint(hf_voicechat_11b, wrapper_dir) - return hf_voicechat_11b, wrapper_dir diff --git a/tests/collections/speechlm2/nemo_inference_pipelines/test_engine_selection.py b/tests/collections/speechlm2/nemo_inference_pipelines/test_engine_selection.py index f253e4c4470f..ada03274977c 100644 --- a/tests/collections/speechlm2/nemo_inference_pipelines/test_engine_selection.py +++ b/tests/collections/speechlm2/nemo_inference_pipelines/test_engine_selection.py @@ -18,6 +18,7 @@ from nemo.collections.speechlm2.inference.model_wrappers.engine_selection import ( native_weight_skip_prefixes, + reject_unimplemented_vllm, resolve_engine_types, ) @@ -75,3 +76,18 @@ def test_unusable_engine_selection_is_named(): resolve_engine_types({"llm_engine_type": "other"}) with pytest.raises(ValueError, match="not a config key"): resolve_engine_types({"engine_type": VLLM}) + + +def test_vllm_selection_is_named_as_not_implemented(): + """vllm_omni stays a legal config value so the native loop is already the + combined-form loop; constructing that backend is rejected rather than + silently running native.""" + with pytest.raises(NotImplementedError, match="not implemented in this PR"): + reject_unimplemented_vllm(VLLM, NATIVE) + with pytest.raises(NotImplementedError, match="not implemented in this PR"): + reject_unimplemented_vllm(NATIVE, VLLM) + + from nemo.collections.speechlm2.inference.model_wrappers.backend.vllm.llm import VllmLLM + + with pytest.raises(NotImplementedError, match="not implemented in this PR"): + VllmLLM() diff --git a/tests/collections/speechlm2/nemo_inference_pipelines/test_nemotron_voicechat_pipeline_vllm.py b/tests/collections/speechlm2/nemo_inference_pipelines/test_nemotron_voicechat_pipeline_vllm.py deleted file mode 100644 index e5d188e27b4e..000000000000 --- a/tests/collections/speechlm2/nemo_inference_pipelines/test_nemotron_voicechat_pipeline_vllm.py +++ /dev/null @@ -1,87 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""One live vLLM-Omni pipeline run on nvidia/NVIDIA-NemotronLabs-VoiceChat-11B. - -Skipped when ``vllm_omni`` is not installed. The 11B snapshot is downloaded -into the Hugging Face cache if needed. -""" - -from __future__ import annotations - -import tempfile - -import pytest -import torch - -pytest.importorskip("vllm_omni") - -from nemo.collections.speechlm2.inference.utils.stepprogressbar import StepProgressBar - -MOCK_SYSTEM_PROMPT = "This is a mock prompt for the test" -_VLLM = "vllm_omni" - -_VLLM_DEFAULTS = { - "s2s": { - "llm_engine_type": _VLLM, - "tts_engine_type": _VLLM, - "deterministic": False, - "decode_audio": True, - "system_prompt": MOCK_SYSTEM_PROMPT, - }, - "streaming": { - "chunk_size_in_secs": 0.08, - "buffer_size_in_secs": 71 * 0.08, - }, -} - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") -def test_pipeline_no_crash_vllm( - build_pipeline, hf_voicechat_11b, voicechat_audio_path, voicechat_speaker_name, real_vllm_omni_wrapper -): - """vLLM/vLLM ``pipeline.run()`` on the public 11B: text and audio exist.""" - _, wrapper_dir = real_vllm_omni_wrapper - output_dir = tempfile.mkdtemp(prefix="no-crash-vllm-") - - pipeline = build_pipeline( - hf_voicechat_11b, - voicechat_audio_path, - output_dir, - _VLLM_DEFAULTS, - { - "s2s": { - "speaker_name": voicechat_speaker_name, - "vllm_omni_config": {"wrapper_dir": wrapper_dir}, - } - }, - ) - wrapper = pipeline.s2s_model - assert wrapper.llm_engine_type == _VLLM - assert wrapper.tts_engine_type == _VLLM - - progress_bar = StepProgressBar.from_audio_filepaths( - [voicechat_audio_path], - chunk_size_in_secs=pipeline.chunk_size_in_secs, - pad_audio_to_sec=pipeline.pad_audio_to_sec, - pad_silence_ratio=pipeline.pad_silence_ratio, - pad_audio_by_sec=pipeline.pad_audio_by_sec, - ) - result = pipeline.run([voicechat_audio_path], progress_bar=progress_bar) - assert result is not None - assert len(result) == 1 - output = result[0] - assert output.token_text is not None and output.token_text.numel() > 0 - audio = output.audio_buffer - assert audio is not None and audio.numel() > 0 diff --git a/tests/collections/speechlm2/nemo_inference_pipelines/test_text_sampling.py b/tests/collections/speechlm2/nemo_inference_pipelines/test_text_sampling.py index bf9aebc32982..1fe93f08a793 100644 --- a/tests/collections/speechlm2/nemo_inference_pipelines/test_text_sampling.py +++ b/tests/collections/speechlm2/nemo_inference_pipelines/test_text_sampling.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pytest import torch from nemo.collections.speechlm2.inference.model_wrappers.text_sampling import sample_text_token @@ -61,86 +60,3 @@ def fail_multinomial(*_args, **_kwargs): special_token_ids={1}, special_ids_tensor=torch.tensor([1]), ).tolist() == [1] - - -def test_vllm_processor_reuses_shared_sampler_and_skips_prefill_history(): - pytest.importorskip("vllm") - from nemo.collections.speechlm2.inference.vllm_omni.nemotron_duplex_h.sampling import ( - SharedTextRequestSampler, - SharedTextSamplingState, - ) - - logits = torch.tensor([0.2, 1.3, 0.9, 0.5]) - params = { - "top_p": 0.9, - "repetition_penalty": 1.2, - "temperature": 0.75, - "special_token_ids": {0}, - } - - torch.manual_seed(17) - expected = sample_text_token( - logits.unsqueeze(0), - torch.tensor([[2, 2]], dtype=torch.long), - 2, - special_ids_tensor=torch.tensor([0]), - **params, - ) - - processor = SharedTextRequestSampler( - history_skip=1, - state=SharedTextSamplingState( - sample_count=3, - tokens=[2, 2], - ), - **params, - ) - torch.manual_seed(17) - forced_logits = processor([], logits.clone()) - - assert int(forced_logits.argmax().item()) == int(expected[0].item()) - assert torch.isfinite(forced_logits).sum().item() == 1 - - -def test_vllm_processor_history_survives_segment_readmission(): - """A vLLM request is re-admitted per streaming segment; the sampling - history has to outlive that or the repetition penalty resets mid-stream.""" - pytest.importorskip("vllm") - from types import SimpleNamespace - - from vllm import SamplingParams - - from nemo.collections.speechlm2.inference.vllm_omni.nemotron_duplex_h.sampling import ( - SHARED_TEXT_SAMPLING_ARG, - SharedTextSamplingLogitsProcessor, - ) - - adapter = SharedTextSamplingLogitsProcessor( - SimpleNamespace(scheduler_config=SimpleNamespace(max_num_seqs=1)), - torch.device("cpu"), - False, - ) - params = SamplingParams( - temperature=0.0, - extra_args={ - SHARED_TEXT_SAMPLING_ARG: { - "top_p": 1.0, - "temperature": 1.0, - "repetition_penalty": 1.0, - "special_token_ids": [0], - "history_skip": 1, - "history_key": "stream-1", - } - }, - ) - - first_segment = adapter.new_req_logits_processor(params) - assert first_segment is not None - first_segment([], torch.tensor([0.0, 2.0, 1.0])) - first_segment([], torch.tensor([0.0, 1.0, 2.0])) - - resumed_segment = adapter.new_req_logits_processor(params) - assert resumed_segment is not None - assert resumed_segment.state is first_segment.state - assert resumed_segment.state.sample_count == 2 - assert resumed_segment.state.tokens == [2] diff --git a/tests/collections/speechlm2/nemo_inference_pipelines/test_vllm_omni_checkpoint.py b/tests/collections/speechlm2/nemo_inference_pipelines/test_vllm_omni_checkpoint.py deleted file mode 100644 index e5d6dc9f8d4f..000000000000 --- a/tests/collections/speechlm2/nemo_inference_pipelines/test_vllm_omni_checkpoint.py +++ /dev/null @@ -1,112 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Wrapper-checkpoint assembly on the filesystem. - -A real vLLM/vLLM ``pipeline.run()`` lives in -``test_nemotron_voicechat_pipeline_vllm.py``. -""" - -from types import SimpleNamespace - -import pytest - -from nemo.collections.speechlm2.inference.vllm_omni.checkpoint import build_wrapper_checkpoint - - -def _stub_source_and_partial_wrapper(tmp_path, ready_component): - """A stub source checkpoint plus a wrapper with one component already built.""" - source = tmp_path / "source" - source.mkdir() - (source / "config.json").write_text("{}") - (source / "model.safetensors").write_bytes(b"source") - - wrapper = tmp_path / "wrapper" - (wrapper / ready_component).mkdir(parents=True) - (wrapper / ready_component / "config.json").write_text("{}") - (wrapper / ready_component / "model.safetensors").write_bytes(b"ready") - (wrapper / "config.json").write_text('{"model_type": "nemotron_voicechat"}') - return source, wrapper - - -@pytest.mark.parametrize("component", ["nemotron", "eartts"]) -def test_wrapper_checkpoint_converts_only_the_requested_component(tmp_path, component): - """Reusing a ready component must not drag the other one along.""" - source, wrapper = _stub_source_and_partial_wrapper(tmp_path, component) - - result = build_wrapper_checkpoint( - str(source), - str(wrapper), - include_nemotron=component == "nemotron", - include_eartts=component == "eartts", - ) - - assert result == str(wrapper) - assert not (wrapper / ("eartts" if component == "nemotron" else "nemotron")).exists() - # Records which source it came from, which is what makes it verifiable. - assert (wrapper / ".nemo_source.json").is_file() - - -def test_wrapper_checkpoint_refuses_to_extend_an_unverified_partial_wrapper(tmp_path): - """A wrapper with no ``.nemo_source.json`` could have come from any source, - so adding a second component to it might silently mix two checkpoints. - """ - source, wrapper = _stub_source_and_partial_wrapper(tmp_path, "nemotron") - - with pytest.raises(ValueError, match="Cannot safely add a component"): - build_wrapper_checkpoint(str(source), str(wrapper), include_nemotron=False, include_eartts=True) - - -def test_converter_applies_voicechat_special_token_overrides(): - """The converted config must carry VoiceChat's BOS/EOS/PAD, not the LLM - backbone's; otherwise conversion succeeds with incorrect system-prompt - prefill token IDs. - """ - from nemo.collections.speechlm2.inference.vllm_omni.scripts.convert_duplex_stt_checkpoint import ( - _apply_source_special_tokens, - ) - - class FakeTokenizer: - def __init__(self): - self.vocab = {"": 0, "": 1, "": 2, "": 12} - - def get_vocab(self): - return self.vocab - - def add_special_tokens(self, values): - for name, token in values.items(): - setattr(self, name, token) - return 0 - - def convert_tokens_to_ids(self, token): - return self.vocab[token] - - config = SimpleNamespace(bos_token_id=1, eos_token_id=12, pad_token_id=0) - source = { - "model": { - "stt": { - "model": { - "override_tokens": { - "bos_token": "", - "eos_token": "", - "pad_token": "", - } - } - } - } - } - - _apply_source_special_tokens(config, FakeTokenizer(), source) - - assert (config.bos_token_id, config.eos_token_id, config.pad_token_id) == (1, 2, 12) diff --git a/tests/collections/speechlm2/nemo_inference_pipelines/test_vllm_omni_eartts_cfg.py b/tests/collections/speechlm2/nemo_inference_pipelines/test_vllm_omni_eartts_cfg.py deleted file mode 100644 index 69c1a969b71e..000000000000 --- a/tests/collections/speechlm2/nemo_inference_pipelines/test_vllm_omni_eartts_cfg.py +++ /dev/null @@ -1,289 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""CPU-only tests for vLLM-Omni EarTTS classifier-free guidance. - -Guidance arithmetic and MaskGIT trajectory sharing are numerical properties -that an end-to-end run cannot localise, so they are checked directly against -the real model definition. Requires vllm-omni to be installed; skipped -otherwise rather than run against a stubbed vLLM, which would only test the -stubs. -""" - -from types import SimpleNamespace - -import pytest -import torch -from torch import nn - - -@pytest.fixture(scope="module") -def eartts(): - pytest.importorskip("vllm_omni") - from nemo.collections.speechlm2.inference.vllm_omni.eartts import eartts as eartts_module - - return eartts_module - - -def test_unconditional_embedding_replaces_only_text_branch(eartts): - config = SimpleNamespace( - hidden_size=2, - emb_vocab_size=4, - codebook_size=3, - latent_size=2, - num_quantizers=2, - use_gated_fusion_for_text_audio=False, - use_audio_prompt_frozen_projection=False, - ) - embedding = eartts.EarTTSInputEmbedding(config) - with torch.no_grad(): - for rvq in embedding.rvq_embs: - rvq.weight.zero_() - embedding.rvq_embs[0].weight[0] = torch.tensor([1.0, 0.0]) - embedding.rvq_embs[1].weight[0] = torch.tensor([0.0, 2.0]) - embedding.embed_code.weight.copy_(torch.eye(2)) - embedding.embed_subword.embed_subwords.weight.zero_() - embedding.embed_subword.embed_subwords.weight[1] = torch.tensor([3.0, 4.0]) - embedding.bos_emb.zero_() - embedding.null_emb.copy_(torch.tensor([9.0, 10.0])) - - kwargs = dict( - acoustic_tokens=torch.zeros(2, 2, dtype=torch.long), - text_tokens=torch.ones(2, dtype=torch.long), - text_mask=torch.ones(2, dtype=torch.long), - bos_mask=torch.zeros(2, dtype=torch.long), - speaker_latent=torch.zeros(2, 2), - ) - original = embedding(**kwargs) - explicit_no_cfg = embedding( - **kwargs, - cfg_is_uncond=torch.zeros(2, dtype=torch.bool), - ) - guided_input = embedding( - **kwargs, - cfg_is_uncond=torch.tensor([False, True]), - ) - - torch.testing.assert_close(original, explicit_no_cfg, rtol=0, atol=0) - torch.testing.assert_close(guided_input[0], torch.tensor([4.0, 6.0])) - torch.testing.assert_close(guided_input[1], torch.tensor([10.0, 12.0])) - torch.testing.assert_close( - guided_input[1] - embedding.null_emb, - torch.tensor([1.0, 2.0]), - ) - assert dict(embedding.named_parameters())["null_emb"] is embedding.null_emb - - -def test_cfg_rows_are_ordered_and_guided_after_mlp(eartts): - hidden = torch.tensor([[5.0], [10.0], [3.0], [20.0]]) - enabled = torch.ones(4, dtype=torch.bool) - is_uncond = torch.tensor([True, False, True, False]) - pair_id = torch.tensor([20, 10, 10, 20]) - scale = torch.tensor([1.5, 2.0, 2.0, 1.5]) - valid = torch.ones(4, dtype=torch.bool) - - ( - ordered, - ordered_is_uncond, - ordered_scale, - active, - partner, - conditional_rep, - inverse, - ) = eartts._prepare_cfg_sampling_batch( - hidden, - enabled, - is_uncond, - pair_id, - scale, - valid, - ) - torch.testing.assert_close( - ordered, - torch.tensor([[10.0], [20.0], [3.0], [5.0]]), - ) - assert ordered_is_uncond.tolist() == [False, False, True, True] - assert active.tolist() == [True, True, True, True] - assert partner.tolist() == [2, 3, 0, 1] - assert conditional_rep.tolist() == [0, 1, 0, 1] - - guided = eartts._apply_cfg_after_mlp( - ordered, - ordered_is_uncond, - ordered_scale, - active, - partner, - ) - torch.testing.assert_close( - guided, - torch.tensor([[24.0], [42.5], [24.0], [42.5]]), - ) - assert torch.equal(ordered[inverse], hidden) - - # An incomplete or disabled batch must fall straight through: no - # reordering, nothing active, and guidance a no-op. - plain = torch.randn(3, 4) - ordered, roles, scales, active, partner, _, inverse = eartts._prepare_cfg_sampling_batch( - plain, - cfg_enabled=torch.tensor([True, True, False]), - cfg_is_uncond=torch.tensor([False, True, False]), - cfg_pair_id=torch.tensor([7, 7, -1]), - cfg_scale=torch.ones(3), - valid=torch.ones(3, dtype=torch.bool), - ) - assert torch.equal(ordered, plain) - assert not active.any() - assert torch.equal(inverse, torch.arange(3)) - assert torch.equal(eartts._apply_cfg_after_mlp(ordered, roles, scales, active, partner), plain) - - -def test_maskgit_shares_one_code_trajectory_per_pair(eartts): - config = SimpleNamespace( - num_quantizers=2, - codebook_size=8, - noise_scale=0.7, - num_iter=2, - exponent=3.0, - latent_size=4, - hidden_size=4, - intermediate_size=8, - mog_num_layers=0, - mog_num_predictions=4, - mog_low_rank=None, - top_p_or_k=None, - mog_min_log_std=-4.0, - mog_eps=1e-6, - ) - sampler = eartts.MaskGITSampler(config) - torch.manual_seed(3) - with torch.no_grad(): - for parameter in sampler.parameters(): - parameter.normal_(mean=0.0, std=0.2) - - hidden = torch.randn(4, 4) - enabled = torch.ones(4, dtype=torch.bool) - roles = torch.tensor([True, False, True, False]) - pairs = torch.tensor([2, 1, 1, 2]) - scales = torch.full((4,), 1.25) - valid = torch.ones(4, dtype=torch.bool) - torch.manual_seed(11) - codes = sampler(hidden, enabled, roles, pairs, scales, valid) - - assert torch.equal(codes[0], codes[3]) - assert torch.equal(codes[1], codes[2]) - - no_cfg_hidden = torch.randn(3, 4) - torch.manual_seed(17) - implicit_no_cfg = sampler(no_cfg_hidden) - torch.manual_seed(17) - explicit_no_cfg = sampler( - no_cfg_hidden, - cfg_enabled=torch.zeros(3, dtype=torch.bool), - cfg_is_uncond=torch.zeros(3, dtype=torch.bool), - cfg_pair_id=torch.full((3,), -1, dtype=torch.long), - cfg_scale=torch.zeros(3), - valid=torch.ones(3, dtype=torch.bool), - ) - assert torch.equal(implicit_no_cfg, explicit_no_cfg) - - -def test_client_facing_stage_emits_the_drainable_audio_key(eartts): - """A final AR audio stage must publish codes under ``model_outputs``. - - vLLM-Omni remaps that key onto the drainable ``audio`` modality, so DELTA - streaming empties it every step. Any other key is retained across steps and - merged with ``CONCAT_LAST``, which widens a ``T x num_quantizers`` frame - instead of appending frames to it. - """ - hidden = torch.zeros(1, 4) - codes = torch.tensor([[3, 5]], dtype=torch.long) - - for single_stage_audio, expected_key in ((True, "model_outputs"), (False, "audio_codes")): - model = object.__new__(eartts.EarTTSForCausalLM) - nn.Module.__init__(model) - model._single_stage_audio = single_stage_audio - model._out_codes = codes.clone() - - output = model.make_omni_output(hidden) - - assert list(output.multimodal_outputs) == [expected_key] - torch.testing.assert_close(output.multimodal_outputs[expected_key], codes) - - stashed = model.postprocess(hidden, output.multimodal_outputs) - torch.testing.assert_close(stashed["last_acoustic_codes"], codes) - - # Per-request CFG metadata lands in model-owned buffers whose addresses - # must stay stable, because CUDA graphs capture them. - model = object.__new__(eartts.EarTTSForCausalLM) - nn.Module.__init__(model) - model.config = SimpleNamespace(guidance_scale=0.5) - model._cfg_enabled = torch.zeros(8, dtype=torch.bool) - model._cfg_is_uncond = torch.zeros(8, dtype=torch.bool) - model._cfg_pair_id = torch.full((8,), -1, dtype=torch.long) - model._cfg_scale = torch.zeros(8) - addresses = tuple( - value.data_ptr() - for value in ( - model._cfg_enabled, - model._cfg_is_uncond, - model._cfg_pair_id, - model._cfg_scale, - ) - ) - - model._write_cfg_state( - start=1, - span_len=2, - info_dict={ - "cfg_enabled": True, - "cfg_role": "cond", - "cfg_pair_id": "request-7", - "cfg_scale": 1.75, - }, - ) - model._write_cfg_state( - start=3, - span_len=1, - info_dict={ - "cfg_enabled": torch.tensor(True), - "cfg_role": ["uncond"], - "cfg_pair_id": "request-7", - "cfg_scale": torch.tensor(1.75), - }, - ) - - assert model._cfg_enabled[1:4].all() - assert model._cfg_is_uncond[1:4].tolist() == [False, False, True] - assert model._cfg_pair_id[1] == model._cfg_pair_id[3] - torch.testing.assert_close(model._cfg_scale[1:4], torch.full((3,), 1.75)) - assert addresses == tuple( - value.data_ptr() - for value in ( - model._cfg_enabled, - model._cfg_is_uncond, - model._cfg_pair_id, - model._cfg_scale, - ) - ) - - with pytest.raises(AssertionError, match="cfg_role"): - model._write_cfg_state( - start=4, - span_len=1, - info_dict={ - "cfg_enabled": True, - "cfg_role": "conditional", - "cfg_pair_id": 4, - }, - ) From 209f3d93019e02d2d420d5fee48fa3ba4f6b06b9 Mon Sep 17 00:00:00 2001 From: Elena Rastorgueva Date: Tue, 1 Sep 2026 19:14:20 +0000 Subject: [PATCH 3/6] docs(speechlm2): keep VoiceChat from_pretrained as a local path Signed-off-by: Elena Rastorgueva --- docs/source/speechlm2/intro.rst | 2 +- docs/source/speechlm2/models.rst | 2 +- docs/source/speechlm2/streaming_inference.rst | 3 ++- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/docs/source/speechlm2/intro.rst b/docs/source/speechlm2/intro.rst index 1de9c44d7896..909afb4d37e5 100644 --- a/docs/source/speechlm2/intro.rst +++ b/docs/source/speechlm2/intro.rst @@ -245,7 +245,7 @@ You can evaluate and run full-duplex inference using the `NemotronVoiceChat` pip from nemo.collections.audio.parts.utils.transforms import resample import nemo.collections.speechlm2 as slm - model = slm.models.NemotronVoiceChat.from_pretrained("nvidia/NVIDIA-NemotronLabs-VoiceChat-11B").eval() + model = slm.models.NemotronVoiceChat.from_pretrained("path/to/pretrained_checkpoint").eval() # Load user audio prompt audio_path = "path/to/user_audio.wav" diff --git a/docs/source/speechlm2/models.rst b/docs/source/speechlm2/models.rst index bb5907363e82..3ff0b9888ea9 100644 --- a/docs/source/speechlm2/models.rst +++ b/docs/source/speechlm2/models.rst @@ -312,7 +312,7 @@ All models in the speechlm2 collection can be instantiated from pretrained check ear_tts_model = slm.models.DuplexEARTTS.from_pretrained("path/to/checkpoint") # Load NemotronVoiceChat (Inference Only) - voicechat_model = slm.models.NemotronVoiceChat.from_pretrained("nvidia/NVIDIA-NemotronLabs-VoiceChat-11B") + voicechat_model = slm.models.NemotronVoiceChat.from_pretrained("path/to/checkpoint") Remote HuggingFace code is disabled by default. If a trusted checkpoint requires custom code, opt in at runtime and pin the repository to a reviewed revision: diff --git a/docs/source/speechlm2/streaming_inference.rst b/docs/source/speechlm2/streaming_inference.rst index 853cbc5e5f84..f728bc4423b3 100644 --- a/docs/source/speechlm2/streaming_inference.rst +++ b/docs/source/speechlm2/streaming_inference.rst @@ -27,7 +27,8 @@ There are two ways to use the pipeline: generate_step(frames) │ ├─ incremental agent audio + text - └─ incremental user ASR text (when the checkpoint has an ASR head) + ├─ incremental user ASR text (when the checkpoint has an ASR head) + └─ [EXPERIMENTAL] incremental function head output (when the checkpoint has a function call head, as in NVIDIA-NemotronLabs-VoiceChat-11B) Each audio file passed to ``run()`` is treated as one continuous audio stream. ``run()`` accumulates the per-step outputs for each stream and writes final audio/text From 6e82f8319fde85839a939161fbf2ab5db292f169 Mon Sep 17 00:00:00 2001 From: Elena Rastorgueva Date: Tue, 1 Sep 2026 20:29:10 +0000 Subject: [PATCH 4/6] refactor(speechlm2): use StreamingEncoder CUDA graphs for perception Replace the VoiceChat-specific encoder capture with encoder.set_streaming_cuda_graphs so subsequent cache-aware steps share the ASR helper. Signed-off-by: Elena Rastorgueva --- docs/source/speechlm2/streaming_inference.rst | 5 + .../conf/s2s_streaming.yaml | 2 +- .../model_wrappers/perception_cache.py | 335 ++---------------- 3 files changed, 42 insertions(+), 300 deletions(-) diff --git a/docs/source/speechlm2/streaming_inference.rst b/docs/source/speechlm2/streaming_inference.rst index f728bc4423b3..0abd612a9ccc 100644 --- a/docs/source/speechlm2/streaming_inference.rst +++ b/docs/source/speechlm2/streaming_inference.rst @@ -218,6 +218,11 @@ S2S Model Settings (``s2s``) * - ``use_perception_cache`` - ``true`` - Cache-aware streaming for the perception encoder. + * - ``use_perception_cudagraph`` + - ``true`` + - CUDA-graph replay of the perception encoder step via + ``encoder.set_streaming_cuda_graphs``. Requires ``use_perception_cache``. + First chunk stays eager; subsequent chunks capture after a short warmup. * - ``use_llm_cache`` - ``false`` - Reuse the native LLM KV cache instead of replaying history. NemotronH diff --git a/examples/speechlm2/nemo_inference_pipelines/conf/s2s_streaming.yaml b/examples/speechlm2/nemo_inference_pipelines/conf/s2s_streaming.yaml index d49efe7782ca..cab10543b9a7 100644 --- a/examples/speechlm2/nemo_inference_pipelines/conf/s2s_streaming.yaml +++ b/examples/speechlm2/nemo_inference_pipelines/conf/s2s_streaming.yaml @@ -79,7 +79,7 @@ s2s: # Inference settings # ======================== use_perception_cache: true # Enable cache-aware streaming for perception encoder - use_perception_cudagraph: true # Enable CUDA graph-accelerated perception encoder + use_perception_cudagraph: true # encoder.set_streaming_cuda_graphs on the perception encoder use_llm_cache: false # Keep the LLM KV cache across steps (native engine only). # False replays the whole history each step: O(n^2) and slow, # and is the portable compatibility default. diff --git a/nemo/collections/speechlm2/inference/model_wrappers/perception_cache.py b/nemo/collections/speechlm2/inference/model_wrappers/perception_cache.py index 9ebc68f79c1c..c8516e23dc87 100644 --- a/nemo/collections/speechlm2/inference/model_wrappers/perception_cache.py +++ b/nemo/collections/speechlm2/inference/model_wrappers/perception_cache.py @@ -15,9 +15,14 @@ """ Cache-aware perception encoder for streaming S2S inference. -Provides incremental mel-spectrogram encoding with optional CUDA graph -acceleration, so that only new audio needs to be processed each step -instead of re-encoding the entire buffer. +Provides incremental mel-spectrogram encoding so that only new audio needs +to be processed each step instead of re-encoding the entire buffer. + +Optional CUDA-graph replay of the encoder step is +``StreamingEncoder.set_streaming_cuda_graphs``. Subsequent chunks pass +``keep_all_outputs=False`` so that helper can capture; for +``att_context_style=chunked_limited`` that flag is a no-op versus ``True`` +(``valid_out_len`` already equals lookahead + 1). """ import copy @@ -46,55 +51,14 @@ def is_initialized(self) -> bool: return None not in [self.cache_last_channel, self.cache_last_time, self.cache_last_channel_len] -@dataclass -class PerceptionCUDAGraphState: - """State for CUDA graph-accelerated perception encoder. - - Holds separate graphs for first chunk (different size) and subsequent chunks. - Also holds static buffers for inputs/outputs to enable graph replay. - """ - - # CUDA graphs - graph_first: torch.cuda.CUDAGraph | None = None - graph_subsequent: torch.cuda.CUDAGraph | None = None - - # Static input buffers (for copying data before graph replay) - static_mel_first: torch.Tensor | None = None - static_mel_subsequent: torch.Tensor | None = None - static_mel_len_first: torch.Tensor | None = None - static_mel_len_subsequent: torch.Tensor | None = None - - # Static cache input buffers - static_cache_channel_in: torch.Tensor | None = None - static_cache_time_in: torch.Tensor | None = None - static_cache_channel_len_in: torch.Tensor | None = None - - # Static output buffers (results are written here during replay) - static_encoded_first: torch.Tensor | None = None - static_encoded_subsequent: torch.Tensor | None = None - static_encoded_len_first: torch.Tensor | None = None - static_encoded_len_subsequent: torch.Tensor | None = None - - # Static cache output buffers - SEPARATE for first and subsequent graphs - # (each graph writes to its own output tensors during replay) - static_cache_channel_out_first: torch.Tensor | None = None - static_cache_time_out_first: torch.Tensor | None = None - static_cache_channel_len_out_first: torch.Tensor | None = None - static_cache_channel_out_subsequent: torch.Tensor | None = None - static_cache_time_out_subsequent: torch.Tensor | None = None - static_cache_channel_len_out_subsequent: torch.Tensor | None = None - - def is_captured(self) -> bool: - """Check if graphs have been captured.""" - return self.graph_first is not None and self.graph_subsequent is not None - - class PerceptionCacheManager: - """Manages cache-aware streaming perception encoding with optional CUDA graphs. + """Manages cache-aware streaming perception encoding. - This class encapsulates all perception cache setup, CUDA graph capture, - and the incremental encoding step. It is created by the inference wrapper - when ``use_perception_cache=True``. + Encapsulates preprocessor setup, encoder-cache state, and the incremental + encoding step. Created by the inference wrapper when + ``use_perception_cache=True``. When ``use_cudagraph=True``, attaches + ``encoder.set_streaming_cuda_graphs`` so subsequent encoder steps replay + from a CUDA graph; adapter and projection stay eager. """ def __init__(self, model, device: torch.device, dtype: torch.dtype, use_cudagraph: bool = False): @@ -108,7 +72,6 @@ def __init__(self, model, device: torch.device, dtype: torch.dtype, use_cudagrap self.subsampling_factor = None self.input_features = None self.sampling_frames = None - self.cudagraph_state: PerceptionCUDAGraphState | None = None def setup(self) -> bool: """Setup cache-aware streaming for the perception encoder. @@ -155,185 +118,11 @@ def setup(self) -> bool: logging.info(f" Subsampling factor: {self.subsampling_factor}") if self.use_cudagraph: - logging.info(" Setting up CUDA graphs for perception encoder...") - self.capture_cudagraphs() - logging.info(" CUDA graphs captured") + encoder.set_streaming_cuda_graphs(enabled=True) + logging.info(" Streaming encoder CUDA graphs enabled (set_streaming_cuda_graphs)") return True - def capture_cudagraphs(self): - """Capture CUDA graphs for perception encoder with both chunk sizes. - - Note: "chunk" in the streaming encoder config (chunk_size, shift_size, etc.) - follows NeMo's cache-aware streaming encoder API and is measured in - mel-spectrogram time-steps, not audio samples or seconds. - """ - encoder = self.model.stt_model.perception.encoder - perception = self.model.stt_model.perception - streaming_cfg = self.streaming_cfg - - if isinstance(streaming_cfg.chunk_size, list): - chunk_size_first = streaming_cfg.chunk_size[0] - chunk_size_subsequent = streaming_cfg.chunk_size[1] - else: - chunk_size_first = streaming_cfg.chunk_size - chunk_size_subsequent = streaming_cfg.chunk_size - - if isinstance(streaming_cfg.pre_encode_cache_size, list): - pre_encode_cache_first = streaming_cfg.pre_encode_cache_size[0] - pre_encode_cache_subsequent = streaming_cfg.pre_encode_cache_size[1] - else: - pre_encode_cache_first = streaming_cfg.pre_encode_cache_size - pre_encode_cache_subsequent = streaming_cfg.pre_encode_cache_size - - mel_len_first = chunk_size_first + pre_encode_cache_first - mel_len_subsequent = chunk_size_subsequent + pre_encode_cache_subsequent - - logging.info(f" CUDA graph mel lengths: first={mel_len_first}, subsequent={mel_len_subsequent}") - - cache_last_channel, cache_last_time, cache_last_channel_len = encoder.get_initial_cache_state(batch_size=1) - - state = PerceptionCUDAGraphState() - - state.static_mel_first = torch.zeros( - (1, self.input_features, mel_len_first), dtype=torch.float32, device=self.device - ) - state.static_mel_subsequent = torch.zeros( - (1, self.input_features, mel_len_subsequent), dtype=torch.float32, device=self.device - ) - state.static_mel_len_first = torch.tensor([mel_len_first], dtype=torch.long, device=self.device) - state.static_mel_len_subsequent = torch.tensor([mel_len_subsequent], dtype=torch.long, device=self.device) - - if cache_last_channel is not None: - state.static_cache_channel_in = cache_last_channel.clone() - if cache_last_time is not None: - state.static_cache_time_in = cache_last_time.clone() - if cache_last_channel_len is not None: - state.static_cache_channel_len_in = cache_last_channel_len.clone() - - logging.info(" Warming up encoder for CUDA graph capture...") - # PyTorch recommends a few eager warmup iterations before CUDA graph - # capture on a side stream; its example uses three iterations: - # https://pytorch.org/docs/stable/notes/cuda.html#cuda-graphs - warmup_stream = torch.cuda.Stream(device=self.device) - warmup_stream.wait_stream(torch.cuda.current_stream(self.device)) - with torch.cuda.stream(warmup_stream), torch.no_grad(): - for _ in range(3): - _ = encoder.cache_aware_stream_step( - processed_signal=state.static_mel_first, - processed_signal_length=state.static_mel_len_first, - cache_last_channel=( - state.static_cache_channel_in.clone() if state.static_cache_channel_in is not None else None - ), - cache_last_time=( - state.static_cache_time_in.clone() if state.static_cache_time_in is not None else None - ), - cache_last_channel_len=( - state.static_cache_channel_len_in.clone() - if state.static_cache_channel_len_in is not None - else None - ), - keep_all_outputs=True, - drop_extra_pre_encoded=0, - ) - _ = encoder.cache_aware_stream_step( - processed_signal=state.static_mel_subsequent, - processed_signal_length=state.static_mel_len_subsequent, - cache_last_channel=( - state.static_cache_channel_in.clone() if state.static_cache_channel_in is not None else None - ), - cache_last_time=( - state.static_cache_time_in.clone() if state.static_cache_time_in is not None else None - ), - cache_last_channel_len=( - state.static_cache_channel_len_in.clone() - if state.static_cache_channel_len_in is not None - else None - ), - keep_all_outputs=True, - drop_extra_pre_encoded=streaming_cfg.drop_extra_pre_encoded, - ) - torch.cuda.current_stream(self.device).wait_stream(warmup_stream) - - # Capture graph for FIRST chunk - logging.info(f" Capturing CUDA graph for first chunk (mel_len={mel_len_first})...") - state.graph_first = torch.cuda.CUDAGraph() - - if state.static_cache_channel_in is not None: - state.static_cache_channel_in.copy_(cache_last_channel) - if state.static_cache_time_in is not None: - state.static_cache_time_in.copy_(cache_last_time) - if state.static_cache_channel_len_in is not None: - state.static_cache_channel_len_in.copy_(cache_last_channel_len) - - with torch.cuda.graph(state.graph_first): - ( - encoded_first, - encoded_len_first, - cache_channel_out_first, - cache_time_out_first, - cache_channel_len_out_first, - ) = encoder.cache_aware_stream_step( - processed_signal=state.static_mel_first, - processed_signal_length=state.static_mel_len_first, - cache_last_channel=state.static_cache_channel_in, - cache_last_time=state.static_cache_time_in, - cache_last_channel_len=state.static_cache_channel_len_in, - keep_all_outputs=True, - drop_extra_pre_encoded=0, - ) - encoded_adapted_first, _ = perception.modality_adapter( - audio_signal=encoded_first, length=encoded_len_first - ) - encoded_chunk_first = perception.proj(encoded_adapted_first.transpose(1, 2)) - - state.static_encoded_first = encoded_chunk_first - state.static_encoded_len_first = encoded_len_first - state.static_cache_channel_out_first = cache_channel_out_first - state.static_cache_time_out_first = cache_time_out_first - state.static_cache_channel_len_out_first = cache_channel_len_out_first - - # Capture graph for SUBSEQUENT chunks - logging.info(f" Capturing CUDA graph for subsequent chunks (mel_len={mel_len_subsequent})...") - state.graph_subsequent = torch.cuda.CUDAGraph() - - if state.static_cache_channel_in is not None: - state.static_cache_channel_in.copy_(cache_last_channel) - if state.static_cache_time_in is not None: - state.static_cache_time_in.copy_(cache_last_time) - if state.static_cache_channel_len_in is not None: - state.static_cache_channel_len_in.copy_(cache_last_channel_len) - - with torch.cuda.graph(state.graph_subsequent): - ( - encoded_subsequent, - encoded_len_subsequent, - cache_channel_out_subsequent, - cache_time_out_subsequent, - cache_channel_len_out_subsequent, - ) = encoder.cache_aware_stream_step( - processed_signal=state.static_mel_subsequent, - processed_signal_length=state.static_mel_len_subsequent, - cache_last_channel=state.static_cache_channel_in, - cache_last_time=state.static_cache_time_in, - cache_last_channel_len=state.static_cache_channel_len_in, - keep_all_outputs=True, - drop_extra_pre_encoded=streaming_cfg.drop_extra_pre_encoded, - ) - encoded_adapted_subsequent, _ = perception.modality_adapter( - audio_signal=encoded_subsequent, length=encoded_len_subsequent - ) - encoded_chunk_subsequent = perception.proj(encoded_adapted_subsequent.transpose(1, 2)) - - state.static_encoded_subsequent = encoded_chunk_subsequent - state.static_encoded_len_subsequent = encoded_len_subsequent - state.static_cache_channel_out_subsequent = cache_channel_out_subsequent - state.static_cache_time_out_subsequent = cache_time_out_subsequent - state.static_cache_channel_len_out_subsequent = cache_channel_len_out_subsequent - - self.cudagraph_state = state - logging.info(" CUDA graphs captured successfully") - def get_initial_state(self, batch_size: int = 1) -> PerceptionCacheState: """Get initial cache state for perception encoder.""" encoder = self.model.stt_model.perception.encoder @@ -483,79 +272,27 @@ def step( chunk_lengths = torch.tensor([mel_chunk.shape[-1]], dtype=torch.long, device=self.device) - if self.use_cudagraph and self.cudagraph_state is not None and self.cudagraph_state.is_captured(): - graph_state = self.cudagraph_state - - if is_first_sub_step: - graph_state.static_mel_first.copy_(mel_chunk) - else: - graph_state.static_mel_subsequent.copy_(mel_chunk) - - if graph_state.static_cache_channel_in is not None and cache_last_channel is not None: - graph_state.static_cache_channel_in.copy_(cache_last_channel) - if graph_state.static_cache_time_in is not None and cache_last_time is not None: - graph_state.static_cache_time_in.copy_(cache_last_time) - if graph_state.static_cache_channel_len_in is not None and cache_last_channel_len is not None: - graph_state.static_cache_channel_len_in.copy_(cache_last_channel_len) - - if is_first_sub_step: - graph_state.graph_first.replay() - encoded_chunk = graph_state.static_encoded_first.clone() - cache_last_channel = ( - graph_state.static_cache_channel_out_first.clone() - if graph_state.static_cache_channel_out_first is not None - else None - ) - cache_last_time = ( - graph_state.static_cache_time_out_first.clone() - if graph_state.static_cache_time_out_first is not None - else None - ) - cache_last_channel_len = ( - graph_state.static_cache_channel_len_out_first.clone() - if graph_state.static_cache_channel_len_out_first is not None - else None - ) - else: - graph_state.graph_subsequent.replay() - encoded_chunk = graph_state.static_encoded_subsequent.clone() - cache_last_channel = ( - graph_state.static_cache_channel_out_subsequent.clone() - if graph_state.static_cache_channel_out_subsequent is not None - else None - ) - cache_last_time = ( - graph_state.static_cache_time_out_subsequent.clone() - if graph_state.static_cache_time_out_subsequent is not None - else None - ) - cache_last_channel_len = ( - graph_state.static_cache_channel_len_out_subsequent.clone() - if graph_state.static_cache_channel_len_out_subsequent is not None - else None - ) + # keep_all_outputs=False so set_streaming_cuda_graphs can capture + # subsequent steps. For chunked_limited VoiceChat this is a no-op vs + # True: valid_out_len already equals lookahead + 1. + ( + encoded, + encoded_len, + cache_last_channel, + cache_last_time, + cache_last_channel_len, + ) = encoder.cache_aware_stream_step( + processed_signal=mel_chunk, + processed_signal_length=chunk_lengths, + cache_last_channel=cache_last_channel, + cache_last_time=cache_last_time, + cache_last_channel_len=cache_last_channel_len, + keep_all_outputs=False, + drop_extra_pre_encoded=drop_extra_pre_encoded, + ) - else: - ( - encoded, - encoded_len, - cache_last_channel, - cache_last_time, - cache_last_channel_len, - ) = encoder.cache_aware_stream_step( - processed_signal=mel_chunk, - processed_signal_length=chunk_lengths, - cache_last_channel=cache_last_channel, - cache_last_time=cache_last_time, - cache_last_channel_len=cache_last_channel_len, - keep_all_outputs=True, - drop_extra_pre_encoded=drop_extra_pre_encoded, - ) - - modality_adapter = perception.modality_adapter - encoded_adapted, _ = modality_adapter(audio_signal=encoded, length=encoded_len) - - encoded_chunk = perception.proj(encoded_adapted.transpose(1, 2)) + encoded_adapted, _ = perception.modality_adapter(audio_signal=encoded, length=encoded_len) + encoded_chunk = perception.proj(encoded_adapted.transpose(1, 2)) encoded_chunks.append(encoded_chunk) From fbd58d67486e6d2d067726c9fd9bbb4cbf1a4120 Mon Sep 17 00:00:00 2001 From: Elena Rastorgueva Date: Tue, 1 Sep 2026 21:35:23 +0000 Subject: [PATCH 5/6] refactor(speechlm2): keep EarTTS sampling checkpoint-owned Remove redundant streaming overrides so native and converted EarTTS use the sampling values carried by the model checkpoint. Signed-off-by: Elena Rastorgueva --- docs/source/speechlm2/streaming_inference.rst | 10 +++------- .../nemo_inference_pipelines/conf/s2s_streaming.yaml | 8 -------- .../inference/model_wrappers/config_overrides.py | 8 +------- .../nemotron_voicechat_inference_wrapper.py | 1 - .../nemo_inference_pipelines/test_config_overrides.py | 10 +--------- 5 files changed, 5 insertions(+), 32 deletions(-) diff --git a/docs/source/speechlm2/streaming_inference.rst b/docs/source/speechlm2/streaming_inference.rst index 0abd612a9ccc..c9e38db45e65 100644 --- a/docs/source/speechlm2/streaming_inference.rst +++ b/docs/source/speechlm2/streaming_inference.rst @@ -648,11 +648,6 @@ reported at load time rather than silently doing nothing. - Applied inside EarTTS on both paths. The converted EarTTS always substitutes codec silence on EOS and has no flag, so it cannot honour ``false``. - * - ``inference_top_p_or_k``, ``inference_noise_scale``, ``inference_guidance_scale`` - - yes - - no - - The vLLM EarTTS takes sampling from the converted checkpoint and - ``vllm_omni_config`` instead. * - ``deterministic`` - yes - no @@ -699,8 +694,9 @@ EarTTS classifier-free guidance uses two explicit requests in the same engine. They have independent KV caches, but a custom scheduler advances them in lockstep. The unconditional request replaces text conditioning with the checkpoint's ``null_emb``; the MaskGIT sampler applies the native guidance -formula and returns only the conditional stream's codes. Configure it with -``vllm_omni_config.guidance_enabled`` and ``guidance_scale``. +formula and returns only the conditional stream's codes. Whether guidance is +enabled can be overridden with ``vllm_omni_config.guidance_enabled``; its scale +comes from the converted TTS checkpoint. Perception, the audio codec and tokenization stay on PyTorch in every engine pairing. diff --git a/examples/speechlm2/nemo_inference_pipelines/conf/s2s_streaming.yaml b/examples/speechlm2/nemo_inference_pipelines/conf/s2s_streaming.yaml index cab10543b9a7..a972c97d8d8f 100644 --- a/examples/speechlm2/nemo_inference_pipelines/conf/s2s_streaming.yaml +++ b/examples/speechlm2/nemo_inference_pipelines/conf/s2s_streaming.yaml @@ -59,7 +59,6 @@ s2s: nemotron_dtype: float32 # dtype for the converted Nemotron LLM checkpoint eartts_precompute_batch_size: 256 # batch size for baking out the subword lookup table guidance_enabled: null # null -> converted eartts/config.json enable_guidance - guidance_scale: null # null -> converted eartts/config.json guidance_scale log_stats: false stage_init_timeout: 600 # seconds to wait for both stage children to come up step_timeout: 60.0 # per-step timeout on the synchronous side (seconds) @@ -120,17 +119,10 @@ s2s: inference_bos_boost: null # Boost agent text BOS logit inference_eos_boost: null # Boost agent text EOS logit - # EarTTS behaviour. Applied by EarTTS internally, so these affect the TTS - # component selected by tts_engine_type. inference_force_speech_silence_on_eos: null # null -> checkpoint value (model default: true). # Substitutes codec silence as the acoustic input of # the step whose text token is EOS. vllm_omni always # does this and cannot honour false. - inference_top_p_or_k: null # null -> checkpoint value (model default: 0.8). native TTS only - inference_noise_scale: null # null -> checkpoint value (model default: 0.8). native TTS only - inference_guidance_scale: null # null -> checkpoint value (model default: 0.5). native TTS only; - # the vllm_omni TTS reads vllm_omni_config.guidance_scale, which - # itself falls back to the converted eartts/config.json value. system_prompt: ??? diff --git a/nemo/collections/speechlm2/inference/model_wrappers/config_overrides.py b/nemo/collections/speechlm2/inference/model_wrappers/config_overrides.py index 47007043df91..4dc1bf15311a 100644 --- a/nemo/collections/speechlm2/inference/model_wrappers/config_overrides.py +++ b/nemo/collections/speechlm2/inference/model_wrappers/config_overrides.py @@ -69,9 +69,6 @@ # DuplexEARTTS.infer_codes_one_step (flag-gated, defaults to True) and the # vLLM EarTTS preprocess (unconditional). See VLLM_FORCES_TRUE. "inference_force_speech_silence_on_eos", - "inference_top_p_or_k", - "inference_noise_scale", - "inference_guidance_scale", ) COMPONENT_OF = {**{key: LLM for key in LLM_KEYS}, **{key: TTS for key in TTS_KEYS}} @@ -83,9 +80,6 @@ # The run is correct, but the setting does nothing: warn. VLLM_IGNORES = { - "inference_top_p_or_k": (TTS, "read by DuplexEARTTS._get_generation_config"), - "inference_noise_scale": (TTS, "read by DuplexEARTTS._get_generation_config"), - "inference_guidance_scale": (TTS, "read by DuplexEARTTS._get_generation_config"), "use_llm_cache": (LLM, "vLLM always keeps a paged KV cache"), "use_tts_torch_compile": (TTS, "vLLM compiles inside the engine"), "use_tts_subword_cache": ( @@ -95,7 +89,7 @@ } # These boolean flags request an optimization only when enabled. Their false -# values are no-ops, unlike numeric sampling values such as a zero noise scale. +# values are no-ops. _IGNORED_ENABLE_FLAGS = frozenset({"use_llm_cache", "use_tts_torch_compile", "use_tts_subword_cache"}) # vLLM does this unconditionally: it can honour True but not False. Warn only diff --git a/nemo/collections/speechlm2/inference/model_wrappers/nemotron_voicechat_inference_wrapper.py b/nemo/collections/speechlm2/inference/model_wrappers/nemotron_voicechat_inference_wrapper.py index 03d82bc8444a..0a638c4d2d16 100644 --- a/nemo/collections/speechlm2/inference/model_wrappers/nemotron_voicechat_inference_wrapper.py +++ b/nemo/collections/speechlm2/inference/model_wrappers/nemotron_voicechat_inference_wrapper.py @@ -191,7 +191,6 @@ def __init__(self, model_cfg: DictConfig): self.omni_wrapper_dir: str | None = None self.omni_speaker_latent: torch.Tensor | None = None self.omni_guidance_enabled = True - self.omni_guidance_scale = 0.5 # Sampling parameters (defaults match s2s_streaming.yaml) self.top_p = float(model_cfg.get("top_p", 0.5)) diff --git a/tests/collections/speechlm2/nemo_inference_pipelines/test_config_overrides.py b/tests/collections/speechlm2/nemo_inference_pipelines/test_config_overrides.py index 91bc59a1e6d5..141f1ca04ea9 100644 --- a/tests/collections/speechlm2/nemo_inference_pipelines/test_config_overrides.py +++ b/tests/collections/speechlm2/nemo_inference_pipelines/test_config_overrides.py @@ -63,16 +63,14 @@ def test_overrides_land_in_the_config_their_consumer_reads(model, warnings): { "inference_user_pad_boost": 0.8, "force_turn_taking": True, - "inference_top_p_or_k": 0.7, }, llm_engine_type="native", tts_engine_type="native", ) - # DuplexSTTModel reads its own cfg; DuplexEARTTS reads its own. + # DuplexSTTModel reads its own cfg. assert model.stt_model.cfg["inference_user_pad_boost"] == 0.8 assert model.stt_model.cfg["force_turn_taking"] is True - assert model.tts_model.cfg["inference_top_p_or_k"] == 0.7 # Untouched by this call, and still reported as the effective value. assert model.stt_model.cfg["inference_pad_boost"] == 1.5 assert effective["inference_pad_boost"] == 1.5 @@ -84,12 +82,6 @@ def test_overrides_land_in_the_config_their_consumer_reads(model, warnings): [ # Boosts and turn-taking work on both backends, so they stay quiet. ({"inference_user_pad_boost": 0.8, "inference_pad_boost": 0.3, "force_turn_taking": True}, "vllm_omni", None), - # vLLM EarTTS takes sampling from the converted checkpoint instead. - ({"inference_noise_scale": 0.9}, "vllm_omni", "inference_noise_scale"), - ({"inference_noise_scale": 0.0}, "vllm_omni", "inference_noise_scale"), - ({"inference_top_p_or_k": 0.0}, "vllm_omni", "inference_top_p_or_k"), - ({"inference_guidance_scale": 0.0}, "vllm_omni", "inference_guidance_scale"), - ({"inference_noise_scale": 0.9}, "native", None), # It forces codec silence on EOS unconditionally: True is honoured, # False cannot be. ({"inference_force_speech_silence_on_eos": False}, "vllm_omni", "force_speech_silence"), From e77e1ac127a63150989a7dcd4b49bb9eb0b4cabd Mon Sep 17 00:00:00 2001 From: Elena Rastorgueva Date: Tue, 1 Sep 2026 22:59:10 +0000 Subject: [PATCH 6/6] test(speechlm2): trim redundant VoiceChat coverage Keep distinct config, engine, and tiny-model parity checks while avoiding repeated cases and a second real-checkpoint test. Signed-off-by: Elena Rastorgueva --- .../test_config_overrides.py | 3 -- .../test_engine_selection.py | 8 ---- ...test_nemotron_voicechat_pipeline_parity.py | 48 +++++-------------- 3 files changed, 12 insertions(+), 47 deletions(-) diff --git a/tests/collections/speechlm2/nemo_inference_pipelines/test_config_overrides.py b/tests/collections/speechlm2/nemo_inference_pipelines/test_config_overrides.py index 141f1ca04ea9..bf9844a8da16 100644 --- a/tests/collections/speechlm2/nemo_inference_pipelines/test_config_overrides.py +++ b/tests/collections/speechlm2/nemo_inference_pipelines/test_config_overrides.py @@ -80,8 +80,6 @@ def test_overrides_land_in_the_config_their_consumer_reads(model, warnings): @pytest.mark.parametrize( ("overrides", "tts_engine_type", "expected"), [ - # Boosts and turn-taking work on both backends, so they stay quiet. - ({"inference_user_pad_boost": 0.8, "inference_pad_boost": 0.3, "force_turn_taking": True}, "vllm_omni", None), # It forces codec silence on EOS unconditionally: True is honoured, # False cannot be. ({"inference_force_speech_silence_on_eos": False}, "vllm_omni", "force_speech_silence"), @@ -112,7 +110,6 @@ def test_a_backend_reports_exactly_the_keys_it_ignores(model, warnings, override ({"use_llm_cache": True}, ("vllm_omni", "native"), "use_llm_cache"), ({"use_llm_cache": True}, ("native", "native"), None), ({"use_tts_torch_compile": True}, ("native", "vllm_omni"), "use_tts_torch_compile"), - ({"use_tts_subword_cache": True}, ("native", "vllm_omni"), "use_tts_subword_cache"), # A falsy value is already a no-op, so it needs no warning. ({"use_tts_torch_compile": False}, ("native", "vllm_omni"), None), ], diff --git a/tests/collections/speechlm2/nemo_inference_pipelines/test_engine_selection.py b/tests/collections/speechlm2/nemo_inference_pipelines/test_engine_selection.py index ada03274977c..b53720f1f309 100644 --- a/tests/collections/speechlm2/nemo_inference_pipelines/test_engine_selection.py +++ b/tests/collections/speechlm2/nemo_inference_pipelines/test_engine_selection.py @@ -35,12 +35,6 @@ {"stt_model.llm.", "tts_model.tts_model."}, set(), ), - ( - {"llm_engine_type": NATIVE, "tts_engine_type": VLLM}, - (NATIVE, VLLM), - {"tts_model.tts_model."}, - {"stt_model.llm."}, - ), ( {"tts_engine_type": VLLM}, (NATIVE, VLLM), @@ -84,8 +78,6 @@ def test_vllm_selection_is_named_as_not_implemented(): silently running native.""" with pytest.raises(NotImplementedError, match="not implemented in this PR"): reject_unimplemented_vllm(VLLM, NATIVE) - with pytest.raises(NotImplementedError, match="not implemented in this PR"): - reject_unimplemented_vllm(NATIVE, VLLM) from nemo.collections.speechlm2.inference.model_wrappers.backend.vllm.llm import VllmLLM diff --git a/tests/collections/speechlm2/nemo_inference_pipelines/test_nemotron_voicechat_pipeline_parity.py b/tests/collections/speechlm2/nemo_inference_pipelines/test_nemotron_voicechat_pipeline_parity.py index 991c53ecae2b..3089fbd5d66c 100644 --- a/tests/collections/speechlm2/nemo_inference_pipelines/test_nemotron_voicechat_pipeline_parity.py +++ b/tests/collections/speechlm2/nemo_inference_pipelines/test_nemotron_voicechat_pipeline_parity.py @@ -14,10 +14,10 @@ """Offline vs. incremental inference parity tests for NemotronVoiceChat. -``test_parity_tiny_model`` (cache × prompt) and -``test_parity_tiny_function_model_without_asr`` run on random-weight models. -``test_parity`` does one pass on ``nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`` -(downloaded into the Hugging Face cache if needed). +``test_parity_tiny_model`` covers the baseline and prompt-plus-cache paths; +``test_parity_tiny_function_model_without_asr`` covers function-token feedback. +Both use random-weight checkpoints so the suite can exercise the parity +invariant without loading the public 11B. Run from the NeMo repo root (use ``-s`` to see live progress):: @@ -260,15 +260,16 @@ def _build_parity_pipeline( @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") -@pytest.mark.parametrize("use_llm_cache", [False, True], ids=["no_cache", "llm_cache"]) -@pytest.mark.parametrize("system_prompt", [None, MOCK_SYSTEM_PROMPT], ids=["no_prompt", "prompt"]) -def test_parity_tiny_model(build_pipeline, tiny_model_artifacts, use_llm_cache, system_prompt): +@pytest.mark.parametrize( + ("system_prompt", "use_llm_cache"), + [(None, False), (MOCK_SYSTEM_PROMPT, True)], + ids=["no_prompt-no_cache", "prompt-llm_cache"], +) +def test_parity_tiny_model(build_pipeline, tiny_model_artifacts, system_prompt, use_llm_cache): """Offline/incremental parity with a tiny random-weight model. - Running both cache settings against the same offline reference is what - pins the invariant that the native KV cache is a speed path and not a - different model: if either matched offline and the other did not, one of - these cases would fail. + The two cases cover the simplest path and the opposite prompt-plus-cache + corner without paying for every combination. """ model_dir, audio_path, _ = tiny_model_artifacts pipeline = _build_parity_pipeline( @@ -296,28 +297,3 @@ def test_parity_tiny_function_model_without_asr(build_pipeline, tiny_function_mo report = run_parity_check(pipeline, audio_path, system_prompt=MOCK_SYSTEM_PROMPT) assert report["asr_token_comparison"]["match"] is None assert_parity(report, strict=True, atol=0.0) - - -# --------------------------------------------------------------------------- -# Public 11B — one load, not the cache×prompt matrix -# --------------------------------------------------------------------------- - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") -def test_parity(build_pipeline, hf_voicechat_11b, voicechat_audio_path, voicechat_speaker_name): - """Offline/incremental parity on ``nvidia/NVIDIA-NemotronLabs-VoiceChat-11B``.""" - pipeline = _build_parity_pipeline( - build_pipeline, - hf_voicechat_11b, - voicechat_audio_path, - tempfile.mkdtemp(prefix="parity-11b-"), - { - "s2s": { - "system_prompt": MOCK_SYSTEM_PROMPT, - "speaker_name": voicechat_speaker_name, - } - }, - ) - report = run_parity_check(pipeline, voicechat_audio_path, system_prompt=MOCK_SYSTEM_PROMPT) - assert report["asr_token_comparison"]["match"] is None - assert_parity(report, strict=True, atol=0.0)