From 9825998fd9a379ff0bcf1f6a0c540230bfb02017 Mon Sep 17 00:00:00 2001
From: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Date: Mon, 31 Aug 2026 01:09:08 -0400
Subject: [PATCH 1/5] Add raw-audio streaming Sortformer session

Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
---
 examples/speaker_tasks/diarization/README.md  |  16 ++
 .../asr/models/sortformer_diar_models.py      |  13 ++
 .../asr/parts/utils/streaming_sortformer.py   | 193 ++++++++++++++++++
 .../test_diar_sortformer_models.py            |  61 ++++++
 4 files changed, 283 insertions(+)
 create mode 100644 nemo/collections/asr/parts/utils/streaming_sortformer.py

diff --git a/examples/speaker_tasks/diarization/README.md b/examples/speaker_tasks/diarization/README.md
index 886abb73b8bd..6bd9d20187e1 100644
--- a/examples/speaker_tasks/diarization/README.md
+++ b/examples/speaker_tasks/diarization/README.md
@@ -86,6 +86,22 @@ diar_model.sortformer_modules.spkcache_update_period = 300
 predicted_segments = diar_model.diarize(audio="/path/to/audio.wav", batch_size=1)
 ```
 
+For a live mono waveform, create one session per audio stream and push float tensors at the model sample rate. A step
+can return no frames until the configured chunk and right context are available. Mark the last chunk with
+``is_final=True`` to flush the tail, or call ``reset()`` before reusing the session for a new stream. The session
+disables dither and feature padding while extracting each chunk, then applies the checkpoint's feature normalization
+over the complete model input window.
+
+```python
+import torch
+
+session = diar_model.create_streaming_session()
+for audio_chunk in audio_stream:
+    new_speaker_probabilities = session.diarize_step(audio_chunk)
+tail_speaker_probabilities = session.diarize_step(torch.empty(0), is_final=True)
+session.reset()
+```
+
 Diarization Error Rate (DER) with post-processing — all evaluations include overlapping speech:
 
 | Dataset | Collar | 30.4s latency | 10.0s latency | 1.04s latency | 0.32s latency |
diff --git a/nemo/collections/asr/models/sortformer_diar_models.py b/nemo/collections/asr/models/sortformer_diar_models.py
index 5ce5c31a3c3f..79cee4805fe0 100644
--- a/nemo/collections/asr/models/sortformer_diar_models.py
+++ b/nemo/collections/asr/models/sortformer_diar_models.py
@@ -614,6 +614,19 @@ def process_signal(self, audio_signal, audio_signal_length):
             torch.cuda.empty_cache()
         return processed_signal, processed_signal_length
 
+    def create_streaming_session(self):
+        """Create an independent high-level raw-audio streaming session.
+
+        The returned session accepts arbitrarily sized mono waveform chunks through ``diarize_step()`` and owns all
+        preprocessing buffers and Sortformer cache state. Create one session per concurrent audio stream.
+
+        Returns:
+            SortformerStreamingSession: A new session bound to this model.
+        """
+        from nemo.collections.asr.parts.utils.streaming_sortformer import SortformerStreamingSession
+
+        return SortformerStreamingSession(self)
+
     def forward(
         self,
         audio_signal,
diff --git a/nemo/collections/asr/parts/utils/streaming_sortformer.py b/nemo/collections/asr/parts/utils/streaming_sortformer.py
new file mode 100644
index 000000000000..9621e986d1f3
--- /dev/null
+++ b/nemo/collections/asr/parts/utils/streaming_sortformer.py
@@ -0,0 +1,193 @@
+# 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.
+
+import copy
+import math
+from typing import TYPE_CHECKING
+
+import torch
+
+from nemo.collections.asr.parts.preprocessing.features import normalize_batch
+
+if TYPE_CHECKING:
+    from nemo.collections.asr.models.sortformer_diar_models import SortformerEncLabelModel
+
+__all__ = ['SortformerStreamingSession']
+
+
+class SortformerStreamingSession:
+    """Stateful raw-audio session for a streaming Sortformer model.
+
+    A session owns the feature and speaker-cache state for one mono audio stream. Incoming waveform chunks may have
+    arbitrary lengths. The session waits until a complete model chunk and its configured right context are available,
+    then returns only the newly committed speaker probabilities. Call :meth:`reset` before reusing a session for a new
+    stream.
+
+    Args:
+        model: Streaming ``SortformerEncLabelModel`` in evaluation mode.
+    """
+
+    def __init__(self, model: 'SortformerEncLabelModel'):
+        if not model.streaming_mode:
+            raise ValueError("SortformerStreamingSession requires a model with streaming_mode=True")
+        if model.training:
+            raise ValueError("SortformerStreamingSession requires an evaluation model; call model.eval() first")
+
+        self.model = model
+        self.device = model.device
+        self._normalization = model.preprocessor.featurizer.normalize
+        self._preprocessor = copy.deepcopy(model.preprocessor).to(self.device).eval()
+        self._preprocessor.featurizer.normalize = None
+        self._preprocessor.featurizer.dither = 0.0
+        self._preprocessor.featurizer.pad_to = 0
+
+        self._hop_length = self._preprocessor.hop_length
+        self._n_fft = self._preprocessor.featurizer.n_fft
+        self._stft_margin_frames = math.ceil((self._n_fft // 2 + 1) / self._hop_length) + 1
+        self._chunk_frames = model.sortformer_modules.chunk_len * model.encoder.subsampling_factor
+        self._left_context_frames = model.sortformer_modules.chunk_left_context * model.encoder.subsampling_factor
+        self._right_context_frames = model.sortformer_modules.chunk_right_context * model.encoder.subsampling_factor
+        self.reset()
+
+    @torch.inference_mode()
+    def diarize_step(self, audio_chunk: torch.Tensor, is_final: bool = False) -> torch.Tensor:
+        """Consume a mono waveform chunk and return newly committed speaker probabilities.
+
+        Args:
+            audio_chunk: Float waveform with shape ``(num_samples,)`` or ``(1, num_samples)`` at the model sample rate.
+            is_final: Flush the remaining audio and close the stream. No more audio can be supplied until ``reset``.
+
+        Returns:
+            Tensor with shape ``(1, new_frames, num_speakers)``. ``new_frames`` can be zero while the session waits
+            for a complete chunk and its right context.
+        """
+        if self._finalized:
+            raise RuntimeError("This streaming session is finalized; call reset() before supplying more audio")
+        if not isinstance(is_final, bool):
+            raise TypeError(f"is_final must be a boolean, got {type(is_final).__name__}")
+
+        audio_chunk = self._validate_audio_chunk(audio_chunk)
+        if audio_chunk.numel() > 0:
+            self._audio_buffer = torch.cat([self._audio_buffer, audio_chunk])
+            self._received_samples += audio_chunk.numel()
+
+        available_frames = self._available_feature_frames(is_final=is_final)
+        emitted = []
+        while self._next_feature_frame < available_frames:
+            central_end = min(self._next_feature_frame + self._chunk_frames, available_frames)
+            if not is_final and central_end + self._right_context_frames > available_frames:
+                break
+
+            feature_start = max(0, self._next_feature_frame - self._left_context_frames)
+            feature_end = min(central_end + self._right_context_frames, available_frames)
+            processed_signal = self._extract_feature_window(feature_start, feature_end)
+            processed_signal_length = torch.tensor([feature_end - feature_start], dtype=torch.long, device=self.device)
+            empty_preds = processed_signal.new_zeros((1, 0, self.model.sortformer_modules.n_spk))
+            self.streaming_state, chunk_preds = self.model.forward_streaming_step(
+                processed_signal=processed_signal.transpose(1, 2),
+                processed_signal_length=processed_signal_length,
+                streaming_state=self.streaming_state,
+                total_preds=empty_preds,
+                left_offset=self._next_feature_frame - feature_start,
+                right_offset=feature_end - central_end,
+            )
+            emitted.append(chunk_preds)
+            self._next_feature_frame = central_end
+
+        self._compact_audio_buffer()
+
+        if is_final:
+            self._finalized = True
+
+        if emitted:
+            return torch.cat(emitted, dim=1)
+        return torch.zeros(
+            (1, 0, self.model.sortformer_modules.n_spk),
+            dtype=next(self.model.parameters()).dtype,
+            device=self.device,
+        )
+
+    def reset(self) -> None:
+        """Clear buffered audio and model state so this session can process a new stream."""
+        self.streaming_state = self.model.sortformer_modules.init_streaming_state(
+            batch_size=1,
+            async_streaming=self.model.async_streaming,
+            device=self.device,
+        )
+        self._audio_buffer = torch.empty(0, dtype=torch.float32, device=self.device)
+        self._audio_buffer_start = 0
+        self._received_samples = 0
+        self._next_feature_frame = 0
+        self._finalized = False
+
+    def _available_feature_frames(self, is_final: bool) -> int:
+        sample_count = torch.tensor(self._received_samples, device=self.device)
+        offline_frames = int(self._preprocessor.featurizer.get_seq_len(sample_count).item())
+        if is_final:
+            return max(0, offline_frames)
+
+        stable_samples = self._received_samples - self._n_fft // 2
+        if stable_samples < 0:
+            return 0
+        stable_frames = stable_samples // self._hop_length + 1
+        return max(0, min(offline_frames, stable_frames))
+
+    def _compact_audio_buffer(self) -> None:
+        first_needed_frame = max(
+            0,
+            self._next_feature_frame - self._left_context_frames - self._stft_margin_frames,
+        )
+        first_needed_sample = first_needed_frame * self._hop_length
+        drop_samples = first_needed_sample - self._audio_buffer_start
+        if drop_samples > 0:
+            self._audio_buffer = self._audio_buffer[drop_samples:].clone()
+            self._audio_buffer_start = first_needed_sample
+
+    def _extract_feature_window(self, feature_start: int, feature_end: int) -> torch.Tensor:
+        segment_start_frame = max(0, feature_start - self._stft_margin_frames)
+        segment_start_sample = segment_start_frame * self._hop_length
+        segment_end_sample = min(
+            self._received_samples,
+            (feature_end + self._stft_margin_frames) * self._hop_length,
+        )
+        buffer_start = segment_start_sample - self._audio_buffer_start
+        buffer_end = segment_end_sample - self._audio_buffer_start
+        audio_signal = self._audio_buffer[buffer_start:buffer_end].unsqueeze(0)
+        audio_signal_length = torch.tensor([audio_signal.shape[1]], dtype=torch.long, device=self.device)
+        features, _ = self._preprocessor(input_signal=audio_signal, length=audio_signal_length)
+
+        local_start = feature_start - segment_start_frame
+        local_end = local_start + feature_end - feature_start
+        if features.shape[2] < local_end:
+            raise RuntimeError(
+                "Streaming preprocessor returned fewer feature frames than required: "
+                f"needed {local_end}, got {features.shape[2]}"
+            )
+        features = features[:, :, local_start:local_end]
+        feature_length = torch.tensor([features.shape[2]], dtype=torch.long, device=self.device)
+        if self._normalization:
+            features, _, _ = normalize_batch(features, feature_length, self._normalization)
+        return features
+
+    def _validate_audio_chunk(self, audio_chunk: torch.Tensor) -> torch.Tensor:
+        if not isinstance(audio_chunk, torch.Tensor):
+            raise TypeError(f"audio_chunk must be a torch.Tensor, got {type(audio_chunk).__name__}")
+        if audio_chunk.ndim == 2 and audio_chunk.shape[0] == 1:
+            audio_chunk = audio_chunk.squeeze(0)
+        if audio_chunk.ndim != 1:
+            raise ValueError(
+                "audio_chunk must contain one mono stream with shape (num_samples,) or (1, num_samples); "
+                f"got {tuple(audio_chunk.shape)}"
+            )
+        return audio_chunk.detach().to(device=self.device, dtype=torch.float32)
diff --git a/tests/collections/speaker_tasks/test_diar_sortformer_models.py b/tests/collections/speaker_tasks/test_diar_sortformer_models.py
index e437c46948f3..c028bde10712 100644
--- a/tests/collections/speaker_tasks/test_diar_sortformer_models.py
+++ b/tests/collections/speaker_tasks/test_diar_sortformer_models.py
@@ -267,6 +267,67 @@ def test_constructor(self, sortformer_model):
         instance2 = SortformerEncLabelModel.from_config_dict(confdict)
         assert isinstance(instance2, SortformerEncLabelModel)
 
+    @pytest.mark.unit
+    def test_raw_audio_streaming_session_is_independent_of_input_chunking(self):
+        model = _create_sortformer_model().eval()
+        model.streaming_mode = True
+        model.sortformer_modules.chunk_len = 2
+        model.sortformer_modules.chunk_left_context = 1
+        model.sortformer_modules.chunk_right_context = 1
+        model._check_streaming_parameters()
+        audio = torch.randn(8193)
+
+        one_shot_session = model.create_streaming_session()
+        one_shot_outputs = [
+            one_shot_session.diarize_step(audio),
+            one_shot_session.diarize_step(torch.empty(0), is_final=True),
+        ]
+        one_shot_preds = torch.cat(one_shot_outputs, dim=1)
+
+        chunked_session = model.create_streaming_session()
+        chunked_outputs = []
+        start = 0
+        for chunk_size in (17, 503, 1600, 81, 2999, 123, 2870):
+            end = min(start + chunk_size, audio.numel())
+            chunked_outputs.append(chunked_session.diarize_step(audio[start:end]))
+            start = end
+            if start == audio.numel():
+                break
+        if start < audio.numel():
+            chunked_outputs.append(chunked_session.diarize_step(audio[start:]))
+        chunked_outputs.append(chunked_session.diarize_step(torch.empty(0), is_final=True))
+        chunked_preds = torch.cat(chunked_outputs, dim=1)
+
+        assert one_shot_preds.shape[1] > 0
+        assert chunked_session._audio_buffer.untyped_storage().nbytes() < audio.untyped_storage().nbytes()
+        torch.testing.assert_close(chunked_preds, one_shot_preds)
+
+    @pytest.mark.unit
+    def test_raw_audio_streaming_session_reset_and_validation(self):
+        offline_model = _create_sortformer_model().eval()
+        with pytest.raises(ValueError, match="streaming_mode=True"):
+            offline_model.create_streaming_session()
+
+        model = _create_sortformer_model().eval()
+        model.streaming_mode = True
+        model.sortformer_modules.chunk_len = 2
+        model.sortformer_modules.chunk_left_context = 1
+        model.sortformer_modules.chunk_right_context = 1
+        model._check_streaming_parameters()
+        audio = torch.randn(4097)
+        session = model.create_streaming_session()
+
+        first_preds = session.diarize_step(audio, is_final=True)
+        with pytest.raises(RuntimeError, match="finalized"):
+            session.diarize_step(torch.empty(0))
+        session.reset()
+        second_preds = session.diarize_step(audio.unsqueeze(0), is_final=True)
+
+        torch.testing.assert_close(second_preds, first_preds)
+        with pytest.raises(ValueError, match="one mono stream"):
+            session.reset()
+            session.diarize_step(torch.randn(2, 100))
+
     @pytest.mark.unit
     @pytest.mark.parametrize(
         "stacking_factor, feature_shape, input_lengths, expected_encoded_lengths",

From 163deb06075cc8c514c6744bd10b7df03c751ab7 Mon Sep 17 00:00:00 2001
From: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Date: Wed, 2 Sep 2026 11:29:52 -0400
Subject: [PATCH 2/5] Address batched streaming session review

Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
---
 examples/speaker_tasks/diarization/README.md  |  24 +-
 .../asr/models/sortformer_diar_models.py      |  28 +-
 .../asr/parts/utils/sortformer_utils.py       | 285 +++++++++++++++++-
 .../asr/parts/utils/streaming_sortformer.py   | 193 ------------
 .../test_diar_sortformer_models.py            |  88 ++++--
 5 files changed, 373 insertions(+), 245 deletions(-)
 delete mode 100644 nemo/collections/asr/parts/utils/streaming_sortformer.py

diff --git a/examples/speaker_tasks/diarization/README.md b/examples/speaker_tasks/diarization/README.md
index 6bd9d20187e1..bba4d0e1a1f8 100644
--- a/examples/speaker_tasks/diarization/README.md
+++ b/examples/speaker_tasks/diarization/README.md
@@ -86,19 +86,25 @@ diar_model.sortformer_modules.spkcache_update_period = 300
 predicted_segments = diar_model.diarize(audio="/path/to/audio.wav", batch_size=1)
 ```
 
-For a live mono waveform, create one session per audio stream and push float tensors at the model sample rate. A step
-can return no frames until the configured chunk and right context are available. Mark the last chunk with
-``is_final=True`` to flush the tail, or call ``reset()`` before reusing the session for a new stream. The session
-disables dither and feature padding while extracting each chunk, then applies the checkpoint's feature normalization
-over the complete model input window.
+For live mono waveforms, create a session with a fixed number of streams and push a padded float tensor at the model
+sample rate together with its valid sample lengths. Each row owns independent preprocessing buffers and asynchronous
+Sortformer cache state. A step can return no frames for a row until its configured chunk and right context are
+available. Use a per-row ``is_final`` mask to flush streams independently; finalized rows remain in the batch with a
+zero input length while other rows continue. Call ``reset()`` before reusing the session. The session disables dither
+and feature padding while extracting each chunk, then applies the checkpoint's feature normalization over each
+complete model input window.
 
 ```python
 import torch
 
-session = diar_model.create_streaming_session()
-for audio_chunk in audio_stream:
-    new_speaker_probabilities = session.diarize_step(audio_chunk)
-tail_speaker_probabilities = session.diarize_step(torch.empty(0), is_final=True)
+session = diar_model.create_streaming_session(batch_size=2)
+for audio_batch, audio_lengths, final_mask in audio_stream:
+    probabilities, probability_lengths = session.diarize_step(
+        audio_batch,
+        audio_chunk_lengths=audio_lengths,
+        is_final=final_mask,
+    )
+    stream_0_probabilities = probabilities[0, : probability_lengths[0]]
 session.reset()
 ```
 
diff --git a/nemo/collections/asr/models/sortformer_diar_models.py b/nemo/collections/asr/models/sortformer_diar_models.py
index 79cee4805fe0..b73be195afcb 100644
--- a/nemo/collections/asr/models/sortformer_diar_models.py
+++ b/nemo/collections/asr/models/sortformer_diar_models.py
@@ -614,18 +614,21 @@ def process_signal(self, audio_signal, audio_signal_length):
             torch.cuda.empty_cache()
         return processed_signal, processed_signal_length
 
-    def create_streaming_session(self):
-        """Create an independent high-level raw-audio streaming session.
+    def create_streaming_session(self, batch_size: int = 1):
+        """Create an independent high-level raw-audio streaming session for a fixed batch of streams.
 
-        The returned session accepts arbitrarily sized mono waveform chunks through ``diarize_step()`` and owns all
-        preprocessing buffers and Sortformer cache state. Create one session per concurrent audio stream.
+        The returned session accepts arbitrarily sized mono waveform chunks through ``diarize_step()`` and owns the
+        per-stream preprocessing buffers and batched asynchronous Sortformer cache state.
+
+        Args:
+            batch_size: Fixed number of independent audio streams owned by the session.
 
         Returns:
             SortformerStreamingSession: A new session bound to this model.
         """
-        from nemo.collections.asr.parts.utils.streaming_sortformer import SortformerStreamingSession
+        from nemo.collections.asr.parts.utils.sortformer_utils import SortformerStreamingSession
 
-        return SortformerStreamingSession(self)
+        return SortformerStreamingSession(self, batch_size=batch_size)
 
     def forward(
         self,
@@ -950,6 +953,7 @@ def forward_streaming_step(
         drop_extra_pre_encoded=0,
         left_offset=0,
         right_offset=0,
+        async_streaming=None,
     ):
         """
         One-step forward pass for diarization inference in streaming mode.
@@ -978,6 +982,7 @@ def forward_streaming_step(
             drop_extra_pre_encoded (int): Number of leading pre-encoded frames to discard before streaming updates.
             left_offset (int): left offset for the current chunk
             right_offset (int): right offset for the current chunk
+            async_streaming (Optional[bool]): Override the model-level asynchronous streaming setting for this call.
 
         Returns:
             streaming_state (SortformerStreamingState):
@@ -987,6 +992,11 @@ def forward_streaming_step(
                 Tensor containing the updated total predicted speaker activity probabilities.
                 Shape: (batch_size, cumulative pred length, num_speakers)
         """
+        if async_streaming is None:
+            async_streaming = self.async_streaming
+        elif not isinstance(async_streaming, bool):
+            raise TypeError(f"async_streaming must be a boolean or None, got {type(async_streaming).__name__}")
+
         chunk_pre_encode_embs, chunk_pre_encode_lengths = self._call_pre_encode(
             processed_signal, processed_signal_length
         )
@@ -995,7 +1005,7 @@ def forward_streaming_step(
             chunk_pre_encode_embs = chunk_pre_encode_embs[:, drop_extra_pre_encoded:, :]
             chunk_pre_encode_lengths = chunk_pre_encode_lengths - drop_extra_pre_encoded
 
-        if self.async_streaming:
+        if async_streaming:
             output_length = None
             if self.async_pad_to_max:
                 output_length = (
@@ -1036,7 +1046,7 @@ def forward_streaming_step(
             spkcache_fifo_chunk_preds = self.sortformer_modules.downsample_preds(
                 high_resolution_preds, self.upsample_factor
             ).detach()
-            if not self.async_streaming and streaming_state.spk_perm is not None:
+            if not async_streaming and streaming_state.spk_perm is not None:
                 inv_spk_perm = torch.stack(
                     [
                         torch.argsort(streaming_state.spk_perm[batch_index])
@@ -1053,7 +1063,7 @@ def forward_streaming_step(
         spkcache_fifo_chunk_preds = self.sortformer_modules.apply_mask_to_preds(
             spkcache_fifo_chunk_preds, spkcache_fifo_chunk_fc_encoder_lengths
         )
-        if self.async_streaming:
+        if async_streaming:
             saved_spkcache_lengths = streaming_state.spkcache_lengths.clone()
             saved_fifo_lengths = streaming_state.fifo_lengths.clone()
             streaming_state, chunk_preds = self.sortformer_modules.streaming_update_async(
diff --git a/nemo/collections/asr/parts/utils/sortformer_utils.py b/nemo/collections/asr/parts/utils/sortformer_utils.py
index 5437b416575f..474019cfdb8c 100644
--- a/nemo/collections/asr/parts/utils/sortformer_utils.py
+++ b/nemo/collections/asr/parts/utils/sortformer_utils.py
@@ -12,21 +12,304 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import copy
 import logging
+import math
 import os
 import time
 from functools import wraps
 from pathlib import Path
 from tempfile import NamedTemporaryFile
-from typing import TYPE_CHECKING, Dict, List, Optional
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
 
 import torch
 from omegaconf import open_dict
 
+from nemo.collections.asr.parts.preprocessing.features import normalize_batch
+
 if TYPE_CHECKING:
     from nemo.collections.asr.models import SortformerEncLabelModel
 
 
+class SortformerStreamingSession:
+    """Stateful raw-audio session for a fixed batch of streaming Sortformer inputs.
+
+    Each row owns independent waveform buffering, progress, finalization, and speaker-cache state. Incoming waveform
+    chunks may have different valid lengths, and the session uses Sortformer's asynchronous streaming update so idle,
+    active, and finalized rows can coexist in one batch.
+
+    Args:
+        model: Streaming ``SortformerEncLabelModel`` in evaluation mode.
+        batch_size: Fixed number of independent audio streams owned by the session.
+    """
+
+    def __init__(self, model: "SortformerEncLabelModel", batch_size: int = 1):
+        if not model.streaming_mode:
+            raise ValueError("SortformerStreamingSession requires a model with streaming_mode=True")
+        if model.training:
+            raise ValueError("SortformerStreamingSession requires an evaluation model; call model.eval() first")
+        if not isinstance(batch_size, int) or isinstance(batch_size, bool) or batch_size < 1:
+            raise ValueError(f"batch_size must be a positive integer, got {batch_size}")
+
+        self.model = model
+        self.batch_size = batch_size
+        self.device = model.device
+        self._normalization = model.preprocessor.featurizer.normalize
+        self._preprocessor = copy.deepcopy(model.preprocessor).to(self.device).eval()
+        self._preprocessor.featurizer.normalize = None
+        self._preprocessor.featurizer.dither = 0.0
+        self._preprocessor.featurizer.pad_to = 0
+
+        self._hop_length = self._preprocessor.hop_length
+        self._n_fft = self._preprocessor.featurizer.n_fft
+        self._stft_margin_frames = math.ceil((self._n_fft // 2 + 1) / self._hop_length) + 1
+        self._chunk_frames = model.sortformer_modules.chunk_len * model.encoder.subsampling_factor
+        self._left_context_frames = model.sortformer_modules.chunk_left_context * model.encoder.subsampling_factor
+        self._right_context_frames = model.sortformer_modules.chunk_right_context * model.encoder.subsampling_factor
+        self.reset()
+
+    @torch.inference_mode()
+    def diarize_step(
+        self,
+        audio_chunks: torch.Tensor,
+        audio_chunk_lengths: Optional[torch.Tensor] = None,
+        is_final: Union[bool, torch.Tensor] = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Consume one raw-audio chunk per stream and return newly committed speaker probabilities.
+
+        Args:
+            audio_chunks: Float waveforms with shape ``(batch_size, max_num_samples)``. A one-dimensional tensor is
+                also accepted when ``batch_size=1``.
+            audio_chunk_lengths: Valid sample count for each padded row, with shape ``(batch_size,)``. If omitted,
+                every row uses the complete waveform width.
+            is_final: Boolean finalization mask with shape ``(batch_size,)`` or one boolean applied to every row.
+                Finalized rows can remain in later calls with a zero audio length while other rows continue.
+
+        Returns:
+            padded_probabilities: Newly committed probabilities with shape
+                ``(batch_size, max_new_frames, num_speakers)``.
+            probability_lengths: Valid output frames for each row, with shape ``(batch_size,)``.
+        """
+        audio_chunks, audio_chunk_lengths = self._validate_audio_chunks(audio_chunks, audio_chunk_lengths)
+        final_mask = self._validate_final_mask(is_final)
+        input_lengths = audio_chunk_lengths.tolist()
+        final_flags = final_mask.tolist()
+
+        for stream_index, chunk_length in enumerate(input_lengths):
+            if self._finalized[stream_index] and chunk_length > 0:
+                raise RuntimeError(
+                    f"Cannot supply audio to finalized stream {stream_index}; call reset() before reusing the session"
+                )
+            if chunk_length > 0:
+                chunk = audio_chunks[stream_index, :chunk_length]
+                self._audio_buffers[stream_index] = torch.cat([self._audio_buffers[stream_index], chunk])
+                self._received_samples[stream_index] += chunk_length
+
+        emitted = [[] for _ in range(self.batch_size)]
+        while True:
+            ready_groups = self._get_ready_groups(final_flags)
+            if not ready_groups:
+                break
+
+            for (left_offset, right_offset), requests in ready_groups.items():
+                processed_signal, processed_signal_length = self._extract_feature_batch(requests)
+                empty_preds = processed_signal.new_zeros((self.batch_size, 0, self.model.sortformer_modules.n_spk))
+                self.streaming_state, chunk_preds = self.model.forward_streaming_step(
+                    processed_signal=processed_signal,
+                    processed_signal_length=processed_signal_length,
+                    streaming_state=self.streaming_state,
+                    total_preds=empty_preds,
+                    left_offset=left_offset,
+                    right_offset=right_offset,
+                    async_streaming=True,
+                )
+
+                for stream_index, _, _, central_end in requests:
+                    committed_feature_frames = central_end - self._next_feature_frames[stream_index]
+                    output_length = math.ceil(committed_feature_frames / self.model.output_subsampling_factor)
+                    if output_length > chunk_preds.shape[1]:
+                        raise RuntimeError(
+                            "Streaming model returned fewer prediction frames than required: "
+                            f"needed {output_length}, got {chunk_preds.shape[1]}"
+                        )
+                    emitted[stream_index].append(chunk_preds[stream_index, :output_length])
+                    self._next_feature_frames[stream_index] = central_end
+
+        self._compact_audio_buffers()
+        for stream_index, is_final_stream in enumerate(final_flags):
+            if is_final_stream:
+                self._finalized[stream_index] = True
+
+        return self._pad_emitted_outputs(emitted)
+
+    def reset(self) -> None:
+        """Clear every stream's buffered audio and initialize a fresh batched asynchronous model state."""
+        self.streaming_state = self.model.sortformer_modules.init_streaming_state(
+            batch_size=self.batch_size,
+            async_streaming=True,
+            device=self.device,
+        )
+        self._audio_buffers = [torch.empty(0, dtype=torch.float32, device=self.device) for _ in range(self.batch_size)]
+        self._audio_buffer_starts = [0] * self.batch_size
+        self._received_samples = [0] * self.batch_size
+        self._next_feature_frames = [0] * self.batch_size
+        self._finalized = [False] * self.batch_size
+
+    def _get_ready_groups(self, final_flags: List[bool]):
+        ready_groups = {}
+        for stream_index in range(self.batch_size):
+            if self._finalized[stream_index]:
+                continue
+            available_frames = self._available_feature_frames(stream_index, is_final=final_flags[stream_index])
+            next_frame = self._next_feature_frames[stream_index]
+            if next_frame >= available_frames:
+                continue
+
+            central_end = min(next_frame + self._chunk_frames, available_frames)
+            if not final_flags[stream_index] and central_end + self._right_context_frames > available_frames:
+                continue
+
+            feature_start = max(0, next_frame - self._left_context_frames)
+            feature_end = min(central_end + self._right_context_frames, available_frames)
+            offsets = (next_frame - feature_start, feature_end - central_end)
+            ready_groups.setdefault(offsets, []).append((stream_index, feature_start, feature_end, central_end))
+        return ready_groups
+
+    def _available_feature_frames(self, stream_index: int, is_final: bool) -> int:
+        received_samples = self._received_samples[stream_index]
+        sample_count = torch.tensor(received_samples, device=self.device)
+        offline_frames = int(self._preprocessor.featurizer.get_seq_len(sample_count).item())
+        if is_final:
+            return max(0, offline_frames)
+
+        stable_samples = received_samples - self._n_fft // 2
+        if stable_samples < 0:
+            return 0
+        stable_frames = stable_samples // self._hop_length + 1
+        return max(0, min(offline_frames, stable_frames))
+
+    def _compact_audio_buffers(self) -> None:
+        for stream_index in range(self.batch_size):
+            first_needed_frame = max(
+                0,
+                self._next_feature_frames[stream_index] - self._left_context_frames - self._stft_margin_frames,
+            )
+            first_needed_sample = first_needed_frame * self._hop_length
+            drop_samples = first_needed_sample - self._audio_buffer_starts[stream_index]
+            if drop_samples > 0:
+                self._audio_buffers[stream_index] = self._audio_buffers[stream_index][drop_samples:].clone()
+                self._audio_buffer_starts[stream_index] = first_needed_sample
+
+    def _extract_feature_batch(self, requests):
+        audio_segments = []
+        audio_lengths = []
+        local_feature_ranges = []
+        for stream_index, feature_start, feature_end, _ in requests:
+            segment_start_frame = max(0, feature_start - self._stft_margin_frames)
+            segment_start_sample = segment_start_frame * self._hop_length
+            segment_end_sample = min(
+                self._received_samples[stream_index],
+                (feature_end + self._stft_margin_frames) * self._hop_length,
+            )
+            buffer_start = segment_start_sample - self._audio_buffer_starts[stream_index]
+            buffer_end = segment_end_sample - self._audio_buffer_starts[stream_index]
+            audio_segment = self._audio_buffers[stream_index][buffer_start:buffer_end]
+            audio_segments.append(audio_segment)
+            audio_lengths.append(audio_segment.shape[0])
+            local_start = feature_start - segment_start_frame
+            local_feature_ranges.append((local_start, local_start + feature_end - feature_start))
+
+        padded_audio = torch.nn.utils.rnn.pad_sequence(audio_segments, batch_first=True)
+        audio_lengths = torch.tensor(audio_lengths, dtype=torch.long, device=self.device)
+        features, feature_lengths = self._preprocessor(input_signal=padded_audio, length=audio_lengths)
+
+        feature_windows = []
+        window_lengths = []
+        for request_index, (local_start, local_end) in enumerate(local_feature_ranges):
+            if feature_lengths[request_index] < local_end:
+                raise RuntimeError(
+                    "Streaming preprocessor returned fewer feature frames than required: "
+                    f"needed {local_end}, got {feature_lengths[request_index].item()}"
+                )
+            window = features[request_index, :, local_start:local_end].transpose(0, 1)
+            feature_windows.append(window)
+            window_lengths.append(window.shape[0])
+
+        active_features = torch.nn.utils.rnn.pad_sequence(feature_windows, batch_first=True).transpose(1, 2)
+        active_lengths = torch.tensor(window_lengths, dtype=torch.long, device=self.device)
+        if self._normalization:
+            active_features, _, _ = normalize_batch(active_features, active_lengths, self._normalization)
+        active_features = active_features.transpose(1, 2)
+
+        batch_features = active_features.new_zeros(
+            (self.batch_size, active_features.shape[1], active_features.shape[2])
+        )
+        batch_lengths = torch.zeros(self.batch_size, dtype=torch.long, device=self.device)
+        for request_index, (stream_index, _, _, _) in enumerate(requests):
+            feature_length = window_lengths[request_index]
+            batch_features[stream_index, :feature_length] = active_features[request_index, :feature_length]
+            batch_lengths[stream_index] = feature_length
+        return batch_features, batch_lengths
+
+    def _pad_emitted_outputs(self, emitted):
+        num_speakers = self.model.sortformer_modules.n_spk
+        dtype = next(self.model.parameters()).dtype
+        row_outputs = [
+            torch.cat(row, dim=0) if row else torch.zeros((0, num_speakers), dtype=dtype, device=self.device)
+            for row in emitted
+        ]
+        output_lengths = torch.tensor([row.shape[0] for row in row_outputs], dtype=torch.long, device=self.device)
+        max_output_length = max((row.shape[0] for row in row_outputs), default=0)
+        padded_outputs = torch.zeros(
+            (self.batch_size, max_output_length, num_speakers), dtype=dtype, device=self.device
+        )
+        for stream_index, row in enumerate(row_outputs):
+            padded_outputs[stream_index, : row.shape[0]] = row
+        return padded_outputs, output_lengths
+
+    def _validate_audio_chunks(self, audio_chunks, audio_chunk_lengths):
+        if not isinstance(audio_chunks, torch.Tensor):
+            raise TypeError(f"audio_chunks must be a torch.Tensor, got {type(audio_chunks).__name__}")
+        if audio_chunks.ndim == 1 and self.batch_size == 1:
+            audio_chunks = audio_chunks.unsqueeze(0)
+        if audio_chunks.ndim != 2 or audio_chunks.shape[0] != self.batch_size:
+            raise ValueError(
+                f"audio_chunks must have batch dimension {self.batch_size} and shape "
+                f"({self.batch_size}, max_num_samples); got {tuple(audio_chunks.shape)}"
+            )
+        audio_chunks = audio_chunks.detach().to(device=self.device, dtype=torch.float32)
+
+        if audio_chunk_lengths is None:
+            audio_chunk_lengths = torch.full(
+                (self.batch_size,), audio_chunks.shape[1], dtype=torch.long, device=self.device
+            )
+        elif not isinstance(audio_chunk_lengths, torch.Tensor):
+            raise TypeError(
+                f"audio_chunk_lengths must be a torch.Tensor or None, got {type(audio_chunk_lengths).__name__}"
+            )
+        elif audio_chunk_lengths.shape != (self.batch_size,):
+            raise ValueError(
+                f"audio_chunk_lengths must have shape ({self.batch_size},), got {tuple(audio_chunk_lengths.shape)}"
+            )
+        elif torch.is_floating_point(audio_chunk_lengths) or audio_chunk_lengths.dtype == torch.bool:
+            raise TypeError("audio_chunk_lengths must contain integers")
+        else:
+            audio_chunk_lengths = audio_chunk_lengths.to(device=self.device, dtype=torch.long)
+
+        if torch.any(audio_chunk_lengths < 0) or torch.any(audio_chunk_lengths > audio_chunks.shape[1]):
+            raise ValueError(f"audio_chunk_lengths must be between 0 and {audio_chunks.shape[1]} samples")
+        return audio_chunks, audio_chunk_lengths
+
+    def _validate_final_mask(self, is_final):
+        if isinstance(is_final, bool):
+            return torch.full((self.batch_size,), is_final, dtype=torch.bool, device=self.device)
+        if not isinstance(is_final, torch.Tensor):
+            raise TypeError(f"is_final must be a boolean or torch.Tensor, got {type(is_final).__name__}")
+        if is_final.dtype != torch.bool or is_final.shape != (self.batch_size,):
+            raise ValueError(f"is_final must be boolean with shape ({self.batch_size},), got {tuple(is_final.shape)}")
+        return is_final.to(device=self.device)
+
+
 def configure_output_subsampling_factor(
     diar_model: "SortformerEncLabelModel",
     output_subsampling_factor: Optional[int],
diff --git a/nemo/collections/asr/parts/utils/streaming_sortformer.py b/nemo/collections/asr/parts/utils/streaming_sortformer.py
deleted file mode 100644
index 9621e986d1f3..000000000000
--- a/nemo/collections/asr/parts/utils/streaming_sortformer.py
+++ /dev/null
@@ -1,193 +0,0 @@
-# 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.
-
-import copy
-import math
-from typing import TYPE_CHECKING
-
-import torch
-
-from nemo.collections.asr.parts.preprocessing.features import normalize_batch
-
-if TYPE_CHECKING:
-    from nemo.collections.asr.models.sortformer_diar_models import SortformerEncLabelModel
-
-__all__ = ['SortformerStreamingSession']
-
-
-class SortformerStreamingSession:
-    """Stateful raw-audio session for a streaming Sortformer model.
-
-    A session owns the feature and speaker-cache state for one mono audio stream. Incoming waveform chunks may have
-    arbitrary lengths. The session waits until a complete model chunk and its configured right context are available,
-    then returns only the newly committed speaker probabilities. Call :meth:`reset` before reusing a session for a new
-    stream.
-
-    Args:
-        model: Streaming ``SortformerEncLabelModel`` in evaluation mode.
-    """
-
-    def __init__(self, model: 'SortformerEncLabelModel'):
-        if not model.streaming_mode:
-            raise ValueError("SortformerStreamingSession requires a model with streaming_mode=True")
-        if model.training:
-            raise ValueError("SortformerStreamingSession requires an evaluation model; call model.eval() first")
-
-        self.model = model
-        self.device = model.device
-        self._normalization = model.preprocessor.featurizer.normalize
-        self._preprocessor = copy.deepcopy(model.preprocessor).to(self.device).eval()
-        self._preprocessor.featurizer.normalize = None
-        self._preprocessor.featurizer.dither = 0.0
-        self._preprocessor.featurizer.pad_to = 0
-
-        self._hop_length = self._preprocessor.hop_length
-        self._n_fft = self._preprocessor.featurizer.n_fft
-        self._stft_margin_frames = math.ceil((self._n_fft // 2 + 1) / self._hop_length) + 1
-        self._chunk_frames = model.sortformer_modules.chunk_len * model.encoder.subsampling_factor
-        self._left_context_frames = model.sortformer_modules.chunk_left_context * model.encoder.subsampling_factor
-        self._right_context_frames = model.sortformer_modules.chunk_right_context * model.encoder.subsampling_factor
-        self.reset()
-
-    @torch.inference_mode()
-    def diarize_step(self, audio_chunk: torch.Tensor, is_final: bool = False) -> torch.Tensor:
-        """Consume a mono waveform chunk and return newly committed speaker probabilities.
-
-        Args:
-            audio_chunk: Float waveform with shape ``(num_samples,)`` or ``(1, num_samples)`` at the model sample rate.
-            is_final: Flush the remaining audio and close the stream. No more audio can be supplied until ``reset``.
-
-        Returns:
-            Tensor with shape ``(1, new_frames, num_speakers)``. ``new_frames`` can be zero while the session waits
-            for a complete chunk and its right context.
-        """
-        if self._finalized:
-            raise RuntimeError("This streaming session is finalized; call reset() before supplying more audio")
-        if not isinstance(is_final, bool):
-            raise TypeError(f"is_final must be a boolean, got {type(is_final).__name__}")
-
-        audio_chunk = self._validate_audio_chunk(audio_chunk)
-        if audio_chunk.numel() > 0:
-            self._audio_buffer = torch.cat([self._audio_buffer, audio_chunk])
-            self._received_samples += audio_chunk.numel()
-
-        available_frames = self._available_feature_frames(is_final=is_final)
-        emitted = []
-        while self._next_feature_frame < available_frames:
-            central_end = min(self._next_feature_frame + self._chunk_frames, available_frames)
-            if not is_final and central_end + self._right_context_frames > available_frames:
-                break
-
-            feature_start = max(0, self._next_feature_frame - self._left_context_frames)
-            feature_end = min(central_end + self._right_context_frames, available_frames)
-            processed_signal = self._extract_feature_window(feature_start, feature_end)
-            processed_signal_length = torch.tensor([feature_end - feature_start], dtype=torch.long, device=self.device)
-            empty_preds = processed_signal.new_zeros((1, 0, self.model.sortformer_modules.n_spk))
-            self.streaming_state, chunk_preds = self.model.forward_streaming_step(
-                processed_signal=processed_signal.transpose(1, 2),
-                processed_signal_length=processed_signal_length,
-                streaming_state=self.streaming_state,
-                total_preds=empty_preds,
-                left_offset=self._next_feature_frame - feature_start,
-                right_offset=feature_end - central_end,
-            )
-            emitted.append(chunk_preds)
-            self._next_feature_frame = central_end
-
-        self._compact_audio_buffer()
-
-        if is_final:
-            self._finalized = True
-
-        if emitted:
-            return torch.cat(emitted, dim=1)
-        return torch.zeros(
-            (1, 0, self.model.sortformer_modules.n_spk),
-            dtype=next(self.model.parameters()).dtype,
-            device=self.device,
-        )
-
-    def reset(self) -> None:
-        """Clear buffered audio and model state so this session can process a new stream."""
-        self.streaming_state = self.model.sortformer_modules.init_streaming_state(
-            batch_size=1,
-            async_streaming=self.model.async_streaming,
-            device=self.device,
-        )
-        self._audio_buffer = torch.empty(0, dtype=torch.float32, device=self.device)
-        self._audio_buffer_start = 0
-        self._received_samples = 0
-        self._next_feature_frame = 0
-        self._finalized = False
-
-    def _available_feature_frames(self, is_final: bool) -> int:
-        sample_count = torch.tensor(self._received_samples, device=self.device)
-        offline_frames = int(self._preprocessor.featurizer.get_seq_len(sample_count).item())
-        if is_final:
-            return max(0, offline_frames)
-
-        stable_samples = self._received_samples - self._n_fft // 2
-        if stable_samples < 0:
-            return 0
-        stable_frames = stable_samples // self._hop_length + 1
-        return max(0, min(offline_frames, stable_frames))
-
-    def _compact_audio_buffer(self) -> None:
-        first_needed_frame = max(
-            0,
-            self._next_feature_frame - self._left_context_frames - self._stft_margin_frames,
-        )
-        first_needed_sample = first_needed_frame * self._hop_length
-        drop_samples = first_needed_sample - self._audio_buffer_start
-        if drop_samples > 0:
-            self._audio_buffer = self._audio_buffer[drop_samples:].clone()
-            self._audio_buffer_start = first_needed_sample
-
-    def _extract_feature_window(self, feature_start: int, feature_end: int) -> torch.Tensor:
-        segment_start_frame = max(0, feature_start - self._stft_margin_frames)
-        segment_start_sample = segment_start_frame * self._hop_length
-        segment_end_sample = min(
-            self._received_samples,
-            (feature_end + self._stft_margin_frames) * self._hop_length,
-        )
-        buffer_start = segment_start_sample - self._audio_buffer_start
-        buffer_end = segment_end_sample - self._audio_buffer_start
-        audio_signal = self._audio_buffer[buffer_start:buffer_end].unsqueeze(0)
-        audio_signal_length = torch.tensor([audio_signal.shape[1]], dtype=torch.long, device=self.device)
-        features, _ = self._preprocessor(input_signal=audio_signal, length=audio_signal_length)
-
-        local_start = feature_start - segment_start_frame
-        local_end = local_start + feature_end - feature_start
-        if features.shape[2] < local_end:
-            raise RuntimeError(
-                "Streaming preprocessor returned fewer feature frames than required: "
-                f"needed {local_end}, got {features.shape[2]}"
-            )
-        features = features[:, :, local_start:local_end]
-        feature_length = torch.tensor([features.shape[2]], dtype=torch.long, device=self.device)
-        if self._normalization:
-            features, _, _ = normalize_batch(features, feature_length, self._normalization)
-        return features
-
-    def _validate_audio_chunk(self, audio_chunk: torch.Tensor) -> torch.Tensor:
-        if not isinstance(audio_chunk, torch.Tensor):
-            raise TypeError(f"audio_chunk must be a torch.Tensor, got {type(audio_chunk).__name__}")
-        if audio_chunk.ndim == 2 and audio_chunk.shape[0] == 1:
-            audio_chunk = audio_chunk.squeeze(0)
-        if audio_chunk.ndim != 1:
-            raise ValueError(
-                "audio_chunk must contain one mono stream with shape (num_samples,) or (1, num_samples); "
-                f"got {tuple(audio_chunk.shape)}"
-            )
-        return audio_chunk.detach().to(device=self.device, dtype=torch.float32)
diff --git a/tests/collections/speaker_tasks/test_diar_sortformer_models.py b/tests/collections/speaker_tasks/test_diar_sortformer_models.py
index c028bde10712..805833d1ea61 100644
--- a/tests/collections/speaker_tasks/test_diar_sortformer_models.py
+++ b/tests/collections/speaker_tasks/test_diar_sortformer_models.py
@@ -268,39 +268,58 @@ def test_constructor(self, sortformer_model):
         assert isinstance(instance2, SortformerEncLabelModel)
 
     @pytest.mark.unit
-    def test_raw_audio_streaming_session_is_independent_of_input_chunking(self):
+    def test_raw_audio_streaming_session_batches_independent_streams(self):
         model = _create_sortformer_model().eval()
         model.streaming_mode = True
         model.sortformer_modules.chunk_len = 2
         model.sortformer_modules.chunk_left_context = 1
         model.sortformer_modules.chunk_right_context = 1
         model._check_streaming_parameters()
-        audio = torch.randn(8193)
-
-        one_shot_session = model.create_streaming_session()
-        one_shot_outputs = [
-            one_shot_session.diarize_step(audio),
-            one_shot_session.diarize_step(torch.empty(0), is_final=True),
-        ]
-        one_shot_preds = torch.cat(one_shot_outputs, dim=1)
-
-        chunked_session = model.create_streaming_session()
-        chunked_outputs = []
-        start = 0
-        for chunk_size in (17, 503, 1600, 81, 2999, 123, 2870):
-            end = min(start + chunk_size, audio.numel())
-            chunked_outputs.append(chunked_session.diarize_step(audio[start:end]))
-            start = end
-            if start == audio.numel():
-                break
-        if start < audio.numel():
-            chunked_outputs.append(chunked_session.diarize_step(audio[start:]))
-        chunked_outputs.append(chunked_session.diarize_step(torch.empty(0), is_final=True))
-        chunked_preds = torch.cat(chunked_outputs, dim=1)
-
-        assert one_shot_preds.shape[1] > 0
-        assert chunked_session._audio_buffer.untyped_storage().nbytes() < audio.untyped_storage().nbytes()
-        torch.testing.assert_close(chunked_preds, one_shot_preds)
+        audio = [torch.randn(8193), torch.randn(5001)]
+
+        reference_preds = []
+        for signal in audio:
+            session = model.create_streaming_session(batch_size=1)
+            preds, pred_lengths = session.diarize_step(
+                signal.unsqueeze(0),
+                audio_chunk_lengths=torch.tensor([signal.numel()]),
+                is_final=torch.tensor([True]),
+            )
+            reference_preds.append(preds[0, : pred_lengths[0]])
+
+        session = model.create_streaming_session(batch_size=2)
+        emitted = [[], []]
+        offsets = [0, 0]
+        step_sizes = [(17, 503), (1600, 81), (2999, 4417), (3577, 0)]
+        for step_index, sizes in enumerate(step_sizes):
+            chunks = []
+            lengths = []
+            final = []
+            for stream_index, size in enumerate(sizes):
+                end = min(offsets[stream_index] + size, audio[stream_index].numel())
+                chunks.append(audio[stream_index][offsets[stream_index] : end])
+                offsets[stream_index] = end
+                lengths.append(chunks[-1].numel())
+                final.append(end == audio[stream_index].numel())
+            padded_audio = torch.nn.utils.rnn.pad_sequence(chunks, batch_first=True)
+            preds, pred_lengths = session.diarize_step(
+                padded_audio,
+                audio_chunk_lengths=torch.tensor(lengths),
+                is_final=torch.tensor(final),
+            )
+            for stream_index in range(2):
+                emitted[stream_index].append(preds[stream_index, : pred_lengths[stream_index]])
+
+            if step_index == 2:
+                assert final == [False, True]
+
+        for stream_index in range(2):
+            batched_preds = torch.cat(emitted[stream_index])
+            torch.testing.assert_close(batched_preds, reference_preds[stream_index])
+            assert (
+                session._audio_buffers[stream_index].untyped_storage().nbytes()
+                < audio[stream_index].untyped_storage().nbytes()
+            )
 
     @pytest.mark.unit
     def test_raw_audio_streaming_session_reset_and_validation(self):
@@ -315,18 +334,21 @@ def test_raw_audio_streaming_session_reset_and_validation(self):
         model.sortformer_modules.chunk_right_context = 1
         model._check_streaming_parameters()
         audio = torch.randn(4097)
-        session = model.create_streaming_session()
+        session = model.create_streaming_session(batch_size=1)
 
-        first_preds = session.diarize_step(audio, is_final=True)
-        with pytest.raises(RuntimeError, match="finalized"):
-            session.diarize_step(torch.empty(0))
+        first_preds, first_lengths = session.diarize_step(audio, is_final=True)
+        with pytest.raises(RuntimeError, match="finalized stream 0"):
+            session.diarize_step(torch.ones(1))
         session.reset()
-        second_preds = session.diarize_step(audio.unsqueeze(0), is_final=True)
+        second_preds, second_lengths = session.diarize_step(audio.unsqueeze(0), is_final=torch.tensor([True]))
 
+        assert torch.equal(second_lengths, first_lengths)
         torch.testing.assert_close(second_preds, first_preds)
-        with pytest.raises(ValueError, match="one mono stream"):
+        with pytest.raises(ValueError, match="batch dimension"):
             session.reset()
             session.diarize_step(torch.randn(2, 100))
+        with pytest.raises(ValueError, match="positive integer"):
+            model.create_streaming_session(batch_size=0)
 
     @pytest.mark.unit
     @pytest.mark.parametrize(

From 19a05241ad699495d51dc1f5cd4c16c90e0dc876 Mon Sep 17 00:00:00 2001
From: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Date: Wed, 2 Sep 2026 16:10:01 -0400
Subject: [PATCH 3/5] Add per-stream speaker limits to Sortformer sessions

Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
---
 .../asr/models/sortformer_diar_models.py      | 19 ++++-
 .../asr/modules/sortformer_modules.py         | 60 ++++++++++++++-
 .../asr/parts/utils/sortformer_utils.py       | 14 +++-
 .../test_diar_sortformer_models.py            | 74 +++++++++++++++++++
 4 files changed, 160 insertions(+), 7 deletions(-)

diff --git a/nemo/collections/asr/models/sortformer_diar_models.py b/nemo/collections/asr/models/sortformer_diar_models.py
index b73be195afcb..aade71a8e092 100644
--- a/nemo/collections/asr/models/sortformer_diar_models.py
+++ b/nemo/collections/asr/models/sortformer_diar_models.py
@@ -17,7 +17,7 @@
 import os
 import random
 from collections import OrderedDict
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
 import torch
@@ -614,7 +614,11 @@ def process_signal(self, audio_signal, audio_signal_length):
             torch.cuda.empty_cache()
         return processed_signal, processed_signal_length
 
-    def create_streaming_session(self, batch_size: int = 1):
+    def create_streaming_session(
+        self,
+        batch_size: int = 1,
+        max_speakers: Optional[Union[int, Sequence[int], torch.Tensor]] = None,
+    ):
         """Create an independent high-level raw-audio streaming session for a fixed batch of streams.
 
         The returned session accepts arbitrarily sized mono waveform chunks through ``diarize_step()`` and owns the
@@ -622,13 +626,15 @@ def create_streaming_session(self, batch_size: int = 1):
 
         Args:
             batch_size: Fixed number of independent audio streams owned by the session.
+            max_speakers: Number of enabled speaker channels for every stream. A scalar applies to the complete batch;
+                a sequence or tensor supplies one value per row. By default, every model speaker channel is enabled.
 
         Returns:
             SortformerStreamingSession: A new session bound to this model.
         """
         from nemo.collections.asr.parts.utils.sortformer_utils import SortformerStreamingSession
 
-        return SortformerStreamingSession(self, batch_size=batch_size)
+        return SortformerStreamingSession(self, batch_size=batch_size, max_speakers=max_speakers)
 
     def forward(
         self,
@@ -1059,6 +1065,13 @@ def forward_streaming_step(
                         for batch_index in range(high_resolution_preds.shape[0])
                     ]
                 )
+            high_resolution_preds = self.sortformer_modules.apply_max_speakers_mask(
+                high_resolution_preds, streaming_state.max_speakers
+            )
+
+        spkcache_fifo_chunk_preds = self.sortformer_modules.apply_max_speakers_mask(
+            spkcache_fifo_chunk_preds, streaming_state.max_speakers
+        )
 
         spkcache_fifo_chunk_preds = self.sortformer_modules.apply_mask_to_preds(
             spkcache_fifo_chunk_preds, spkcache_fifo_chunk_fc_encoder_lengths
diff --git a/nemo/collections/asr/modules/sortformer_modules.py b/nemo/collections/asr/modules/sortformer_modules.py
index b6a13b0c5578..5020a2add638 100644
--- a/nemo/collections/asr/modules/sortformer_modules.py
+++ b/nemo/collections/asr/modules/sortformer_modules.py
@@ -14,7 +14,8 @@
 
 import math
 from dataclasses import dataclass
-from typing import List, Optional, Tuple
+from numbers import Integral
+from typing import List, Optional, Sequence, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -45,6 +46,7 @@ class StreamingSortformerState:
         spk_perm (torch.Tensor): Speaker permutation information for the speaker cache
         mean_sil_emb (torch.Tensor): Mean silence embedding
         n_sil_frames (torch.Tensor): Number of silence frames
+        max_speakers (torch.Tensor): Per-row number of enabled speaker channels, or ``None`` for all channels.
     """
 
     spkcache = None  # Speaker cache to store embeddings from start
@@ -57,6 +59,7 @@ class StreamingSortformerState:
     spk_perm = None
     mean_sil_emb = None
     n_sil_frames = None
+    max_speakers = None
 
     def to(self, device):
         """
@@ -85,6 +88,8 @@ def to(self, device):
             self.mean_sil_emb = self.mean_sil_emb.to(device)
         if self.n_sil_frames is not None:
             self.n_sil_frames = self.n_sil_frames.to(device)
+        if self.max_speakers is not None:
+            self.max_speakers = self.max_speakers.to(device)
 
 
 class SortformerModules(NeuralModule, Exportable):
@@ -485,7 +490,13 @@ def concat_and_pad(embs: List[torch.Tensor], lengths: List[torch.Tensor], output
         output = flat_output[: batch_size * sig_length].view(batch_size, sig_length, emb_dim)
         return output, total_lengths
 
-    def init_streaming_state(self, batch_size: int = 1, async_streaming: bool = False, device: torch.device = None):
+    def init_streaming_state(
+        self,
+        batch_size: int = 1,
+        async_streaming: bool = False,
+        device: torch.device = None,
+        max_speakers: Optional[Union[int, Sequence[int], torch.Tensor]] = None,
+    ):
         """
         Initializes StreamingSortformerState with empty tensors or zero-valued tensors.
 
@@ -493,6 +504,8 @@ def init_streaming_state(self, batch_size: int = 1, async_streaming: bool = Fals
             batch_size (int): Batch size for tensors in streaming state
             async_streaming (bool): True for asynchronous update, False for synchronous update
             device (torch.device): Device for tensors in streaming state
+            max_speakers (Optional[Union[int, Sequence[int], torch.Tensor]]): Number of enabled speaker channels for
+                every row. A scalar applies to the complete batch; a sequence or tensor supplies one value per row.
 
         Returns:
             streaming_state (SortformerStreamingState): initialized streaming state
@@ -510,8 +523,51 @@ def init_streaming_state(self, batch_size: int = 1, async_streaming: bool = Fals
             streaming_state.fifo = torch.zeros((batch_size, 0, self.fc_d_model), device=device)
         streaming_state.mean_sil_emb = torch.zeros((batch_size, self.fc_d_model), device=device)
         streaming_state.n_sil_frames = torch.zeros((batch_size,), dtype=torch.long, device=device)
+        streaming_state.max_speakers = self._normalize_max_speakers(max_speakers, batch_size, device)
         return streaming_state
 
+    def _normalize_max_speakers(self, max_speakers, batch_size, device):
+        """Validate a scalar or per-row speaker limit and return it as a device tensor."""
+        if max_speakers is None:
+            return None
+        if isinstance(max_speakers, bool):
+            raise TypeError("max_speakers must be an integer, sequence of integers, or integer tensor")
+        if isinstance(max_speakers, Integral):
+            max_speakers = torch.full((batch_size,), int(max_speakers), dtype=torch.long, device=device)
+        elif isinstance(max_speakers, torch.Tensor):
+            if (
+                max_speakers.dtype == torch.bool
+                or torch.is_floating_point(max_speakers)
+                or torch.is_complex(max_speakers)
+            ):
+                raise TypeError("max_speakers must contain integers")
+            if max_speakers.ndim == 0:
+                max_speakers = max_speakers.expand(batch_size)
+            elif max_speakers.ndim != 1 or max_speakers.numel() != batch_size:
+                raise ValueError(f"max_speakers must contain one value per batch row; expected {batch_size}")
+            max_speakers = max_speakers.to(device=device, dtype=torch.long)
+        else:
+            if not isinstance(max_speakers, Sequence) or isinstance(max_speakers, (str, bytes)):
+                raise TypeError("max_speakers must be an integer, sequence of integers, or integer tensor")
+            if len(max_speakers) != batch_size:
+                raise ValueError(f"max_speakers must contain one value per batch row; expected {batch_size}")
+            if any(isinstance(value, bool) or not isinstance(value, Integral) for value in max_speakers):
+                raise TypeError("max_speakers must contain integers")
+            max_speakers = torch.tensor(max_speakers, dtype=torch.long, device=device)
+
+        if torch.any(max_speakers < 1) or torch.any(max_speakers > self.n_spk):
+            raise ValueError(f"max_speakers values must be between 1 and {self.n_spk}")
+        return max_speakers
+
+    @staticmethod
+    def apply_max_speakers_mask(predictions, max_speakers):
+        """Zero speaker channels at or above each row's configured limit."""
+        if max_speakers is None:
+            return predictions
+        speaker_indices = torch.arange(predictions.shape[2], device=predictions.device).view(1, 1, -1)
+        enabled_speakers = speaker_indices < max_speakers.view(-1, 1, 1)
+        return predictions.masked_fill(~enabled_speakers, 0.0)
+
     @staticmethod
     def apply_mask_to_preds(spkcache_fifo_chunk_preds, spkcache_fifo_chunk_fc_encoder_lengths):
         """
diff --git a/nemo/collections/asr/parts/utils/sortformer_utils.py b/nemo/collections/asr/parts/utils/sortformer_utils.py
index 474019cfdb8c..fa5a6485a30a 100644
--- a/nemo/collections/asr/parts/utils/sortformer_utils.py
+++ b/nemo/collections/asr/parts/utils/sortformer_utils.py
@@ -20,7 +20,7 @@
 from functools import wraps
 from pathlib import Path
 from tempfile import NamedTemporaryFile
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple, Union
 
 import torch
 from omegaconf import open_dict
@@ -41,9 +41,16 @@ class SortformerStreamingSession:
     Args:
         model: Streaming ``SortformerEncLabelModel`` in evaluation mode.
         batch_size: Fixed number of independent audio streams owned by the session.
+        max_speakers: Number of enabled speaker channels for every stream. A scalar applies to the complete batch; a
+            sequence or tensor supplies one value per row. By default, every model speaker channel is enabled.
     """
 
-    def __init__(self, model: "SortformerEncLabelModel", batch_size: int = 1):
+    def __init__(
+        self,
+        model: "SortformerEncLabelModel",
+        batch_size: int = 1,
+        max_speakers: Optional[Union[int, Sequence[int], torch.Tensor]] = None,
+    ):
         if not model.streaming_mode:
             raise ValueError("SortformerStreamingSession requires a model with streaming_mode=True")
         if model.training:
@@ -53,6 +60,7 @@ def __init__(self, model: "SortformerEncLabelModel", batch_size: int = 1):
 
         self.model = model
         self.batch_size = batch_size
+        self._max_speakers = max_speakers
         self.device = model.device
         self._normalization = model.preprocessor.featurizer.normalize
         self._preprocessor = copy.deepcopy(model.preprocessor).to(self.device).eval()
@@ -148,7 +156,9 @@ def reset(self) -> None:
             batch_size=self.batch_size,
             async_streaming=True,
             device=self.device,
+            max_speakers=self._max_speakers,
         )
+        self._max_speakers = self.streaming_state.max_speakers
         self._audio_buffers = [torch.empty(0, dtype=torch.float32, device=self.device) for _ in range(self.batch_size)]
         self._audio_buffer_starts = [0] * self.batch_size
         self._received_samples = [0] * self.batch_size
diff --git a/tests/collections/speaker_tasks/test_diar_sortformer_models.py b/tests/collections/speaker_tasks/test_diar_sortformer_models.py
index 805833d1ea61..6ba083a9dc1e 100644
--- a/tests/collections/speaker_tasks/test_diar_sortformer_models.py
+++ b/tests/collections/speaker_tasks/test_diar_sortformer_models.py
@@ -321,6 +321,80 @@ def test_raw_audio_streaming_session_batches_independent_streams(self):
                 < audio[stream_index].untyped_storage().nbytes()
             )
 
+    @pytest.mark.unit
+    @pytest.mark.parametrize("high_resolution", [False, True])
+    def test_raw_audio_streaming_session_applies_per_row_max_speakers_before_state_update(self, high_resolution):
+        model = _create_sortformer_model(high_resolution=high_resolution).eval()
+        model.streaming_mode = True
+        model.sortformer_modules.chunk_len = 2
+        model.sortformer_modules.chunk_left_context = 1
+        model.sortformer_modules.chunk_right_context = 1
+        model._check_streaming_parameters()
+        captured_state_predictions = []
+        streaming_update_async = model.sortformer_modules.streaming_update_async
+
+        def capture_streaming_update(**kwargs):
+            captured_state_predictions.append(kwargs["preds"].clone())
+            return streaming_update_async(**kwargs)
+
+        model.sortformer_modules.streaming_update_async = capture_streaming_update
+        session = model.create_streaming_session(batch_size=2, max_speakers=[2, 4])
+        audio = torch.nn.utils.rnn.pad_sequence([torch.randn(8193), torch.randn(5001)], batch_first=True)
+        predictions, prediction_lengths = session.diarize_step(
+            audio,
+            audio_chunk_lengths=torch.tensor([8193, 5001]),
+            is_final=torch.tensor([True, True]),
+        )
+
+        assert torch.count_nonzero(predictions[0, : prediction_lengths[0], 2:]) == 0
+        assert torch.count_nonzero(predictions[1, : prediction_lengths[1], 2:]) > 0
+        assert captured_state_predictions
+        for state_predictions in captured_state_predictions:
+            assert torch.count_nonzero(state_predictions[0, :, 2:]) == 0
+        assert any(
+            torch.count_nonzero(state_predictions[1, :, 2:]) > 0 for state_predictions in captured_state_predictions
+        )
+        assert torch.count_nonzero(session.streaming_state.fifo_preds[0, :, 2:]) == 0
+        assert torch.count_nonzero(session.streaming_state.spkcache_preds[0, :, 2:]) == 0
+        assert session.streaming_state.max_speakers.tolist() == [2, 4]
+        session.reset()
+        assert session.streaming_state.max_speakers.tolist() == [2, 4]
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "max_speakers, error_type, error_match",
+        [
+            ([2], ValueError, "one value per batch row"),
+            ([2, 5], ValueError, "between 1 and 4"),
+            ([2, 1.5], TypeError, "must contain integers"),
+            (torch.tensor([2.0, 4.0]), TypeError, "must contain integers"),
+            (True, TypeError, "must be an integer"),
+        ],
+    )
+    def test_raw_audio_streaming_session_rejects_invalid_max_speakers(self, max_speakers, error_type, error_match):
+        model = _create_sortformer_model().eval()
+        model.streaming_mode = True
+
+        with pytest.raises(error_type, match=error_match):
+            model.create_streaming_session(batch_size=2, max_speakers=max_speakers)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "max_speakers, expected",
+        [
+            (2, [2, 2]),
+            (torch.tensor(3), [3, 3]),
+            (torch.tensor([1, 4]), [1, 4]),
+        ],
+    )
+    def test_raw_audio_streaming_session_normalizes_max_speakers(self, max_speakers, expected):
+        model = _create_sortformer_model().eval()
+        model.streaming_mode = True
+
+        session = model.create_streaming_session(batch_size=2, max_speakers=max_speakers)
+
+        assert session.streaming_state.max_speakers.tolist() == expected
+
     @pytest.mark.unit
     def test_raw_audio_streaming_session_reset_and_validation(self):
         offline_model = _create_sortformer_model().eval()

From 7437574de55b054b7f42056e273110119b9da66f Mon Sep 17 00:00:00 2001
From: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Date: Wed, 2 Sep 2026 17:03:59 -0400
Subject: [PATCH 4/5] Document Sortformer session speaker limits

Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
---
 examples/speaker_tasks/diarization/README.md | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/examples/speaker_tasks/diarization/README.md b/examples/speaker_tasks/diarization/README.md
index bba4d0e1a1f8..46bfae2da905 100644
--- a/examples/speaker_tasks/diarization/README.md
+++ b/examples/speaker_tasks/diarization/README.md
@@ -94,10 +94,14 @@ zero input length while other rows continue. Call ``reset()`` before reusing the
 and feature padding while extracting each chunk, then applies the checkpoint's feature normalization over each
 complete model input window.
 
+Use ``max_speakers`` to enable fewer than the checkpoint's maximum number of speaker channels. A scalar applies to
+every stream, while a sequence or integer tensor supplies one limit per row. Disabled channels are zeroed in returned
+probabilities and excluded before speaker-cache, FIFO, and silence-profile updates.
+
 ```python
 import torch
 
-session = diar_model.create_streaming_session(batch_size=2)
+session = diar_model.create_streaming_session(batch_size=2, max_speakers=[2, 4])
 for audio_batch, audio_lengths, final_mask in audio_stream:
     probabilities, probability_lengths = session.diarize_step(
         audio_batch,

From 01c9c68a2f59321e8b27feffa42e7e0812800c54 Mon Sep 17 00:00:00 2001
From: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Date: Tue, 8 Sep 2026 02:23:54 -0400
Subject: [PATCH 5/5] Preserve speech activity for Sortformer silence profiles

Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
---
 examples/speaker_tasks/diarization/README.md  |  3 +-
 .../asr/models/sortformer_diar_models.py      | 14 ++-
 .../asr/modules/sortformer_modules.py         | 92 +++++++++++++++++--
 .../test_diar_sortformer_models.py            |  6 ++
 .../test_diar_sortformer_modules.py           | 42 +++++++++
 5 files changed, 146 insertions(+), 11 deletions(-)

diff --git a/examples/speaker_tasks/diarization/README.md b/examples/speaker_tasks/diarization/README.md
index 46bfae2da905..d3ec6d248e03 100644
--- a/examples/speaker_tasks/diarization/README.md
+++ b/examples/speaker_tasks/diarization/README.md
@@ -96,7 +96,8 @@ complete model input window.
 
 Use ``max_speakers`` to enable fewer than the checkpoint's maximum number of speaker channels. A scalar applies to
 every stream, while a sequence or integer tensor supplies one limit per row. Disabled channels are zeroed in returned
-probabilities and excluded before speaker-cache, FIFO, and silence-profile updates.
+probabilities and excluded from speaker-cache and FIFO state. Their raw activity is used only to prevent speech frames
+from being added to the running silence profile.
 
 ```python
 import torch
diff --git a/nemo/collections/asr/models/sortformer_diar_models.py b/nemo/collections/asr/models/sortformer_diar_models.py
index aade71a8e092..69c7d88ff3ad 100644
--- a/nemo/collections/asr/models/sortformer_diar_models.py
+++ b/nemo/collections/asr/models/sortformer_diar_models.py
@@ -1069,13 +1069,17 @@ def forward_streaming_step(
                 high_resolution_preds, streaming_state.max_speakers
             )
 
-        spkcache_fifo_chunk_preds = self.sortformer_modules.apply_max_speakers_mask(
-            spkcache_fifo_chunk_preds, streaming_state.max_speakers
-        )
-
         spkcache_fifo_chunk_preds = self.sortformer_modules.apply_mask_to_preds(
             spkcache_fifo_chunk_preds, spkcache_fifo_chunk_fc_encoder_lengths
         )
+        silence_profile_preds = None
+        if streaming_state.max_speakers is not None:
+            # Disabled channels must not enter output/cache state, but their activity still prevents a speech frame
+            # from being folded into the running silence embedding.
+            silence_profile_preds = spkcache_fifo_chunk_preds
+            spkcache_fifo_chunk_preds = self.sortformer_modules.apply_max_speakers_mask(
+                spkcache_fifo_chunk_preds, streaming_state.max_speakers
+            )
         if async_streaming:
             saved_spkcache_lengths = streaming_state.spkcache_lengths.clone()
             saved_fifo_lengths = streaming_state.fifo_lengths.clone()
@@ -1086,6 +1090,7 @@ def forward_streaming_step(
                 preds=spkcache_fifo_chunk_preds,
                 lc=lc_enc,
                 rc=rc_enc,
+                silence_profile_preds=silence_profile_preds,
             )
             if self.high_resolution:
                 max_chunk_len = chunk_pre_encode_embs.shape[1] - lc_enc - rc_enc
@@ -1107,6 +1112,7 @@ def forward_streaming_step(
                 preds=spkcache_fifo_chunk_preds,
                 lc=lc_enc,
                 rc=rc_enc,
+                silence_profile_preds=silence_profile_preds,
             )
             if self.high_resolution:
                 chunk_len = chunk_pre_encode_embs.shape[1] - lc_enc - rc_enc
diff --git a/nemo/collections/asr/modules/sortformer_modules.py b/nemo/collections/asr/modules/sortformer_modules.py
index 5020a2add638..1457692d15c1 100644
--- a/nemo/collections/asr/modules/sortformer_modules.py
+++ b/nemo/collections/asr/modules/sortformer_modules.py
@@ -702,6 +702,8 @@ def _update_async_fifo(
         max_chunk_len,
         max_pop_out_len,
         lc,
+        silence_fifo_preds=None,
+        silence_chunk_preds=None,
     ):
         """
         Pop and retain logical ``[FIFO | chunk]`` frames and mutate the FIFO streaming state.
@@ -723,12 +725,19 @@ def _update_async_fifo(
             max_chunk_len (int): Physical chunk capacity excluding context.
             max_pop_out_len (int): Physical capacity of the rectangular popped-frame buffer.
             lc (int): Left context offset of the current chunk.
+            silence_fifo_preds (Optional[torch.Tensor]): Unmasked predictions for the valid current FIFO region, used
+                only to decide whether popped frames are silent.
+            silence_chunk_preds (Optional[torch.Tensor]): Unmasked predictions for the current chunk region, used only
+                to decide whether popped frames are silent.
 
         Returns:
             pop_out_embs (torch.Tensor): Left-aligned popped embeddings.
                 Shape: (batch_size, max_pop_out_len, emb_dim)
             pop_out_preds (torch.Tensor): Left-aligned predictions for popped embeddings.
                 Shape: (batch_size, max_pop_out_len, n_spk)
+            pop_out_silence_preds (torch.Tensor): Predictions used only for silence classification. This is identical
+                to ``pop_out_preds`` unless unmasked silence-profile predictions were supplied.
+                Shape: (batch_size, max_pop_out_len, n_spk)
             valid_pop_mask (torch.Tensor): Mask identifying valid popped frames.
                 Shape: (batch_size, max_pop_out_len)
         """
@@ -781,12 +790,27 @@ def _update_async_fifo(
         )
         pop_out_embs, updated_fifo = torch.split(gathered_fifo_embs, [max_pop_out_len, max_fifo_len], dim=1)
         pop_out_preds, updated_fifo_preds = torch.split(gathered_fifo_preds, [max_pop_out_len, max_fifo_len], dim=1)
+        pop_out_silence_preds = pop_out_preds
+        if silence_fifo_preds is not None and silence_chunk_preds is not None:
+            silence_fifo_chunk_preds = torch.cat(
+                [
+                    silence_fifo_preds,
+                    silence_chunk_preds,
+                    silence_fifo_preds.new_zeros((batch_size, 1, n_spk)),
+                ],
+                dim=1,
+            )
+            pop_out_silence_preds = torch.gather(
+                silence_fifo_chunk_preds,
+                1,
+                fifo_physical_indices[:, :max_pop_out_len].unsqueeze(-1).expand(-1, -1, n_spk),
+            )
         valid_pop_mask = pop_positions < pop_out_lengths.unsqueeze(1)
 
         streaming_state.fifo = updated_fifo
         streaming_state.fifo_preds = updated_fifo_preds
         streaming_state.fifo_lengths.copy_(new_fifo_lengths)
-        return pop_out_embs, pop_out_preds, valid_pop_mask
+        return pop_out_embs, pop_out_preds, pop_out_silence_preds, valid_pop_mask
 
     def _update_async_silence_profile(self, streaming_state, pop_out_embs, pop_out_preds, valid_pop_mask):
         """
@@ -901,7 +925,16 @@ def _update_async_spkcache(
             streaming_state.spkcache_compressed[idx] = True
         streaming_state.spkcache_lengths.copy_(updated_spkcache_lengths.clamp(max=self.spkcache_len))
 
-    def streaming_update_async(self, streaming_state, chunk, chunk_lengths, preds, lc: int = 0, rc: int = 0):
+    def streaming_update_async(
+        self,
+        streaming_state,
+        chunk,
+        chunk_lengths,
+        preds,
+        lc: int = 0,
+        rc: int = 0,
+        silence_profile_preds=None,
+    ):
         """
         Update the speaker cache and FIFO queue with the chunk of embeddings and speaker predictions.
         Asynchronous version, which means speaker cache, FIFO and chunk may have different lengths within a batch.
@@ -917,6 +950,8 @@ def streaming_update_async(self, streaming_state, chunk, chunk_lengths, preds, l
                 Shape: (batch_size, spkcache_len + fifo_len + lc+chunk_len+rc, num_spks)
             lc (int): Left-context offset. Only ``chunk[:, lc:chunk_len+lc]`` is used to update the state.
             rc (int): Right-context offset excluded from the speaker-cache and FIFO update.
+            silence_profile_preds (Optional[torch.Tensor]): Unmasked speaker predictions used only to classify popped
+                frames as speech or silence. State and returned predictions continue to use ``preds``.
 
         Returns:
             streaming_state (SortformerStreamingState): Current streaming state including speaker cache and FIFO
@@ -948,11 +983,25 @@ def streaming_update_async(self, streaming_state, chunk, chunk_lengths, preds, l
             max_chunk_len,
             lc,
         )
+        silence_fifo_preds = None
+        silence_chunk_preds = None
+        if silence_profile_preds is not None:
+            _, silence_fifo_preds, silence_chunk_preds = self._gather_async_predictions(
+                streaming_state,
+                silence_profile_preds,
+                spkcache_lengths,
+                fifo_lengths,
+                chunk_lengths,
+                max_spkcache_len,
+                max_fifo_len,
+                max_chunk_len,
+                lc,
+            )
 
         pop_out_lengths, new_fifo_lengths = self._compute_async_fifo_pop_lengths(
             spkcache_lengths, fifo_lengths, chunk_lengths, max_fifo_len
         )
-        pop_out_embs, pop_out_preds, valid_pop_mask = self._update_async_fifo(
+        pop_out_embs, pop_out_preds, pop_out_silence_preds, valid_pop_mask = self._update_async_fifo(
             streaming_state,
             chunk,
             current_fifo_preds,
@@ -963,8 +1012,10 @@ def streaming_update_async(self, streaming_state, chunk, chunk_lengths, preds, l
             max_chunk_len,
             max_pop_out_len,
             lc,
+            silence_fifo_preds=silence_fifo_preds,
+            silence_chunk_preds=silence_chunk_preds,
         )
-        self._update_async_silence_profile(streaming_state, pop_out_embs, pop_out_preds, valid_pop_mask)
+        self._update_async_silence_profile(streaming_state, pop_out_embs, pop_out_silence_preds, valid_pop_mask)
         self._update_async_spkcache(
             streaming_state,
             current_spkcache_preds,
@@ -983,7 +1034,15 @@ def streaming_update_async(self, streaming_state, chunk, chunk_lengths, preds, l
 
         return streaming_state, chunk_preds
 
-    def streaming_update(self, streaming_state, chunk, preds, lc: int = 0, rc: int = 0):
+    def streaming_update(
+        self,
+        streaming_state,
+        chunk,
+        preds,
+        lc: int = 0,
+        rc: int = 0,
+        silence_profile_preds=None,
+    ):
         """
         Update the speaker cache and FIFO queue with the chunk of embeddings and speaker predictions.
         Synchronous version, which means speaker cahce, FIFO queue and chunk have same lengths within a batch.
@@ -997,6 +1056,8 @@ def streaming_update(self, streaming_state, chunk, preds, lc: int = 0, rc: int =
                 Shape: (batch_size, spkcache_len + fifo_len + lc+chunk_len+rc, num_spks)
             lc (int): Left-context offset. Only ``chunk[:, lc:chunk_len+lc]`` is used to update the state.
             rc (int): Right-context offset excluded from the speaker-cache and FIFO update.
+            silence_profile_preds (Optional[torch.Tensor]): Unmasked speaker predictions used only to classify popped
+                frames as speech or silence. State and returned predictions continue to use ``preds``.
 
         Returns:
             streaming_state (SortformerStreamingState): current streaming state including speaker cache and FIFO
@@ -1018,10 +1079,26 @@ def streaming_update(self, streaming_state, chunk, preds, lc: int = 0, rc: int =
             preds = torch.stack(
                 [preds[batch_index, :, inv_spk_perm[batch_index]] for batch_index in range(batch_size)]
             )
+            if silence_profile_preds is not None:
+                silence_profile_preds = torch.stack(
+                    [
+                        silence_profile_preds[batch_index, :, inv_spk_perm[batch_index]]
+                        for batch_index in range(batch_size)
+                    ]
+                )
 
         streaming_state.fifo_preds = preds[:, spkcache_len : spkcache_len + fifo_len]
         chunk = chunk[:, lc : chunk_len + lc]
         chunk_preds = preds[:, spkcache_len + fifo_len + lc : spkcache_len + fifo_len + chunk_len + lc]
+        silence_profile_fifo_preds = None
+        if silence_profile_preds is not None:
+            silence_profile_fifo_preds = torch.cat(
+                [
+                    silence_profile_preds[:, spkcache_len : spkcache_len + fifo_len],
+                    silence_profile_preds[:, spkcache_len + fifo_len + lc : spkcache_len + fifo_len + chunk_len + lc],
+                ],
+                dim=1,
+            )
 
         # append chunk to fifo
         streaming_state.fifo = torch.cat([streaming_state.fifo, chunk], dim=1)
@@ -1035,12 +1112,15 @@ def streaming_update(self, streaming_state, chunk, preds, lc: int = 0, rc: int =
 
             pop_out_embs = streaming_state.fifo[:, :pop_out_len]
             pop_out_preds = streaming_state.fifo_preds[:, :pop_out_len]
+            silence_profile_pop_out_preds = (
+                pop_out_preds if silence_profile_fifo_preds is None else silence_profile_fifo_preds[:, :pop_out_len]
+            )
             if not self.use_learnable_sil_emb:
                 streaming_state.mean_sil_emb, streaming_state.n_sil_frames = self._get_silence_profile(
                     streaming_state.mean_sil_emb,
                     streaming_state.n_sil_frames,
                     pop_out_embs,
-                    pop_out_preds,
+                    silence_profile_pop_out_preds,
                 )
             streaming_state.fifo = streaming_state.fifo[:, pop_out_len:]
             streaming_state.fifo_preds = streaming_state.fifo_preds[:, pop_out_len:]
diff --git a/tests/collections/speaker_tasks/test_diar_sortformer_models.py b/tests/collections/speaker_tasks/test_diar_sortformer_models.py
index 6ba083a9dc1e..3c76249eba4d 100644
--- a/tests/collections/speaker_tasks/test_diar_sortformer_models.py
+++ b/tests/collections/speaker_tasks/test_diar_sortformer_models.py
@@ -331,10 +331,12 @@ def test_raw_audio_streaming_session_applies_per_row_max_speakers_before_state_u
         model.sortformer_modules.chunk_right_context = 1
         model._check_streaming_parameters()
         captured_state_predictions = []
+        captured_silence_profile_predictions = []
         streaming_update_async = model.sortformer_modules.streaming_update_async
 
         def capture_streaming_update(**kwargs):
             captured_state_predictions.append(kwargs["preds"].clone())
+            captured_silence_profile_predictions.append(kwargs["silence_profile_preds"].clone())
             return streaming_update_async(**kwargs)
 
         model.sortformer_modules.streaming_update_async = capture_streaming_update
@@ -351,6 +353,10 @@ def capture_streaming_update(**kwargs):
         assert captured_state_predictions
         for state_predictions in captured_state_predictions:
             assert torch.count_nonzero(state_predictions[0, :, 2:]) == 0
+        assert any(
+            torch.count_nonzero(silence_predictions[0, :, 2:]) > 0
+            for silence_predictions in captured_silence_profile_predictions
+        )
         assert any(
             torch.count_nonzero(state_predictions[1, :, 2:]) > 0 for state_predictions in captured_state_predictions
         )
diff --git a/tests/collections/speaker_tasks/test_diar_sortformer_modules.py b/tests/collections/speaker_tasks/test_diar_sortformer_modules.py
index 299b9020957b..bdfcc4d3f1ea 100644
--- a/tests/collections/speaker_tasks/test_diar_sortformer_modules.py
+++ b/tests/collections/speaker_tasks/test_diar_sortformer_modules.py
@@ -2512,6 +2512,48 @@ def test_async_zero_capacity_fifo_masks_ragged_silence_updates(
         assert streaming_state.n_sil_frames.tolist() == list(expected_silence_counts)
         torch.testing.assert_close(streaming_state.mean_sil_emb, torch.tensor(expected_silence_means))
 
+    @pytest.mark.unit
+    @pytest.mark.parametrize("async_streaming", [False, True])
+    def test_speaker_limit_does_not_turn_disabled_channel_activity_into_silence(self, async_streaming):
+        sortformer_modules = SortformerModules(
+            num_spks=4,
+            fc_d_model=2,
+            spkcache_len=8,
+            fifo_len=0,
+            chunk_len=2,
+            spkcache_update_period=2,
+            spkcache_sil_frames_per_spk=0,
+            use_learnable_sil_emb=False,
+        )
+        streaming_state = sortformer_modules.init_streaming_state(
+            batch_size=1, async_streaming=async_streaming, max_speakers=2
+        )
+        chunk = torch.tensor([[[1.0, 3.0], [5.0, 7.0]]])
+        silence_profile_preds = torch.zeros(1, 2, 4)
+        silence_profile_preds[0, 0, 2] = 1.0
+        state_preds = sortformer_modules.apply_max_speakers_mask(silence_profile_preds, streaming_state.max_speakers)
+
+        if async_streaming:
+            streaming_state, chunk_preds = sortformer_modules.streaming_update_async(
+                streaming_state,
+                chunk,
+                torch.tensor([2]),
+                state_preds,
+                silence_profile_preds=silence_profile_preds,
+            )
+        else:
+            streaming_state, chunk_preds = sortformer_modules.streaming_update(
+                streaming_state,
+                chunk,
+                state_preds,
+                silence_profile_preds=silence_profile_preds,
+            )
+
+        assert streaming_state.n_sil_frames.tolist() == [1]
+        torch.testing.assert_close(streaming_state.mean_sil_emb, chunk[:, 1])
+        assert torch.count_nonzero(streaming_state.spkcache_preds[..., 2:]) == 0
+        assert torch.count_nonzero(chunk_preds[..., 2:]) == 0
+
     @pytest.mark.unit
     @pytest.mark.parametrize(
         (