From 5a9d07e9dce3b93d007c1502d1b585c800b4f109 Mon Sep 17 00:00:00 2001 From: naymaraq Date: Mon, 7 Sep 2026 21:47:46 +0400 Subject: [PATCH 01/15] asr_streaming_infer: read audio once before the timed loop so RTFx excludes 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 (cherry picked from commit 22eb62222019c4aeb60af8902ff3b81e0e780ce6) --- .../asr_streaming_infer.py | 8 +++++- .../asr/inference/pipelines/base_pipeline.py | 5 +++- .../streaming/framing/multi_stream.py | 27 ++++++++++++++++--- 3 files changed, 34 insertions(+), 6 deletions(-) diff --git a/examples/asr/asr_streaming_inference/asr_streaming_infer.py b/examples/asr/asr_streaming_inference/asr_streaming_infer.py index 29dbe6b5ace5..6bfb6c5fbba2 100644 --- a/examples/asr/asr_streaming_inference/asr_streaming_infer.py +++ b/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,9 @@ def main(cfg): # Build the pipeline pipeline = PipelineBuilder.build_pipeline(cfg) + # Read the audio once, outside the timed region: RTFx measures the pipeline, not disk I/O. + 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 +107,9 @@ 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()) diff --git a/nemo/collections/asr/inference/pipelines/base_pipeline.py b/nemo/collections/asr/inference/pipelines/base_pipeline.py index 1c65a78d5fe6..442a30a187c6 100644 --- a/nemo/collections/asr/inference/pipelines/base_pipeline.py +++ b/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 = {} diff --git a/nemo/collections/asr/inference/streaming/framing/multi_stream.py b/nemo/collections/asr/inference/streaming/framing/multi_stream.py index 65ba10746ea5..8c4a1aac4827 100644 --- a/nemo/collections/asr/inference/streaming/framing/multi_stream.py +++ b/nemo/collections/asr/inference/streaming/framing/multi_stream.py @@ -124,18 +124,28 @@ def __init__( self._progress_bar = None self.processed_streams = set() - def set_audio_filepaths(self, audio_filepaths: list[str], options: list[RequestOptions]) -> None: + def set_audio_filepaths( + self, + audio_filepaths: list[str], + options: list[RequestOptions], + audio_samples: list[torch.Tensor] | None = None, + ) -> None: """ Set the audio filepaths Args: audio_filepaths (list[str]): The list of audio filepaths options (list[RequestOptions]): The list of options + audio_samples (list[torch.Tensor] | None): Preloaded audio samples, one tensor per filepath. + When given, streams are fed from these tensors and no audio file is read from disk. """ if len(audio_filepaths) != len(options): raise ValueError("audio_filepaths and options must have the same length") + if audio_samples is not None and len(audio_samples) != len(audio_filepaths): + raise ValueError("audio_filepaths and audio_samples must have the same length") self.audio_filepaths = audio_filepaths self.options = options + self.audio_samples = audio_samples self.n_audio_files = len(audio_filepaths) self.total_progress_steps = self.n_audio_files * 2 # One step for adding, one for processing self.sid2filepath = {} @@ -183,7 +193,10 @@ def add_stream(self) -> None: options = self.options[self.stream_id] self.sid2filepath[self.stream_id] = audio_filepath self.elapsed_durations[self.stream_id] = 0.0 - stream.load_audio(audio_filepath, options) + if self.audio_samples is not None: + stream.load_audio(self.audio_samples[self.stream_id], options) + else: + stream.load_audio(audio_filepath, options) # Add the stream to the multi streamer self.multi_streamer.add_stream(stream, stream_id=self.stream_id) @@ -284,14 +297,20 @@ def __init__( ) self.right_pad_features = right_pad_features - def set_audio_filepaths(self, audio_filepaths: list[str], options: list[RequestOptions]) -> None: + def set_audio_filepaths( + self, + audio_filepaths: list[str], + options: list[RequestOptions], + audio_samples: list[torch.Tensor] | None = None, + ) -> None: """ Set the audio filepaths Args: audio_filepaths (list[str]): The list of audio filepaths options (list[RequestOptions]): The list of options + audio_samples (list[torch.Tensor] | None): Preloaded audio samples, one tensor per filepath. """ - self.multi_streamer.set_audio_filepaths(audio_filepaths, options) + self.multi_streamer.set_audio_filepaths(audio_filepaths, options, audio_samples=audio_samples) def set_progress_bar(self, progress_bar: ProgressBar) -> None: """ From 8bd91989e1b90dd7c5fa409a50d6310ec1037f25 Mon Sep 17 00:00:00 2001 From: naymaraq Date: Mon, 7 Sep 2026 21:56:51 +0400 Subject: [PATCH 02/15] manifest_io: sort input longest-first so the continuous-batching tail is short Signed-off-by: naymaraq (cherry picked from commit 12e293f7b1f9397ea8c879173b5563485dfbe59f) --- nemo/collections/asr/inference/utils/manifest_io.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/nemo/collections/asr/inference/utils/manifest_io.py b/nemo/collections/asr/inference/utils/manifest_io.py index a5c579f9bc66..3ca96f18a137 100644 --- a/nemo/collections/asr/inference/utils/manifest_io.py +++ b/nemo/collections/asr/inference/utils/manifest_io.py @@ -50,7 +50,8 @@ def prepare_audio_data( Args: audio_file: (str) Path to the audio file, folder or manifest file per_stream_biasing_defaults: default params for per-stream biasing - sort_by_duration: (bool) If True, sort the audio files by duration from shortest to longest + sort_by_duration: (bool) If True, sort the audio files by duration from longest to shortest, so the + continuous-batching tail (when no files are left to refill the batch) is made of short streams Returns: (list[str], list[dict] | None, list[ASRRequestOptions] | None, dict[str, int]) List of audio filepaths, manifest, options and filepath order @@ -98,7 +99,7 @@ def prepare_audio_data( indices = list(range(len(filepaths))) durations = [librosa.get_duration(path=filepaths[i]) for i in indices] indices_with_durations = list(zip(indices, durations)) - indices_with_durations.sort(key=lambda x: x[1]) + indices_with_durations.sort(key=lambda x: x[1], reverse=True) filepaths = [filepaths[i] for i, duration in indices_with_durations] if manifest is not None: # keep manifest in the same order as filepaths for consistency From 60cee0bf8715342e1206dd735aa4cf9da664534a Mon Sep 17 00:00:00 2001 From: naymaraq Date: Mon, 7 Sep 2026 22:23:48 +0400 Subject: [PATCH 03/15] cache_feature_bufferer: shift the whole batch of feature buffers at once Signed-off-by: naymaraq (cherry picked from commit e8936b0bfe8046f98b0b553a49c1907595cc838f) --- .../buffering/cache_feature_bufferer.py | 35 ++++++++++++------- 1 file changed, 22 insertions(+), 13 deletions(-) diff --git a/nemo/collections/asr/inference/streaming/buffering/cache_feature_bufferer.py b/nemo/collections/asr/inference/streaming/buffering/cache_feature_bufferer.py index dbec26c688e6..ee26660512b1 100644 --- a/nemo/collections/asr/inference/streaming/buffering/cache_feature_bufferer.py +++ b/nemo/collections/asr/inference/streaming/buffering/cache_feature_bufferer.py @@ -148,21 +148,27 @@ def preprocess( right_padding = (features.shape[2] - feature_lens).clamp(min=0).to(torch.long) return features, right_padding - def _update_feature_buffer(self, slot_ids: list[int], feat_chunk: Tensor) -> None: + def _update_feature_buffer(self, slot_ids: Tensor, feat_chunk: Tensor) -> Tensor: """ - Add an extracted feature to `feature_buffer` + Shift the buffers of `slot_ids` left by the chunk length and append the new feature chunk. + + The whole batch is shifted with one gather, one concatenation and one scatter instead of three + copies per slot; every slot of a batch holds one distinct stream and every chunk has the same + length, so this moves exactly the same data as the per-slot loop did. Args: - slot_ids (list[int]): list of slot ids - feat_chunk (Tensor): feature chunk of shape (B, F, T) + slot_ids (Tensor): slot indices of the batch, shape (B,). + feat_chunk (Tensor): feature chunk of shape (B, F, T). + Returns: + (Tensor) the updated feature buffers of the batch, shape (B, F, feature_buffer_len). """ - for i, slot_id in enumerate(slot_ids): - chunk_len = feat_chunk[i].shape[-1] - if chunk_len > self.feature_buffer_len: - raise ValueError(f"feat_chunk ({chunk_len}) longer than buffer ({self.feature_buffer_len})") + chunk_len = feat_chunk.shape[-1] + if chunk_len > self.feature_buffer_len: + raise ValueError(f"feat_chunk ({chunk_len}) longer than buffer ({self.feature_buffer_len})") - shifted = self.feature_buffer[slot_id, :, chunk_len:].clone() - self.feature_buffer[slot_id, :, :-chunk_len].copy_(shifted) - self.feature_buffer[slot_id, :, -chunk_len:].copy_(feat_chunk[i]) + buffers = self.feature_buffer.index_select(0, slot_ids) + buffers = torch.cat([buffers[:, :, chunk_len:], feat_chunk], dim=2) + self.feature_buffer.index_copy_(0, slot_ids, buffers) + return buffers def update(self, frames: list[Frame]) -> tuple[list[Tensor], list[int]]: """ @@ -217,8 +223,11 @@ def update(self, frames: list[Frame]) -> tuple[list[Tensor], list[int]]: right_paddings=right_paddings, expected_feat_len=self.feature_chunk_len + self.plus_one, ) - self._update_feature_buffer(slot_ids=slot_ids, feat_chunk=features[:, :, -self.feature_chunk_len :]) - fbuffers = list(self.feature_buffer[slot_ids].unbind(0)) + slot_ids_tensor = torch.tensor(slot_ids, device=self.device, dtype=torch.long) + buffers = self._update_feature_buffer( + slot_ids=slot_ids_tensor, feat_chunk=features[:, :, -self.feature_chunk_len :] + ) + fbuffers = list(buffers.unbind(0)) if len(slots_to_free) > 0: self.free_slots(slots_to_free) From dd1cc75790282f1bc2667390548e87d6951a6d9a Mon Sep 17 00:00:00 2001 From: naymaraq Date: Mon, 7 Sep 2026 22:35:06 +0400 Subject: [PATCH 04/15] LSTMDropout: unroll the single-timestep eval step with torch.lstm_cell 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 (cherry picked from commit d8b21aa39a7995250b77dc12010b749ecf5873bf) --- nemo/collections/common/parts/rnn.py | 47 ++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/nemo/collections/common/parts/rnn.py b/nemo/collections/common/parts/rnn.py index 9be7277c8b17..8c77f41e73f3 100644 --- a/nemo/collections/common/parts/rnn.py +++ b/nemo/collections/common/parts/rnn.py @@ -228,9 +228,56 @@ def __init__( if 'weight' in name or 'bias' in name: v.data *= float(weights_init_scale) + def _single_step_supported(self) -> bool: + """ + Whether the unrolled single-timestep path may stand in for the cuDNN call. + + Only the plain unidirectional, non-projected, sequence-first case is covered, and only in eval + mode, where the inter-layer dropout of `torch.nn.LSTM` and `self.dropout` are both identities. + """ + lstm = self.lstm + return not lstm.bidirectional and lstm.proj_size == 0 and not lstm.batch_first + + def _single_step( + self, x: torch.Tensor, h: Tuple[torch.Tensor, torch.Tensor] + ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + """ + One timestep of the stacked LSTM, layer by layer, with `torch.lstm_cell`. + + Autoregressive transducer decoding calls this module once per symbol with a sequence length of + one. cuDNN spends about 0.4 ms of host time per such call setting up descriptors for 0.03 ms of + device work; `torch.lstm_cell` is the same recurrence without that setup. Accumulation order + inside the gate GEMM differs from cuDNN, so results agree to within low-precision rounding + rather than bit for bit. + Args: + x: (torch.Tensor) input of shape (1, B, input_size). + h: (tuple) `(h_0, c_0)`, each of shape (num_layers, B, hidden_size). + Returns: + (tuple) output of shape (1, B, hidden_size) and the new `(h_n, c_n)`. + """ + h_0, c_0 = h + inp = x[0] + h_n, c_n = [], [] + for layer in range(self.lstm.num_layers): + h_l, c_l = torch.lstm_cell( + inp, + (h_0[layer], c_0[layer]), + getattr(self.lstm, f"weight_ih_l{layer}"), + getattr(self.lstm, f"weight_hh_l{layer}"), + getattr(self.lstm, f"bias_ih_l{layer}"), + getattr(self.lstm, f"bias_hh_l{layer}"), + ) + h_n.append(h_l) + c_n.append(c_l) + inp = h_l + return inp.unsqueeze(0), (torch.stack(h_n), torch.stack(c_n)) + def forward( self, x: torch.Tensor, h: Optional[Tuple[torch.Tensor, torch.Tensor]] = None ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + if not self.training and h is not None and x.shape[0] == 1 and self._single_step_supported(): + return self._single_step(x, h) + x, h = self.lstm(x, h) if self.dropout: From 0d5f191ed0a1df9dbd5539b62dc89e28787e7ea6 Mon Sep 17 00:00:00 2001 From: naymaraq Date: Tue, 8 Sep 2026 15:49:01 +0400 Subject: [PATCH 05/15] Review: bound the audio preloading and drop hardware-specific numbers 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 --- .../asr_streaming_infer.py | 16 ++++++++++------ .../asr_streaming_inference/buffered_ctc.yaml | 1 + .../asr_streaming_inference/buffered_rnnt.yaml | 1 + .../asr_streaming_inference/buffered_salm.yaml | 1 + .../asr_streaming_inference/cache_aware_ctc.yaml | 1 + .../cache_aware_rnnt.yaml | 1 + nemo/collections/common/parts/rnn.py | 9 +++++---- 7 files changed, 20 insertions(+), 10 deletions(-) diff --git a/examples/asr/asr_streaming_inference/asr_streaming_infer.py b/examples/asr/asr_streaming_inference/asr_streaming_infer.py index 6bfb6c5fbba2..7909c5b27ed5 100644 --- a/examples/asr/asr_streaming_inference/asr_streaming_infer.py +++ b/examples/asr/asr_streaming_inference/asr_streaming_infer.py @@ -93,8 +93,15 @@ def main(cfg): # Build the pipeline pipeline = PipelineBuilder.build_pipeline(cfg) - # Read the audio once, outside the timed region: RTFx measures the pipeline, not disk I/O. - audio_samples = [read_audio(path, target_sr=cfg.streaming.sample_rate, mono=True) for path in audio_filepaths] + # 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() @@ -107,9 +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, audio_samples=audio_samples - ) + 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()) @@ -119,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)") diff --git a/examples/asr/conf/asr_streaming_inference/buffered_ctc.yaml b/examples/asr/conf/asr_streaming_inference/buffered_ctc.yaml index b139b5fa30f1..acab9b259af3 100644 --- a/examples/asr/conf/asr_streaming_inference/buffered_ctc.yaml +++ b/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 diff --git a/examples/asr/conf/asr_streaming_inference/buffered_rnnt.yaml b/examples/asr/conf/asr_streaming_inference/buffered_rnnt.yaml index 24a344d40cab..4d2dc75ee8e2 100644 --- a/examples/asr/conf/asr_streaming_inference/buffered_rnnt.yaml +++ b/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 diff --git a/examples/asr/conf/asr_streaming_inference/buffered_salm.yaml b/examples/asr/conf/asr_streaming_inference/buffered_salm.yaml index 75e1912fb84d..66070c19a0a0 100644 --- a/examples/asr/conf/asr_streaming_inference/buffered_salm.yaml +++ b/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 diff --git a/examples/asr/conf/asr_streaming_inference/cache_aware_ctc.yaml b/examples/asr/conf/asr_streaming_inference/cache_aware_ctc.yaml index 7f4ed69152ec..cb2704c82a96 100644 --- a/examples/asr/conf/asr_streaming_inference/cache_aware_ctc.yaml +++ b/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 diff --git a/examples/asr/conf/asr_streaming_inference/cache_aware_rnnt.yaml b/examples/asr/conf/asr_streaming_inference/cache_aware_rnnt.yaml index 823d26f52acc..5a1aca7b361d 100644 --- a/examples/asr/conf/asr_streaming_inference/cache_aware_rnnt.yaml +++ b/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 diff --git a/nemo/collections/common/parts/rnn.py b/nemo/collections/common/parts/rnn.py index 8c77f41e73f3..b4679739cda5 100644 --- a/nemo/collections/common/parts/rnn.py +++ b/nemo/collections/common/parts/rnn.py @@ -245,10 +245,11 @@ def _single_step( One timestep of the stacked LSTM, layer by layer, with `torch.lstm_cell`. Autoregressive transducer decoding calls this module once per symbol with a sequence length of - one. cuDNN spends about 0.4 ms of host time per such call setting up descriptors for 0.03 ms of - device work; `torch.lstm_cell` is the same recurrence without that setup. Accumulation order - inside the gate GEMM differs from cuDNN, so results agree to within low-precision rounding - rather than bit for bit. + one. Each such call pays cuDNN's fixed per-call setup — descriptors and workspace — for a single + timestep of work, and on a small batch that setup, not the arithmetic, is what the step costs; + `torch.lstm_cell` runs the same recurrence without it. How much this saves depends on the GPU, + the driver and the size of the prediction network. Accumulation order inside the gate GEMM + differs from cuDNN, so results agree to within low-precision rounding rather than bit for bit. Args: x: (torch.Tensor) input of shape (1, B, input_size). h: (tuple) `(h_0, c_0)`, each of shape (num_layers, B, hidden_size). From 79d5f5335ca627a849b7d481dce095a666880b6b Mon Sep 17 00:00:00 2001 From: naymaraq Date: Tue, 8 Sep 2026 15:57:19 +0400 Subject: [PATCH 06/15] rnn: document LSTMDropout.forward for pylint 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 --- nemo/collections/common/parts/rnn.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/nemo/collections/common/parts/rnn.py b/nemo/collections/common/parts/rnn.py index b4679739cda5..78ecd8096958 100644 --- a/nemo/collections/common/parts/rnn.py +++ b/nemo/collections/common/parts/rnn.py @@ -276,6 +276,15 @@ def _single_step( def forward( self, x: torch.Tensor, h: Optional[Tuple[torch.Tensor, torch.Tensor]] = None ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + """ + Run the stacked LSTM over `x`, taking the unrolled path for a single inference timestep. + + Args: + x: (torch.Tensor) input of shape (T, B, input_size). + h: (tuple | None) optional initial `(h_0, c_0)`, each of shape (num_layers, B, hidden_size). + Returns: + (tuple) output of shape (T, B, hidden_size) and the final `(h_n, c_n)`. + """ if not self.training and h is not None and x.shape[0] == 1 and self._single_step_supported(): return self._single_step(x, h) From f2d36eaca818bb3f78e7fc73a51d919aba784c63 Mon Sep 17 00:00:00 2001 From: naymaraq Date: Tue, 8 Sep 2026 19:49:24 +0400 Subject: [PATCH 07/15] LSTMDropout: keep the cuDNN path while tracing 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 --- nemo/collections/common/parts/rnn.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/nemo/collections/common/parts/rnn.py b/nemo/collections/common/parts/rnn.py index 78ecd8096958..fcd5faa0f605 100644 --- a/nemo/collections/common/parts/rnn.py +++ b/nemo/collections/common/parts/rnn.py @@ -234,7 +234,12 @@ def _single_step_supported(self) -> bool: Only the plain unidirectional, non-projected, sequence-first case is covered, and only in eval mode, where the inter-layer dropout of `torch.nn.LSTM` and `self.dropout` are both identities. + Tracing keeps the cuDNN path whatever the shapes: `torch.lstm_cell` lowers to + `aten::_thnn_fused_lstm_cell`, which the ONNX exporter has no symbolic for, and a traced graph + should hold the general recurrence rather than one unrolled timestep of it. """ + if torch.jit.is_tracing(): + return False lstm = self.lstm return not lstm.bidirectional and lstm.proj_size == 0 and not lstm.batch_first From a6cde0533d9b17e271f3aa3d832cafd92cf3b478 Mon Sep 17 00:00:00 2001 From: naymaraq Date: Tue, 8 Sep 2026 23:23:11 +0400 Subject: [PATCH 08/15] LSTMDropout: keep the cuDNN path under autocast 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 --- nemo/collections/common/parts/rnn.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/nemo/collections/common/parts/rnn.py b/nemo/collections/common/parts/rnn.py index fcd5faa0f605..e176cb496c26 100644 --- a/nemo/collections/common/parts/rnn.py +++ b/nemo/collections/common/parts/rnn.py @@ -237,9 +237,14 @@ def _single_step_supported(self) -> bool: Tracing keeps the cuDNN path whatever the shapes: `torch.lstm_cell` lowers to `aten::_thnn_fused_lstm_cell`, which the ONNX exporter has no symbolic for, and a traced graph should hold the general recurrence rather than one unrolled timestep of it. + Autocast keeps it too: `torch.nn.LSTM` casts to float16 under autocast whatever the autocast + dtype is, while `torch.lstm_cell` follows that dtype, so the two paths would return states of + different dtypes and a later `batch_copy_states` into a cached state would fail. """ if torch.jit.is_tracing(): return False + if torch.is_autocast_enabled(self.lstm.weight_ih_l0.device.type): + return False lstm = self.lstm return not lstm.bidirectional and lstm.proj_size == 0 and not lstm.batch_first From c3a3b54a1069960e56e78b95fcfd9791e3939714 Mon Sep 17 00:00:00 2001 From: naymaraq Date: Mon, 7 Sep 2026 23:29:28 +0400 Subject: [PATCH 09/15] cache-aware pipeline: compile the encoder layers with torch.compile 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 (cherry picked from commit f3ba9824f58fdd81af7a72547f7ddccb67c687ac) --- .../cache_aware_asr_inference_wrapper.py | 24 +++++++++++++++++++ .../pipelines/cache_aware_rnnt_pipeline.py | 13 ++++++++++ 2 files changed, 37 insertions(+) diff --git a/nemo/collections/asr/inference/model_wrappers/cache_aware_asr_inference_wrapper.py b/nemo/collections/asr/inference/model_wrappers/cache_aware_asr_inference_wrapper.py index 559e70a3876e..4f7f438affd3 100644 --- a/nemo/collections/asr/inference/model_wrappers/cache_aware_asr_inference_wrapper.py +++ b/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,29 @@ 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: + """ + Wrap each encoder layer in ``torch.compile`` so inductor fuses the layer's elementwise work. + + Cache-aware streaming spends far more host time launching kernels than the device spends + running them, so fusing a layer into fewer, larger kernels is what shortens the step. Mode is + the default one, not ``reduce-overhead``: that mode replays through CUDA graphs, which this + pipeline deliberately does not use. Shapes are marked dynamic because the batch shrinks as + streams finish, and a static compile would recompile on every new width. + + 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) number of layers compiled. + """ + layers = getattr(self.asr_model.encoder, "layers", None) + if layers is None: + return 0 + for i, layer in enumerate(layers): + layers[i] = torch.compile(layer, dynamic=True) + return len(layers) + def stream_step(self, *args, **kwargs) -> Any: """ Executes a single streaming step. diff --git a/nemo/collections/asr/inference/pipelines/cache_aware_rnnt_pipeline.py b/nemo/collections/asr/inference/pipelines/cache_aware_rnnt_pipeline.py index 28bafc986d2a..71e84f4fc748 100644 --- a/nemo/collections/asr/inference/pipelines/cache_aware_rnnt_pipeline.py +++ b/nemo/collections/asr/inference/pipelines/cache_aware_rnnt_pipeline.py @@ -90,8 +90,21 @@ def __init__( self.init_text_processor(cfg, itn_model) self.init_nmt_model(nmt_model) self.init_decoding_computer() + self.init_compiled_encoder(cfg) super().__init__() + def init_compiled_encoder(self, cfg: DictConfig) -> None: + """ + Compile the encoder layers unless the encoder is running its own CUDA-graph path. + Args: + cfg: (DictConfig) Configuration parameters. + """ + if cfg.asr.get("use_cuda_graphs", False): + return + compiled = self.asr_model.compile_encoder_layers() + if compiled: + logging.info(f"Compiled {compiled} encoder layers with torch.compile") + def init_decoding_computer(self) -> None: """Initialize ``decoding_computer``.""" self.decoding_computer = None From b9174d8a63ad5a5d39299352e0532d3bf7726e36 Mon Sep 17 00:00:00 2001 From: naymaraq Date: Tue, 8 Sep 2026 00:27:33 +0400 Subject: [PATCH 10/15] cache-aware pipeline: compile the encoder as one graph with the chunk 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 (cherry picked from commit 3b6740034d7320aaf8955f6546f1e3fb94d02f89) --- .../cache_aware_asr_inference_wrapper.py | 34 +++++++++++++------ .../pipelines/cache_aware_rnnt_pipeline.py | 5 ++- 2 files changed, 25 insertions(+), 14 deletions(-) diff --git a/nemo/collections/asr/inference/model_wrappers/cache_aware_asr_inference_wrapper.py b/nemo/collections/asr/inference/model_wrappers/cache_aware_asr_inference_wrapper.py index 4f7f438affd3..1e16516ae43f 100644 --- a/nemo/collections/asr/inference/model_wrappers/cache_aware_asr_inference_wrapper.py +++ b/nemo/collections/asr/inference/model_wrappers/cache_aware_asr_inference_wrapper.py @@ -143,26 +143,38 @@ def set_streaming_cuda_graphs(self, enabled: bool = True) -> None: def compile_encoder_layers(self) -> int: """ - Wrap each encoder layer in ``torch.compile`` so inductor fuses the layer's elementwise work. + 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 a layer into fewer, larger kernels is what shortens the step. Mode is - the default one, not ``reduce-overhead``: that mode replays through CUDA graphs, which this - pipeline deliberately does not use. Shapes are marked dynamic because the batch shrinks as - streams finish, and a static compile would recompile on every new width. + 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) number of layers compiled. + (int) 1 when the encoder was compiled, 0 when it has no compilable body. """ - layers = getattr(self.asr_model.encoder, "layers", None) - if layers is None: + encoder = self.asr_model.encoder + if not hasattr(encoder, "forward_internal"): return 0 - for i, layer in enumerate(layers): - layers[i] = torch.compile(layer, dynamic=True) - return len(layers) + 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: """ diff --git a/nemo/collections/asr/inference/pipelines/cache_aware_rnnt_pipeline.py b/nemo/collections/asr/inference/pipelines/cache_aware_rnnt_pipeline.py index 71e84f4fc748..acca33a9a285 100644 --- a/nemo/collections/asr/inference/pipelines/cache_aware_rnnt_pipeline.py +++ b/nemo/collections/asr/inference/pipelines/cache_aware_rnnt_pipeline.py @@ -101,9 +101,8 @@ def init_compiled_encoder(self, cfg: DictConfig) -> None: """ if cfg.asr.get("use_cuda_graphs", False): return - compiled = self.asr_model.compile_encoder_layers() - if compiled: - logging.info(f"Compiled {compiled} encoder layers with torch.compile") + if self.asr_model.compile_encoder_layers(): + logging.info("Compiled the encoder with torch.compile") def init_decoding_computer(self) -> None: """Initialize ``decoding_computer``.""" From 1d95cc4a542e78da64c0218f3853b373df48d95b Mon Sep 17 00:00:00 2001 From: naymaraq Date: Tue, 8 Sep 2026 01:58:43 +0400 Subject: [PATCH 11/15] context manager: drop the identity gather and rebuild no slot tensors 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 (cherry picked from commit 3553d32b7183e5da40dbe59f6c3ad9ba7fa71575) --- .../asr/inference/utils/context_manager.py | 54 +++++++++++++------ 1 file changed, 39 insertions(+), 15 deletions(-) diff --git a/nemo/collections/asr/inference/utils/context_manager.py b/nemo/collections/asr/inference/utils/context_manager.py index 577292cc49bf..705d4274ce2f 100644 --- a/nemo/collections/asr/inference/utils/context_manager.py +++ b/nemo/collections/asr/inference/utils/context_manager.py @@ -89,6 +89,27 @@ def reset(self) -> None: self.cache_last_channel_len, # B ) = self.cache_aware_model.get_initial_cache_state(self.num_slots) self.device = self.cache_last_channel.device + # reusable buffers for the slot ids of the active batch, so no tensor is built per step + self._slot_ids_host = torch.empty(self.num_slots, dtype=torch.long, pin_memory=True) + self._slot_ids_dev = torch.empty(self.num_slots, dtype=torch.long, device=self.device) + self._slot_ids_list: list[int] = [] + self._slot_ids_contiguous = False + + def _set_active_slots(self, slot_ids: list[int]) -> Tensor: + """ + Records the slots of the active batch in the reusable buffers and returns them as a tensor. + Args: + slot_ids: slot index per element of the active batch, in batch order + Returns: + (Tensor) the slot ids on the cache device + """ + self._slot_ids_list = slot_ids + self._slot_ids_contiguous = slot_ids == list(range(len(slot_ids))) + host = self._slot_ids_host[: len(slot_ids)] + host.copy_(torch.tensor(slot_ids, dtype=torch.long)) + dev = self._slot_ids_dev[: len(slot_ids)] + dev.copy_(host, non_blocking=True) + return dev def _reset_slots(self, slot_ids: list[int]) -> None: """ @@ -96,7 +117,7 @@ def _reset_slots(self, slot_ids: list[int]) -> None: Args: slot_ids: list of slot indices to reset """ - if self.cache_disabled: + if self.cache_disabled or len(slot_ids) == 0: return slot_ids_tensor = torch.tensor(slot_ids, device=self.device, dtype=torch.long) @@ -124,19 +145,19 @@ def update_cache(self, stream_ids: list[int], new_context: CacheAwareContext, ma return slot_ids_list = [self.streamidx2slotidx[sid] for sid in stream_ids] - slot_ids = torch.tensor(slot_ids_list, device=self.device, dtype=torch.long) - tgt_slot_ids = torch.tensor( - [mapping[sid] for sid in slot_ids_list], - device=self.device, - dtype=torch.long, - ) + # `mapping` sends each slot to its position in the batch, so the gather it would drive is the + # identity and is skipped; new_context is already in batch order. + num = len(slot_ids_list) + + if slot_ids_list == self._slot_ids_list: + slot_ids = self._slot_ids_dev[:num] + else: + slot_ids = torch.tensor(slot_ids_list, device=self.device, dtype=torch.long) # In-place copy along batch/slot dimension - self.cache_last_channel.index_copy_(1, slot_ids, new_context.cache_last_channel.index_select(1, tgt_slot_ids)) - self.cache_last_time.index_copy_(1, slot_ids, new_context.cache_last_time.index_select(1, tgt_slot_ids)) - self.cache_last_channel_len.index_copy_( - 0, slot_ids, new_context.cache_last_channel_len.index_select(0, tgt_slot_ids) - ) + self.cache_last_channel.index_copy_(1, slot_ids, new_context.cache_last_channel) + self.cache_last_time.index_copy_(1, slot_ids, new_context.cache_last_time) + self.cache_last_channel_len.index_copy_(0, slot_ids, new_context.cache_last_channel_len) def reset_slots(self, stream_ids: list[int], eos_flags: list[bool]) -> None: """ @@ -182,9 +203,12 @@ def get_context(self, stream_ids: list[int]) -> tuple[CacheAwareContext, dict]: # get the cache for the particular stream_ids slot_ids = [self.streamidx2slotidx[stream_id] for stream_id in stream_ids] - cache_last_channel = self.cache_last_channel[:, slot_ids, :, :] - cache_last_time = self.cache_last_time[:, slot_ids, :, :] - cache_last_channel_len = self.cache_last_channel_len[slot_ids] + slot_ids_dev = self._set_active_slots(slot_ids) + # a gather, not a view of the slot prefix: the encoder is sensitive to the layout of the cache + # it reads, and a strided view changes which bf16 kernels it picks + cache_last_channel = self.cache_last_channel.index_select(1, slot_ids_dev) + cache_last_time = self.cache_last_time.index_select(1, slot_ids_dev) + cache_last_channel_len = self.cache_last_channel_len.index_select(0, slot_ids_dev) # create a context object context = CacheAwareContext( From ed6e4ff257de4d334a1d8583a35bc86274a10f47 Mon Sep 17 00:00:00 2001 From: naymaraq Date: Tue, 8 Sep 2026 02:10:18 +0400 Subject: [PATCH 12/15] cache aware rnnt: one step for a batch mixing last and non-last chunks 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 (cherry picked from commit ca7ce65c73cbe432b6dfc2b412912552af8db4a5) --- .../cache_aware_rnnt_inference_wrapper.py | 29 ++++--- .../pipelines/cache_aware_rnnt_pipeline.py | 80 +++++++------------ 2 files changed, 51 insertions(+), 58 deletions(-) diff --git a/nemo/collections/asr/inference/model_wrappers/cache_aware_rnnt_inference_wrapper.py b/nemo/collections/asr/inference/model_wrappers/cache_aware_rnnt_inference_wrapper.py index fe36fa711b0e..a70d45ed2811 100644 --- a/nemo/collections/asr/inference/model_wrappers/cache_aware_rnnt_inference_wrapper.py +++ b/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,12 @@ 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 +167,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 +180,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 +212,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 +288,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 +301,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]. diff --git a/nemo/collections/asr/inference/pipelines/cache_aware_rnnt_pipeline.py b/nemo/collections/asr/inference/pipelines/cache_aware_rnnt_pipeline.py index acca33a9a285..c24d0c0f0fd5 100644 --- a/nemo/collections/asr/inference/pipelines/cache_aware_rnnt_pipeline.py +++ b/nemo/collections/asr/inference/pipelines/cache_aware_rnnt_pipeline.py @@ -308,7 +308,7 @@ def _streaming_step( context, previous_hypotheses: list[Hypothesis | None], drop_extra_pre_encoded: int, - keep_all_outputs: bool, + keep_all_outputs: bool | Tensor, prompt_vectors: Tensor | None, ) -> tuple[list[Hypothesis], object]: """ @@ -437,7 +437,7 @@ def cache_aware_transcribe_step( features: list[Tensor], right_paddings: list[int], ready_state_ids: set, - keep_all_outputs: bool = False, + keep_all_outputs: bool | Tensor = False, ) -> None: """ Cache Aware Transcribe Step @@ -456,7 +456,8 @@ def cache_aware_transcribe_step( features: (list[Tensor]) List of feature buffers. right_paddings: (list[int] | None) List of right paddings. ready_state_ids: (set) Set of ready state IDs. - 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, so one call can serve last and non-last chunks together. """ feature_buffers, feature_buffer_lens = self.preprocess(features, right_paddings) @@ -524,6 +525,24 @@ def cache_aware_transcribe_step( if eos: state.reset_beam_decoding_state_() + def _keep_mask(self, is_last: list[bool]) -> bool | Tensor: + """ + The keep_all_outputs argument for a batch of streams. + + A batch that is all last or all non-last keeps the cheaper scalar form, which lets the encoder + trim the right context itself; only a mixed batch needs the per-stream vector, which is what + lets one call serve it instead of two. + Args: + is_last: (list[bool]) whether each stream of the batch is on its last chunk + Returns: + (bool | Tensor) scalar when the batch is uniform, otherwise a bool vector of shape [B] + """ + if all(is_last): + return True + if not any(is_last): + return False + return torch.tensor(is_last, device=self.device, dtype=torch.bool) + def transcribe_step_for_feature_buffers(self, fbuffers: list[FeatureBuffer]) -> None: """ Transcribes the feature buffers in a streaming manner. @@ -534,30 +553,13 @@ def transcribe_step_for_feature_buffers(self, fbuffers: list[FeatureBuffer]) -> """ ready_state_ids = set() - final_fbuffers, final_features = [], [] - nonfinal_fbuffers, nonfinal_features = [], [] - final_right_paddings = [] - - for fbuffer in fbuffers: - feature = fbuffer.features - right_padding = max(0, self.expected_feature_buffer_len - fbuffer.valid_size) - - if fbuffer.is_last: - final_fbuffers.append(fbuffer) - final_features.append(feature) - final_right_paddings.append(right_padding) - else: - nonfinal_fbuffers.append(fbuffer) - nonfinal_features.append(feature) + features = [fbuffer.features for fbuffer in fbuffers] + right_paddings = [max(0, self.expected_feature_buffer_len - fbuffer.valid_size) for fbuffer in fbuffers] + is_last = [fbuffer.is_last for fbuffer in fbuffers] - if len(nonfinal_fbuffers) > 0: + if len(fbuffers) > 0: self.cache_aware_transcribe_step( - nonfinal_fbuffers, nonfinal_features, None, ready_state_ids, keep_all_outputs=False - ) - - if len(final_fbuffers) > 0: - self.cache_aware_transcribe_step( - final_fbuffers, final_features, final_right_paddings, ready_state_ids, keep_all_outputs=True + fbuffers, features, right_paddings, ready_state_ids, keep_all_outputs=self._keep_mask(is_last) ) if len(ready_state_ids) > 0: @@ -580,30 +582,10 @@ def transcribe_step_for_frames(self, frames: list[Frame]) -> None: # streams that contains multiple frames if len(all_fbuffers) > 0: - final_frames, final_fbuffers = [], [] - nonfinal_frames, nonfinal_fbuffers = [], [] - final_right_paddings = [] - - for jdx, bfeature in enumerate(all_fbuffers): - bframe = frames[jdx] - - if bframe.is_last: - final_frames.append(bframe) - final_fbuffers.append(bfeature) - final_right_paddings.append(right_paddings[jdx]) - else: - nonfinal_frames.append(bframe) - nonfinal_fbuffers.append(bfeature) - - if len(nonfinal_frames) > 0: - self.cache_aware_transcribe_step( - nonfinal_frames, nonfinal_fbuffers, None, ready_state_ids, keep_all_outputs=False - ) - - if len(final_frames) > 0: - self.cache_aware_transcribe_step( - final_frames, final_fbuffers, final_right_paddings, ready_state_ids, keep_all_outputs=True - ) + is_last = [frame.is_last for frame in frames] + self.cache_aware_transcribe_step( + frames, all_fbuffers, right_paddings, ready_state_ids, keep_all_outputs=self._keep_mask(is_last) + ) # post-process the ready states if len(ready_state_ids) > 0: From db7871d2423271e430f96fa5a1cf944049a55752 Mon Sep 17 00:00:00 2001 From: naymaraq Date: Tue, 8 Sep 2026 02:19:40 +0400 Subject: [PATCH 13/15] cache aware feature bufferer: keep the batch batched across the step 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 (cherry picked from commit 8b666093ce3cd953afe14445298d929fd72cae29) --- .../pipelines/cache_aware_ctc_pipeline.py | 2 ++ .../pipelines/cache_aware_rnnt_pipeline.py | 33 ++++++++++++++----- .../buffering/cache_feature_bufferer.py | 22 ++++++++----- 3 files changed, 40 insertions(+), 17 deletions(-) diff --git a/nemo/collections/asr/inference/pipelines/cache_aware_ctc_pipeline.py b/nemo/collections/asr/inference/pipelines/cache_aware_ctc_pipeline.py index 2a373ddf87d0..deee07b9c8c3 100644 --- a/nemo/collections/asr/inference/pipelines/cache_aware_ctc_pipeline.py +++ b/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: diff --git a/nemo/collections/asr/inference/pipelines/cache_aware_rnnt_pipeline.py b/nemo/collections/asr/inference/pipelines/cache_aware_rnnt_pipeline.py index c24d0c0f0fd5..dc16135fc25c 100644 --- a/nemo/collections/asr/inference/pipelines/cache_aware_rnnt_pipeline.py +++ b/nemo/collections/asr/inference/pipelines/cache_aware_rnnt_pipeline.py @@ -279,15 +279,31 @@ def get_sep(self) -> str: """Return the separator for the text processor.""" return self.sep - def preprocess(self, buffers: list[Tensor], right_paddings: list[int] | None = None) -> tuple[Tensor, Tensor]: + def preprocess( + self, + buffers: list[Tensor] | Tensor, + right_paddings: list[int] | Tensor | None = None, + ) -> tuple[Tensor, Tensor]: """ Preprocess the feature buffers by stacking them and computing the lengths Args: - buffers: (list[Tensor]) List of feature buffers. - right_paddings: (list[int] | None) List of right paddings. + buffers: (list[Tensor] | Tensor) List of feature buffers, or one batched (B, F, T) tensor. + right_paddings: (list[int] | Tensor | None) Right paddings, per stream or batched. Returns: (tuple[Tensor, Tensor]) Processed feature buffers and their lengths. """ + if isinstance(buffers, Tensor): + # already batched by the bufferer: no stacking, and no host round trip for the lengths + feature_buffers = drop_trailing_features(buffers, self.expected_feature_buffer_len) + feature_buffer_lens = torch.full( + (feature_buffers.shape[0],), feature_buffers.shape[2], device=self.device, dtype=torch.long + ) + if right_paddings is not None: + if not isinstance(right_paddings, Tensor): + right_paddings = torch.tensor(right_paddings, device=self.device) + feature_buffer_lens = feature_buffer_lens - right_paddings + return feature_buffers.to(self.device), feature_buffer_lens + feature_buffers = [f_buffer.unsqueeze_(0) for f_buffer in buffers] # Trim to expected feature buffer length (safeguard for external feature buffer inputs) feature_buffers = [ @@ -295,7 +311,8 @@ def preprocess(self, buffers: list[Tensor], right_paddings: list[int] | None = N ] feature_buffer_lens = torch.tensor([f_buffer.shape[2] for f_buffer in feature_buffers], device=self.device) if right_paddings is not None: - right_paddings = torch.tensor(right_paddings, device=feature_buffer_lens.device) + if not isinstance(right_paddings, Tensor): + right_paddings = torch.tensor(right_paddings, device=feature_buffer_lens.device) feature_buffer_lens = feature_buffer_lens - right_paddings feature_buffers = torch.cat(feature_buffers).to(self.device) return feature_buffers, feature_buffer_lens @@ -434,8 +451,8 @@ def run_greedy_decoder(self, state: CacheAwareRNNTStreamingState, request: Reque def cache_aware_transcribe_step( self, requests: list[Request], - features: list[Tensor], - right_paddings: list[int], + features: list[Tensor] | Tensor, + right_paddings: list[int] | Tensor, ready_state_ids: set, keep_all_outputs: bool | Tensor = False, ) -> None: @@ -453,8 +470,8 @@ def cache_aware_transcribe_step( 8. Update the ready states to indicate that the state is ready for text post-processing Args: requests: (list[Request]) List of requests (frames or feature buffers) to transcribe. - features: (list[Tensor]) List of feature buffers. - right_paddings: (list[int] | None) List of right paddings. + features: (list[Tensor] | Tensor) Feature buffers, per stream or batched (B, F, T). + right_paddings: (list[int] | Tensor | None) Right paddings, per stream or batched. ready_state_ids: (set) Set of ready state IDs. keep_all_outputs: (bool | Tensor) Whether to keep all outputs; a bool vector of shape [B] keeps them per stream, so one call can serve last and non-last chunks together. diff --git a/nemo/collections/asr/inference/streaming/buffering/cache_feature_bufferer.py b/nemo/collections/asr/inference/streaming/buffering/cache_feature_bufferer.py index ee26660512b1..9118ace9befc 100644 --- a/nemo/collections/asr/inference/streaming/buffering/cache_feature_bufferer.py +++ b/nemo/collections/asr/inference/streaming/buffering/cache_feature_bufferer.py @@ -170,17 +170,21 @@ def _update_feature_buffer(self, slot_ids: Tensor, feat_chunk: Tensor) -> Tensor self.feature_buffer.index_copy_(0, slot_ids, buffers) return buffers - def update(self, frames: list[Frame]) -> tuple[list[Tensor], list[int]]: + def update(self, frames: list[Frame]) -> tuple[Tensor, Tensor]: """ Update the feature bufferers with the new frames. + + The buffers and the right paddings stay batched: unbinding them here only to have the caller + stack them again costs a copy per stream, and reading the paddings back to Python costs a + device synchronization on every step. Args: frames (list[Frame]): list of frames with length equal to batch size Returns: - tuple[list[Tensor], list[int]]: feature buffers and right paddings + tuple[Tensor, Tensor]: feature buffers of shape (B, F, T) and right paddings of shape (B,) """ - # if there are no frames, return empty lists + # if there are no frames, return empty batches if len(frames) == 0: - return [], [] + return torch.empty(0, device=self.device), torch.empty(0, dtype=torch.long, device=self.device) # if the stream_id is new, we need to assign a slot to it slot_ids, slots_to_reset, slots_to_free = [], [], [] @@ -203,11 +207,13 @@ def update(self, frames: list[Frame]) -> tuple[list[Tensor], list[int]]: if len(slots_to_reset) > 0: self.reset_slots(slots_to_reset) - right_paddings = torch.zeros(len(frames), dtype=torch.long, device=self.device) + # one transfer for the whole batch instead of a scalar copy per stream + right_paddings = torch.tensor( + [frame.size - frame.valid_size for frame in frames], dtype=torch.long, device=self.device + ) audio_buffers = [] for i, frame in enumerate(frames): slot_id = slot_ids[i] - right_paddings[i] = frame.size - frame.valid_size self.audio_bufferers[slot_id].update(frame) buffer = self.audio_bufferers[slot_id].sample_buffer @@ -227,9 +233,7 @@ def update(self, frames: list[Frame]) -> tuple[list[Tensor], list[int]]: buffers = self._update_feature_buffer( slot_ids=slot_ids_tensor, feat_chunk=features[:, :, -self.feature_chunk_len :] ) - fbuffers = list(buffers.unbind(0)) - if len(slots_to_free) > 0: self.free_slots(slots_to_free) - return fbuffers, right_paddings.tolist() + return buffers, right_paddings From f865039c1a47826961952209ac61d0d83f6eb778 Mon Sep 17 00:00:00 2001 From: naymaraq Date: Wed, 9 Sep 2026 01:01:15 +0400 Subject: [PATCH 14/15] cache_aware_rnnt_inference_wrapper: black formatting 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 --- .../model_wrappers/cache_aware_rnnt_inference_wrapper.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/nemo/collections/asr/inference/model_wrappers/cache_aware_rnnt_inference_wrapper.py b/nemo/collections/asr/inference/model_wrappers/cache_aware_rnnt_inference_wrapper.py index a70d45ed2811..4d449db8a849 100644 --- a/nemo/collections/asr/inference/model_wrappers/cache_aware_rnnt_inference_wrapper.py +++ b/nemo/collections/asr/inference/model_wrappers/cache_aware_rnnt_inference_wrapper.py @@ -145,9 +145,7 @@ def encoder_step( 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) - ) + 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] From 6dadd333d2640924429c79834f6b26c45af9f6da Mon Sep 17 00:00:00 2001 From: naymaraq Date: Thu, 10 Sep 2026 16:59:50 +0400 Subject: [PATCH 15/15] fix(asr): keep the fused subsampling geometry as ints for torch.compile `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 Signed-off-by: naymaraq --- .../asr/parts/triton/subsampling.py | 48 ++++++++++++------- .../collections/asr/test_fast_subsampling.py | 20 ++++++++ 2 files changed, 50 insertions(+), 18 deletions(-) diff --git a/nemo/collections/asr/parts/triton/subsampling.py b/nemo/collections/asr/parts/triton/subsampling.py index 0524b5070e84..481052205a77 100644 --- a/nemo/collections/asr/parts/triton/subsampling.py +++ b/nemo/collections/asr/parts/triton/subsampling.py @@ -71,6 +71,18 @@ TAPS = tl.constexpr(KERNEL.value * KERNEL.value) PAD = tl.constexpr(1 << (max(TAPS.value, WINDOW_SIZE.value) - 1).bit_length()) # a power of two +# The same geometry as plain ints, for the host code below. A ``tl.constexpr`` implements +# ``__index__``, so eager host code can reshape and slice with the wrappers directly, but Dynamo +# does not honour it: under ``torch.compile`` a wrapper used as a slice bound raises TypeError. +# The kernels keep the wrappers; everything outside them uses these. +_KERNEL = KERNEL.value +_STRIDE = STRIDE.value +_NUM_BINS = NUM_BINS.value +_WINDOW_COLS = WINDOW_COLS.value +_WINDOW_SIZE = WINDOW_SIZE.value +_TAPS = TAPS.value +_PAD = PAD.value + def _forward_configs(): return [ @@ -373,18 +385,18 @@ def _backward_kernel( def _downsampled_length(length, pad_total): """Output extent of one strided convolution stage, floor mode. Takes ints or int tensors.""" - return (length + pad_total - KERNEL.value) // STRIDE.value + 1 + return (length + pad_total - _KERNEL) // _STRIDE + 1 def _as_window(row): """Read a flat PAD-wide row back as the KERNEL x WINDOW_COLS window it holds.""" - return row[:, :WINDOW_SIZE].reshape(-1, KERNEL, WINDOW_COLS) + return row[:, :_WINDOW_SIZE].reshape(-1, _KERNEL, _WINDOW_COLS) def _as_row(window): """Flatten a window into the PAD-wide row the kernel loads; PAD is the next power of two.""" - row = window.new_zeros(window.shape[0], PAD) - row[:, :WINDOW_SIZE] = window.reshape(window.shape[0], WINDOW_SIZE) + row = window.new_zeros(window.shape[0], _PAD) + row[:, :_WINDOW_SIZE] = window.reshape(window.shape[0], _WINDOW_SIZE) return row @@ -403,11 +415,11 @@ def _build_bin_taps(depth_weight): zeros stand in for a slice, which the flattened window tile cannot express. """ channels = depth_weight.shape[0] - taps = depth_weight.reshape(channels, KERNEL, KERNEL) - bin0 = depth_weight.new_zeros(channels, KERNEL, WINDOW_COLS) - bin1 = depth_weight.new_zeros(channels, KERNEL, WINDOW_COLS) - bin0[..., :KERNEL] = taps - bin1[..., STRIDE:] = taps + taps = depth_weight.reshape(channels, _KERNEL, _KERNEL) + bin0 = depth_weight.new_zeros(channels, _KERNEL, _WINDOW_COLS) + bin1 = depth_weight.new_zeros(channels, _KERNEL, _WINDOW_COLS) + bin0[..., :_KERNEL] = taps + bin1[..., _STRIDE:] = taps return _as_row(bin0), _as_row(bin1) @@ -438,7 +450,7 @@ def forward( def grid(meta): return ( - triton.cdiv(out_time, meta["TIME_ROWS"]) * triton.cdiv(out_freq, NUM_BINS), + triton.cdiv(out_time, meta["TIME_ROWS"]) * triton.cdiv(out_freq, _NUM_BINS), triton.cdiv(channels, meta["CHANNEL_BLOCK"]), batch_size, ) @@ -489,14 +501,14 @@ def backward(ctx, grad_output): (batch_size, channels, mel_freq, relu_out_freq, out_time, out_freq, pad_start) = ctx.shapes grad_output = grad_output.contiguous() device = grad_output.device - grad_conv_weight = torch.zeros((channels, PAD), device=device, dtype=torch.float32) + grad_conv_weight = torch.zeros((channels, _PAD), device=device, dtype=torch.float32) grad_conv_bias = torch.zeros((channels,), device=device, dtype=torch.float32) - acc_bin0 = torch.zeros((channels, PAD), device=device, dtype=torch.float32) - acc_bin1 = torch.zeros((channels, PAD), device=device, dtype=torch.float32) + acc_bin0 = torch.zeros((channels, _PAD), device=device, dtype=torch.float32) + acc_bin1 = torch.zeros((channels, _PAD), device=device, dtype=torch.float32) grad_depth_bias = torch.zeros((channels,), device=device, dtype=torch.float32) def grid(meta): - tiles = batch_size * triton.cdiv(out_time, meta["TIME_ROWS"]) * triton.cdiv(out_freq, NUM_BINS) + tiles = batch_size * triton.cdiv(out_time, meta["TIME_ROWS"]) * triton.cdiv(out_freq, _NUM_BINS) return triton.cdiv(channels, meta["CHANNEL_BLOCK"]), min(meta["TILE_SPLITS"], tiles) _backward_kernel[grid]( @@ -528,15 +540,15 @@ def grid(meta): grad_output.stride(2), ) # Undo `_build_bin_taps`: each bin accumulated into the columns it read. - grad_depth_weight = (_as_window(acc_bin0)[..., :KERNEL] + _as_window(acc_bin1)[..., STRIDE:]).reshape( - channels, TAPS + grad_depth_weight = (_as_window(acc_bin0)[..., :_KERNEL] + _as_window(acc_bin1)[..., _STRIDE:]).reshape( + channels, _TAPS ) conv_w_dtype, conv_b_dtype, depth_w_dtype, depth_b_dtype = ctx.param_dtypes return ( None, # mel - grad_conv_weight[:, :TAPS].view(channels, 1, KERNEL, KERNEL).to(conv_w_dtype), + grad_conv_weight[:, :_TAPS].view(channels, 1, _KERNEL, _KERNEL).to(conv_w_dtype), grad_conv_bias.to(conv_b_dtype), - grad_depth_weight.view(channels, 1, KERNEL, KERNEL).to(depth_w_dtype), + grad_depth_weight.view(channels, 1, _KERNEL, _KERNEL).to(depth_w_dtype), grad_depth_bias.to(depth_b_dtype), None, # mel_lengths None, # relu_out_lengths diff --git a/tests/collections/asr/test_fast_subsampling.py b/tests/collections/asr/test_fast_subsampling.py index f41821267ce0..64d7d7245fd7 100644 --- a/tests/collections/asr/test_fast_subsampling.py +++ b/tests/collections/asr/test_fast_subsampling.py @@ -382,3 +382,23 @@ def test_state_dict_is_identical_with_and_without_triton(): without_triton.load_state_dict(with_triton.state_dict()) with_triton.load_state_dict(without_triton.state_dict()) + + +@pytest.mark.unit +@pytest.mark.skipif(not CUDA_TRITON_AVAILABLE, reason="CUDA and Triton are required") +def test_bin_taps_survive_dynamo(): + """The geometry constants must reach Dynamo as ints; a ``tl.constexpr`` is not one to it. + + Eager code may index with the wrappers, since ``tl.constexpr`` implements ``__index__``, but + ``torch.compile`` traces rather than runs, and a wrapper used as a slice bound raises + ``TypeError: slice indices must be integers``. A caller that compiles the encoder puts this + function under Dynamo, so the host code holds the plain ints. + """ + from nemo.collections.asr.parts.triton.subsampling import _build_bin_taps + + depth_weight = torch.arange(CONV_CHANNELS * 9, dtype=torch.float32, device="cuda").reshape(CONV_CHANNELS, 9) + + eager = _build_bin_taps(depth_weight) + compiled = torch.compile(_build_bin_taps, backend="eager")(depth_weight) + + assert all(torch.equal(before, after) for before, after in zip(eager, compiled))