Cache-aware Streaming ASR: up to 5x faster encoder with CUDA graphs

docs/source/asr/featured_models.rst

@@ -68,6 +68,16 @@ Streaming models trained with limited right context for real-time inference with
 * Simulation script: ``examples/asr/asr_cache_aware_streaming/speech_to_text_cache_aware_streaming_infer.py``
 * Supports multiple look-aheads with ``att_context_size`` lists
 
+**CUDA graphs for the streaming encoder step (inference):** the steady-state streaming step has
+static shapes, so it can be captured once into a CUDA graph and replayed with a single kernel
+launch per chunk instead of ~10^3 launches, which substantially reduces per-step latency for
+low-latency streaming. For the covered non-autocast configurations the graph path is expected to
+preserve eager execution semantics (unit tests assert exact equality against eager); non-uniform
+steps (first/last chunk) automatically fall back to eager. Enable with
+``asr_model.encoder.set_streaming_cuda_graphs(True)`` or ``use_cuda_graphs=true`` in the
+simulation script above (see
+:class:`~nemo.collections.asr.parts.submodules.streaming_encoder_cuda_graphs.CudaGraphsStreamingEncoderStep`).
+
 Configs: ``examples/asr/conf/fastconformer/cache_aware_streaming/``
 
 .. _RNNT-Prompt_model:

examples/asr/asr_cache_aware_streaming/speech_to_text_cache_aware_streaming_infer.py

@@ -172,6 +172,12 @@ class TranscriptionConfig:
     )
     matmul_precision: str = "high"  # Literal["highest", "high", "medium"]
 
+    # Capture the steady-state encoder streaming step into a CUDA graph and replay it with a
+    # single kernel launch per chunk (removes the per-step host launch overhead). For the covered
+    # non-autocast configs the graph path preserves eager semantics; requires CUDA and non-uniform
+    # steps fall back to eager.
+    use_cuda_graphs: bool = False
+
     # Decoding strategy for CTC models
     ctc_decoding: CTCDecodingConfig = field(default_factory=CTCDecodingConfig)
     # Decoding strategy for RNNT models
@@ -392,6 +398,14 @@ def main(cfg: TranscriptionConfig):
             chunk_size=cfg.chunk_size, left_chunks=cfg.left_chunks, shift_size=shift_size
         )
 
+    # Enable CUDA-graph replay for the encoder streaming step (single launch per steady-state chunk)
+    if cfg.use_cuda_graphs:
+        if device.type != "cuda":
+            raise ValueError("use_cuda_graphs=true requires a CUDA device")
+        if not hasattr(asr_model.encoder, "set_streaming_cuda_graphs"):
+            raise ValueError("Model encoder does not support CUDA graphs for the streaming step.")
+        asr_model.encoder.set_streaming_cuda_graphs(enabled=True)
+
     # In streaming, offline normalization is not feasible as we don't have access to the whole audio at the beginning
     # When online_normalization is enabled, the normalization of the input features (mel-spectrograms) are done per step
     # It is suggested to train the streaming models without any normalization in the input features.

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

@@ -126,6 +126,19 @@ def setup_streaming_params(self, chunk_size: int, shift_size: int) -> None:
         """
         self.asr_model.encoder.setup_streaming_params(chunk_size=chunk_size, shift_size=shift_size)
 
+    def set_streaming_cuda_graphs(self, enabled: bool = True) -> None:
+        """
+        Enable or disable CUDA-graph replay for the encoder streaming step (inference only).
+        Steady-state chunks are captured once into a CUDA graph and replayed with a single
+        kernel launch; for the covered non-autocast configurations the graph path is expected to
+        preserve eager execution semantics, and non-steady-state steps (including under CUDA
+        autocast) fall back to eager. See
+        :class:`~nemo.collections.asr.parts.submodules.streaming_encoder_cuda_graphs.CudaGraphsStreamingEncoderStep`.
+        Args:
+            enabled: (bool) whether to enable CUDA graphs for the encoder streaming step.
+        """
+        self.asr_model.encoder.set_streaming_cuda_graphs(enabled=enabled)
+
     def stream_step(self, *args, **kwargs) -> Any:
         """
         Executes a single streaming step.

nemo/collections/asr/models/asr_model.py

@@ -181,11 +181,16 @@ def maybe_enable_cuda_graphs(self, force_reinit=False) -> bool:
         # check that self.decoding.decoding exists and is instance of WithOptionalCudaGraphs
         if isinstance(getattr(getattr(self, "decoding", None), "decoding", None), WithOptionalCudaGraphs):
             state_changed |= self.decoding.decoding.maybe_enable_cuda_graphs()
-        if state_changed:
-            logging.info(
-                f"CUDA graphs enabled for {type(self).__name__}::{type(self.decoding).__name__}::"
-                f"{type(self.decoding.decoding).__name__}"
-            )
+            if state_changed:
+                logging.info(
+                    f"CUDA graphs enabled for {type(self).__name__}::{type(self.decoding).__name__}::"
+                    f"{type(self.decoding.decoding).__name__}"
+                )
+        # streaming-encoder CUDA graphs (attached via encoder.set_streaming_cuda_graphs)
+        encoder_graphs = getattr(getattr(self, "encoder", None), "_stream_step_cuda_graphs", None)
+        if isinstance(encoder_graphs, WithOptionalCudaGraphs) and encoder_graphs.maybe_enable_cuda_graphs():
+            logging.info(f"CUDA graphs enabled for {type(self).__name__}::encoder streaming step")
+            state_changed = True
         return state_changed
 
     def disable_cuda_graphs(self) -> bool:
