chore: import upstream snapshot with attribution

nemo/collections/asr/modules/parallel_expert_encoder.py

@@ -0,0 +1,717 @@
+# 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.
+
+"""Parallel Expert Speech Encoder.
+
+Runs a Sortformer speaker-diarization expert and an ASR Conformer encoder on the
+same mel input, then fuses their outputs (LayerNorm + sinusoidal speaker-kernel +
+ADD). Expects un-normalised mels; the ASR branch re-applies ``normalize_batch``
+internally. I/O matches :class:`ConformerEncoder` (drop-in). Only self-contained PE
+bundles (inline ``asr_encoder_cfg`` + ``diarization_model_cfg`` in
+``model_config.yaml``) are supported.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import math
+import os
+import tarfile
+from typing import List, Optional, Union
+
+import torch
+import torch.distributed as dist
+from lightning.pytorch import Trainer
+from omegaconf import DictConfig, OmegaConf
+from torch import nn
+from tqdm import tqdm
+
+from nemo.collections.asr.modules.conformer_encoder import ConformerEncoder
+from nemo.collections.asr.parts.preprocessing.features import normalize_batch
+from nemo.core.classes import ModelPT
+from nemo.core.classes.common import PretrainedModelInfo
+from nemo.core.classes.module import freeze, unfreeze
+from nemo.utils import logging
+from nemo.utils.decorators import experimental
+
+__all__ = [
+    'ParallelExpertEncoder',
+    'ParallelExpertEncoderPT',
+]
+
+
+@contextlib.contextmanager
+def _default_dtype(dtype: torch.dtype):
+    """Temporarily set the global default float dtype.
+
+    Makes ``SortformerModules.init_streaming_state`` allocate its dtype-less
+    speaker-cache / FIFO buffers in the diarizer's dtype, avoiding fp32/bf16 mismatch.
+    """
+    prev = torch.get_default_dtype()
+    if dtype == prev or not dtype.is_floating_point:
+        yield
+        return
+    torch.set_default_dtype(dtype)
+    try:
+        yield
+    finally:
+        torch.set_default_dtype(prev)
+
+
+@contextlib.contextmanager
+def _disable_dist_feature_sync():
+    """Temporarily make ``torch.distributed`` look uninitialized.
+
+    Skips the cross-rank ``all_reduce`` in ``SortformerEncLabelModel.forward_streaming``,
+    which is unnecessary and unsafe for single-recording inference (e.g. a vLLM worker).
+    The original ``dist.is_initialized`` is always restored.
+    """
+    if not (hasattr(dist, "is_initialized") and dist.is_initialized()):
+        yield
+        return
+    orig_is_initialized = dist.is_initialized
+    dist.is_initialized = lambda: False
+    try:
+        yield
+    finally:
+        dist.is_initialized = orig_is_initialized
+
+
+def _clone_config(config: Optional[DictConfig]) -> Optional[DictConfig]:
+    """Deep-copy a ``DictConfig`` without resolving interpolations.
+
+    ``from_config_dict`` mutates its input in place, so sub-target builders get a copy.
+    """
+    if config is None:
+        return None
+    return OmegaConf.create(OmegaConf.to_container(config, resolve=False))
+
+
+@experimental
+class ParallelExpertEncoderPT(ModelPT):
+    """ModelPT shell so a :class:`ParallelExpertEncoder` can be saved/restored as a
+    ``.nemo`` archive (inline ``asr_encoder_cfg`` + ``diarization_model_cfg``).
+    """
+
+    def __init__(self, cfg: DictConfig, trainer: Optional[Trainer] = None):
+        super().__init__(cfg=cfg, trainer=trainer)
+        self.encoder = ParallelExpertEncoder(
+            asr_encoder_cfg=self._cfg.get('asr_encoder_cfg', None),
+            diarization_model_cfg=self._cfg.get('diarization_model_cfg', None),
+            asr_normalize_type=self._cfg.get('asr_normalize_type', None),
+            freeze_diar=self._cfg.get('freeze_diar', True),
+            freeze_asr=self._cfg.get('freeze_asr', False),
+            online_inference_length=self._cfg.get('online_inference_length', 500),
+            chunk_left_context=self._cfg.get('chunk_left_context', 50),
+            chunk_right_context=self._cfg.get('chunk_right_context', 50),
+            diar_fifo_len=self._cfg.get('diar_fifo_len', 40),
+            diar_spkcache_update_period=self._cfg.get('diar_spkcache_update_period', 300),
+            diar_spkcache_len=self._cfg.get('diar_spkcache_len', 188),
+        )
+
+    @classmethod
+    def list_available_models(cls) -> List[PretrainedModelInfo]:
+        return []
+
+    def setup_training_data(self, train_data_config: Union[DictConfig, dict]):
+        pass
+
+    def setup_validation_data(self, val_data_config: Union[DictConfig, dict]):
+        pass
+
+    @staticmethod
+    def is_pe_nemo(nemo_path: str) -> bool:
+        """Detect whether a ``.nemo`` archive is a :class:`ParallelExpertEncoderPT` bundle.
+
+        Reads only ``model_config.yaml`` and checks its ``target:``.
+
+        Args:
+            nemo_path (str): Path to a ``.nemo`` archive.
+
+        Returns:
+            ``True`` if ``target`` ends with ``ParallelExpertEncoderPT``, else ``False``.
+        """
+        if not (isinstance(nemo_path, str) and nemo_path.endswith('.nemo') and os.path.isfile(nemo_path)):
+            return False
+        try:
+            with tarfile.open(nemo_path, mode='r') as tf:
+                for member in tf.getmembers():
+                    if os.path.basename(member.name) == 'model_config.yaml':
+                        fobj = tf.extractfile(member)
+                        if fobj is None:
+                            return False
+                        cfg = OmegaConf.create(fobj.read().decode('utf-8'))
+                        return str(cfg.get('target', '')).endswith('ParallelExpertEncoderPT')
+        except (tarfile.TarError, OSError) as exc:
+            logging.warning("[ParallelExpertEncoder] Could not inspect %s: %s", nemo_path, exc)
+            return False
+        return False
+
+    @classmethod
+    def load_from_nemo(
+        cls,
+        model_path_or_name: str,
+        *,
+        map_location: Union[str, torch.device] = 'cpu',
+        strict: bool = True,
+    ) -> ParallelExpertEncoder:
+        """Load a self-contained PE bundle and return its inner encoder.
+
+        Follows the standard NeMo :class:`~nemo.core.classes.common.Model`
+        convention for resolving a checkpoint reference:
+
+        * a local ``.nemo`` file is restored with :meth:`ModelPT.restore_from`;
+        * otherwise ``model_path_or_name`` is treated as a pretrained model
+          identifier -- a HuggingFace Hub repo id (``{repo}/{name}``) or an NGC
+          alias -- and resolved with :meth:`Model.from_pretrained`, which
+          downloads/caches the ``.nemo`` (honouring the HuggingFace cache and
+          ``HF_HUB_OFFLINE``, so a prefetched cache works on offline nodes).
+
+        This mirrors ``speechlm2.parts.pretrained.load_pretrained_nemo`` so PE
+        bundles load uniformly from local files or model cards.
+
+        Args:
+            model_path_or_name (str): Local ``.nemo`` path or pretrained model id.
+            map_location (str | torch.device): Device to map weights onto.
+            strict (bool): Enforce exact state-dict match.
+
+        Returns:
+            The restored :class:`ParallelExpertEncoder`.
+        """
+        if (
+            isinstance(model_path_or_name, str)
+            and model_path_or_name.endswith('.nemo')
+            and os.path.isfile(model_path_or_name)
+        ):
+            bundle = cls.restore_from(
+                restore_path=model_path_or_name,
+                map_location=map_location,
+                strict=strict,
+            )
+        else:
+            bundle = cls.from_pretrained(
+                model_name=model_path_or_name,
+                map_location=map_location,
+                strict=strict,
+            )
+        return bundle.encoder
+
+    @classmethod
+    def save_to_nemo(
+        cls,
+        encoder: ParallelExpertEncoder,
+        output_nemo_path: str,
+        *,
+        template_bundle_path: str,
+    ) -> None:
+        """Save ``encoder`` as a self-contained PE ``.nemo``, reusing ``model_config.yaml``
+        from ``template_bundle_path``.
+
+        The template must describe the same architecture (``d_model``, ``n_spk``);
+        mismatches raise :class:`ValueError` fail-fast.
+
+        Args:
+            encoder (ParallelExpertEncoder): The encoder whose weights are persisted.
+            output_nemo_path (str): Destination ``.nemo`` path.
+            template_bundle_path (str): Existing PE ``.nemo`` whose ``model_config.yaml`` is reused.
+        """
+        if not isinstance(encoder, ParallelExpertEncoder):
+            raise TypeError(f"save_to_nemo expects a ParallelExpertEncoder, " f"got {type(encoder).__name__}")
+        if not os.path.isfile(template_bundle_path):
+            raise FileNotFoundError(f"template_bundle_path does not exist: {template_bundle_path}")
+
+        template_cfg: Optional[DictConfig] = None
+        with tarfile.open(template_bundle_path, mode='r') as tf:
+            for member in tf.getmembers():
+                if os.path.basename(member.name) == 'model_config.yaml':
+                    fobj = tf.extractfile(member)
+                    if fobj is not None:
+                        template_cfg = OmegaConf.create(fobj.read().decode('utf-8'))
+                    break
+        if template_cfg is None:
+            raise RuntimeError(f"Could not read 'model_config.yaml' from template bundle: {template_bundle_path}")
+
+        tmpl_asr = template_cfg.get('asr_encoder_cfg', None)
+        tmpl_diar = template_cfg.get('diarization_model_cfg', None)
+        if tmpl_asr in (None, {}, '') or tmpl_diar in (None, {}, ''):
+            raise ValueError(
+                f"Template bundle {template_bundle_path} is not self-contained "
+                "(asr_encoder_cfg / diarization_model_cfg missing); it cannot be "
+                "used as a save template."
+            )
+
+        tmpl_d_model = int(tmpl_asr.get('d_model', -1))
+        tmpl_n_spk = int(tmpl_diar.get('sortformer_modules', {}).get('num_spks', -1))
+        enc_d_model = int(encoder.d_model)
+        enc_n_spk = int(encoder.n_spk)
+        if tmpl_d_model != enc_d_model:
+            raise ValueError(
+                f"Template asr_encoder_cfg.d_model={tmpl_d_model} does not match "
+                f"encoder.d_model={enc_d_model}; the saved bundle would fail "
+                "strict reload."
+            )
+        if tmpl_n_spk != enc_n_spk:
+            raise ValueError(
+                f"Template diarization_model_cfg.sortformer_modules.num_spks="
+                f"{tmpl_n_spk} does not match encoder.n_spk={enc_n_spk}; the "
+                "saved bundle would fail strict reload."
+            )
+
+        # Fresh PT shell from the template cfg to reuse NeMo's save_to; swap in encoder.
+        shell = cls(cfg=template_cfg, trainer=None)
+        shell.encoder = encoder
+        # Pin `_cfg` to the verbatim template so save_to round-trips it exactly.
+        shell._cfg = template_cfg
+
+        shell.save_to(output_nemo_path)
+        logging.info(
+            "[ParallelExpertEncoder] Saved PE bundle to %s using template config from %s",
+            output_nemo_path,
+            template_bundle_path,
+        )
+
+
+@experimental
+class ParallelExpertEncoder(nn.Module):
+    """Sortformer-diarizer + ASR Conformer encoder; I/O identical to :class:`ConformerEncoder`.
+
+    Reconstructed from inline configs in the PE bundle's ``model_config.yaml``.
+
+    Args:
+        asr_encoder_cfg (DictConfig): Inline config for the ASR-side :class:`ConformerEncoder`.
+        diarization_model_cfg (DictConfig): Inline config for the :class:`SortformerEncLabelModel`.
+        asr_normalize_type (str, optional): Normalization replayed on the ASR branch. Defaults to ``per_feature``.
+        freeze_diar (bool): Freeze the Sortformer parameters. Defaults to ``True``.
+        freeze_asr (bool): Freeze the wrapped ASR ConformerEncoder. Defaults to ``False``.
+        online_inference_length (int): Online-inference window in encoder output frames
+            (default ``500`` ~= 40s); ``<= 0`` disables it.
+        chunk_left_context (int): Left context (output frames) per online window, shared by
+            both branches. Default ``50``.
+        chunk_right_context (int): Right context (output frames) per online window, shared by
+            both branches. Default ``50``.
+        diar_fifo_len (int): Sortformer streaming ``fifo_len``. Default ``40``.
+        diar_spkcache_update_period (int): Sortformer streaming ``spkcache_update_period``. Default ``300``.
+        diar_spkcache_len (int): Sortformer streaming ``spkcache_len``. Default ``188``.
+    """
+
+    def __init__(
+        self,
+        asr_encoder_cfg: DictConfig,
+        diarization_model_cfg: DictConfig,
+        asr_normalize_type: Optional[str] = None,
+        freeze_diar: bool = True,
+        freeze_asr: bool = False,
+        online_inference_length: int = 500,
+        chunk_left_context: int = 50,
+        chunk_right_context: int = 50,
+        diar_fifo_len: int = 40,
+        diar_spkcache_update_period: int = 300,
+        diar_spkcache_len: int = 188,
+    ):
+        super().__init__()
+
+        # Lazy import: SortformerEncLabelModel imports from asr.modules (circular).
+        from nemo.collections.asr.models.sortformer_diar_models import SortformerEncLabelModel
+
+        if asr_encoder_cfg is None or diarization_model_cfg is None:
+            raise ValueError(
+                "ParallelExpertEncoder requires both `asr_encoder_cfg` and "
+                "`diarization_model_cfg`; self-contained PE bundles supply "
+                "these inline in their model_config.yaml."
+            )
+
+        self.asr_encoder = ConformerEncoder.from_config_dict(_clone_config(asr_encoder_cfg))
+        if not isinstance(self.asr_encoder, ConformerEncoder):
+            raise TypeError(
+                f"Expected `asr_encoder_cfg._target_` to instantiate a "
+                f"ConformerEncoder, got {type(self.asr_encoder).__name__} instead."
+            )
+        self.asr_normalize_type = asr_normalize_type or 'per_feature'
+        self._feat_in = self.asr_encoder._feat_in
+
+        self.diarization_model = SortformerEncLabelModel.from_config_dict(_clone_config(diarization_model_cfg))
+
+        self.freeze_diar = freeze_diar
+        self.freeze_asr = freeze_asr
+
+        # Long-form / online inference configuration.
+        self.online_inference_length = int(online_inference_length)
+        # Overlap-and-trim context (output frames) shared by both branches.
+        self.chunk_left_context = max(0, int(chunk_left_context))
+        self.chunk_right_context = max(0, int(chunk_right_context))
+        # Online-inference window + context in input mel frames (constant per session).
+        self.chunk_feat_len = self.online_inference_length * self.asr_encoder.subsampling_factor
+        self.left_ctx_feat_len = self.chunk_left_context * self.asr_encoder.subsampling_factor
+        self.right_ctx_feat_len = self.chunk_right_context * self.asr_encoder.subsampling_factor
+        self.diar_fifo_len = int(diar_fifo_len)
+        self.diar_spkcache_update_period = int(diar_spkcache_update_period)
+        self.diar_spkcache_len = int(diar_spkcache_len)
+
+        self.n_spk = int(self.diarization_model.sortformer_modules.n_spk)
+        self.asr_d_model = self.asr_encoder.d_model
+
+        self.asr_norm = nn.LayerNorm(self.asr_d_model)
+        self.diar_norm = nn.LayerNorm(self.n_spk)
+        self.register_buffer(
+            "diar_kernel",
+            self._build_sinusoid_position_encoding(self.n_spk, self.asr_d_model),
+            persistent=False,
+        )
+
+        if self.freeze_diar:
+            self.diarization_model.eval()
+            for p in self.diarization_model.parameters():
+                p.requires_grad = False
+        if self.freeze_asr:
+            self.asr_encoder.eval()
+            for p in self.asr_encoder.parameters():
+                p.requires_grad = False
+
+    def train(self, mode: bool = True) -> "ParallelExpertEncoder":
+        """Set training mode, but keep frozen sub-branches in eval.
+
+        The parent ``model.train()`` recurses into every sub-module, which would re-enable
+        dropout / BatchNorm stat updates in a frozen branch. This re-asserts ``eval()`` on
+        the frozen Sortformer (and ASR encoder) so their outputs stay deterministic.
+
+        Args:
+            mode (bool): Whether to set training mode (``True``) or eval mode (``False``).
+
+        Returns:
+            ParallelExpertEncoder: ``self``, matching ``nn.Module.train``d.
+        """
+        super().train(mode)
+        if self.freeze_diar:
+            self.diarization_model.eval()
+        if self.freeze_asr:
+            self.asr_encoder.eval()
+        return self
+
+    # ConformerEncoder-compatible properties (drop-in for SALM perception).
+    @property
+    def d_model(self) -> int:
+        return self.asr_d_model
+
+    @property
+    def subsampling_factor(self) -> int:
+        return self.asr_encoder.subsampling_factor
+
+    @property
+    def pre_encode(self):
+        return self.asr_encoder.pre_encode
+
+    # freeze/unfreeze parity (plain nn.Module re-exposing the standalone helpers).
+    def freeze(self) -> None:
+        freeze(self)
+
+    def unfreeze(self, partial: bool = False) -> None:
+        unfreeze(self, partial=partial)
+
+    # Fusion helpers
+    @staticmethod
+    def _build_sinusoid_position_encoding(max_position: int, embedding_dim: int) -> torch.Tensor:
+        """Mirror of ``MSEncDecMultiTaskModel.get_sinusoid_position_encoding``."""
+        position = torch.arange(max_position, dtype=torch.float32).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, embedding_dim, 2, dtype=torch.float32) * -(math.log(10000.0) / embedding_dim)
+        )
+        pe = torch.zeros(max_position, embedding_dim, dtype=torch.float32)
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        return pe
+
+    @staticmethod
+    def _align_diar_frames(spk_targets: torch.Tensor, target_len: int) -> torch.Tensor:
+        """Pad-by-repeat or truncate ``spk_targets`` to ``target_len`` along time."""
+        cur_len = spk_targets.shape[1]
+        if cur_len < target_len:
+            last = spk_targets[:, -1:, :]
+            spk_targets = torch.cat([spk_targets, last.repeat(1, target_len - cur_len, 1)], dim=1)
+        elif cur_len > target_len:
+            spk_targets = spk_targets[:, :target_len, :]
+        return spk_targets
+
+    @staticmethod
+    def _match_module_io(tensor: torch.Tensor, module: nn.Module) -> torch.Tensor:
+        """Cast ``tensor`` to ``module``'s parameter device & dtype (mels arrive fp32, experts run bf16).
+
+        Args:
+            tensor (Tensor): Input to align (e.g. mel features).
+            module (nn.Module): Module whose first parameter sets the target device/dtype.
+
+        Returns:
+            ``tensor`` moved to the module's device/dtype, or unchanged if it has no parameters.
+        """
+        param = next(module.parameters(), None)
+        if param is None:
+            return tensor
+        return tensor.to(device=param.device, dtype=param.dtype)
+
+    def _fuse_diar_and_asr(self, asr_encoded: torch.Tensor, spk_targets: torch.Tensor) -> torch.Tensor:
+        """Fuse ASR states with speaker-activity preds (LayerNorm + sinusoidal kernel + ADD).
+
+        Args:
+            asr_encoded (Tensor): ASR encoder output. Shape ``(B, D, T_asr)``.
+            spk_targets (Tensor): Speaker-activity predictions. Shape ``(B, T_diar, n_spk)``.
+
+        Returns:
+            Fused encoder output. Shape ``(B, D, T_asr)``.
+        """
+        asr_enc_states = asr_encoded.transpose(1, 2)  # (B, T, D)
+        spk_targets = self._align_diar_frames(spk_targets, asr_enc_states.shape[1]).to(asr_enc_states.dtype)
+
+        asr_enc_states = self.asr_norm(asr_enc_states)
+        spk_targets = self.diar_norm(spk_targets)
+        speaker_infusion = torch.matmul(spk_targets, self.diar_kernel.to(spk_targets.dtype))
+        fused = speaker_infusion + asr_enc_states
+
+        return fused.transpose(1, 2)  # (B, D, T)
+
+    # Forward — identical signature to ConformerEncoder.forward
+    def forward(
+        self,
+        audio_signal,
+        length,
+        spk_targets=None,
+    ):
+        """Encode ``audio_signal``, optionally fusing diarization.
+
+        Dispatches to :meth:`_forward` (offline) or :meth:`_forward_online` (long-form,
+        inference-only, when the input exceeds one window).
+
+        Args:
+            audio_signal (Tensor): Un-normalised mel features. Shape ``(B, feat_in, n_frames)``.
+            length (Tensor): Per-sample feature lengths. Shape ``(B,)``.
+            spk_targets (Tensor, optional): ``(B, T, n_spk)`` speaker-activity override (RTTM/oracle);
+                when ``None`` the wrapped Sortformer is run.
+
+        Returns:
+            Tuple ``(outputs, encoded_lengths)`` with ``outputs`` of shape ``(B, D, T_asr)``.
+        """
+        if spk_targets is not None:
+            use_online = False
+        elif self.online_inference_length > 0 and not self.training:
+            # Even if spk_targets is None, use offline if audio is short enough
+            use_online = audio_signal.shape[-1] > self.chunk_feat_len
+        else:
+            use_online = False
+
+        if use_online:
+            return self._forward_online(audio_signal=audio_signal, length=length, spk_targets=spk_targets)
+
+        return self._forward(
+            audio_signal=audio_signal,
+            length=length,
+            spk_targets=spk_targets,
+        )
+
+    def _forward(
+        self,
+        audio_signal,
+        length,
+        spk_targets=None,
+    ):
+        """Offline (non-chunked) forward pass. See :meth:`forward` for argument semantics."""
+        if spk_targets is None:
+            # Cast fp32 mels to the diarizer's device/dtype before its conv subsampling.
+            diar_signal = self._match_module_io(audio_signal, self.diarization_model)
+            diar_length = length.to(device=diar_signal.device)
+            with torch.set_grad_enabled(not self.freeze_diar):
+                emb_seq, emb_seq_length = self.diarization_model.frontend_encoder(
+                    processed_signal=diar_signal,
+                    processed_signal_length=diar_length,
+                    bypass_pre_encode=False,
+                )
+                spk_targets = self.diarization_model.forward_infer(
+                    emb_seq=emb_seq,
+                    emb_seq_length=emb_seq_length,
+                )
+
+        if self.asr_normalize_type:
+            asr_audio_signal, _, _ = normalize_batch(
+                audio_signal,
+                length,
+                normalize_type=self.asr_normalize_type,
+            )
+        else:
+            asr_audio_signal = audio_signal
+        # Cast fp32 mels to the ASR encoder's device/dtype before its conv subsampling.
+        asr_audio_signal = self._match_module_io(asr_audio_signal, self.asr_encoder)
+        asr_length = length.to(device=asr_audio_signal.device)
+
+        with torch.set_grad_enabled(not self.freeze_asr):
+            asr_encoded, asr_encoded_len = self.asr_encoder(
+                audio_signal=asr_audio_signal,
+                length=asr_length,
+            )
+
+        if spk_targets is not None:
+            outputs = self._fuse_diar_and_asr(asr_encoded, spk_targets)
+        else:
+            outputs = asr_encoded
+
+        return outputs, asr_encoded_len
+
+    def _forward_online(self, audio_signal, length, spk_targets=None):
+        """Long-form online inference: a lock-step loop over fixed windows.
+
+        Walks the recording in non-overlapping windows of ``online_inference_length``
+        output frames. Both experts run on the same context-extended slice
+        ``[stt - left : end + right]`` (differing only in normalization): the ASR
+        encoder uses overlap-and-trim, while the streaming Sortformer carries its
+        speaker-cache / FIFO state across windows and trims context internally.
+        Per-window diar outputs are aligned to the ASR frame count, then both buffers
+        are concatenated and fused once.
+
+        Args:
+            audio_signal (Tensor): Un-normalised mel features. Shape ``(B, feat_in, n_frames)``.
+            length (Tensor): Per-sample feature lengths. Shape ``(B,)``.
+            spk_targets (Tensor, optional): ``(B, T, n_spk)`` override; when given, only ASR is chunked.
+
+        Returns:
+            Tuple ``(outputs, encoded_lengths)`` with ``outputs`` of shape ``(B, D, T_asr)``.
+        """
+        total_feat_len = min(audio_signal.shape[-1], int(length.max().item()))
+        num_chunks = max(1, math.ceil(total_feat_len / self.chunk_feat_len))
+
+        # Normalise the whole utterance once (not per chunk) to match offline stats.
+        if self.asr_normalize_type:
+            asr_audio_signal, _, _ = normalize_batch(
+                audio_signal,
+                length,
+                normalize_type=self.asr_normalize_type,
+            )
+        else:
+            asr_audio_signal = audio_signal
+
+        # Match the ASR encoder's device/dtype (mels arrive fp32, encoder runs bf16).
+        asr_audio_signal = self._match_module_io(asr_audio_signal, self.asr_encoder)
+        length = length.to(device=asr_audio_signal.device)
+
+        run_streaming_diar = spk_targets is None
+        if run_streaming_diar:
+            streaming_state, stream_dtype, diar_audio_signal, diar_length = self._init_streaming_diar(
+                audio_signal,
+                length,
+                batch_size=audio_signal.shape[0],
+            )
+            n_spk = self.diarization_model.sortformer_modules.n_spk
+            total_preds = torch.zeros(
+                (diar_audio_signal.shape[0], 0, n_spk),
+                device=diar_audio_signal.device,
+                dtype=stream_dtype,
+            )
+
+        asr_chunks: List[torch.Tensor] = []
+        diar_chunks: List[torch.Tensor] = []
+        asr_encoded_len = torch.zeros_like(length)
+
+        for chunk_idx in tqdm(
+            range(num_chunks),
+            total=num_chunks,
+            desc="PEE online inference",
+            disable=getattr(self, '_suppress_online_pbar', False),
+        ):
+            stt = chunk_idx * self.chunk_feat_len
+            end = min(stt + self.chunk_feat_len, total_feat_len)
+
+            # Shared context-extended window (input mel frames) for both branches.
+            enc_stt = max(stt - self.left_ctx_feat_len, 0)
+            enc_end = min(end + self.right_ctx_feat_len, total_feat_len)
+            left_offset = stt - enc_stt
+            right_offset = enc_end - end
+
+            asr_chunk = asr_audio_signal[:, :, enc_stt:enc_end]
+            chunk_length = (length - enc_stt).clamp(min=0, max=enc_end - enc_stt)
+            with torch.set_grad_enabled(not self.freeze_asr):
+                enc_ctx, _ = self.asr_encoder(audio_signal=asr_chunk, length=chunk_length)
+            # Trim context off in output-frame space using rounded cumulative positions.
+            left_drop = left_offset // self.subsampling_factor
+            core_len = round(end / self.subsampling_factor) - round(stt / self.subsampling_factor)
+            core_len = max(0, min(core_len, enc_ctx.shape[-1] - left_drop))
+            enc_chunk = enc_ctx[:, :, left_drop : left_drop + core_len]
+            asr_chunks.append(enc_chunk)
+            asr_encoded_len += core_len
+            align_target = enc_chunk.shape[-1]
+
+            # Diar branch: stream the same window; Sortformer trims context internally.
+            if run_streaming_diar:
+                prev_len = total_preds.shape[1]
+                diar_chunk = diar_audio_signal[:, :, enc_stt:enc_end].transpose(1, 2)  # (B, t, feat_in)
+                diar_chunk_length = (diar_length - enc_stt).clamp(min=0, max=enc_end - enc_stt)
+                with (
+                    torch.set_grad_enabled(not self.freeze_diar),
+                    _disable_dist_feature_sync(),
+                    _default_dtype(stream_dtype),
+                ):
+                    streaming_state, total_preds = self.diarization_model.forward_streaming_step(
+                        processed_signal=diar_chunk,
+                        processed_signal_length=diar_chunk_length,
+                        streaming_state=streaming_state,
+                        total_preds=total_preds,
+                        left_offset=left_offset,
+                        right_offset=right_offset,
+                    )
+                diar_raw = total_preds[:, prev_len:]
+                # Newly emitted frames, aligned to the ASR chunk (frame-parallel).
+                new_preds = self._align_diar_frames(diar_raw, align_target)
+                diar_chunks.append(new_preds)
+
+        asr_encoded = torch.cat(asr_chunks, dim=2)  # (B, D, T_asr)
+        if run_streaming_diar:
+            spk_targets = torch.cat(diar_chunks, dim=1)  # (B, T_asr, n_spk)
+
+        if spk_targets is not None:
+            outputs = self._fuse_diar_and_asr(asr_encoded, spk_targets)
+        else:
+            outputs = asr_encoded
+
+        return outputs, asr_encoded_len
+
+    def _init_streaming_diar(self, audio_signal: torch.Tensor, length: torch.Tensor, batch_size: int):
+        """Configure the wrapped Sortformer for streaming and build its initial state.
+
+        Args:
+            audio_signal (Tensor): Input mel features. Shape ``(B, feat_in, n_frames)``.
+            length (Tensor): Per-sample feature lengths. Shape ``(B,)``.
+            batch_size (int): Batch size for the streaming state.
+
+        Returns:
+            ``(streaming_state, stream_dtype, diar_audio_signal, diar_length)`` cast onto
+            the diarizer's device & dtype.
+        """
+        sm = self.diarization_model.sortformer_modules
+        sm.chunk_len = self.online_inference_length
+        sm.fifo_len = self.diar_fifo_len
+        sm.spkcache_update_period = self.diar_spkcache_update_period
+        sm.spkcache_len = self.diar_spkcache_len
+        sm._check_streaming_parameters()
+
+        diar_param = next(self.diarization_model.parameters(), None)
+        if diar_param is not None:
+            self.diarization_model.to(diar_param.device)
+            diar_device, stream_dtype = diar_param.device, diar_param.dtype
+        else:
+            diar_device, stream_dtype = audio_signal.device, torch.get_default_dtype()
+
+        diar_audio_signal = audio_signal.to(device=diar_device, dtype=stream_dtype)
+        diar_length = length.to(device=diar_device)
+
+        with _disable_dist_feature_sync(), _default_dtype(stream_dtype):
+            streaming_state = sm.init_streaming_state(
+                batch_size=batch_size,
+                async_streaming=self.diarization_model.async_streaming,
+                device=diar_device,
+            )
+        return streaming_state, stream_dtype, diar_audio_signal, diar_length