Core updates for Streaming Sortformer inference

[ipmedenn]

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 ?

This PR introduces a large set of mostly inference-side updates for the upcoming 8-speaker Sortformer model release.

The main new feature is high-resolution output. Previously, Sortformer prediction resolution was tied to the encoder resolution, which is 80 ms in the standard model configurations. High-resolution mode decouples the two: the encoder continues operating at its configured resolution, while a trainable subpixel upsampling layer produces predictions at a native 10 ms resolution. These predictions can then be returned at any configured multiple of 10 ms, such as 20, 30, 80, or 160 ms. Previously released models retain their existing output resolution and behavior by default.

To make training practical for higher speaker counts, PIL and ATS targets are now generated using Hungarian assignment instead of exhaustive permutation search.

Asynchronous inference has been substantially refactored and vectorized, improving efficiency and fixing the ragged-batch correctness issues described in #16002.

The PR also adds an optional learnable silence embedding for speaker-cache compression; improves the inference script with vectorized binarization and post-processing, profiling, prediction-tensor caching, and optional encoder compilation; and fixes streaming ONNX export.

Collection: ASR / Speaker diarization

Changelog

Output resolution

• Add high_resolution mode with a trainable subpixel upsampling layer that produces predictions at the preprocessor frame resolution (10 ms in the standard configurations).
• Add output_subsampling_factor to control the returned prediction frame step independently of the encoder subsampling factor.
• Align both legacy and Lhotse diarization targets with the configured output resolution.
• Use the configured output resolution in diarization post-processing.
• Reject streaming configurations where chunk boundaries would split an output-downsampling window.
• Ensure the Sortformer output resolution matches the ASR encoder resolution inside ParallelExpertEncoder.

Sortformer model and training

• Replace exhaustive PIL/ATS speaker-permutation search with Hungarian assignment.
• Add configurable arrival-time tolerance to ATS target generation.
• Support different numbers of reference speakers and prediction streams.
• Retain reference speakers beyond the model capacity during evaluation so they are counted as false negatives.
• Add an optional learnable silence embedding for speaker-cache compression; the running mean remains the default.
• Retain the legacy hidden_to_spks module for checkpoint compatibility while freezing its unused parameters.
• Allow transformer_encoder to be omitted from new model configurations.
• Remove the need for find_unused_parameters=True, allowing standard DDP.

Streaming correctness and asynchronous inference

• Track speaker-cache compression state and refresh cached predictions from the current forward pass before the first compression in both synchronous and asynchronous streaming.
• Base FIFO eviction on each sample's actual chunk length in ragged asynchronous batches.
• Vectorize asynchronous state updates, eliminating per-row Python loops and scalar device-to-host synchronization.
• Vectorize state concatenation and high-resolution chunk-prediction extraction for ragged asynchronous batches.
• Flush the remaining FIFO when an asynchronous row is finalized with a zero-length chunk.
• Add optional async_pad_to_max behavior to keep encoder input shapes fixed at their maximum capacity.
• Add optional async_desync_updates behavior to randomize initial FIFO eviction and stagger speaker-cache updates across batch rows.
• Add regression coverage verifying that a short stream is unaffected by longer streams in the same batch.

Inference utilities and export

• Replace the boolean save_preds_tensors option with the explicit out_preds_tensors cache path. Validate existing cache metadata and write new caches atomically; use overwrite_preds_tensors=true to replace an existing cache.
• Add detailed inference profiling, including timings for preprocessing, streaming-state concatenation, encoding, cache compression, and state updates.
• Change the default inference precision used by e2e_diarize_speech.py from FP32 to BF16. Set precision=32 to retain the previous behavior.
• Vectorize diarization binarization and post-processing, substantially reducing post-processing overhead.
• Add optional torch.compile support for the encoder through compile_encoder.
• Add model-sized default input examples for streaming export.
• Allow callers to provide an explicit input example for streaming export.
• Keep the concatenation width fixed during ONNX tracing while allowing logical cache, FIFO, and chunk lengths to vary at runtime.
• Add an ONNX Runtime regression test covering empty, partially filled, and fully populated streaming states.

Unit tests

• high-resolution training, inference, and streaming;
• output-length calculation and target-resolution alignment;
• models configured without transformer_encoder;
• legacy checkpoint loading;
• learnable silence embeddings;
• Hungarian PIL/ATS assignment at higher speaker counts and with different numbers of reference and predicted speakers;
• evaluation with more reference speakers than the model output capacity;
• ParallelExpertEncoder output-resolution alignment;
• prediction-cache validation;
• first-compression correctness and vectorized asynchronous state updates for ragged batches;
• asynchronous row finalization and zero-length chunk flushing;
• fixed-capacity asynchronous encoder padding;
• vectorized high-resolution asynchronous chunk-prediction extraction;
• streaming ONNX export with empty, partially filled, and fully populated runtime states.

Usage

Model configuration

Existing Sortformer model configurations remain backward-compatible. When high_resolution is omitted or set to false, the native prediction resolution remains tied to the encoder subsampling factor, and output_subsampling_factor defaults to that same value.

For a standard 10 ms preprocessor frame step, high-resolution training with 10 ms returned predictions can be configured as:

model:
  high_resolution: true
  output_subsampling_factor: 1

In high-resolution mode, output_subsampling_factor is expressed in preprocessor frames. With a 10 ms preprocessor frame step, values of 1, 2, 3, 8, and 16 correspond to returned prediction steps of 10, 20, 30, 80, and 160 ms, respectively.

By default, speaker-cache compression uses the running mean silence embedding. A learnable silence embedding can be enabled with:

model:
  sortformer_modules:
    use_learnable_sil_emb: true

For new architectures that do not require the secondary transformer encoder, the transformer_encoder section may be omitted. Previously released model configurations containing this section continue to instantiate it normally.

Streaming output constraints

Streaming output downsampling is applied independently to each chunk, so chunk boundaries must align with complete downsampling windows. Configurations whose chunk length is incompatible with output_subsampling_factor are rejected with a descriptive error.

Inference script options

The inference script e2e_diarize_speech.py provides the following asynchronous-streaming, caching, and performance controls:

precision: bf16
out_preds_tensors: null
overwrite_preds_tensors: false
async_streaming: false
async_pad_to_max: false
async_desync_updates: false
compile_encoder: false
profile_inference: true

async_streaming: true enables asynchronous batched state updates; async_pad_to_max and async_desync_updates only affect this mode.

async_pad_to_max: true keeps encoder inputs at a fixed maximum shape, trading additional padded computation for stable input shapes.

async_desync_updates: true randomizes the initial FIFO eviction to stagger speaker-cache updates across batch rows.

out_preds_tensors specifies a prediction-cache path and enables both loading and saving of prediction tensors. Existing caches are validated before reuse. Set overwrite_preds_tensors: true to recompute and replace an existing cache.

The default inference precision is BF16. Set precision: 32 to retain the previous FP32 behavior.

Setting compile_encoder: true applies torch.compile(..., dynamic=True) to the encoder only; it does not compile the complete Sortformer model.

Detailed inference profiling is enabled by default and can be disabled with profile_inference: false.

ONNX export

Streaming export uses a model-sized default input example, and callers may provide an explicit input example when different physical capacities are required.

The exported ONNX graph has a fixed batch size and fixed physical cache, FIFO, and chunk capacities. Their logical lengths may vary independently for each batch row at runtime, including empty, partially filled, and fully populated states.

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? Necessary unit tests were added.
• Did you add or update any necessary documentation? New configuration fields include inline 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

• Fixes #16002.

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