Speed up the cache-aware streaming RNNT pipeline - #16231
Open
naymaraq wants to merge 16 commits into
Open
Conversation
…cludes disk I/O pipeline.run() and the request streamer accept preloaded audio_samples; the entry script loads every file at the streaming sample rate before warmup and passes the tensors in, so no file is read from disk inside the timed region. Signed-off-by: naymaraq <dkaramyan@nvidia.com> (cherry picked from commit 22eb62222019c4aeb60af8902ff3b81e0e780ce6)
… is short Signed-off-by: naymaraq <dkaramyan@nvidia.com> (cherry picked from commit 12e293f7b1f9397ea8c879173b5563485dfbe59f)
Signed-off-by: naymaraq <dkaramyan@nvidia.com> (cherry picked from commit e8936b0bfe8046f98b0b553a49c1907595cc838f)
Transducer decoding calls the prediction network once per symbol with sequence length one; cuDNN spends ~0.4 ms of host time per call for ~0.03 ms of device work. Signed-off-by: naymaraq <dkaramyan@nvidia.com> (cherry picked from commit d8b21aa39a7995250b77dc12010b749ecf5873bf)
Preloaded audio stays in host memory for the whole run, which a large manifest cannot afford. preload_audio (true by default, in all five inference configs) turns the preloading off, and the entry script logs how much host memory it is about to take before it loads anything. With preload_audio false the pipeline reads from disk during the run as before, and the reported RTFx includes that disk I/O. Both paths produce identical predictions. The LSTMDropout single-step docstring and commit message quoted per-call host and device times measured on one GPU. Reworded to say what is true anywhere: the call pays cuDNN's fixed per-call setup for a single timestep of work, and how much torch.lstm_cell saves depends on the GPU, the driver and the model size. Also applies black to asr_streaming_infer.py. Signed-off-by: naymaraq <dkaramyan@nvidia.com>
The unrolled single-step branch made forward longer than pylint's docstring-min-length, so C0116 now fires on it under .pylintrc.other, which covers nemo/collections/common. Adds the missing docstring; no behaviour change. Signed-off-by: naymaraq <dkaramyan@nvidia.com>
naymaraq
marked this pull request as ready for review
September 8, 2026 12:03
Collaborator
Author
|
/ok to test 79d5f53 |
The unrolled single-timestep path calls torch.lstm_cell, which lowers to aten::_thnn_fused_lstm_cell on CUDA. The ONNX exporter has no symbolic for that operator, so exporting an RNNT model failed with UnsupportedOperatorError once the prediction network took the unrolled path during tracing (tests/collections/asr/test_asr_exportables.py::TestExportable::test_EncDecRNNTModel_export_to_onnx). A traced graph should hold the general recurrence rather than one unrolled timestep of it in any case, so _single_step_supported now returns False under torch.jit.is_tracing(). Eager decoding is unchanged, and the measured RTFx numbers, which come from eager runs, still stand. Signed-off-by: naymaraq <dkaramyan@nvidia.com>
Collaborator
Author
|
/ok to test f2d36ea |
torch.nn.LSTM casts to float16 under autocast whatever the autocast dtype is, while torch.lstm_cell follows that dtype. With bfloat16 autocast the unrolled single-timestep path therefore returned bfloat16 states where the library path returns float16, and on CPU it ignored autocast and returned float32. Nothing failed at the LSTM itself; the mismatch surfaced later, when RNNTDecoder.batch_copy_states wrote the new state into the cached one: RuntimeError: Index put requires the source and destination dtypes match, got BFloat16 for the destination and Half for the source. which failed tests/collections/asr/test_asr_interctc_models.py ::TestInterCTCLoss::test_forward for EncDecHybridRNNTCTCModel. _single_step_supported now returns False when autocast is enabled, so under autocast the module behaves exactly as it did before. The unrolled path still applies where it was measured: the cache-aware configs set use_amp=false, and the CUDA graph path requires it. Signed-off-by: naymaraq <dkaramyan@nvidia.com>
Collaborator
Author
|
/ok to test a6cde05 |
The run is launch bound, three times more host time than device time, so fusing each conformer layer into fewer kernels shortens the step. Default mode, not reduce-overhead, so no CUDA graphs; dynamic shapes, since the batch shrinks as streams finish. Skipped when the encoder's own CUDA-graph path is enabled. Signed-off-by: naymaraq <dkaramyan@nvidia.com> (cherry picked from commit f3ba9824f58fdd81af7a72547f7ddccb67c687ac)
… length pinned Compiling forward_internal once lets inductor fuse across layer boundaries. The chunk's feature length is pinned with mark_static because the bufferer fixes it; left symbolic it makes the subsampling output length an expression the shape solver cannot divide by, which it reports once per layer per batch width. Signed-off-by: naymaraq <dkaramyan@nvidia.com> (cherry picked from commit 3b6740034d7320aaf8955f6546f1e3fb94d02f89)
… per step update_cache built tgt_slot_ids from the mapping get_context returned, but that mapping always sends each slot to its own position in the batch, so the index_select it drove was a full copy of the cache into an identically ordered tensor on every step. It is removed; new_context is already in batch order. The slot ids of the active batch now live in reusable pinned host and device buffers filled once per step by get_context, so update_cache reuses them instead of building two more tensors, and _reset_slots returns before launching three index_fill_ when no stream has ended. Predictions are unchanged: gathers, not views, are handed to the encoder, since a strided view of the slot prefix changes which bf16 kernels it picks. Signed-off-by: naymaraq <dkaramyan@nvidia.com> (cherry picked from commit 3553d32b7183e5da40dbe59f6c3ad9ba7fa71575)
transcribe_step_for_frames and transcribe_step_for_feature_buffers split the batch by is_last and called cache_aware_transcribe_step twice, so every step on which a stream ended paid for two encoder and two decoder calls. keep_all_outputs now also accepts a bool vector of shape [B]. Given one, the encoder keeps every output and encoder_step clamps the lengths per stream instead, which the decoder already honours, so one call serves the whole batch. A uniform batch still passes the scalar and takes exactly the path it did before, encoder-side trimming included; only a mixed batch builds the vector. Signed-off-by: naymaraq <dkaramyan@nvidia.com> (cherry picked from commit ca7ce65c73cbe432b6dfc2b412912552af8db4a5)
…boundary update unbound its feature buffer into one tensor per stream and read the right paddings back to the host, and the pipeline's preprocess then stacked the buffers again with a cat and rebuilt the paddings as a device tensor. The unbind and the cat cost a copy per stream, and the tolist forced a device synchronization on every step of a launch-bound run. update now returns the batched buffers and the paddings as they already are, and preprocess takes either form. The per-stream scalar assignment that filled the paddings becomes one transfer. The CTC pipeline still works per stream, so it restores the list and the ints at its call site and is unchanged. Signed-off-by: naymaraq <dkaramyan@nvidia.com> (cherry picked from commit 8b666093ce3cd953afe14445298d929fd72cae29)
Contributor
|
[🤖]: Hi @naymaraq 👋, We wanted to let you know that a CICD pipeline for this PR just finished successfully. So it might be time to merge this PR or get some approvals. |
The torch.where call added by "one step for a batch mixing last and non-last chunks" was split over three lines where it fits on one at the project's line length of 119, so black --check failed on the file. Whitespace only. Signed-off-by: naymaraq <dkaramyan@nvidia.com>
Collaborator
Author
|
/ok to test f865039 |
`subsampling.py` states its kernel geometry as `tl.constexpr` wrappers, which `@triton.jit` requires of a global a kernel reads. The host helpers then used those wrappers where an int belongs: `_build_bin_taps` reshapes and slices with `KERNEL` and `STRIDE`, `_as_row` and `_as_window` with `PAD` and `WINDOW_SIZE`. Running that works, because `tl.constexpr` implements `__index__`. Tracing it does not: Dynamo carries the wrapper as an opaque object, so a wrapper used as a slice bound raises `TypeError: slice indices must be integers or None or have an __index__ method`. Any caller that wraps the encoder in `torch.compile` puts `_build_bin_taps` under Dynamo and hits it, before a single Triton kernel is launched. Unwrap the geometry once, next to the wrappers the kernels keep, and use the plain ints in every host helper, the backward pass and both grid lambdas. This is what the comment above the constexpr block already prescribes, and what `_downsampled_length` already did. Behaviour is unchanged: the taps are bit-identical before and after, in eager and under `torch.compile`. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: naymaraq <dkaramyan@nvidia.com>
Collaborator
Author
|
/ok to test c3bb6a7 |
Contributor
|
[🤖]: Hi @naymaraq 👋, We wanted to let you know that a CICD pipeline for this PR just finished successfully. So it might be time to merge this PR or get some approvals. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Important
The
Update branchbutton must only be pressed in very rare occassions.An outdated branch is never blocking the merge of a PR.
Please reach out to the automation team before pressing that button.
What does this PR do ?
Speed up the cache-aware streaming RNNT inference.
Collection: [ASR]
Changelog
Nine ideas, one commit each, in the order they stack. Four came from one autoresearch campaign and
five from two later ones; each was accepted on its own A/B benchmark before it was picked here.
asr_streaming_infer.py,base_pipeline.py,multi_stream.py, the five inference YAMLs(A1 fix-rtfx): read every audio file once before the timed loop and pass the tensors into
pipeline.run(), so no file is read from disk inside the measured region. This changes what RTFxreports, not how fast the pipeline is. Preloaded audio stays in host memory for the whole run, so
the script logs how much that is before loading, and
preload_audio: falsestreams from disk asbefore for a manifest that does not fit, with the reported RTFx then including disk I/O.
manifest_io.py(A2 longest-first): sort the input manifest longest-first, so the tail of acontinuous-batching run — where no streams are left to refill the batch — is made of short streams.
cache_feature_bufferer.py(A3 batched-feat-shift): shift the whole batch of feature bufferswith one gather, one shift and one scatter instead of a Python loop over slots (about 190 kernel
launches per step on a batch of 64).
rnn.py,LSTMDropout(A4 lstm-cell): unroll the single-timestep eval step withtorch.lstm_cellinstead of the cuDNN call, for unidirectional, non-projected, sequence-firstLSTMs with identity dropout. Transducer decoding calls the prediction network once per symbol with
sequence length one, where cuDNN's fixed per-call setup dominates a single timestep of work on a
small batch; how much this saves depends on the GPU, the driver and the size of the prediction
network. Training, tracing and autocast are unchanged: the unrolled path is skipped in all three.
cache_aware_asr_inference_wrapper.py,cache_aware_ctc_pipeline.py(B1 compile-layers):compile the cache-aware encoder layers with
torch.compile.cache_aware_asr_inference_wrapper.py,cache_aware_rnnt_pipeline.py(B2 compile-encoder-static-len): compile the encoder as one graph with the chunk length pinned,
so the shape solver sees a static length.
context_manager.py(B3 ctx-manager): drop the identity gather and stop rebuilding theno-slot tensors on every step.
cache_aware_rnnt_pipeline.py,cache_aware_rnnt_inference_wrapper.py(B4 merged-step): runone step for a batch that mixes last and non-last chunks, instead of splitting the batch in two.
cache_feature_bufferer.py,cache_aware_rnnt_pipeline.py,cache_aware_asr_inference_wrapper.py(B5 batch-plumbing): keep the batch batched across the feature-bufferer step boundary.
Usage
No API change: the same entry point, the same YAML, the same flags.
Before / After
Full LibriSpeech test-other (2939 utterances, 5.34 h),
nvidia/nemotron-speech-streaming-en-0.6b,cache-aware RNNT, greedy without LM, no CUDA graphs, no end-of-utterance, bfloat16,
num_slots=256,one RTX 5000 Ada (unshared), 1 warmup and 3 measured repeats per point, RTFx reported as the median
over repeats. WER is scored outside the pipeline with
--norm-mode lowercase_punctuation.The branch as a whole
Both sides measured in one session, the baseline re-run on upstream
de26b36at the same time as thecandidate.
What each idea contributes
Two ladders, measured against different bases, so read them separately rather than as one chain.
A1–A4, each rung against the one below, measured in an earlier session against upstream
de26b36:RTFx before → after at that rung, with the change over the rung below.
de26b36B1–B5, each rung against the one below, measured in this session on top of A1–A4:
◇ marks a rung whose measured ranges overlap: the sign is what was measured, but the effect is not
separated from run-to-run variation at that batch size.
Three things these rows say that the totals do not:
the pipeline does the same work at the same speed. Its rung is the size of the disk I/O that used
to be counted.
of a continuous-batching run, which matters more as the batch grows. It is kept because the stack
is faster with it at the batch sizes this model is served at, not because it helps everywhere.
size, and B1 and B2 fade as the batch grows — B2 is within noise at 256.
Peak GPU memory
Peak device memory in use during the run, MiB, at each rung of the B ladder. The A1–A4 rungs are
omitted: they were measured in the earlier session and none of them moved memory by more than the
8 MiB seen between the two baselines below.
de26b36A second benchmark: Earnings-22
The stack was developed and tuned entirely on LibriSpeech test-other, so it was also run end to end
on Earnings-22, which it was never tuned against.
The cache-aware CTC path
Two of the commits touch code the CTC pipeline shares — B1 edits
cache_aware_ctc_pipeline.pyandcache_aware_asr_inference_wrapper.py, B5 edits the feature bufferer and the same wrapper — and nobenchmark above exercises that path. It was therefore run end to end as well, at the
cache_aware_ctc.yamldefaults: the default modelstt_en_fastconformer_hybrid_large_streaming_multi,CTC decoding,
num_slots=1024and end-of-utterance enabled at 800 ms, none of which the RNNTmeasurements above use. Full tier, batch size 256, upstream
de26b36against the tip of this branch.GitHub Actions CI
The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.
Trusted PRs run automatically through copy-pr-bot. For an untrusted PR, a maintainer can trigger CI by commenting
/ok to test <head-sha>; repeat this after a new push if the PR remains untrusted.Before your PR is "Ready for review"
Pre checks:
PR Type:
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.
Additional Information