diff --git a/examples/asr/asr_streaming_inference/asr_streaming_infer.py b/examples/asr/asr_streaming_inference/asr_streaming_infer.py index 29dbe6b5ace5..7909c5b27ed5 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,16 @@ def main(cfg): # Build the pipeline pipeline = PipelineBuilder.build_pipeline(cfg) + # Read the audio once, outside the timed region, so RTFx measures the pipeline and not disk I/O. + # Every file then stays in host memory for the whole run: set preload_audio=false for a manifest + # that does not fit, and the pipeline reads from disk as before, with the RTFx below including that. + data_dur, durations = calculate_duration(audio_filepaths) + audio_samples = None + if cfg.get("preload_audio", True): + needed_gib = data_dur * cfg.streaming.sample_rate * 4 / 1024**3 # mono float32 samples + logging.info(f"Preloading {len(audio_filepaths)} audio files into host memory, about {needed_gib:.1f} GiB") + audio_samples = [read_audio(path, target_sr=cfg.streaming.sample_rate, mono=True) for path in audio_filepaths] + # Warmup and run the pipeline timer = SimpleTimer() measurements = [] @@ -103,7 +114,7 @@ def main(cfg): progress_bar = TQDMProgressBar() timer.reset() timer.start(device=pipeline.device) - output = pipeline.run(audio_filepaths, progress_bar=progress_bar, options=options) + output = pipeline.run(audio_filepaths, progress_bar=progress_bar, options=options, audio_samples=audio_samples) timer.stop(pipeline.device) if run_step >= cfg.warmup_steps: measurements.append(timer.total_sec()) @@ -113,7 +124,6 @@ def main(cfg): logging.warning( "RTFx measurement enabled, but warmup_steps=0. At least one warmup step is recommended to measure RTFx." ) - data_dur, durations = calculate_duration(audio_filepaths) exec_dur = sum(measurements) / len(measurements) rtfx = data_dur / exec_dur if exec_dur > 0 else float('inf') logging.info(f"RTFx: {rtfx:.2f} ({data_dur:.2f}s / {exec_dur:.2f}s)") 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/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..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 @@ -16,6 +16,7 @@ from typing import Any +import torch from torch import Tensor from nemo.collections.asr.inference.model_wrappers.asr_inference_wrapper import ASRInferenceWrapper @@ -140,6 +141,41 @@ def set_streaming_cuda_graphs(self, enabled: bool = True) -> None: """ self.asr_model.encoder.set_streaming_cuda_graphs(enabled=enabled) + def compile_encoder_layers(self) -> int: + """ + Compile the encoder body with ``torch.compile`` so inductor fuses its elementwise work. + + Cache-aware streaming spends far more host time launching kernels than the device spends + running them, so fusing the encoder into fewer, larger kernels is what shortens the step. + The whole body is compiled as one graph rather than layer by layer, which lets inductor fuse + across layer boundaries. Mode is the default one, not ``reduce-overhead``: that mode replays + through CUDA graphs, which this pipeline deliberately does not use. The batch width stays + dynamic, since it shrinks as streams finish, but the chunk's feature length is pinned; see + the note in the wrapper below. + + Compilation happens on the first call and therefore inside the warmup iteration. Fusion + changes the order of low-precision arithmetic, so decoded text can differ from eager in the + last bits; a run that must match eager exactly should not call this. + Returns: + (int) 1 when the encoder was compiled, 0 when it has no compilable body. + """ + encoder = self.asr_model.encoder + if not hasattr(encoder, "forward_internal"): + return 0 + compiled = torch.compile(encoder.forward_internal, dynamic=True) + + def forward_internal(audio_signal, length, bypass_pre_encode=False, **kwargs): + if not bypass_pre_encode: + # A cache-aware chunk always carries the bufferer's feature length, so pinning that + # dimension states a fact about the workload rather than constraining it. Left + # symbolic, the subsampling output length becomes an expression the shape solver + # cannot divide by, and it gives up once per layer per distinct batch width. + torch._dynamo.mark_static(audio_signal, 2) + return compiled(audio_signal, length, bypass_pre_encode=bypass_pre_encode, **kwargs) + + encoder.forward_internal = forward_internal + return 1 + def stream_step(self, *args, **kwargs) -> Any: """ Executes a single streaming step. 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..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 @@ -94,7 +94,7 @@ def encoder_step( processed_signal_length: Tensor, context: CacheAwareContext, drop_extra_pre_encoded: int | None, - keep_all_outputs: bool, + keep_all_outputs: bool | Tensor, drop_left_context: int | None = None, valid_out_len: int | None = None, prompt_vectors: Tensor | None = None, @@ -106,13 +106,17 @@ def encoder_step( processed_signal_length: (Tensor) input signal length tensor. context: (CacheAwareContext) context object. drop_extra_pre_encoded: (int | None) number of extra pre-encoded frames to drop. - keep_all_outputs: (bool) whether to keep all outputs or not. + keep_all_outputs: (bool | Tensor) whether to keep all outputs; a bool vector of shape [B] + keeps them per stream, which serves a batch mixing last and non-last chunks. drop_left_context: (int | None) number of left context frames to drop. valid_out_len: (int | None) number of valid output frames. prompt_vectors: (Tensor | None) per-stream one-hot language prompts of shape [B, num_prompts]. Returns: (tuple[Tensor, Tensor, CacheAwareContext]) encoder output, encoder output lengths, and new context. """ + # a boolean vector asks for a mixed batch: the encoder keeps every output and the per-stream + # right context is trimmed afterwards by clamping the lengths, which the decoder honours + per_sample_keep = isinstance(keep_all_outputs, Tensor) ( encoded, encoded_len, @@ -125,7 +129,7 @@ def encoder_step( cache_last_channel=context.cache_last_channel, cache_last_time=context.cache_last_time, cache_last_channel_len=context.cache_last_channel_len, - keep_all_outputs=keep_all_outputs, + keep_all_outputs=True if per_sample_keep else keep_all_outputs, drop_extra_pre_encoded=drop_extra_pre_encoded, ) new_context = CacheAwareContext( @@ -139,7 +143,10 @@ def encoder_step( encoded = encoded[:, :, drop_left_context:] encoded_len = encoded_len - drop_left_context - if valid_out_len and not keep_all_outputs: + if valid_out_len and per_sample_keep: + # drop right context per stream: the streams that are not last keep valid_out_len frames + encoded_len = torch.where(keep_all_outputs, encoded_len, torch.full_like(encoded_len, valid_out_len)) + elif valid_out_len and not keep_all_outputs: # drop right context if any encoded = encoded[:, :, :valid_out_len] encoded_len = torch.ones_like(encoded_len) * valid_out_len @@ -158,7 +165,7 @@ def execute_step( context: CacheAwareContext, previous_hypotheses: list[Hypothesis] | None, drop_extra_pre_encoded: int | None, - keep_all_outputs: bool, + keep_all_outputs: bool | Tensor, drop_left_context: int | None = None, valid_out_len: int | None = None, prompt_vectors: Tensor | None = None, @@ -171,7 +178,8 @@ def execute_step( context: (CacheAwareContext) context object. previous_hypotheses: (list[Hypothesis] | None) list of previous hypotheses for RNNT decoding. drop_extra_pre_encoded: (int | None) number of extra pre-encoded frames to drop. - keep_all_outputs: (bool) whether to keep all outputs or not. + keep_all_outputs: (bool | Tensor) whether to keep all outputs; a bool vector of shape [B] + keeps them per stream, which serves a batch mixing last and non-last chunks. drop_left_context: (int | None) number of left context frames to drop. valid_out_len: (int | None) number of valid output frames. prompt_vectors: (Tensor | None) Optional prompt vectors of shape [B, num_prompts]. @@ -202,7 +210,7 @@ def malsd_stream_step( processed_signal_length: Tensor, context: CacheAwareContext, drop_extra_pre_encoded: int | None, - keep_all_outputs: bool, + keep_all_outputs: bool | Tensor, drop_left_context: int | None = None, valid_out_len: int | None = None, prompt_vectors: Tensor | None = None, @@ -278,7 +286,7 @@ def stream_step( context: CacheAwareContext = None, previous_hypotheses: list[Hypothesis] | None = None, drop_extra_pre_encoded: int | None = None, - keep_all_outputs: bool = False, + keep_all_outputs: bool | Tensor = False, drop_left_context: int | None = None, valid_out_len: int | None = None, prompt_vectors: Tensor | None = None, @@ -291,7 +299,8 @@ def stream_step( context: (CacheAwareContext) context object. previous_hypotheses: (list[Hypothesis] | None) list of previous hypotheses for RNNT decoding. drop_extra_pre_encoded: (int | None) number of extra pre-encoded frames to drop. - keep_all_outputs: (bool) whether to keep all outputs or not. + keep_all_outputs: (bool | Tensor) whether to keep all outputs; a bool vector of shape [B] + keeps them per stream, which serves a batch mixing last and non-last chunks. drop_left_context: (int | None) number of left context frames to drop. valid_out_len: (int | None) number of valid output frames. prompt_vectors: (Tensor | None) Optional prompt vectors of shape [B, num_prompts]. 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/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 28bafc986d2a..dc16135fc25c 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,20 @@ 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 + if self.asr_model.compile_encoder_layers(): + logging.info("Compiled the encoder with torch.compile") + def init_decoding_computer(self) -> None: """Initialize ``decoding_computer``.""" self.decoding_computer = None @@ -267,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 = [ @@ -283,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 @@ -296,7 +325,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]: """ @@ -422,10 +451,10 @@ 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 = False, + keep_all_outputs: bool | Tensor = False, ) -> None: """ Cache Aware Transcribe Step @@ -441,10 +470,11 @@ 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) 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) @@ -512,6 +542,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. @@ -522,30 +570,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 = [] + 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] - 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) - - 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: @@ -568,30 +599,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: 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..9118ace9befc 100644 --- a/nemo/collections/asr/inference/streaming/buffering/cache_feature_bufferer.py +++ b/nemo/collections/asr/inference/streaming/buffering/cache_feature_bufferer.py @@ -148,33 +148,43 @@ 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]]: + 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 = [], [], [] @@ -197,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 @@ -217,10 +229,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 :] + ) if len(slots_to_free) > 0: self.free_slots(slots_to_free) - return fbuffers, right_paddings.tolist() + return buffers, right_paddings 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: """ 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( 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 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/nemo/collections/common/parts/rnn.py b/nemo/collections/common/parts/rnn.py index 9be7277c8b17..e176cb496c26 100644 --- a/nemo/collections/common/parts/rnn.py +++ b/nemo/collections/common/parts/rnn.py @@ -228,9 +228,76 @@ 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. + 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 + + 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. 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). + 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]]: + """ + 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) + x, h = self.lstm(x, h) if self.dropout: 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))