[speechlm2] Add streaming inference pipeline for NemotronVoiceChat

[erastorgueva-nv]

Important

The Update branch button must only be pressed in very rare occassions.
An outdated branch is never blocking the merge of a PR.
Please reach out to the automation team before pressing that button.

What does this PR do ?

Add a streaming (real-time, chunk-by-chunk) inference pipeline for NemotronVoiceChat,
following the same architecture as the NeMo ASR Inference Pipelines.

Collection: speechlm2

Changelog

• Add StreamingS2SPipeline with generate_step() API for both batch file processing and server integration
• Add NemotronVoicechatInferenceWrapper with infer_one_step() implementing perception → LLM → TTS → codec decode
• Add perception cache with optional CUDA graph support for cache-aware streaming encoding
• Add S2SPipelineBuilder factory and Hydra config (s2s_streaming.yaml) for easy setup
• Add state management: slot-based S2SContextManager for decode state lifecycle, S2SStreamingState for output accumulation
• Add s2s_streaming_infer.py entry script for batch inference on files/manifests
• Extend DuplexSTTModel: KV cache support for Nemotron hybrid Mamba/Attention (with monkey-patches for upstream HF bugs), save_pretrained with tokenizer export, function head, ASR logit boosts, cache_position forwarding
• Speed up model loading: meta device initialization, skip codec silence token computation when codec has random weights
• Fix sampling: nan/inf check before top-p filtering, vectorized repetition penalty
• Fix byte-level BPE decoding and BOS/EOS preservation in text output
• Refactor tests: shared conftest.py fixtures, offline-vs-streaming parity test, no-crash config sweep
• Add docs: streaming_inference.rst with architecture, config reference, and server integration guide

Modifications to more general code - FYI @kevinhu-nv @Edresson

• Extend NemotronVoiceChat: from_pretrained supports loading from HF-format checkpoint with llm_artifacts/
• Extend EarTTSModel: vectorized depth-sum, precomputed RVQ schedule, optional torch.compile, subword cache
• Patches: _patch_nemotron_cache_bugs and _patch_nemotron_block_forward methods in DuplexSTTModel are patching bugs in the HF Nemotron model code so we can get the KV caching to work. The patches seem to work for me, though I wonder if we can use more up-to-date code that doesn't have the patches.

Usage

python examples/speechlm2/nemo_inference_pipelines/s2s_streaming_infer.py \
    audio_file=/path/to/audio.wav \
    s2s.model_path=/path/to/checkpoint \
    s2s.speaker_name="<name>" \
    s2s.engine_type="native" \
    streaming.chunk_size_in_secs=0.08 \
    streaming.buffer_size_in_secs=1.68

from nemo.collections.speechlm2.inference import S2SPipelineBuilder

pipeline = S2SPipelineBuilder.build_pipeline(cfg)
output = pipeline.run(audio_filepaths, options=options)

GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

• New Feature
• Bugfix
• Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

• Related to # (issue)

[pzelasko]

Finalized my first review pass :)

[pzelasko]

Is it possible to support different top_p/temperature/repetition_penalty in different examples in a batch without using a for loop over them?

Is there a strong motivation to support that? Or could we expect the user to define session-level decoding parameters and have easier batching?

[erastorgueva-nv]

Yes, I think it will be possible. We can implement it once we support B > 1?
I think doing "session-level decoding parameters" would introduce an uncomfortable constraint - I would rather have it so that each stream can have its own sampling parameters

[pzelasko]

Is this compatible with transformers v5 or v4? Or both? We already bumped NeMo to v5

[erastorgueva-nv]

The meta device initialization is handled by PyTorch's torch.device('meta') context manager, which is independent of transformers. The transformers APIs used here (AutoConfig.from_pretrained, AutoModelForCausalLM.from_config/from_pretrained) are stable across v4 and v5.

[pzelasko]

Since so much logic is dedicated to Nemotron v2 in VoiceChat, shouldn't we test against Nemotron v2 as well? Or are you concerned it will take very long time to load in CI?

[erastorgueva-nv]

Exactly, I used tiny llama to keep model loading relatively quick.
Note that there are some tests in this PR that use the full voicechat checkpoint. They are only run if you specify the path to the local checkpoint. Once the checkpoint is public, I plan to make the test use that public checkpoint. So the test can easily run on CI, not just locally. And I believe that will satisfy the requirement of making sure we do test with nemotron v2

[pzelasko]

OK makes sense. If you want to add proper regression/functional tests, follow the pattern in tests/functional_tests and also please add a nightly E2E test in tests/e2e_tests

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[pzelasko]

Thanks, great work, it's greatly improved!
Left a few comments, my main concern remains hardcoding batch_size==1 in some asserts, while the entire pipeline looks capable of processing batch_size>1 at least for an offline evaluation (online inference is a separate story)

[pzelasko]

Quick Start should live at the top of this page

[pzelasko]

flagging - is this still true?

[pzelasko]

It should explain that it's not the typical padding used for batching, but extra silence following user's recording in order to allow the model to generate response.

[pzelasko]

why is it called model.py and not llm.py like in vllm/ directory?

[pzelasko]

Can you explain where the 10ms is coming from? I thought cache-aware has only "regular" chunks of 80ms

[pzelasko]

Is this still a limitation? The code seems capable of bs>1 inference

[pzelasko]

In that case could we split this up into a library module (here) and a script (the best place seems to be examples/speechlm2/voicechat/to_vllm.py)?

This repo doesn't have scripts-as-importable-modules-living-in-library pattern (I'm sure it has its own name)

[pzelasko]

Can we split up this script in the same way as the one above (ideally make to_vllm.py just work with both types of checkpoints)

[pzelasko]

This seems like a dangerous constant to hardcode, can't we pass model.token_equivalent_duration directly (not sure if this property has been exposed to voicechat class but it's in AudioPerceptionModule and some SALM and S2S classes IIRC)

[pzelasko]

OK makes sense. If you want to add proper regression/functional tests, follow the pattern in tests/functional_tests and also please add a nightly E2E test in tests/e2e_tests