@@ -194,11 +199,16 @@ def disable_cuda_graphs(self) -> bool:
         # check that self.decoding.decoding exists and is instance of WithOptionalCudaGraphs
         if isinstance(getattr(getattr(self, "decoding", None), "decoding", None), WithOptionalCudaGraphs):
             state_changed = self.decoding.decoding.disable_cuda_graphs()
-        if state_changed:
-            logging.info(
-                f"CUDA graphs disabled for {type(self).__name__}::{type(self.decoding).__name__}::"
-                f"{type(self.decoding.decoding).__name__}"
-            )
+            if state_changed:
+                logging.info(
+                    f"CUDA graphs disabled for {type(self).__name__}::{type(self.decoding).__name__}::"
+                    f"{type(self.decoding.decoding).__name__}"
+                )
+        # streaming-encoder CUDA graphs (attached via encoder.set_streaming_cuda_graphs)
+        encoder_graphs = getattr(getattr(self, "encoder", None), "_stream_step_cuda_graphs", None)
+        if isinstance(encoder_graphs, WithOptionalCudaGraphs) and encoder_graphs.disable_cuda_graphs():
+            logging.info(f"CUDA graphs disabled for {type(self).__name__}::encoder streaming step")
+            state_changed = True
         return state_changed
 
     def on_train_epoch_start(self) -> None:

nemo/collections/asr/parts/mixins/streaming.py

@@ -13,9 +13,17 @@
 # limitations under the License.
 
 from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from nemo.collections.asr.parts.submodules.streaming_encoder_cuda_graphs import CudaGraphsStreamingEncoderStep
 
 
 class StreamingEncoder(ABC):
+    # Optional CUDA-graph accelerator for `cache_aware_stream_step`, attached by
+    # `set_streaming_cuda_graphs`. Kept as a plain (non-module) attribute.
+    _stream_step_cuda_graphs = None
+
     @abstractmethod
     def setup_streaming_params(
         self,
@@ -37,6 +45,41 @@ def to_numpy(tensor):
             return None
         return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
 
+    def set_streaming_cuda_graphs(
+        self, enabled: bool = True, warmup_steps: int = 3, max_graphs: int = 8
+    ) -> Optional["CudaGraphsStreamingEncoderStep"]:
+        """Enable or disable CUDA-graph replay for `cache_aware_stream_step` (inference only).
+
+        When enabled, steady-state streaming steps are captured once into a
+        :class:`torch.cuda.CUDAGraph` and replayed with a single kernel launch, removing the
+        per-step host launch overhead of the eager encoder. Non-uniform steps (first step,
+        final step with ``keep_all_outputs=True``) automatically run eager. For the covered
+        non-autocast configurations the graph path is expected to preserve eager execution
+        semantics (the tests assert ``torch.equal`` against eager). See
+        :class:`~nemo.collections.asr.parts.submodules.streaming_encoder_cuda_graphs.CudaGraphsStreamingEncoderStep`.
+
+        Args:
+            enabled: attach (True) or remove (False) the CUDA-graph accelerator.
+            warmup_steps: eager calls per unique step shape before capturing it.
+            max_graphs: maximum number of distinct captured graphs kept alive.
+
+        Returns:
+            The attached ``CudaGraphsStreamingEncoderStep`` helper, or None when disabling.
+        """
+        if not enabled:
+            helper = self._stream_step_cuda_graphs
+            if helper is not None:
+                helper.disable_cuda_graphs()
+            self._stream_step_cuda_graphs = None
+            return None
+        # local import: avoid a circular import at module load time
+        from nemo.collections.asr.parts.submodules.streaming_encoder_cuda_graphs import CudaGraphsStreamingEncoderStep
+
+        self._stream_step_cuda_graphs = CudaGraphsStreamingEncoderStep(
+            self, warmup_steps=warmup_steps, max_graphs=max_graphs
+        )
+        return self._stream_step_cuda_graphs
+
     def cache_aware_stream_step(
         self,
         processed_signal,
@@ -47,6 +90,39 @@ def cache_aware_stream_step(
         keep_all_outputs=True,
         drop_extra_pre_encoded=None,
         bypass_pre_encode=False,
+    ):
+        if self._stream_step_cuda_graphs is not None:
+            return self._stream_step_cuda_graphs.stream_step(
+                processed_signal,
+                processed_signal_length=processed_signal_length,
+                cache_last_channel=cache_last_channel,
+                cache_last_time=cache_last_time,
+                cache_last_channel_len=cache_last_channel_len,
+                keep_all_outputs=keep_all_outputs,
+                drop_extra_pre_encoded=drop_extra_pre_encoded,
+                bypass_pre_encode=bypass_pre_encode,
+            )
+        return self._cache_aware_stream_step_impl(
+            processed_signal,
+            processed_signal_length=processed_signal_length,
+            cache_last_channel=cache_last_channel,
+            cache_last_time=cache_last_time,
+            cache_last_channel_len=cache_last_channel_len,
+            keep_all_outputs=keep_all_outputs,
+            drop_extra_pre_encoded=drop_extra_pre_encoded,
+            bypass_pre_encode=bypass_pre_encode,
+        )
+
+    def _cache_aware_stream_step_impl(
+        self,
+        processed_signal,
+        processed_signal_length=None,
+        cache_last_channel=None,
+        cache_last_time=None,
+        cache_last_channel_len=None,
+        keep_all_outputs=True,
+        drop_extra_pre_encoded=None,
+        bypass_pre_encode=False,
     ):
         if self.streaming_cfg is None:
             self.setup_streaming_params()