Speed up the cache-aware streaming RNNT pipeline

examples/asr/asr_streaming_inference/asr_streaming_infer.py

@@ -48,6 +48,7 @@
 
 from nemo.collections.asr.inference.factory.pipeline_builder import PipelineBuilder
 
+from nemo.collections.asr.inference.utils.audio_io import read_audio
 from nemo.collections.asr.inference.utils.manifest_io import calculate_duration, dump_output, prepare_audio_data
 from nemo.collections.asr.inference.utils.pipeline_eval import (
     calculate_asr_laal,
@@ -92,6 +93,16 @@ def main(cfg):
     # Build the pipeline
     pipeline = PipelineBuilder.build_pipeline(cfg)
 
+    # Read the audio once, outside the timed region, so RTFx measures the pipeline and not disk I/O.
+    # Every file then stays in host memory for the whole run: set preload_audio=false for a manifest
+    # that does not fit, and the pipeline reads from disk as before, with the RTFx below including that.
+    data_dur, durations = calculate_duration(audio_filepaths)
+    audio_samples = None
+    if cfg.get("preload_audio", True):
+        needed_gib = data_dur * cfg.streaming.sample_rate * 4 / 1024**3  # mono float32 samples
+        logging.info(f"Preloading {len(audio_filepaths)} audio files into host memory, about {needed_gib:.1f} GiB")
+        audio_samples = [read_audio(path, target_sr=cfg.streaming.sample_rate, mono=True) for path in audio_filepaths]
+
     # Warmup and run the pipeline
     timer = SimpleTimer()
     measurements = []
@@ -103,7 +114,7 @@ def main(cfg):
         progress_bar = TQDMProgressBar()
         timer.reset()
         timer.start(device=pipeline.device)
-        output = pipeline.run(audio_filepaths, progress_bar=progress_bar, options=options)
+        output = pipeline.run(audio_filepaths, progress_bar=progress_bar, options=options, audio_samples=audio_samples)
         timer.stop(pipeline.device)
         if run_step >= cfg.warmup_steps:
             measurements.append(timer.total_sec())
@@ -113,7 +124,6 @@ def main(cfg):
         logging.warning(
             "RTFx measurement enabled, but warmup_steps=0. At least one warmup step is recommended to measure RTFx."
         )
-    data_dur, durations = calculate_duration(audio_filepaths)
     exec_dur = sum(measurements) / len(measurements)
     rtfx = data_dur / exec_dur if exec_dur > 0 else float('inf')
     logging.info(f"RTFx: {rtfx:.2f} ({data_dur:.2f}s / {exec_dur:.2f}s)")

examples/asr/conf/asr_streaming_inference/buffered_ctc.yaml

@@ -83,6 +83,7 @@ streaming:
 # ========================
 # Pipeline settings
 # ========================
+preload_audio: true                           # Read all audio into host memory before the timed loop, so RTFx measures inference and not disk I/O; set false for a manifest that does not fit in memory
 matmul_precision: high                        # Matrix multiplication precision: highest, high, medium
 log_level: 20                                 # Logging level: 0 (NOTSET), 10 (DEBUG), 20 (INFO), 30 (WARNING), 40 (ERROR), 50 (CRITICAL)
 pipeline_type: buffered                       # Pipeline type: buffered, cache_aware

examples/asr/conf/asr_streaming_inference/buffered_rnnt.yaml

@@ -115,6 +115,7 @@ streaming:
 # ========================
 # Pipeline settings
 # ========================
+preload_audio: true                           # Read all audio into host memory before the timed loop, so RTFx measures inference and not disk I/O; set false for a manifest that does not fit in memory
 matmul_precision: high                     # Matrix multiplication precision: highest, high, medium
 log_level: 20                              # Logging level: 0 (NOTSET), 10 (DEBUG), 20 (INFO), 30 (WARNING), 40 (ERROR), 50 (CRITICAL)
 pipeline_type: buffered                    # Pipeline type: buffered, cache_aware

examples/asr/conf/asr_streaming_inference/buffered_salm.yaml

@@ -69,6 +69,7 @@ streaming:
 # ========================
 # Pipeline settings
 # ========================
+preload_audio: true                           # Read all audio into host memory before the timed loop, so RTFx measures inference and not disk I/O; set false for a manifest that does not fit in memory
 matmul_precision: high                        # Matrix multiplication precision: highest, high, medium
 log_level: 20                                 # Logging level: 0 (NOTSET), 10 (DEBUG), 20 (INFO), 30 (WARNING), 40 (ERROR), 50 (CRITICAL)
 pipeline_type: buffered                       # Pipeline type: buffered, cache_aware

examples/asr/conf/asr_streaming_inference/cache_aware_ctc.yaml

@@ -84,6 +84,7 @@ streaming:
 # ========================
 # Pipeline settings
 # ========================
+preload_audio: true                           # Read all audio into host memory before the timed loop, so RTFx measures inference and not disk I/O; set false for a manifest that does not fit in memory
 matmul_precision: high                        # Matrix multiplication precision: highest, high, medium
 log_level: 20                                 # Logging level: 0 (NOTSET), 10 (DEBUG), 20 (INFO), 30 (WARNING), 40 (ERROR), 50 (CRITICAL)
 pipeline_type: cache_aware                    # Pipeline type: buffered, cache_aware

examples/asr/conf/asr_streaming_inference/cache_aware_rnnt.yaml

@@ -135,6 +135,7 @@ streaming:
 # ========================
 # Pipeline settings
 # ========================
+preload_audio: true                           # Read all audio into host memory before the timed loop, so RTFx measures inference and not disk I/O; set false for a manifest that does not fit in memory
 matmul_precision: high                        # Matrix multiplication precision: highest, high, medium
 log_level: 20                                 # Logging level: 0 (NOTSET), 10 (DEBUG), 20 (INFO), 30 (WARNING), 40 (ERROR), 50 (CRITICAL)
 pipeline_type: cache_aware                    # Pipeline type: buffered, cache_aware

nemo/collections/asr/inference/model_wrappers/cache_aware_asr_inference_wrapper.py

@@ -16,6 +16,7 @@
 
 from typing import Any
 
+import torch
 from torch import Tensor
 
 from nemo.collections.asr.inference.model_wrappers.asr_inference_wrapper import ASRInferenceWrapper
@@ -140,6 +141,41 @@ def set_streaming_cuda_graphs(self, enabled: bool = True) -> None:
         """
         self.asr_model.encoder.set_streaming_cuda_graphs(enabled=enabled)
 
+    def compile_encoder_layers(self) -> int:
+        """
+        Compile the encoder body with ``torch.compile`` so inductor fuses its elementwise work.
+
+        Cache-aware streaming spends far more host time launching kernels than the device spends
+        running them, so fusing the encoder into fewer, larger kernels is what shortens the step.
+        The whole body is compiled as one graph rather than layer by layer, which lets inductor fuse
+        across layer boundaries. Mode is the default one, not ``reduce-overhead``: that mode replays
+        through CUDA graphs, which this pipeline deliberately does not use. The batch width stays
+        dynamic, since it shrinks as streams finish, but the chunk's feature length is pinned; see
+        the note in the wrapper below.
+
+        Compilation happens on the first call and therefore inside the warmup iteration. Fusion
+        changes the order of low-precision arithmetic, so decoded text can differ from eager in the
+        last bits; a run that must match eager exactly should not call this.
+        Returns:
+            (int) 1 when the encoder was compiled, 0 when it has no compilable body.
+        """
+        encoder = self.asr_model.encoder
+        if not hasattr(encoder, "forward_internal"):
+            return 0
+        compiled = torch.compile(encoder.forward_internal, dynamic=True)
+
+        def forward_internal(audio_signal, length, bypass_pre_encode=False, **kwargs):
+            if not bypass_pre_encode:
+                # A cache-aware chunk always carries the bufferer's feature length, so pinning that
+                # dimension states a fact about the workload rather than constraining it. Left
+                # symbolic, the subsampling output length becomes an expression the shape solver
+                # cannot divide by, and it gives up once per layer per distinct batch width.
+                torch._dynamo.mark_static(audio_signal, 2)
+            return compiled(audio_signal, length, bypass_pre_encode=bypass_pre_encode, **kwargs)
+
+        encoder.forward_internal = forward_internal
+        return 1
+
     def stream_step(self, *args, **kwargs) -> Any:
         """
         Executes a single streaming step.

nemo/collections/asr/inference/model_wrappers/cache_aware_rnnt_inference_wrapper.py

@@ -94,7 +94,7 @@ def encoder_step(
         processed_signal_length: Tensor,
         context: CacheAwareContext,
         drop_extra_pre_encoded: int | None,
-        keep_all_outputs: bool,
+        keep_all_outputs: bool | Tensor,
         drop_left_context: int | None = None,
         valid_out_len: int | None = None,
         prompt_vectors: Tensor | None = None,
@@ -106,13 +106,17 @@ def encoder_step(
             processed_signal_length: (Tensor) input signal length tensor.
             context: (CacheAwareContext) context object.
             drop_extra_pre_encoded: (int | None) number of extra pre-encoded frames to drop.
-            keep_all_outputs: (bool) whether to keep all outputs or not.
+            keep_all_outputs: (bool | Tensor) whether to keep all outputs; a bool vector of shape [B]
+                keeps them per stream, which serves a batch mixing last and non-last chunks.
             drop_left_context: (int | None) number of left context frames to drop.
             valid_out_len: (int | None) number of valid output frames.
             prompt_vectors: (Tensor | None) per-stream one-hot language prompts of shape [B, num_prompts].
         Returns:
             (tuple[Tensor, Tensor, CacheAwareContext]) encoder output, encoder output lengths, and new context.
         """
+        # a boolean vector asks for a mixed batch: the encoder keeps every output and the per-stream
+        # right context is trimmed afterwards by clamping the lengths, which the decoder honours
+        per_sample_keep = isinstance(keep_all_outputs, Tensor)
         (
             encoded,
             encoded_len,
@@ -125,7 +129,7 @@ def encoder_step(
             cache_last_channel=context.cache_last_channel,
             cache_last_time=context.cache_last_time,
             cache_last_channel_len=context.cache_last_channel_len,
-            keep_all_outputs=keep_all_outputs,
+            keep_all_outputs=True if per_sample_keep else keep_all_outputs,
             drop_extra_pre_encoded=drop_extra_pre_encoded,
         )
         new_context = CacheAwareContext(
@@ -139,7 +143,10 @@ def encoder_step(
             encoded = encoded[:, :, drop_left_context:]
             encoded_len = encoded_len - drop_left_context
 
-        if valid_out_len and not keep_all_outputs:
+        if valid_out_len and per_sample_keep:
+            # drop right context per stream: the streams that are not last keep valid_out_len frames
+            encoded_len = torch.where(keep_all_outputs, encoded_len, torch.full_like(encoded_len, valid_out_len))
+        elif valid_out_len and not keep_all_outputs:
             # drop right context if any
             encoded = encoded[:, :, :valid_out_len]
             encoded_len = torch.ones_like(encoded_len) * valid_out_len
@@ -158,7 +165,7 @@ def execute_step(
         context: CacheAwareContext,
         previous_hypotheses: list[Hypothesis] | None,
         drop_extra_pre_encoded: int | None,
-        keep_all_outputs: bool,
+        keep_all_outputs: bool | Tensor,
         drop_left_context: int | None = None,
         valid_out_len: int | None = None,
         prompt_vectors: Tensor | None = None,
@@ -171,7 +178,8 @@ def execute_step(
             context: (CacheAwareContext) context object.
             previous_hypotheses: (list[Hypothesis] | None) list of previous hypotheses for RNNT decoding.
             drop_extra_pre_encoded: (int | None) number of extra pre-encoded frames to drop.
-            keep_all_outputs: (bool) whether to keep all outputs or not.
+            keep_all_outputs: (bool | Tensor) whether to keep all outputs; a bool vector of shape [B]
+                keeps them per stream, which serves a batch mixing last and non-last chunks.
             drop_left_context: (int | None) number of left context frames to drop.
             valid_out_len: (int | None) number of valid output frames.
             prompt_vectors: (Tensor | None) Optional prompt vectors of shape [B, num_prompts].
@@ -202,7 +210,7 @@ def malsd_stream_step(
         processed_signal_length: Tensor,
         context: CacheAwareContext,
         drop_extra_pre_encoded: int | None,
-        keep_all_outputs: bool,
+        keep_all_outputs: bool | Tensor,
         drop_left_context: int | None = None,
         valid_out_len: int | None = None,
         prompt_vectors: Tensor | None = None,
@@ -278,7 +286,7 @@ def stream_step(
         context: CacheAwareContext = None,
         previous_hypotheses: list[Hypothesis] | None = None,
         drop_extra_pre_encoded: int | None = None,
-        keep_all_outputs: bool = False,
+        keep_all_outputs: bool | Tensor = False,
         drop_left_context: int | None = None,
         valid_out_len: int | None = None,
         prompt_vectors: Tensor | None = None,
@@ -291,7 +299,8 @@ def stream_step(
             context: (CacheAwareContext) context object.
             previous_hypotheses: (list[Hypothesis] | None) list of previous hypotheses for RNNT decoding.
             drop_extra_pre_encoded: (int | None) number of extra pre-encoded frames to drop.
-            keep_all_outputs: (bool) whether to keep all outputs or not.
+            keep_all_outputs: (bool | Tensor) whether to keep all outputs; a bool vector of shape [B]
+                keeps them per stream, which serves a batch mixing last and non-last chunks.
             drop_left_context: (int | None) number of left context frames to drop.
             valid_out_len: (int | None) number of valid output frames.
             prompt_vectors: (Tensor | None) Optional prompt vectors of shape [B, num_prompts].

nemo/collections/asr/inference/pipelines/base_pipeline.py

@@ -589,6 +589,7 @@ def run(
         audio_filepaths: list[str],
         options: list[ASRRequestOptions] | None = None,
         progress_bar: ProgressBar | None = None,
+        audio_samples: list[Tensor] | None = None,
     ) -> dict:
         """
         Orchestrates reading from audio_filepaths in a streaming manner,
@@ -597,6 +598,8 @@ def run(
             audio_filepaths (list[str]): List of audio filepaths to transcribe.
             options (list[ASRRequestOptions] | None): List of RequestOptions for each stream.
             progress_bar (ProgressBar | None): Progress bar to show the progress. Default is None.
+            audio_samples (list[Tensor] | None): Preloaded audio samples, one tensor per filepath, already at
+                the pipeline's sample rate. When given, no audio file is read from disk during the run.
         Returns:
             dict: A dictionary containing transcriptions and segments for each stream.
         """
@@ -611,7 +614,7 @@ def run(
             raise ValueError("options must be the same length as audio_filepaths")
 
         request_generator = self.get_request_generator()
-        request_generator.set_audio_filepaths(audio_filepaths, options)
+        request_generator.set_audio_filepaths(audio_filepaths, options, audio_samples=audio_samples)
         request_generator.set_progress_bar(progress_bar)
 
         pipeline_output = {}

nemo/collections/asr/inference/pipelines/cache_aware_ctc_pipeline.py

@@ -338,6 +338,8 @@ def transcribe_step_for_frames(self, frames: list[Frame]) -> None:
             frames: (list[Frame]) List of frames to transcribe.
         """
         all_fbuffers, right_paddings = self.bufferer.update(frames)
+        # the bufferer keeps the batch together for the RNNT pipeline; this one still works per stream
+        all_fbuffers, right_paddings = list(all_fbuffers.unbind(0)), right_paddings.tolist()
 
         ready_state_ids = set()
         if len(all_fbuffers) > 0: