Add per-stream speaker limits to Sortformer sessions

[tianyi-zhang-02]

Important

This PR is intentionally stacked on #16174. Both branches include current main (de26b3696); please do not use the Update branch button.

What does this PR do?

Add a fixed per-session or per-row speaker-channel limit to raw-audio Sortformer streaming sessions.

Collection: ASR / Speaker Tasks

Changelog

• Add max_speakers to create_streaming_session() as a scalar, per-row sequence, or integer tensor.
• Validate limits against the checkpoint's speaker capacity and keep them in StreamingSortformerState across reset/device moves.
• Zero disabled channels before returned output, FIFO updates, speaker-cache scoring, and compression.
• Use the original all-channel activity only as a silence veto. This keeps a frame where a disabled channel fired out of mean_sil_emb without allowing that channel to mutate cache/output state.
• Preserve existing behavior when max_speakers=None.
• Cover heterogeneous batches, low/high-resolution output, cache/FIFO state, reset, finalized rows, and invalid inputs.

Usage

session = model.create_streaming_session(batch_size=2, max_speakers=[2, 4])
predictions, prediction_lengths = session.diarize_step(
    audio_chunks,
    audio_chunk_lengths=audio_chunk_lengths,
    is_final=is_final,
)

Motivation and semantics

Output-only masking is insufficient because disabled predictions can affect later chunks through FIFO and speaker-cache state. Masking before every state update fixes that, but it must not turn a speech frame into apparent silence. The implementation therefore separates the two concerns: masked predictions drive output/cache state, while pre-mask activity is consulted only when deciding whether a popped frame is truly silent.

Closes #15711.

Dependency

This is stacked on #16174, which introduces the raw-audio SortformerStreamingSession. The current head includes #16174 at 58357154c and the feature/fix at b07e61e90.

Testing

• Deterministic negative control for the issue author's finding: mask-only behavior incorrectly counted 2 silence frames with mean [3, 5]; the fixed path counts only the true silence frame with mean [5, 7].
• pytest -q tests/collections/speaker_tasks/test_diar_sortformer_models.py tests/collections/speaker_tasks/test_diar_sortformer_modules.py tests/collections/speaker_tasks/utils/test_sortformer_utils.py --cpu — 228 passed.
• The new silence-profile regression covers both synchronous and asynchronous state updates and asserts that disabled cache/output channels remain exactly zero.
• pre-commit run --files ... on the five changed files.
• git diff --check.
• Current head b07e61e90 passed 13/13 named CUDA checks on RTX 4090 / SM89 with Python 3.13.14, PyTorch 2.11.0+cu130, and CUDA 13.0. The deterministic sync/async controls reproduced the mask-only pollution and verified the corrected silence mean on CUDA.
• The same run restored the official nvidia/diar_streaming_sortformer_4spk-v2.1 checkpoint and exercised a heterogeneous [2, 4] batch with staggered chunks, early row-0 finalization, and later zero-length row-0 input. Across 16 asynchronous updates, row 0's disabled output/FIFO/cache channels stayed exact zero, row 1 retained channels 2/3, pre-mask activity reached the dedicated silence path, and all inspected state/update tensors were on cuda:0.
• The bounded run completed in 11.156 s with 930.2 MiB peak allocated and 1190.0 MiB peak reserved CUDA memory. No install, model download, or RunPod was used.

The official-checkpoint smoke used synthetic audio and did not itself produce a disabled-only frame above the silence threshold. The triggering-state proof comes from the deterministic CUDA controls; the reporter's 8.1% to 6.7% DER measurement is not claimed here because their audio and downstream pipeline are unavailable.

Before your PR is "Ready for review"

Pre checks:

• Read and followed the contributor guidelines.
• Added regression and edge-case tests for the new behavior.
• Updated the Streaming Sortformer usage documentation.
• This change does not affect an optional-install component.

PR Type:

• New Feature
• Bugfix
• Documentation

Additional Information

• Related issue: support max_speakers configuration in diar_streaming_sortformer_4spk-v2.1 #15711
• Base API: Add raw-audio streaming sessions for Sortformer #16174

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