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

[hamuzhan]

What does this PR do ?

This is an opt-in, off-by-default speedup for the cache-aware FastConformer streaming encoder
steady-state path. Non-steady-state steps, CUDA autocast, capture failures and unseen shapes all
run eager, the number of retained graphs is bounded, and the tests cover the autocast fallback,
cache-buffer aliasing, the train/val lifecycle and the max_graphs limit.

Adds an opt-in CUDA graph path for the cache-aware streaming encoder step
(StreamingEncoder.cache_aware_stream_step). This targets any cache-aware FastConformer streaming
encoder that goes through StreamingEncoder.cache_aware_stream_step (Parakeet / Nemotron Speech /
stt_*_fastconformer_*_streaming_*); it is not a generic path for all ASR encoders. Low-latency
streaming inference of these models spends most of its time waiting on the host, not computing.
Each streaming step launches around 1.5k small kernels, so the GPU is idle for most of the step
while the host enqueues them.

The steady-state streaming step has fully static shapes. The caches are allocated at full size up
front and only the cache_last_channel_len values change between steps. That means the whole step
can be captured once into a torch.cuda.CUDAGraph and replayed with a single launch. The existing
RNNT/TDT label-looping decoder (GreedyBatchedLabelLoopingComputerBase and the WithOptionalCudaGraphs
infrastructure) already applies this idea to decoding; this PR reuses the same WithOptionalCudaGraphs
enable/disable and fallback design for the encoder step. The encoder is simpler because it has no
data-dependent control flow, so plain CUDA graphs are enough (no conditional nodes, no cuda-python
dependency).

You turn it on with asr_model.encoder.set_streaming_cuda_graphs(True) or use_cuda_graphs=true in
the cache-aware streaming inference script. For the covered non-autocast configurations the graph
path is expected to preserve eager execution semantics, and the unit tests assert torch.equal
(not allclose) against eager. Steps that are not part of the steady state (the first chunk, the
final keep_all_outputs=True chunk, or any shape we have not seen) just run eager.

A few design notes:

• Stable output buffers. The captured step copies its outputs into buffers owned by the regular
  caching allocator. Without this, the decoder's own CUDA graph aliased the encoder graph's private
  memory pool and corrupted decoding, so it is needed for correctness.
• Safe teardown. Graphs are dropped only after a full device sync. Tearing down a graph pool
  while replays are still in flight is not safe.
• WithOptionalCudaGraphs. Implemented so the existing ASRModel train/val hooks disable and
  re-enable the encoder graphs the same way they already do for the decoder (graphs do not release
  memory during training).
• Graph key. A step is captured/replayed only when the input shape, dtype, device, the
  streaming params that change the computation (att_context_size, last_channel_cache_size,
  valid_out_len, drop_extra_pre_encoded) and keep_all_outputs all match. The
  cache_last_channel_len value (which changes during the cache-fill phase) is a graph input, not
  part of the key: its shape is fixed, so it is just copied into the static input buffer each step.
• CUDA autocast runs eager. A captured graph fixes the dtype/kernel/workspace choices at capture
  time, so replaying a non-autocast graph under autocast (or the reverse) can silently return
  wrong-precision outputs. We could add the autocast state to the graph key, but a wrong key there
  is a silent-correctness bug, so the graph path is skipped whenever CUDA autocast is active
  (cache-aware streaming models run in float32 anyway). There is a test for this. The autocast check
  handles both the device-specific and the no-argument torch.is_autocast_enabled signatures for
  PyTorch compatibility.
• Eager fallback if capture ever fails; after a failure the mode is set to no-graphs so the
  same shape is not retried on every step (unless a mode is forced for testing).

Results on nvidia/nemotron-3.5-asr-streaming-0.6b (GH200, fp32/TF32), steady-state step p50:

att_context_size (chunk)	batch	eager	graphed	speedup	real-time streams (est.)
[56, 0] (80 ms)	1	22.5 ms	4.4 ms	5.14x	4 -> 18
[56, 0] (80 ms)	32	27.0 ms	8.8 ms	3.08x	95 -> 292
[56, 3] (320 ms)	8	28.9 ms	8.4 ms	3.45x	89 -> 305
[56, 3] (320 ms)	32	31.0 ms	13.3 ms	2.32x	331 -> 767
[56, 13] (1120 ms)	32	39.4 ms	23.2 ms	1.69x	911 -> 1543

The win is biggest in the low-latency, low-batch case (the voice-agent use case), where launch
overhead dominates the step. End-to-end on the documented script ([56,0], batch 10): 9.87 s ->
3.08 s (3.2x).

real-time streams (est.) is a rough capacity estimate, not a full-pipeline benchmark. It is
batch * chunk_duration / step_p50 using the per-step encoder+decoder time measured above (chunk
duration = (right_context + 1) * 80 ms); it does not model audio I/O, tokenization or host-side
scheduling at high concurrency, so treat it as an upper bound that scales with the step speedup.

Memory overhead (extra peak device memory retained by the graph pool plus static input/output
buffers vs eager, encoder-only, GH200, same model):

att_context_size	batch	eager peak	graphed peak	overhead
[56, 0]	1	2651 MB	2697 MB	+46 MB
[56, 0]	32	3317 MB	3757 MB	+440 MB
[56, 3]	32	3325 MB	3763 MB	+438 MB
[56, 13]	32	3363 MB	3800 MB	+437 MB

The overhead scales with batch (one step of intermediate activations) and is essentially
independent of the look-ahead. Each captured graph retains its pool until disable_cuda_graphs()
or the helper is detached, so at most max_graphs (default 8) pools can be live at once.

Correctness: speech_to_text_cache_aware_streaming_infer.py ... use_cuda_graphs=true compare_vs_offline=true reports 0 differences streaming-vs-offline, identical WER, and transcripts
that match the eager run exactly (checked at [56,13] and [56,0]). Tests cover:

• graphed vs eager encoder outputs, exact (torch.equal), including the cache-fill phase and the
  final keep_all_outputs=True step;
• the static input cache buffers and the stable output cache buffers are distinct (a replay does
  not overwrite its own inputs);
• a new graph is captured after an att_context_size switch (no stale-graph reuse), and the
  max_graphs limit is respected (further keys stay eager);
• eager fallback in training mode, under CUDA autocast, and in forced no-graphs mode;
• disable/re-enable lifecycle, the recursive WithOptionalCudaGraphs toggle, and that the real
  ASRModel.disable_cuda_graphs() / maybe_enable_cuda_graphs() also toggle the encoder graphs;
• no-op without CUDA, and argument validation.

Collection: ASR

Changelog

• Add CudaGraphsStreamingEncoderStep (nemo/collections/asr/parts/submodules/streaming_encoder_cuda_graphs.py):
  keyed capture/replay of the streaming encoder step with warmup, stable output buffers, eager
  fallback, and a forced-mode hook for testing; implements WithOptionalCudaGraphs.
• Add StreamingEncoder.set_streaming_cuda_graphs() and route cache_aware_stream_step through the
  optional graph path (nemo/collections/asr/parts/mixins/streaming.py).
• Extend ASRModel.maybe_enable_cuda_graphs() / disable_cuda_graphs() to also toggle the encoder
  streaming graphs, so the training/validation hooks manage them like the decoder graphs
  (nemo/collections/asr/models/asr_model.py).
• Add a use_cuda_graphs flag to speech_to_text_cache_aware_streaming_infer.py.
• Add CacheAwareASRInferenceWrapper.set_streaming_cuda_graphs() delegation.
• Document the feature in the cache-aware streaming Conformer section of the ASR docs.
• Add tests/collections/asr/test_streaming_encoder_cuda_graphs.py (14 tests): graphed-vs-eager
  correctness for the supported non-autocast CUDA inference path (cache-fill phase, final
  keep_all_outputs fallback), max_graphs behavior, input/output cache-buffer aliasing,
  model-level train/val hooks, and CUDA autocast fallback.

Usage

import nemo.collections.asr as nemo_asr

asr_model = nemo_asr.models.ASRModel.from_pretrained("nvidia/nemotron-3.5-asr-streaming-0.6b").eval()
asr_model.encoder.set_streaming_cuda_graphs(True)  # opt-in; CUDA autocast falls back to eager
# ... then run the usual cache-aware streaming loop (conformer_stream_step / the inference wrapper)

From the simulation script:

python examples/asr/asr_cache_aware_streaming/speech_to_text_cache_aware_streaming_infer.py \
    pretrained_name=nvidia/nemotron-3.5-asr-streaming-0.6b \
    dataset_manifest=<manifest.json> \
    att_context_size="[56,0]" target_lang=en-US \
    use_cuda_graphs=true

GitHub Actions CI

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

• New Feature
• Bugfix
• Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

cc @nithinraok (NeMo core / ASR).

Additional Information

• The feature is opt-in and off by default, so there is no behavior change unless
  set_streaming_cuda_graphs(True) or use_cuda_graphs=true is set. It no-ops without CUDA.
• This extends the existing decoder CUDA-graph work to the encoder side: it reuses the
  WithOptionalCudaGraphs interface and the same enable/disable + graceful-fallback pattern used by
  the label-looping decoders, so the train/val lifecycle behaves identically for both.

Limitations / scope:

• CUDA inference in eval mode only; training and validation disable it via the existing hooks.
• First chunk, final keep_all_outputs=True chunk, unseen shapes, active autocast, and capture
  failures fall back to eager.
• Graph pools are retained until disable_cuda_graphs() / detach (see the memory table above).
• Wired through the encoder setter and the cache-aware streaming script; the streaming Pipeline API
  config wiring is left for a follow-up.
• Not validated for multi-GPU / model-parallel encoders; the graph is keyed by device index and a
  device change re-captures, but DDP inference has not been exercised here.

Follow-ups (separate PRs, not in here): wiring the flag through the streaming Pipeline API config;
cutting the attention op count (SDPA plus positional-embedding precompute); a cache ring-buffer.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[hamuzhan]

@nithinraok

[hamuzhan]

@artbataev

[artbataev]

@hamuzhan thank you so much for your contribution! The speedup is truly impressive!
I'm a bit swamped this week, but I'll make time to review the PR by the end of the week.

[hamuzhan]

@artbataev Thank you for taking the time to review it. Looking forward to your feedback! Have a great week!