chore: import upstream snapshot with attribution

tests/collections/asr/decoding/conftest.py

@@ -0,0 +1,135 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from pathlib import Path
+
+import pytest
+
+from nemo.collections.asr.models import ASRModel
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest, write_manifest
+from tests.collections.asr.decoding.utils import make_preprocessor_deterministic, preserve_decoding_cfg_and_cpu_device
+
+CHECKPOINTS_PATH = Path("/home/TestData/asr")
+
+
+@pytest.fixture(scope="session")
+def an4_val_manifest_corrected(tmp_path_factory, test_data_dir):
+    """
+    Correct an4_val manifest audio filepaths, e.g.,
+    "tests/data/asr/test/an4/wav/an440-mjgm-b.wav" -> test_data_dir / "test/an4/wav/an440-mjgm-b.wav"
+    """
+    an4_val_manifest_orig_path = Path(test_data_dir) / "asr/an4_val.json"
+    an4_val_manifest_corrected_path = tmp_path_factory.mktemp("manifests") / "an4_val_corrected.json"
+    an4_val_records = read_manifest(an4_val_manifest_orig_path)
+    for record in an4_val_records:
+        record["audio_filepath"] = record["audio_filepath"].replace(
+            "tests/data/asr", str(an4_val_manifest_orig_path.resolve().parent)
+        )
+    write_manifest(an4_val_manifest_corrected_path, an4_val_records)
+    return an4_val_manifest_corrected_path
+
+
+@pytest.fixture(scope="session")
+def an4_train_manifest_corrected(tmp_path_factory, test_data_dir):
+    """
+    Correct an4_train manifest audio filepaths, e.g.,
+    "tests/data/asr/test/an4/wav/an440-mjgm-b.wav" -> test_data_dir / "test/an4/wav/an440-mjgm-b.wav"
+    """
+    an4_train_manifest_orig_path = Path(test_data_dir) / "asr/an4_train.json"
+    an4_train_manifest_corrected_path = tmp_path_factory.mktemp("manifests") / "an4_train_corrected.json"
+    an4_train_records = read_manifest(an4_train_manifest_orig_path)
+    for record in an4_train_records:
+        record["audio_filepath"] = record["audio_filepath"].replace(
+            "tests/data/asr", str(an4_train_manifest_orig_path.resolve().parent)
+        )
+    write_manifest(an4_train_manifest_corrected_path, an4_train_records)
+    return an4_train_manifest_corrected_path
+
+
+@pytest.fixture(scope="package")
+def _stt_en_conformer_transducer_small_raw():
+    if CHECKPOINTS_PATH.exists():
+        model = ASRModel.restore_from(
+            str(CHECKPOINTS_PATH / "stt_en_conformer_transducer_small.nemo"), map_location="cpu"
+        )
+    else:
+        model_name = "stt_en_conformer_transducer_small"
+        model = ASRModel.from_pretrained(model_name, map_location="cpu")
+    make_preprocessor_deterministic(model)
+    return model
+
+
+@pytest.fixture(scope="package")
+def _stt_en_fastconformer_transducer_large_raw():
+    if CHECKPOINTS_PATH.exists():
+        model = ASRModel.restore_from(
+            str(CHECKPOINTS_PATH / "stt_en_fastconformer_transducer_large.nemo"), map_location="cpu"
+        )
+    else:
+        model_name = "stt_en_fastconformer_transducer_large"
+        model = ASRModel.from_pretrained(model_name, map_location="cpu")
+    make_preprocessor_deterministic(model)
+    return model
+
+
+@pytest.fixture(scope="package")
+def _stt_en_fastconformer_tdt_large_raw():
+    if CHECKPOINTS_PATH.exists():
+        model = ASRModel.restore_from(
+            str(CHECKPOINTS_PATH / "stt_en_fastconformer_tdt_large.nemo"), map_location="cpu"
+        )
+    else:
+        model_name = "nvidia/stt_en_fastconformer_tdt_large"
+        model = ASRModel.from_pretrained(model_name, map_location="cpu")
+    make_preprocessor_deterministic(model)
+    return model
+
+
+@pytest.fixture(scope="package")
+def _canary_180m_flash_raw():
+    model_name = "nvidia/canary-180m-flash"
+    model = ASRModel.from_pretrained(model_name, map_location="cpu")
+    make_preprocessor_deterministic(model)
+    return model
+
+
+@pytest.fixture
+def stt_en_conformer_transducer_small(_stt_en_conformer_transducer_small_raw):
+    """Function-level fixture for model. Guarantees to preserve decoding config and device"""
+    model = _stt_en_conformer_transducer_small_raw
+    with preserve_decoding_cfg_and_cpu_device(model):
+        yield model
+
+
+@pytest.fixture
+def stt_en_fastconformer_transducer_large(_stt_en_fastconformer_transducer_large_raw):
+    """Function-level fixture for model. Guarantees to preserve decoding config and device"""
+    model = _stt_en_fastconformer_transducer_large_raw
+    with preserve_decoding_cfg_and_cpu_device(model):
+        yield model
+
+
+@pytest.fixture
+def stt_en_fastconformer_tdt_large(_stt_en_fastconformer_tdt_large_raw):
+    """Function-level fixture for model. Guarantees to preserve decoding config and device"""
+    model = _stt_en_fastconformer_tdt_large_raw
+    with preserve_decoding_cfg_and_cpu_device(model):
+        yield model
+
+
+@pytest.fixture
+def canary_180m_flash(_canary_180m_flash_raw):
+    """Function-level fixture for model. Guarantees to preserve decoding config and device"""
+    model = _canary_180m_flash_raw
+    with preserve_decoding_cfg_and_cpu_device(model):
+        yield model

tests/collections/asr/decoding/test_batched_beam_decoding.py

@@ -0,0 +1,883 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import glob
+import os
+from pathlib import Path
+
+import pytest
+import torch
+from kaldialign import edit_distance
+from omegaconf import open_dict
+from tqdm import tqdm
+
+from nemo.collections.asr.models import ASRModel
+from nemo.collections.asr.models.ctc_models import EncDecCTCModel
+from nemo.collections.asr.parts.submodules.ctc_beam_decoding import BeamBatchedCTCInfer
+from nemo.collections.asr.parts.submodules.ngram_lm import NGramGPULanguageModel
+from nemo.collections.asr.parts.submodules.rnnt_beam_decoding import BeamBatchedRNNTInfer
+from nemo.collections.asr.parts.submodules.tdt_beam_decoding import BeamBatchedTDTInfer
+from nemo.collections.asr.parts.utils import rnnt_utils
+from nemo.core.utils import numba_utils
+from nemo.core.utils.cuda_python_utils import skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported
+from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
+from tests.collections.asr.decoding.utils import load_audio
+
+RNNT_MODEL = "stt_en_conformer_transducer_small"
+CTC_MODEL = "nvidia/stt_en_conformer_ctc_small"
+TDT_MODEL = "nvidia/stt_en_fastconformer_tdt_large"
+MAX_SAMPLES = 10
+
+DEVICES = [torch.device("cpu")]
+
+if torch.cuda.is_available():
+    DEVICES.append(torch.device('cuda'))
+
+NUMBA_RNNT_LOSS_AVAILABLE = numba_utils.numba_cpu_is_supported(
+    __NUMBA_MINIMUM_VERSION__
+) or numba_utils.numba_cuda_is_supported(__NUMBA_MINIMUM_VERSION__)
+
+
+# available audio filename fixtures
+@pytest.fixture(scope="module")
+def test_audio_filenames(test_data_dir):
+    return tuple(glob.glob(os.path.join(test_data_dir, "asr", "test", "an4", "wav", "*.wav")))
+
+
+# model fixtures
+@pytest.fixture(scope="module")
+def rnnt_model():
+    model = ASRModel.from_pretrained(model_name=RNNT_MODEL, map_location="cpu")
+    model.eval()
+    return model
+
+
+@pytest.fixture(scope="module")
+def tdt_model():
+    model = ASRModel.from_pretrained(model_name=TDT_MODEL, map_location="cpu")
+    model.eval()
+    return model
+
+
+@pytest.fixture(scope="module")
+def ctc_model():
+    model = ASRModel.from_pretrained(model_name=CTC_MODEL, map_location="cpu")
+    model.eval()
+    return model
+
+
+# encoder output fixtures
+@pytest.fixture(scope="module")
+def get_rnnt_encoder_output(rnnt_model, test_audio_filenames):
+    encoder_output, encoded_lengths = get_transducer_model_encoder_output(
+        test_audio_filenames, MAX_SAMPLES, rnnt_model
+    )
+    return encoder_output, encoded_lengths
+
+
+@pytest.fixture(scope="module")
+def get_tdt_encoder_output(tdt_model, test_audio_filenames):
+    encoder_output, encoded_lengths = get_transducer_model_encoder_output(test_audio_filenames, MAX_SAMPLES, tdt_model)
+    return encoder_output, encoded_lengths
+
+
+@pytest.fixture(scope="module")
+def get_ctc_output(ctc_model, test_audio_filenames):
+    encoder_output, encoded_lengths = get_ctc_model_output(test_audio_filenames, MAX_SAMPLES, ctc_model)
+    return encoder_output, encoded_lengths
+
+
+@pytest.fixture(scope="module")
+def kenlm_model_path(tmp_path_factory, test_data_dir):
+    lm_path = Path(test_data_dir) / "asr/kenlm_ngram_lm/parakeet-tdt_ctc-110m-libri-1024.kenlm.tmp.arpa"
+    assert os.path.exists(lm_path), f"LM file not found: {lm_path}"
+    lm_nemo_path = tmp_path_factory.mktemp("lm") / f"{lm_path.name}.nemo"
+    NGramGPULanguageModel.from_file(lm_path, vocab_size=1024).save_to(f"{lm_nemo_path}")
+    return f"{lm_nemo_path}"
+
+
+def get_transducer_model_encoder_output(
+    test_audio_filenames,
+    num_samples: int,
+    model: ASRModel,
+    device: torch.device = torch.device("cpu"),
+    dtype: torch.dtype = torch.float32,
+):
+    audio_filepaths = test_audio_filenames[:num_samples]
+
+    with torch.no_grad():
+        model.preprocessor.featurizer.dither = 0.0
+        model.preprocessor.featurizer.pad_to = 0
+        model.eval()
+
+        all_inputs, all_lengths = [], []
+        for audio_file in tqdm(audio_filepaths, desc="Loading audio files"):
+            audio_tensor, _ = load_audio(audio_file)
+            all_inputs.append(audio_tensor)
+            all_lengths.append(torch.tensor(audio_tensor.shape[0], dtype=torch.int64))
+
+        input_batch = torch.nn.utils.rnn.pad_sequence(all_inputs, batch_first=True).to(device=device, dtype=dtype)
+        length_batch = torch.tensor(all_lengths, dtype=torch.int64).to(device)
+
+        encoded_outputs, encoded_length = model(input_signal=input_batch, input_signal_length=length_batch)
+
+    return encoded_outputs, encoded_length
+
+
+def get_ctc_model_output(
+    test_audio_filenames,
+    num_samples: int,
+    model: ASRModel,
+    device: torch.device = torch.device("cpu"),
+    dtype: torch.dtype = torch.float32,
+):
+    audio_filepaths = test_audio_filenames[:num_samples]
+
+    with torch.no_grad():
+        model.preprocessor.featurizer.dither = 0.0
+        model.preprocessor.featurizer.pad_to = 0
+        model.eval()
+
+        all_inputs, all_lengths = [], []
+        for audio_file in tqdm(audio_filepaths, desc="Loading audio files"):
+            audio_tensor, _ = load_audio(audio_file)
+            all_inputs.append(audio_tensor)
+            all_lengths.append(torch.tensor(audio_tensor.shape[0], dtype=torch.int64))
+
+        input_batch = torch.nn.utils.rnn.pad_sequence(all_inputs, batch_first=True).to(device=device, dtype=dtype)
+        length_batch = torch.tensor(all_lengths, dtype=torch.int64).to(device)
+
+        log_probs, encoded_length, _ = model(input_signal=input_batch, input_signal_length=length_batch)
+
+    return log_probs, encoded_length
+
+
+def print_unit_test_info(strategy, batch_size, beam_size, allow_cuda_graphs, device):
+    print(
+        f"""Beam search algorithm: {strategy},
+                Batch size: {batch_size},
+                Beam size: {beam_size},
+                Cuda Graphs: {allow_cuda_graphs},
+                Decoding device: {device}
+            """
+    )
+
+
+def check_res_best_hyps(num_samples, hyps):
+    assert type(hyps) == list
+    assert type(hyps[0]) == rnnt_utils.Hypothesis
+
+    assert len(hyps) == num_samples
+
+    assert all(
+        [
+            hasattr(hyps[hyp_idx], "y_sequence")
+            and hasattr(hyps[hyp_idx], "score")
+            and hasattr(hyps[hyp_idx], "timestamp")
+            for hyp_idx in range(num_samples)
+        ]
+    )
+
+
+def print_res_best_hyps(hyps):
+    for hyp_idx, hyp in enumerate(hyps):
+        print("Sample: ", hyp_idx)
+        print("Decoded text: ", hyp.text)
+        print("Score: ", hyp.score)
+        print("Transcript", hyp.y_sequence)
+        print("Timesteps", hyp.timestamp)
+        print()
+
+
+def check_res_nbest_hyps(num_samples, batch_nbest_hyps):
+    assert type(batch_nbest_hyps) == list
+    assert type(batch_nbest_hyps[0]) == rnnt_utils.NBestHypotheses
+
+    assert len(batch_nbest_hyps) == num_samples
+
+    for idx in range(num_samples):
+        assert all(
+            [
+                hasattr(batch_nbest_hyps[idx].n_best_hypotheses[hyp_idx], "y_sequence")
+                and hasattr(batch_nbest_hyps[idx].n_best_hypotheses[hyp_idx], "score")
+                and hasattr(batch_nbest_hyps[idx].n_best_hypotheses[hyp_idx], "timestamp")
+                for hyp_idx in range(len(batch_nbest_hyps[idx].n_best_hypotheses))
+            ]
+        )
+
+        # Empty transcript (blank-only beam) is valid; y_sequence and timestamp must stay aligned.
+        assert all(
+            len(batch_nbest_hyps[idx].n_best_hypotheses[hyp_idx].y_sequence)
+            == len(batch_nbest_hyps[idx].n_best_hypotheses[hyp_idx].timestamp)
+            for hyp_idx in range(len(batch_nbest_hyps[idx].n_best_hypotheses))
+        )
+
+
+def print_res_nbest_hyps(batch_nbest_hyps):
+    for batch_idx, nbest_hyps in enumerate(batch_nbest_hyps):
+        print(f"Batch idx: {batch_idx}")
+        for idx, hyp in enumerate(nbest_hyps):
+            print(f"Hyp index: {idx + 1}")
+            print("Text: ", hyp.text)
+            print("Score: ", hyp.score)
+            print("Transcripts: ", hyp.y_sequence)
+            print("Timesteps: ", hyp.timestamp)
+            print()
+
+
+def decode_text_from_hypotheses(hyps, model):
+    if isinstance(model, EncDecCTCModel):
+        return model.decoding.decode_hypothesis(hyps, fold_consecutive=False)
+    else:
+        return model.decoding.decode_hypothesis(hyps)
+
+
+def decode_text_from_nbest_hypotheses(hyps, model):
+    if isinstance(model, EncDecCTCModel):
+        return [
+            model.decoding.decode_hypothesis(nbest_hyp.n_best_hypotheses, fold_consecutive=False) for nbest_hyp in hyps
+        ]
+    else:
+        return [model.decoding.decode_hypothesis(nbest_hyp.n_best_hypotheses) for nbest_hyp in hyps]
+
+
+class TestRNNTDecoding:
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {"search_type": "malsd_batch", "allow_cuda_graphs": False},
+            {"search_type": "malsd_batch", "allow_cuda_graphs": True},
+            {"search_type": "maes_batch", "allow_cuda_graphs": False},
+        ],
+    )
+    @pytest.mark.parametrize("beam_size", [4])
+    @pytest.mark.parametrize("batch_size", [4, 16])
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_rnnt_beam_decoding_return_best_hypothesis(
+        self, test_audio_filenames, rnnt_model, get_rnnt_encoder_output, beam_config, device, batch_size, beam_size
+    ):
+        num_samples = min(batch_size, len(test_audio_filenames))
+        model = rnnt_model.to(device)
+        encoder_output, encoded_lengths = get_rnnt_encoder_output
+        encoder_output, encoded_lengths = encoder_output[:num_samples].to(device), encoded_lengths[:num_samples].to(
+            device
+        )
+
+        vocab_size = model.tokenizer.vocab_size
+        decoding = BeamBatchedRNNTInfer(
+            model.decoder,
+            model.joint,
+            blank_index=vocab_size,
+            beam_size=beam_size,
+            score_norm=True,
+            return_best_hypothesis=True,
+            **beam_config,
+        )
+
+        print_unit_test_info(
+            strategy=beam_config['search_type'],
+            batch_size=batch_size,
+            beam_size=beam_size,
+            allow_cuda_graphs=beam_config.get('allow_cuda_graphs', True),
+            device=device,
+        )
+
+        with torch.no_grad():
+            hyps = decoding(encoder_output=encoder_output, encoded_lengths=encoded_lengths)[0]
+
+            check_res_best_hyps(num_samples, hyps)
+            hyps = decode_text_from_hypotheses(hyps, model)
+            print_res_best_hyps(hyps)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Test is only GPU-based decoding")
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {"search_type": "malsd_batch", "allow_cuda_graphs": False},
+            {"search_type": "malsd_batch", "allow_cuda_graphs": True},
+            {"search_type": "maes_batch", "allow_cuda_graphs": False},
+        ],
+    )
+    @pytest.mark.parametrize("beam_size", [4])
+    @pytest.mark.parametrize("batch_size", [4])
+    def test_rnnt_beam_decoding_return_nbest(
+        self, test_audio_filenames, rnnt_model, get_rnnt_encoder_output, beam_config, device, beam_size, batch_size
+    ):
+        device = torch.device("cuda")
+        num_samples = min(batch_size, len(test_audio_filenames))
+        model = rnnt_model.to(device)
+        encoder_output, encoded_lengths = get_rnnt_encoder_output
+        encoder_output, encoded_lengths = encoder_output[:num_samples].to(device), encoded_lengths[:num_samples].to(
+            device
+        )
+
+        vocab_size = model.tokenizer.vocab_size
+        decoding = BeamBatchedRNNTInfer(
+            model.decoder,
+            model.joint,
+            blank_index=vocab_size,
+            beam_size=beam_size,
+            score_norm=True,
+            return_best_hypothesis=False,
+            **beam_config,
+        )
+
+        print_unit_test_info(
+            strategy=beam_config['search_type'],
+            batch_size=batch_size,
+            beam_size=beam_size,
+            allow_cuda_graphs=beam_config.get('allow_cuda_graphs', True),
+            device=device,
+        )
+
+        with torch.no_grad():
+            batch_nbest_hyps = decoding(encoder_output=encoder_output, encoded_lengths=encoded_lengths)[0]
+
+            check_res_nbest_hyps(num_samples, batch_nbest_hyps)
+            batch_nbest_hyps = decode_text_from_nbest_hypotheses(batch_nbest_hyps, model)
+            print_res_nbest_hyps(batch_nbest_hyps)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Test is only GPU-based decoding")
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {"search_type": "malsd_batch", "allow_cuda_graphs": False},
+            {"search_type": "maes_batch", "allow_cuda_graphs": False},
+            {"search_type": "malsd_batch", "allow_cuda_graphs": True},
+        ],
+    )
+    @pytest.mark.parametrize("batch_size", [4])
+    @pytest.mark.parametrize("beam_size", [4])
+    @pytest.mark.parametrize("pruning_mode", ["late", "early"])
+    @pytest.mark.parametrize("blank_lm_score_mode", ["no_score", "lm_weighted_full"])
+    def test_rnnt_beam_decoding_kenlm(
+        self,
+        kenlm_model_path,
+        test_audio_filenames,
+        rnnt_model,
+        get_rnnt_encoder_output,
+        beam_config,
+        device,
+        batch_size,
+        beam_size,
+        pruning_mode,
+        blank_lm_score_mode,
+    ):
+        device = torch.device("cuda")
+
+        num_samples = min(batch_size, len(test_audio_filenames))
+        model = rnnt_model.to(device)
+        encoder_output, encoded_lengths = get_rnnt_encoder_output
+        encoder_output, encoded_lengths = encoder_output[:num_samples].to(device), encoded_lengths[:num_samples].to(
+            device
+        )
+
+        vocab_size = model.tokenizer.vocab_size
+
+        fusion_models = [NGramGPULanguageModel.from_file(lm_path=kenlm_model_path, vocab_size=vocab_size)]
+        fusion_models_alpha = [0.3]
+
+        decoding = BeamBatchedRNNTInfer(
+            model.decoder,
+            model.joint,
+            blank_index=vocab_size,
+            beam_size=beam_size,
+            score_norm=True,
+            return_best_hypothesis=True,
+            pruning_mode=pruning_mode,
+            blank_lm_score_mode=blank_lm_score_mode,
+            fusion_models=fusion_models,
+            fusion_models_alpha=fusion_models_alpha,
+            **beam_config,
+        )
+
+        print_unit_test_info(
+            strategy=beam_config['search_type'],
+            batch_size=batch_size,
+            beam_size=beam_size,
+            allow_cuda_graphs=beam_config.get('allow_cuda_graphs', True),
+            device=device,
+        )
+
+        with torch.no_grad():
+            hyps = decoding(encoder_output=encoder_output, encoded_lengths=encoded_lengths)[0]
+
+            check_res_best_hyps(num_samples, hyps)
+            hyps = decode_text_from_hypotheses(hyps, model)
+            print_res_best_hyps(hyps)
+
+
+class TestTDTDecoding:
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {"search_type": "malsd_batch", "allow_cuda_graphs": False},
+            {"search_type": "malsd_batch", "allow_cuda_graphs": True},
+        ],
+    )
+    @pytest.mark.parametrize("beam_size", [4])
+    @pytest.mark.parametrize("batch_size", [4, 16])
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_tdt_beam_decoding_return_best_hypothesis(
+        self, test_audio_filenames, tdt_model, get_tdt_encoder_output, beam_config, device, batch_size, beam_size
+    ):
+        num_samples = min(batch_size, len(test_audio_filenames))
+        model = tdt_model.to(device)
+        encoder_output, encoded_lengths = get_tdt_encoder_output
+        encoder_output, encoded_lengths = encoder_output[:num_samples].to(device), encoded_lengths[:num_samples].to(
+            device
+        )
+
+        model_config = model.to_config_dict()
+        durations = list(model_config["model_defaults"]["tdt_durations"])
+
+        vocab_size = model.tokenizer.vocab_size
+        decoding = BeamBatchedTDTInfer(
+            model.decoder,
+            model.joint,
+            blank_index=vocab_size,
+            durations=durations,
+            beam_size=beam_size,
+            score_norm=True,
+            return_best_hypothesis=True,
+            **beam_config,
+        )
+
+        print_unit_test_info(
+            strategy=beam_config['search_type'],
+            batch_size=batch_size,
+            beam_size=beam_size,
+            allow_cuda_graphs=beam_config.get('allow_cuda_graphs', True),
+            device=device,
+        )
+
+        with torch.no_grad():
+            hyps = decoding(encoder_output=encoder_output, encoded_lengths=encoded_lengths)[0]
+
+            check_res_best_hyps(num_samples, hyps)
+            hyps = decode_text_from_hypotheses(hyps, model)
+            print_res_best_hyps(hyps)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Test is only GPU-based decoding")
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {"search_type": "malsd_batch", "allow_cuda_graphs": False},
+            {"search_type": "malsd_batch", "allow_cuda_graphs": True},
+        ],
+    )
+    @pytest.mark.parametrize("beam_size", [4])
+    @pytest.mark.parametrize("batch_size", [4])
+    def test_tdt_beam_decoding_return_nbest(
+        self, test_audio_filenames, tdt_model, get_tdt_encoder_output, beam_config, device, beam_size, batch_size
+    ):
+        device = torch.device("cuda")
+        num_samples = min(batch_size, len(test_audio_filenames))
+        model = tdt_model.to(device)
+        encoder_output, encoded_lengths = get_tdt_encoder_output
+        encoder_output, encoded_lengths = encoder_output[:num_samples].to(device), encoded_lengths[:num_samples].to(
+            device
+        )
+
+        model_config = model.to_config_dict()
+        durations = list(model_config["model_defaults"]["tdt_durations"])
+
+        vocab_size = model.tokenizer.vocab_size
+        decoding = BeamBatchedTDTInfer(
+            model.decoder,
+            model.joint,
+            blank_index=vocab_size,
+            durations=durations,
+            beam_size=beam_size,
+            score_norm=True,
+            return_best_hypothesis=False,
+            **beam_config,
+        )
+
+        print_unit_test_info(
+            strategy=beam_config['search_type'],
+            batch_size=batch_size,
+            beam_size=beam_size,
+            allow_cuda_graphs=beam_config.get('allow_cuda_graphs', True),
+            device=device,
+        )
+
+        with torch.no_grad():
+            batch_nbest_hyps = decoding(encoder_output=encoder_output, encoded_lengths=encoded_lengths)[0]
+
+            check_res_nbest_hyps(num_samples, batch_nbest_hyps)
+            batch_nbest_hyps = decode_text_from_nbest_hypotheses(batch_nbest_hyps, model)
+            print_res_nbest_hyps(batch_nbest_hyps)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Test is only GPU-based decoding")
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {"search_type": "malsd_batch", "allow_cuda_graphs": False},
+            {"search_type": "malsd_batch", "allow_cuda_graphs": True},
+        ],
+    )
+    @pytest.mark.parametrize("batch_size", [4])
+    @pytest.mark.parametrize("beam_size", [4])
+    @pytest.mark.parametrize("pruning_mode", ["late", "early"])
+    @pytest.mark.parametrize("blank_lm_score_mode", ["lm_weighted_full", "no_score"])
+    def test_tdt_beam_decoding_kenlm(
+        self,
+        kenlm_model_path,
+        test_audio_filenames,
+        tdt_model,
+        get_tdt_encoder_output,
+        beam_config,
+        device,
+        batch_size,
+        beam_size,
+        pruning_mode,
+        blank_lm_score_mode,
+    ):
+        device = torch.device("cuda")
+
+        num_samples = min(batch_size, len(test_audio_filenames))
+        model = tdt_model.to(device)
+        encoder_output, encoded_lengths = get_tdt_encoder_output
+        encoder_output, encoded_lengths = encoder_output[:num_samples].to(device), encoded_lengths[:num_samples].to(
+            device
+        )
+
+        model_config = model.to_config_dict()
+        durations = list(model_config["model_defaults"]["tdt_durations"])
+
+        vocab_size = model.tokenizer.vocab_size
+
+        fusion_models = [NGramGPULanguageModel.from_file(lm_path=kenlm_model_path, vocab_size=vocab_size)]
+        fusion_models_alpha = [0.3]
+
+        decoding = BeamBatchedTDTInfer(
+            model.decoder,
+            model.joint,
+            blank_index=vocab_size,
+            durations=durations,
+            beam_size=beam_size,
+            score_norm=True,
+            return_best_hypothesis=True,
+            pruning_mode=pruning_mode,
+            blank_lm_score_mode=blank_lm_score_mode,
+            fusion_models=fusion_models,
+            fusion_models_alpha=fusion_models_alpha,
+            **beam_config,
+        )
+
+        print_unit_test_info(
+            strategy=beam_config['search_type'],
+            batch_size=batch_size,
+            beam_size=beam_size,
+            allow_cuda_graphs=beam_config.get('allow_cuda_graphs', True),
+            device=device,
+        )
+
+        with torch.no_grad():
+            hyps = decoding(encoder_output=encoder_output, encoded_lengths=encoded_lengths)[0]
+
+            check_res_best_hyps(num_samples, hyps)
+            hyps = decode_text_from_hypotheses(hyps, model)
+            print_res_best_hyps(hyps)
+
+
+class TestTransducerCudaGraphBeamDecoding:
+    """
+    Tests CudaGraphs implementations from Transducer models (RNN-T and TDT)
+    """
+
+    @pytest.mark.with_downloads
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA decoder can run only on CUDA")
+    @pytest.mark.parametrize("force_mode", ["no_graphs", "no_while_loops", "full_graph"])
+    @pytest.mark.parametrize("model_type", ["rnnt", "tdt"])
+    def test_stated_stateless(self, test_audio_filenames, rnnt_model, tdt_model, model_type, force_mode: str):
+        """
+        Compares pure Pytorch and with three modes of statefull implementations for double floating point precision.
+            1. Pure pytorch, but statefull implementation: no_graphs
+            2. With CudaGrpahs: no_while_loops and full_graph.
+        """
+        if force_mode == "full_graph":
+            skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported()
+
+        batch_size = 16
+        device = torch.device("cuda")
+        model = rnnt_model.to(device) if model_type == "rnnt" else tdt_model.to(device)
+        decoding_config = copy.deepcopy(model.cfg.decoding)
+
+        with open_dict(decoding_config):
+            decoding_config["strategy"] = "malsd_batch"
+            decoding_config["beam"]["beam_size"] = 4
+            decoding_config["beam"]["return_best_hypothesis"] = False
+            decoding_config["beam"]["allow_cuda_graphs"] = False
+
+        model.change_decoding_strategy(decoding_config)
+
+        actual_hypotheses = model.transcribe(test_audio_filenames, batch_size=batch_size, num_workers=None)
+        actual_transcripts = [[hyp.text for hyp in actual_beam] for actual_beam in actual_hypotheses]
+        actual_scores = [[hyp.score for hyp in actual_beam] for actual_beam in actual_hypotheses]
+        actual_timestamps = [[hyp.timestamp for hyp in actual_beam] for actual_beam in actual_hypotheses]
+
+        # transcribe with use implementation with cuda graphs
+        decoding_config["beam"]["allow_cuda_graphs"] = True
+        model.change_decoding_strategy(decoding_config)
+        model.decoding.decoding.decoding_computer.force_cuda_graphs_mode(mode=force_mode)
+
+        cudagraph_hypotheses = model.transcribe(test_audio_filenames, batch_size=batch_size, num_workers=None)
+        cudagraph_transcripts = [[hyp.text for hyp in cudagraphs_beam] for cudagraphs_beam in cudagraph_hypotheses]
+        cudagraph_scores = [[hyp.score for hyp in cudagraph_beam] for cudagraph_beam in cudagraph_hypotheses]
+        cudagraph_timestamps = [[hyp.timestamp for hyp in cudagraph_beam] for cudagraph_beam in cudagraph_hypotheses]
+
+        for batch_idx in range(min(batch_size, len(test_audio_filenames))):
+            assert len(actual_transcripts[batch_idx]) == len(cudagraph_transcripts[batch_idx])
+            assert cudagraph_scores[batch_idx] == pytest.approx(
+                actual_scores[batch_idx], abs=1e-2
+            ), f"Scores mismatch for batch_idx {batch_idx}"
+            assert (
+                cudagraph_timestamps[batch_idx] == actual_timestamps[batch_idx]
+            ), f"Timestamps mismatch for batch_idx {batch_idx}"
+
+            for actual, fast in zip(actual_transcripts[batch_idx], cudagraph_transcripts[batch_idx]):
+                ref_words = actual.split()
+                hyp_words = fast.split()
+                wer = edit_distance(ref_words, hyp_words)['total'] / max(len(ref_words), 1)
+
+                assert wer <= 1e-3, "Cuda graph beam decoder should match original decoder implementation."
+
+                if actual != fast:
+                    print("Erroneous samples in batch:", batch_idx)
+                    print("Original transcript:", actual)
+                    print("New transcript:", fast)
+
+    @pytest.mark.with_downloads
+    @pytest.mark.skipif(
+        not (torch.cuda.is_available() and torch.cuda.is_bf16_supported()), reason="CUDA decoder can run only on CUDA"
+    )
+    @pytest.mark.parametrize("model_type", ["rnnt", "tdt"])
+    def test_stated_stateless_bf16(self, test_audio_filenames, rnnt_model, tdt_model, model_type):
+        """
+        Checks that we are able to run without errors all decodings in bfloat16.
+        Computational errors accumulate, so just checking if algorithms run without errors
+        """
+        batch_size = 16
+        device = torch.device("cuda")
+        model = rnnt_model.to(device) if model_type == "rnnt" else tdt_model.to(device)
+        decoding_config = copy.deepcopy(model.cfg.decoding)
+
+        # checking pytorch implementation
+        with open_dict(decoding_config):
+            decoding_config["strategy"] = "malsd_batch"
+            decoding_config["beam"]["beam_size"] = 4
+            decoding_config["beam"]["return_best_hypothesis"] = False
+            decoding_config["beam"]["allow_cuda_graphs"] = False
+
+        model.change_decoding_strategy(decoding_config)
+
+        with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
+            model.transcribe(test_audio_filenames, batch_size=batch_size, num_workers=None)
+
+        modes = ["no_graphs", "no_while_loops", "full_graph"]
+        for force_mode in modes:
+            if force_mode == "full_graph":
+                skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported()
+
+            # transcribe with use implementation with cuda graphs
+            decoding_config["beam"]["allow_cuda_graphs"] = True
+            model.change_decoding_strategy(decoding_config)
+            model.decoding.decoding.decoding_computer.force_cuda_graphs_mode(mode=force_mode)
+
+            with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
+                model.transcribe(test_audio_filenames, batch_size=batch_size, num_workers=None)
+
+
+class TestCTCDecoding:
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {"allow_cuda_graphs": False},
+            {"allow_cuda_graphs": True},
+        ],
+    )
+    @pytest.mark.parametrize("beam_size", [4])
+    @pytest.mark.parametrize("batch_size", [4, 16])
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_ctc_beam_decoding_return_best_hypothesis(
+        self, test_audio_filenames, ctc_model, get_ctc_output, beam_config, device, batch_size, beam_size
+    ):
+        num_samples = min(batch_size, len(test_audio_filenames))
+        model = ctc_model.to(device)
+        log_probs, encoded_lengths = get_ctc_output
+        log_probs, encoded_lengths = log_probs[:num_samples].to(device), encoded_lengths[:num_samples].to(device)
+
+        vocab_size = model.tokenizer.vocab_size
+        decoding = BeamBatchedCTCInfer(
+            blank_index=vocab_size,
+            beam_size=beam_size,
+            return_best_hypothesis=True,
+            **beam_config,
+        )
+
+        print_unit_test_info(
+            strategy="beam_batch",
+            batch_size=batch_size,
+            beam_size=beam_size,
+            allow_cuda_graphs=beam_config.get('allow_cuda_graphs', True),
+            device=device,
+        )
+
+        with torch.no_grad():
+            hyps = decoding(decoder_output=log_probs, decoder_lengths=encoded_lengths)[0]
+
+            check_res_best_hyps(num_samples, hyps)
+            hyps = decode_text_from_hypotheses(hyps, model)
+            print_res_best_hyps(hyps)
+
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Test is only GPU-based decoding")
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {"allow_cuda_graphs": False},
+            {"allow_cuda_graphs": True},
+        ],
+    )
+    @pytest.mark.parametrize("beam_size", [4])
+    @pytest.mark.parametrize("batch_size", [4])
+    def test_ctc_beam_decoding_return_nbest(
+        self, test_audio_filenames, ctc_model, get_ctc_output, beam_config, device, beam_size, batch_size
+    ):
+        device = torch.device("cuda")
+        num_samples = min(batch_size, len(test_audio_filenames))
+        model = ctc_model.to(device)
+        log_probs, encoded_lengths = get_ctc_output
+        log_probs, encoded_lengths = log_probs[:num_samples].to(device), encoded_lengths[:num_samples].to(device)
+
+        vocab_size = model.tokenizer.vocab_size
+        decoding = BeamBatchedCTCInfer(
+            blank_index=vocab_size,
+            beam_size=beam_size,
+            return_best_hypothesis=False,
+            **beam_config,
+        )
+
+        print_unit_test_info(
+            strategy="beam_batch",
+            batch_size=batch_size,
+            beam_size=beam_size,
+            allow_cuda_graphs=beam_config.get('allow_cuda_graphs', True),
+            device=device,
+        )
+
+        with torch.no_grad():
+            batch_nbest_hyps = decoding(decoder_output=log_probs, decoder_lengths=encoded_lengths)[0]
+
+            check_res_nbest_hyps(num_samples, batch_nbest_hyps)
+            batch_nbest_hyps = decode_text_from_nbest_hypotheses(batch_nbest_hyps, model)
+            print_res_nbest_hyps(batch_nbest_hyps)
+
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Test is only GPU-based decoding")
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {"allow_cuda_graphs": False, "ngram_lm_alpha": 0.3, "beam_beta": 1.0},
+            {"allow_cuda_graphs": False, "ngram_lm_alpha": 0.3, "beam_beta": 1.0},
+        ],
+    )
+    @pytest.mark.parametrize("batch_size", [4])
+    @pytest.mark.parametrize("beam_size", [4])
+    def test_ctc_beam_decoding_kenlm(
+        self,
+        kenlm_model_path,
+        test_audio_filenames,
+        ctc_model,
+        get_ctc_output,
+        beam_config,
+        device,
+        batch_size,
+        beam_size,
+    ):
+        device = torch.device("cuda")
+        beam_config["ngram_lm_model"] = kenlm_model_path
+
+        num_samples = min(batch_size, len(test_audio_filenames))
+        model = ctc_model.to(device)
+        decoder_output, decoder_lengths = get_ctc_output
+        decoder_output, decoder_lengths = decoder_output[:num_samples].to(device), decoder_lengths[:num_samples].to(
+            device
+        )
+
+        vocab_size = model.tokenizer.vocab_size
+        decoding = BeamBatchedCTCInfer(
+            blank_index=vocab_size,
+            beam_size=beam_size,
+            return_best_hypothesis=True,
+            **beam_config,
+        )
+
+        print_unit_test_info(
+            strategy="beam_batch",
+            batch_size=batch_size,
+            beam_size=beam_size,
+            allow_cuda_graphs=beam_config.get('allow_cuda_graphs', True),
+            device=device,
+        )
+
+        with torch.no_grad():
+            hyps = decoding(decoder_output=decoder_output, decoder_lengths=decoder_lengths)[0]
+
+            check_res_best_hyps(num_samples, hyps)
+            hyps = decode_text_from_hypotheses(hyps, model)
+            print_res_best_hyps(hyps)

tests/collections/asr/decoding/test_batched_hyps.py

@@ -0,0 +1,797 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List
+
+import pytest
+import torch
+
+from nemo.collections.asr.parts.submodules.transducer_decoding.batched_hyps import BatchedHyps
+from nemo.collections.asr.parts.utils.rnnt_utils import batched_hyps_to_hypotheses
+from tests.collections.asr.decoding.utils import avoid_sync_operations
+
+DEVICES: List[torch.device] = [torch.device("cpu")]
+
+if torch.cuda.is_available():
+    DEVICES.append(torch.device("cuda"))
+
+if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+    DEVICES.append(torch.device("mps"))
+
+# blank id that does not collide with any non-blank label used in the "no blank steps" tests
+NON_COLLIDING_BLANK_ID = 1024
+
+
+class TestBatchedHyps:
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_instantiate(self, device: torch.device):
+        hyps = BatchedHyps(batch_size=2, init_length=3, blank_id=NON_COLLIDING_BLANK_ID, device=device)
+        assert torch.is_tensor(hyps.timestamps)
+        # device: for mps device we need to use `type`, not directly compare
+        assert hyps.timestamps.device.type == device.type
+        assert hyps.timestamps.shape == (2, 3)
+        assert hyps.transcript.shape == (2, 3)
+        assert hyps.scores.shape == (2,)
+        assert hyps.current_lengths.tolist() == [0, 0]
+        # optional storage is disabled by default
+        assert hyps.token_durations is None
+        assert hyps.step_confidence is None
+        assert hyps.logits is None
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("batch_size", [-1, 0])
+    def test_instantiate_incorrect_batch_size(self, batch_size):
+        with pytest.raises(ValueError):
+            _ = BatchedHyps(batch_size=batch_size, init_length=3, blank_id=0)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("init_length", [-1, 0])
+    def test_instantiate_incorrect_init_length(self, init_length):
+        with pytest.raises(ValueError):
+            _ = BatchedHyps(batch_size=1, init_length=init_length, blank_id=0)
+
+    @pytest.mark.unit
+    def test_instantiate_with_logits_requires_logits_dim(self):
+        # `with_logits=True` without `logits_dim` is invalid
+        with pytest.raises(ValueError):
+            _ = BatchedHyps(batch_size=1, init_length=3, blank_id=0, with_logits=True, logits_dim=None)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_instantiate_optional_storage(self, device: torch.device):
+        logits_dim = 7
+        hyps = BatchedHyps(
+            batch_size=2,
+            init_length=3,
+            blank_id=0,
+            logits_dim=logits_dim,
+            device=device,
+            with_durations=True,
+            with_step_confidence=True,
+            with_duration_confidence=True,
+            with_logits=True,
+        )
+        assert hyps.token_durations.shape == (2, 3)
+        # duration confidence makes the confidence tensor store a pair (step + duration) per token
+        assert hyps.step_confidence.shape == (2, 3, 2)
+        assert hyps.logits.shape == (2, 3, logits_dim)
+
+        # without duration confidence the confidence tensor is 2d
+        hyps_no_dur_conf = BatchedHyps(
+            batch_size=2,
+            init_length=3,
+            blank_id=0,
+            device=device,
+            with_step_confidence=True,
+            with_duration_confidence=False,
+        )
+        assert hyps_no_dur_conf.step_confidence.shape == (2, 3)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_add_results_masked(self, device: torch.device):
+        # batch of size 2, add label for first utterance only (second is inactive)
+        hyps = BatchedHyps(batch_size=2, init_length=1, blank_id=NON_COLLIDING_BLANK_ID, device=device)
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, False], device=device),
+            labels=torch.tensor([5, 1], device=device),
+            time_indices=torch.tensor([1, 0], device=device),
+            scores=torch.tensor([0.5, 10.0], device=device),
+        )
+        assert hyps.current_lengths.tolist() == [1, 0]
+        assert hyps.transcript.tolist()[0][:1] == [5]
+        assert hyps.timestamps.tolist()[0][:1] == [1]
+        assert hyps.scores.tolist() == pytest.approx([0.5, 0.0])  # inactive score should be ignored!
+        assert hyps.last_nb_timestamp.tolist() == [1, -1]
+        assert hyps.last_nb_timestamp_lasts.tolist() == [1, 0]
+        assert hyps.last_nb_labels.tolist() == [5, -1]
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_add_multiple_results_masked(self, device: torch.device):
+        # batch of size 2, add label for first utterance, then add labels for both utterances
+        hyps = BatchedHyps(batch_size=2, init_length=1, blank_id=NON_COLLIDING_BLANK_ID, device=device)
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, False], device=device),
+            labels=torch.tensor([5, 2], device=device),
+            time_indices=torch.tensor([1, 0], device=device),
+            scores=torch.tensor([0.5, 10.0], device=device),
+        )
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, True], device=device),
+            labels=torch.tensor([2, 4], device=device),
+            time_indices=torch.tensor([1, 2], device=device),
+            scores=torch.tensor([1.0, 1.0], device=device),
+        )
+        assert hyps.current_lengths.tolist() == [2, 1]
+        assert hyps.transcript.tolist()[0][:2] == [5, 2]
+        assert hyps.transcript.tolist()[1][:1] == [4]
+        assert hyps.timestamps.tolist()[0][:2] == [1, 1]
+        assert hyps.timestamps.tolist()[1][:1] == [2]
+        assert hyps.scores.tolist() == pytest.approx([1.5, 1.0])
+        assert hyps.last_nb_timestamp.tolist() == [1, 2]
+        assert hyps.last_nb_timestamp_lasts.tolist() == [2, 1]
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_add_results_masked_no_checks(self, device: torch.device):
+        # `check_lengths=False` must contain no host<->device synchronization (blocking) operations
+        hyps = BatchedHyps(batch_size=2, init_length=4, blank_id=NON_COLLIDING_BLANK_ID, device=device)
+        active_mask = torch.tensor([True, False], device=device)
+        time_indices = torch.tensor([1, 0], device=device)
+        scores = torch.tensor([0.5, 10.0], device=device)
+        labels = torch.tensor([5, 1], device=device)
+        # check there are no blocking operations
+        with avoid_sync_operations(device=device):
+            hyps.add_results_masked_(
+                active_mask=active_mask,
+                labels=labels,
+                time_indices=time_indices,
+                scores=scores,
+                check_lengths=False,
+            )
+        assert hyps.current_lengths.tolist() == [1, 0]
+        assert hyps.transcript.tolist()[0][:1] == [5]
+        assert hyps.timestamps.tolist()[0][:1] == [1]
+        assert hyps.scores.tolist() == pytest.approx([0.5, 0.0])  # inactive score should be ignored!
+        assert hyps.last_nb_timestamp.tolist() == [1, -1]
+        assert hyps.last_nb_timestamp_lasts.tolist() == [1, 0]
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_add_results_masked_reallocates(self, device: torch.device):
+        # init_length is intentionally small; storage must grow transparently when check_lengths=True
+        hyps = BatchedHyps(batch_size=2, init_length=1, blank_id=NON_COLLIDING_BLANK_ID, device=device)
+        for step in range(5):
+            hyps.add_results_masked_(
+                active_mask=torch.tensor([True, True], device=device),
+                labels=torch.tensor([step, step + 10], device=device),
+                time_indices=torch.tensor([step, step], device=device),
+                scores=torch.tensor([1.0, 1.0], device=device),
+                check_lengths=True,
+            )
+        assert hyps._max_length >= 5
+        assert hyps.current_lengths.tolist() == [5, 5]
+        assert hyps.transcript.tolist()[0][:5] == [0, 1, 2, 3, 4]
+        assert hyps.transcript.tolist()[1][:5] == [10, 11, 12, 13, 14]
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_add_results_masked_with_blank_steps(self, device: torch.device):
+        # with_blank_steps=True: blank labels are stored in the transcript, but they do NOT advance
+        # the score / last-non-blank tracking. Single utterance for clarity.
+        blank_id = 0
+        hyps = BatchedHyps(batch_size=1, init_length=2, blank_id=blank_id, device=device, with_blank_steps=True)
+        # (label, time, score): two non-blank tokens at t=0, then blank, then non-blank at t=1, then blank
+        steps = [
+            (3, 0, 1.0),  # non-blank
+            (5, 0, 1.5),  # non-blank, same timestamp
+            (blank_id, 0, 0.1),  # blank
+            (7, 1, 2.0),  # non-blank
+            (blank_id, 1, 0.2),  # blank
+        ]
+        for label, time, score in steps:
+            hyps.add_results_masked_(
+                active_mask=torch.tensor([True], device=device),
+                labels=torch.tensor([label], device=device),
+                time_indices=torch.tensor([time], device=device),
+                scores=torch.tensor([score], device=device),
+            )
+        # all steps (including blanks) are stored
+        assert hyps.current_lengths.tolist() == [5]
+        assert hyps.transcript.tolist()[0][:5] == [3, 5, blank_id, 7, blank_id]
+        assert hyps.timestamps.tolist()[0][:5] == [0, 0, 0, 1, 1]
+        # only non-blank scores accumulate: 1.0 + 1.5 + 2.0
+        assert hyps.scores.tolist() == pytest.approx([4.5])
+        # last-non-blank tracking ignores blanks
+        assert hyps.last_nb_timestamp.tolist() == [1]
+        assert hyps.last_nb_timestamp_lasts.tolist() == [1]
+        assert hyps.last_nb_labels.tolist() == [7]
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_get_data_without_blank_no_blank_steps(self, device: torch.device):
+        # with_blank_steps=False: data is returned as-is (it never contained blanks)
+        hyps = BatchedHyps(batch_size=2, init_length=2, blank_id=NON_COLLIDING_BLANK_ID, device=device)
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, True], device=device),
+            labels=torch.tensor([5, 4], device=device),
+            time_indices=torch.tensor([0, 1], device=device),
+            scores=torch.tensor([1.0, 1.0], device=device),
+        )
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, False], device=device),
+            labels=torch.tensor([2, 0], device=device),
+            time_indices=torch.tensor([1, 0], device=device),
+            scores=torch.tensor([1.0, 0.0], device=device),
+        )
+        lengths, transcript, timestamps, durations, confidence = hyps.get_data_without_blank()
+        # returned objects are the underlying (unmodified) tensors
+        assert lengths is hyps.current_lengths
+        assert transcript is hyps.transcript
+        assert timestamps is hyps.timestamps
+        assert durations is None
+        assert confidence is None
+        assert lengths.tolist() == [2, 1]
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_get_data_without_blank_with_blank_steps(self, device: torch.device):
+        # with_blank_steps=True: blanks are stripped and non-blank order is preserved
+        blank_id = 0
+        hyps = BatchedHyps(batch_size=2, init_length=2, blank_id=blank_id, device=device, with_blank_steps=True)
+        # seq 0: [5, blank, 2, blank]  -> [5, 2]
+        # seq 1: [blank, 4]            -> [4]
+        steps = [
+            # (labels, times, active_mask)
+            ([5, blank_id], [0, 0], [True, True]),
+            ([blank_id, 4], [0, 1], [True, True]),
+            ([2, blank_id], [1, 1], [True, False]),
+            ([blank_id, 0], [1, 0], [True, False]),
+        ]
+        for labels, times, active in steps:
+            hyps.add_results_masked_(
+                active_mask=torch.tensor(active, device=device),
+                labels=torch.tensor(labels, device=device),
+                time_indices=torch.tensor(times, device=device),
+                scores=torch.tensor([1.0, 1.0], device=device),
+            )
+        lengths, transcript, timestamps, durations, confidence = hyps.get_data_without_blank()
+        assert lengths.tolist() == [2, 1]
+        assert transcript[0, :2].tolist() == [5, 2]
+        assert transcript[1, :1].tolist() == [4]
+        assert timestamps[0, :2].tolist() == [0, 1]
+        assert timestamps[1, :1].tolist() == [1]
+        assert durations is None
+        assert confidence is None
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_get_last_labels(self, device: torch.device):
+        hyps = BatchedHyps(batch_size=2, init_length=2, blank_id=NON_COLLIDING_BLANK_ID, device=device)
+        # no labels yet -> pad_id everywhere
+        assert hyps.get_last_labels(pad_id=-1).tolist() == [-1, -1]
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, False], device=device),
+            labels=torch.tensor([5, 1], device=device),
+            time_indices=torch.tensor([0, 0], device=device),
+            scores=torch.tensor([1.0, 0.0], device=device),
+        )
+        assert hyps.get_last_labels(pad_id=-1).tolist() == [5, -1]
+        assert hyps.get_last_labels(pad_id=100).tolist() == [5, 100]
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_clear(self, device: torch.device):
+        hyps = BatchedHyps(batch_size=2, init_length=2, blank_id=NON_COLLIDING_BLANK_ID, device=device)
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, True], device=device),
+            labels=torch.tensor([5, 4], device=device),
+            time_indices=torch.tensor([0, 0], device=device),
+            scores=torch.tensor([1.0, 1.0], device=device),
+        )
+        hyps.clear_()
+        assert hyps.current_lengths.tolist() == [0, 0]
+        assert hyps.scores.tolist() == pytest.approx([0.0, 0.0])
+        assert hyps.transcript.tolist() == [[0, 0], [0, 0]]
+        assert hyps.last_nb_timestamp.tolist() == [-1, -1]
+        assert hyps.last_nb_timestamp_lasts.tolist() == [0, 0]
+        assert hyps.last_nb_labels.tolist() == [-1, -1]
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_clone(self, device: torch.device):
+        logits_dim = 7
+        hyps = BatchedHyps(
+            batch_size=2,
+            init_length=2,
+            blank_id=0,
+            logits_dim=logits_dim,
+            device=device,
+            with_durations=True,
+            with_step_confidence=True,
+            with_logits=True,
+            with_blank_steps=True,
+        )
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, True], device=device),
+            labels=torch.tensor([5, 4], device=device),
+            time_indices=torch.tensor([0, 0], device=device),
+            scores=torch.tensor([1.0, 1.0], device=device),
+            token_durations=torch.tensor([1, 2], device=device),
+            confidence=torch.tensor([0.9, 0.8], device=device),
+            logits=torch.rand((2, logits_dim), device=device),
+        )
+        clone = hyps.clone()
+        # flags carried over
+        assert clone.with_durations and clone.with_step_confidence and clone.with_logits and clone.with_blank_steps
+        assert clone.blank_id == hyps.blank_id
+        # values copied
+        assert clone.current_lengths.tolist() == hyps.current_lengths.tolist()
+        assert torch.equal(clone.transcript, hyps.transcript)
+        assert torch.allclose(clone.logits, hyps.logits)
+        assert torch.allclose(clone.step_confidence, hyps.step_confidence)
+        assert torch.equal(clone.token_durations, hyps.token_durations)
+        # clone is independent of the original
+        hyps.transcript.fill_(0)
+        hyps.scores.fill_(0.0)
+        assert clone.transcript.tolist()[0][:1] == [5]
+        assert clone.scores.tolist() == pytest.approx([1.0, 1.0])
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_merge(self, device: torch.device):
+        # merge two batched hypotheses (basic case: no blank steps, no optional storage)
+        def build(labels_per_step, times_per_step, masks_per_step, scores_per_step):
+            hyps = BatchedHyps(batch_size=2, init_length=4, blank_id=NON_COLLIDING_BLANK_ID, device=device)
+            for labels, times, mask, scores in zip(labels_per_step, times_per_step, masks_per_step, scores_per_step):
+                hyps.add_results_masked_(
+                    active_mask=torch.tensor(mask, device=device),
+                    labels=torch.tensor(labels, device=device),
+                    time_indices=torch.tensor(times, device=device),
+                    scores=torch.tensor(scores, device=device),
+                )
+            return hyps
+
+        # A: seq0=[5, 2], seq1=[4]
+        hyps_a = build(
+            labels_per_step=[[5, 4], [2, 0]],
+            times_per_step=[[0, 0], [1, 0]],
+            masks_per_step=[[True, True], [True, False]],
+            scores_per_step=[[0.5, 0.7], [0.5, 0.0]],
+        )
+        # B: seq0=[7], seq1=[8, 9]
+        hyps_b = build(
+            labels_per_step=[[7, 8], [0, 9]],
+            times_per_step=[[2, 2], [0, 3]],
+            masks_per_step=[[True, True], [False, True]],
+            scores_per_step=[[0.3, 0.3], [0.0, 0.4]],
+        )
+
+        hyps_a.merge_(hyps_b)
+        assert hyps_a.current_lengths.tolist() == [3, 3]
+        assert hyps_a.transcript[0, :3].tolist() == [5, 2, 7]
+        assert hyps_a.transcript[1, :3].tolist() == [4, 8, 9]
+        assert hyps_a.scores.tolist() == pytest.approx([1.3, 1.4])
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_merge_with_logits(self, device: torch.device):
+        # Regression test: merge_ must use dim=1 (sequence axis) for logits scatter/cat,
+        # not dim=-1 (logits-dim axis), which would silently corrupt the data.
+        logits_dim = 5
+        blank_id = NON_COLLIDING_BLANK_ID
+
+        def build_with_logits(labels_per_step, times_per_step, masks_per_step, scores_per_step, logits_per_step):
+            hyps = BatchedHyps(
+                batch_size=2,
+                init_length=4,
+                blank_id=blank_id,
+                logits_dim=logits_dim,
+                device=device,
+                float_dtype=torch.float32,
+                with_logits=True,
+            )
+            for labels, times, mask, scores, logits in zip(
+                labels_per_step, times_per_step, masks_per_step, scores_per_step, logits_per_step
+            ):
+                hyps.add_results_masked_(
+                    active_mask=torch.tensor(mask, device=device),
+                    labels=torch.tensor(labels, device=device),
+                    time_indices=torch.tensor(times, device=device),
+                    scores=torch.tensor(scores, device=device),
+                    logits=logits,
+                )
+            return hyps
+
+        # Fixed logits so we can assert exact values after merge
+        logits_a0 = torch.full((2, logits_dim), 1.0, device=device)  # step for seq0=[5], seq1=[4]
+        logits_a1 = torch.full((2, logits_dim), 2.0, device=device)  # step for seq0=[2], seq1 inactive
+
+        logits_b0 = torch.full((2, logits_dim), 3.0, device=device)  # step for seq0=[7], seq1=[8]
+        logits_b1 = torch.full((2, logits_dim), 4.0, device=device)  # step for seq0 inactive, seq1=[9]
+
+        # A: seq0=[5, 2], seq1=[4]
+        hyps_a = build_with_logits(
+            labels_per_step=[[5, 4], [2, 0]],
+            times_per_step=[[0, 0], [1, 0]],
+            masks_per_step=[[True, True], [True, False]],
+            scores_per_step=[[0.5, 0.7], [0.5, 0.0]],
+            logits_per_step=[logits_a0, logits_a1],
+        )
+        # B: seq0=[7], seq1=[8, 9]
+        hyps_b = build_with_logits(
+            labels_per_step=[[7, 8], [0, 9]],
+            times_per_step=[[2, 2], [0, 3]],
+            masks_per_step=[[True, True], [False, True]],
+            scores_per_step=[[0.3, 0.3], [0.0, 0.4]],
+            logits_per_step=[logits_b0, logits_b1],
+        )
+
+        hyps_a.merge_(hyps_b)
+
+        # Sequence lengths: seq0=2+1=3, seq1=1+2=3
+        assert hyps_a.current_lengths.tolist() == [3, 3]
+
+        # seq0 logits: positions [0,1] from A (value=1 then 2), position [2] from B (value=3)
+        assert torch.allclose(hyps_a.logits[0, 0], torch.full((logits_dim,), 1.0, device=device))
+        assert torch.allclose(hyps_a.logits[0, 1], torch.full((logits_dim,), 2.0, device=device))
+        assert torch.allclose(hyps_a.logits[0, 2], torch.full((logits_dim,), 3.0, device=device))
+
+        # seq1 logits: position [0] from A (value=1), positions [1,2] from B (value=3 then 4)
+        assert torch.allclose(hyps_a.logits[1, 0], torch.full((logits_dim,), 1.0, device=device))
+        assert torch.allclose(hyps_a.logits[1, 1], torch.full((logits_dim,), 3.0, device=device))
+        assert torch.allclose(hyps_a.logits[1, 2], torch.full((logits_dim,), 4.0, device=device))
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_merge_with_logits_triggers_reallocation(self, device: torch.device):
+        # When the combined length exceeds _max_length, merge_ must reallocate logits along dim=1
+        # (sequence axis). With the dim=-1 bug, the reallocation would expand the logits-dim
+        # axis instead, causing a shape mismatch on the subsequent scatter_.
+        logits_dim = 5
+        blank_id = NON_COLLIDING_BLANK_ID
+
+        # Use init_length=1 to force reallocation during merge
+        hyps_a = BatchedHyps(
+            batch_size=1,
+            init_length=1,
+            blank_id=blank_id,
+            logits_dim=logits_dim,
+            device=device,
+            float_dtype=torch.float32,
+            with_logits=True,
+        )
+        hyps_b = BatchedHyps(
+            batch_size=1,
+            init_length=1,
+            blank_id=blank_id,
+            logits_dim=logits_dim,
+            device=device,
+            float_dtype=torch.float32,
+            with_logits=True,
+        )
+
+        logits_a = torch.full((1, logits_dim), 1.0, device=device)
+        logits_b = torch.full((1, logits_dim), 2.0, device=device)
+
+        hyps_a.add_results_masked_(
+            active_mask=torch.tensor([True], device=device),
+            labels=torch.tensor([5], device=device),
+            time_indices=torch.tensor([0], device=device),
+            scores=torch.tensor([1.0], device=device),
+            logits=logits_a,
+        )
+        hyps_b.add_results_masked_(
+            active_mask=torch.tensor([True], device=device),
+            labels=torch.tensor([7], device=device),
+            time_indices=torch.tensor([1], device=device),
+            scores=torch.tensor([1.0], device=device),
+            logits=logits_b,
+        )
+
+        # cur_len=1, other_len=1 -> combined=2 >= init_length=1, so reallocation is triggered
+        hyps_a.merge_(hyps_b)
+
+        assert hyps_a.current_lengths.tolist() == [2]
+        assert hyps_a.logits.shape == (1, hyps_a._max_length, logits_dim)
+        assert torch.allclose(hyps_a.logits[0, 0], torch.full((logits_dim,), 1.0, device=device))
+        assert torch.allclose(hyps_a.logits[0, 1], torch.full((logits_dim,), 2.0, device=device))
+
+
+class TestConvertToHypotheses:
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_convert_no_blank_steps(self, device: torch.device):
+        # with_blank_steps=False: transcript already contains only non-blank labels
+        hyps = BatchedHyps(batch_size=2, init_length=1, blank_id=NON_COLLIDING_BLANK_ID, device=device)
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, False], device=device),
+            labels=torch.tensor([5, 0], device=device),
+            time_indices=torch.tensor([1, 0], device=device),
+            scores=torch.tensor([0.5, 0.0], device=device),
+        )
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, True], device=device),
+            labels=torch.tensor([2, 4], device=device),
+            time_indices=torch.tensor([1, 2], device=device),
+            scores=torch.tensor([1.0, 1.0], device=device),
+        )
+        hypotheses = batched_hyps_to_hypotheses(hyps)
+        assert (hypotheses[0].y_sequence == torch.tensor([5, 2], device="cpu")).all()
+        assert (hypotheses[1].y_sequence == torch.tensor([4], device="cpu")).all()
+        assert hypotheses[0].score == pytest.approx(1.5)
+        assert hypotheses[1].score == pytest.approx(1.0)
+        assert (hypotheses[0].timestamp == torch.tensor([1, 1], device="cpu")).all()
+        assert (hypotheses[1].timestamp == torch.tensor([2], device="cpu")).all()
+        # no blank steps -> no alignments / frame confidence
+        assert hypotheses[0].alignments is None
+        assert hypotheses[1].alignments is None
+        assert hypotheses[0].frame_confidence is None
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_convert_batch_size_arg(self, device: torch.device):
+        # batch_size arg returns only the first `batch_size` hypotheses (CUDA-graph constant batch)
+        hyps = BatchedHyps(batch_size=4, init_length=1, blank_id=NON_COLLIDING_BLANK_ID, device=device)
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, True, True, True], device=device),
+            labels=torch.tensor([5, 4, 3, 2], device=device),
+            time_indices=torch.tensor([0, 0, 0, 0], device=device),
+            scores=torch.tensor([1.0, 1.0, 1.0, 1.0], device=device),
+        )
+        hypotheses = batched_hyps_to_hypotheses(hyps, batch_size=2)
+        assert len(hypotheses) == 2
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_convert_with_blank_steps_strips_blanks(self, device: torch.device):
+        # with_blank_steps=True but no logits/confidence: blanks must be stripped from y_sequence/timestamps,
+        # while alignments are NOT produced (no logits recorded)
+        blank_id = 0
+        hyps = BatchedHyps(batch_size=2, init_length=2, blank_id=blank_id, device=device, with_blank_steps=True)
+        # seq 0: [5, blank, 2, blank] -> [5, 2]
+        # seq 1: [blank, 4]           -> [4]
+        steps = [
+            ([5, blank_id], [0, 0], [True, True], [0.5, 0.1]),
+            ([blank_id, 4], [0, 1], [True, True], [0.1, 1.0]),
+            ([2, blank_id], [1, 1], [True, False], [1.0, 0.0]),
+            ([blank_id, 0], [1, 0], [True, False], [0.1, 0.0]),
+        ]
+        for labels, times, active, scores in steps:
+            hyps.add_results_masked_(
+                active_mask=torch.tensor(active, device=device),
+                labels=torch.tensor(labels, device=device),
+                time_indices=torch.tensor(times, device=device),
+                scores=torch.tensor(scores, device=device),
+            )
+        hypotheses = batched_hyps_to_hypotheses(hyps)
+        assert (hypotheses[0].y_sequence == torch.tensor([5, 2], device="cpu")).all()
+        assert (hypotheses[1].y_sequence == torch.tensor([4], device="cpu")).all()
+        assert (hypotheses[0].timestamp == torch.tensor([0, 1], device="cpu")).all()
+        assert (hypotheses[1].timestamp == torch.tensor([1], device="cpu")).all()
+        # only non-blank scores accumulated
+        assert hypotheses[0].score == pytest.approx(1.5)
+        assert hypotheses[1].score == pytest.approx(1.0)
+        # no logits recorded -> alignments stay None even though blank steps were stored
+        assert hypotheses[0].alignments is None
+        assert hypotheses[1].alignments is None
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_convert_logits_no_alignments_without_blank_steps(self, device: torch.device):
+        # logits are recorded but with_blank_steps=False -> alignments must NOT be produced
+        logits_dim = 7
+        hyps = BatchedHyps(
+            batch_size=2,
+            init_length=2,
+            blank_id=6,
+            logits_dim=logits_dim,
+            device=device,
+            with_logits=True,
+            with_blank_steps=False,
+        )
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, True], device=device),
+            labels=torch.tensor([5, 4], device=device),
+            time_indices=torch.tensor([0, 0], device=device),
+            scores=torch.tensor([1.0, 1.0], device=device),
+            logits=torch.rand((2, logits_dim), device=device),
+        )
+        hypotheses = batched_hyps_to_hypotheses(hyps)
+        assert hypotheses[0].alignments is None
+        assert hypotheses[1].alignments is None
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_convert_with_blank_steps_and_logits_alignments(self, device: torch.device):
+        # Reproduces alignment recovery: with_blank_steps=True + with_logits=True
+        batch_size = 2
+        logits_dim = 7
+        blank_index = 6
+        hyps = BatchedHyps(
+            batch_size=batch_size,
+            init_length=1,
+            blank_id=blank_index,
+            logits_dim=logits_dim,
+            device=device,
+            with_logits=True,
+            with_blank_steps=True,
+        )
+        # sequence 0: [[5, blank], [2, blank]] -> [5, 2]
+        # sequence 1: [[blank   ], [4, blank]] -> [4]
+        # one logits row per (batch, add-call); rows belonging to inactive entries are ignored
+        L = [torch.rand((batch_size, logits_dim), device=device) for _ in range(4)]
+        # call0: seq0=5@t0, seq1=blank@t0
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, True], device=device),
+            labels=torch.tensor([5, blank_index], device=device),
+            time_indices=torch.tensor([0, 0], device=device),
+            scores=torch.tensor([0.5, 0.1], device=device),
+            logits=L[0],
+        )
+        # call1: seq0=blank@t0, seq1=4@t1
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, True], device=device),
+            labels=torch.tensor([blank_index, 4], device=device),
+            time_indices=torch.tensor([0, 1], device=device),
+            scores=torch.tensor([0.1, 1.0], device=device),
+            logits=L[1],
+        )
+        # call2: seq0=2@t1, seq1=blank@t1
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, True], device=device),
+            labels=torch.tensor([2, blank_index], device=device),
+            time_indices=torch.tensor([1, 1], device=device),
+            scores=torch.tensor([1.0, 0.1], device=device),
+            logits=L[2],
+        )
+        # call3: seq0=blank@t1, seq1 inactive
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True, False], device=device),
+            labels=torch.tensor([blank_index, 0], device=device),
+            time_indices=torch.tensor([1, 0], device=device),
+            scores=torch.tensor([0.1, 0.0], device=device),
+            logits=L[3],
+        )
+
+        hypotheses = batched_hyps_to_hypotheses(hyps)
+        assert (hypotheses[0].y_sequence == torch.tensor([5, 2], device="cpu")).all()
+        assert (hypotheses[1].y_sequence == torch.tensor([4], device="cpu")).all()
+        assert hypotheses[0].score == pytest.approx(1.5)
+        assert hypotheses[1].score == pytest.approx(1.0)
+        assert (hypotheses[0].timestamp == torch.tensor([0, 1], device="cpu")).all()
+        assert (hypotheses[1].timestamp == torch.tensor([1], device="cpu")).all()
+
+        # alignments are grouped by timestamp; each entry is a (logits, label) tuple
+        etalon = [
+            [
+                [(L[0][0].cpu(), 5), (L[1][0].cpu(), blank_index)],
+                [(L[2][0].cpu(), 2), (L[3][0].cpu(), blank_index)],
+            ],
+            [
+                [(L[0][1].cpu(), blank_index)],
+                [(L[1][1].cpu(), 4), (L[2][1].cpu(), blank_index)],
+            ],
+        ]
+        for batch_i in range(batch_size):
+            assert len(hypotheses[batch_i].alignments) == len(etalon[batch_i])
+            for t, group_for_timestamp in enumerate(etalon[batch_i]):
+                assert len(hypotheses[batch_i].alignments[t]) == len(group_for_timestamp)
+                for step, (current_logits, label) in enumerate(group_for_timestamp):
+                    assert torch.allclose(hypotheses[batch_i].alignments[t][step][0], current_logits)
+                    assert hypotheses[batch_i].alignments[t][step][1] == label
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_convert_with_durations(self, device: torch.device):
+        # TDT-style: token durations are stored and (with blank steps) stripped together with blanks
+        blank_id = 0
+        hyps = BatchedHyps(
+            batch_size=1,
+            init_length=2,
+            blank_id=blank_id,
+            device=device,
+            with_durations=True,
+            with_blank_steps=True,
+        )
+        # transcript [3, blank, 7, blank] with durations [2, 1, 4, 1] -> [3, 7] with durations [2, 4]
+        steps = [
+            (3, 0, 2, 1.0),
+            (blank_id, 2, 1, 0.1),
+            (7, 3, 4, 2.0),
+            (blank_id, 7, 1, 0.2),
+        ]
+        for label, time, duration, score in steps:
+            hyps.add_results_masked_(
+                active_mask=torch.tensor([True], device=device),
+                labels=torch.tensor([label], device=device),
+                time_indices=torch.tensor([time], device=device),
+                scores=torch.tensor([score], device=device),
+                token_durations=torch.tensor([duration], device=device),
+            )
+        hypotheses = batched_hyps_to_hypotheses(hyps)
+        assert (hypotheses[0].y_sequence == torch.tensor([3, 7], device="cpu")).all()
+        assert (hypotheses[0].timestamp == torch.tensor([0, 3], device="cpu")).all()
+        assert (hypotheses[0].token_duration == torch.tensor([2, 4], device="cpu")).all()
+        assert hypotheses[0].score == pytest.approx(3.0)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_convert_with_step_confidence_no_blank_steps(self, device: torch.device):
+        # with_blank_steps=False: per-token confidence is precomputed, no frame_confidence (no blank steps)
+        hyps = BatchedHyps(
+            batch_size=1,
+            init_length=2,
+            blank_id=NON_COLLIDING_BLANK_ID,
+            device=device,
+            with_step_confidence=True,
+        )
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True], device=device),
+            labels=torch.tensor([5], device=device),
+            time_indices=torch.tensor([0], device=device),
+            scores=torch.tensor([1.0], device=device),
+            confidence=torch.tensor([0.9], device=device),
+        )
+        hyps.add_results_masked_(
+            active_mask=torch.tensor([True], device=device),
+            labels=torch.tensor([2], device=device),
+            time_indices=torch.tensor([1], device=device),
+            scores=torch.tensor([1.0], device=device),
+            confidence=torch.tensor([0.8], device=device),
+        )
+        hypotheses = batched_hyps_to_hypotheses(hyps)
+        assert hypotheses[0].non_blank_step_confidence_precomputed == pytest.approx([0.9, 0.8])
+        assert hypotheses[0].frame_confidence is None
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_convert_with_step_confidence_and_blank_steps(self, device: torch.device):
+        # with_blank_steps=True: frame_confidence is grouped per timestamp (incl. blanks),
+        # while non_blank_step_confidence_precomputed holds only non-blank tokens
+        blank_id = 0
+        hyps = BatchedHyps(
+            batch_size=1,
+            init_length=2,
+            blank_id=blank_id,
+            device=device,
+            with_step_confidence=True,
+            with_blank_steps=True,
+        )
+        # transcript [3, blank, 7], confidence [0.9, 0.5, 0.8], timestamps [0, 0, 1]
+        steps = [
+            (3, 0, 0.9, 1.0),
+            (blank_id, 0, 0.5, 0.1),
+            (7, 1, 0.8, 2.0),
+        ]
+        for label, time, confidence, score in steps:
+            hyps.add_results_masked_(
+                active_mask=torch.tensor([True], device=device),
+                labels=torch.tensor([label], device=device),
+                time_indices=torch.tensor([time], device=device),
+                scores=torch.tensor([score], device=device),
+                confidence=torch.tensor([confidence], device=device),
+            )
+        hypotheses = batched_hyps_to_hypotheses(hyps)
+        # non-blank tokens only
+        assert hypotheses[0].non_blank_step_confidence_precomputed == pytest.approx([0.9, 0.8])
+        # grouped by timestamp: t=0 has 2 steps (token + blank), t=1 has 1 step
+        frame_confidence = hypotheses[0].frame_confidence
+        assert len(frame_confidence) == 2
+        assert len(frame_confidence[0]) == 2
+        assert len(frame_confidence[1]) == 1
+        flat = [float(c) for group in frame_confidence for c in group]
+        assert flat == pytest.approx([0.9, 0.5, 0.8])

tests/collections/asr/decoding/test_biasing_multi_model.py

@@ -0,0 +1,279 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from nemo.collections.asr.parts.context_biasing.biasing_multi_model import (
+    GPUBiasingMultiModel,
+    GPUBiasingMultiModelReference,
+)
+from nemo.collections.asr.parts.context_biasing.boosting_graph_batched import (
+    BoostingTreeModelConfig,
+    GPUBoostingTreeModel,
+)
+from nemo.core.utils.optional_libs import TRITON_AVAILABLE
+
+DEVICES = [torch.device("cpu")]
+if torch.cuda.is_available():
+    DEVICES.append(torch.device("cuda"))
+if hasattr(torch, "mps") and torch.mps.is_available():
+    DEVICES.append(torch.device("mps"))
+
+# Triton only works on CUDA, so only test use_triton=True if Triton is available
+USE_TRITON_OPTIONS = [False, True] if TRITON_AVAILABLE else [False]
+
+
+def create_boosting_model(phrases: list[str], tokenizer, device: torch.device) -> GPUBoostingTreeModel:
+    """Helper to create boosting model from phrases"""
+    cfg = BoostingTreeModelConfig(key_phrases_list=phrases, context_score=1.0)
+    model = GPUBoostingTreeModel.from_config(cfg, tokenizer=tokenizer)
+    return model.to(device)
+
+
+class TestGPUBiasingMultiModel:
+    @pytest.mark.unit
+    @pytest.mark.with_downloads
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_add_models_incremental(self, stt_en_conformer_transducer_small, device: torch.device):
+        """Test adding 2 boosting models one-by-one, verifying arcs and states are correctly merged."""
+        tokenizer = stt_en_conformer_transducer_small.tokenizer
+        vocab_size = tokenizer.vocab_size
+
+        # Create empty multi-model
+        multi_model = GPUBiasingMultiModel(vocab_size=vocab_size).to(device)
+
+        # Initially empty
+        assert multi_model.num_models == 0
+        assert multi_model.has_models() is False
+        assert multi_model.num_states_total == 0
+        assert multi_model.num_arcs_extended_total == 0
+
+        # Create and add first model
+        model1 = create_boosting_model(["hello", "world"], tokenizer, device)
+        model_id1 = multi_model.add_model(model1, alpha=1.0)
+
+        # Verify after first model
+        assert model_id1 == 0
+        assert multi_model.num_models == 1
+        assert multi_model.has_models() is True
+        assert multi_model.model2active[model_id1].item() is True
+        assert multi_model.num_states_total == model1.num_states
+        assert multi_model.num_arcs_extended_total == model1.num_arcs_extended
+        assert multi_model.model2num_states[model_id1].item() == model1.num_states
+        assert multi_model.model2num_arcs_extended[model_id1].item() == model1.num_arcs_extended
+
+        # Create and add second model
+        model2 = create_boosting_model(["test", "one", "two"], tokenizer, device)
+        model_id2 = multi_model.add_model(model2, alpha=1.5)
+
+        # Verify after second model
+        assert model_id2 == 1
+        assert multi_model.num_models == 2
+        assert multi_model.has_models() is True
+        assert multi_model.model2active[model_id1].item() is True
+        assert multi_model.model2active[model_id2].item() is True
+        assert multi_model.num_states_total == model1.num_states + model2.num_states
+        assert multi_model.num_arcs_extended_total == model1.num_arcs_extended + model2.num_arcs_extended
+
+        # Verify offsets
+        assert multi_model.model2states_offset[model_id1].item() == 0
+        assert multi_model.model2states_offset[model_id2].item() == model1.num_states
+        assert multi_model.model2arcs_offset[model_id1].item() == 0
+        assert multi_model.model2arcs_offset[model_id2].item() == model1.num_arcs_extended
+
+        # Verify init states work
+        init_states = multi_model.get_init_states(batch_size=4, bos=True)
+        assert init_states.shape == (4,)
+        assert init_states.device.type == device.type
+
+    @pytest.mark.unit
+    @pytest.mark.with_downloads
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_add_then_remove_model(self, stt_en_conformer_transducer_small, device: torch.device):
+        """Test adding 2 models then removing the first one."""
+        tokenizer = stt_en_conformer_transducer_small.tokenizer
+        vocab_size = tokenizer.vocab_size
+
+        multi_model = GPUBiasingMultiModel(vocab_size=vocab_size).to(device)
+
+        # Add two models
+        model1 = create_boosting_model(["alpha", "beta"], tokenizer, device)
+        model2 = create_boosting_model(["gamma", "delta"], tokenizer, device)
+
+        model_id1 = multi_model.add_model(model1, alpha=1.0)
+        model_id2 = multi_model.add_model(model2, alpha=2.0)
+
+        # Store counts before removal
+        model1_num_states = model1.num_states
+        model1_num_arcs = model1.num_arcs_extended
+        total_states_before = multi_model.num_states_total
+        total_arcs_before = multi_model.num_arcs_extended_total
+
+        assert multi_model.model2active[model_id1].item() is True
+        assert multi_model.model2active[model_id2].item() is True
+
+        # Remove first model
+        multi_model.remove_model(model_id1)
+
+        # Verify removal
+        assert model_id1 in multi_model.free_ids
+        assert multi_model.model2active[model_id1].item() is False
+        assert multi_model.model2active[model_id2].item() is True
+        assert multi_model.model2alpha[model_id1].item() == 0.0
+        assert multi_model.model2alpha[model_id2].item() == 2.0
+
+        # Verify state/arc counts decreased
+        assert multi_model.num_states_total == total_states_before - model1_num_states
+        assert multi_model.num_arcs_extended_total == total_arcs_before - model1_num_arcs
+
+        # Verify model2 offset updated (shifted left)
+        assert multi_model.model2states_offset[model_id2].item() == 0
+        assert multi_model.model2arcs_offset[model_id2].item() == 0
+
+    @pytest.mark.unit
+    @pytest.mark.with_downloads
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_model_id_reuse(self, stt_en_conformer_transducer_small, device):
+        """Test that removed model IDs are reused."""
+        tokenizer = stt_en_conformer_transducer_small.tokenizer
+        vocab_size = tokenizer.vocab_size
+
+        multi_model = GPUBiasingMultiModel(vocab_size=vocab_size).to(device)
+
+        # Add model1 -> id=0
+        model1 = create_boosting_model(["first"], tokenizer, device)
+        model_id1 = multi_model.add_model(model1)
+        assert model_id1 == 0
+
+        # Add model2 -> id=1
+        model2 = create_boosting_model(["second"], tokenizer, device)
+        model_id2 = multi_model.add_model(model2)
+        assert model_id2 == 1
+
+        # Remove model1
+        multi_model.remove_model(model_id1)
+        assert model_id1 in multi_model.free_ids
+
+        # Add model3 -> should reuse id=0
+        model3 = create_boosting_model(["third"], tokenizer, device)
+        model_id3 = multi_model.add_model(model3)
+        assert model_id3 == model_id1  # Reused ID
+        assert model_id1 not in multi_model.free_ids  # No longer free
+
+        # Verify model3 is active
+        assert multi_model.model2active[model_id3].item() is True
+
+    @pytest.mark.unit
+    @pytest.mark.with_downloads
+    @pytest.mark.parametrize("device", DEVICES)
+    @pytest.mark.parametrize("batch_size", [1, 4])
+    @pytest.mark.parametrize("use_triton", USE_TRITON_OPTIONS)
+    @pytest.mark.parametrize("bos", [True, False])
+    def test_advance_matches_reference(
+        self, stt_en_conformer_transducer_small, device: torch.device, batch_size: int, use_triton, bos: bool
+    ):
+        """Verify GPUBiasingMultiModel produces same scores/states as reference implementation."""
+        tokenizer = stt_en_conformer_transducer_small.tokenizer
+        vocab_size = tokenizer.vocab_size
+
+        # Create both implementations
+        multi_model = GPUBiasingMultiModel(vocab_size=vocab_size, use_triton=use_triton).to(device)
+        reference = GPUBiasingMultiModelReference(vocab_size=vocab_size).to(device)
+
+        # Create boosting models with same phrases
+        phrases1 = ["hello world", "test"]
+        phrases2 = ["neural", "network"]
+
+        model1_mm = create_boosting_model(phrases1, tokenizer, device)
+        model1_ref = create_boosting_model(phrases1, tokenizer, device)
+        model2_mm = create_boosting_model(phrases2, tokenizer, device)
+        model2_ref = create_boosting_model(phrases2, tokenizer, device)
+
+        # Add models to both with same alpha values
+        alpha1, alpha2 = 1.0, 1.5
+        model_id1_mm = multi_model.add_model(model1_mm, alpha=alpha1)
+        model_id1_ref = reference.add_model(model1_ref, alpha=alpha1)
+        model_id2_mm = multi_model.add_model(model2_mm, alpha=alpha2)
+        model_id2_ref = reference.add_model(model2_ref, alpha=alpha2)
+
+        assert model_id1_mm == model_id1_ref
+        assert model_id2_mm == model_id2_ref
+
+        # Get initial states
+        states_mm = multi_model.get_init_states(batch_size=batch_size, bos=bos)
+        states_ref = reference.get_init_states(batch_size=batch_size, bos=bos)
+
+        # Create model_ids tensor with alternating models
+        model_ids = torch.tensor(
+            [model_id1_mm if i % 2 == 0 else model_id2_mm for i in range(batch_size)],
+            dtype=torch.long,
+            device=device,
+        )
+
+        # Call advance on both
+        scores_mm, next_states_mm = multi_model.advance(states_mm, model_ids)
+        scores_ref, next_states_ref = reference.advance(states_ref, model_ids)
+
+        # Verify shapes
+        assert scores_mm.shape == (batch_size, vocab_size)
+        assert next_states_mm.shape == (batch_size, vocab_size)
+        assert scores_ref.shape == (batch_size, vocab_size)
+        assert next_states_ref.shape == (batch_size, vocab_size)
+
+        # Verify scores and states match
+        assert torch.allclose(
+            scores_mm, scores_ref, atol=1e-5
+        ), f"Scores mismatch: max diff = {(scores_mm - scores_ref).abs().max()}"
+        assert torch.equal(next_states_mm, next_states_ref), "Next states mismatch"
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_empty_multi_model(self, device: torch.device):
+        """Test behavior of empty multi-model."""
+        vocab_size = 100
+        multi_model = GPUBiasingMultiModel(vocab_size=vocab_size, use_triton=False).to(device)
+
+        # Verify empty state
+        assert multi_model.has_models() is False
+        assert multi_model.num_models == 0
+        assert multi_model.num_states_total == 0
+        assert multi_model.num_arcs_extended_total == 0
+
+        # get_init_states should work and return START_STATE
+        init_states = multi_model.get_init_states(batch_size=4, bos=True)
+        assert init_states.shape == (4,)
+        assert (init_states == GPUBiasingMultiModel.START_STATE).all()
+
+    @pytest.mark.unit
+    @pytest.mark.with_downloads
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_get_alphas(self, stt_en_conformer_transducer_small, device: torch.device):
+        """Per-stream alpha lookup returns model weight or 0 for invalid ids."""
+        tokenizer = stt_en_conformer_transducer_small.tokenizer
+        vocab_size = tokenizer.vocab_size
+        multi_model = GPUBiasingMultiModel(vocab_size=vocab_size).to(device)
+        model1 = create_boosting_model(["hello"], tokenizer, device)
+        model2 = create_boosting_model(["world"], tokenizer, device)
+        model_id1 = multi_model.add_model(model1, alpha=1.0)
+        model_id2 = multi_model.add_model(model2, alpha=2.5)
+
+        model_ids = torch.tensor([model_id1, model_id2, -1, model_id1], device=device, dtype=torch.long)
+        alphas = multi_model.get_alphas(model_ids)
+
+        assert alphas.shape == (4,)
+        assert alphas[0].item() == pytest.approx(1.0)
+        assert alphas[1].item() == pytest.approx(2.5)
+        assert alphas[2].item() == pytest.approx(0.0)
+        assert alphas[3].item() == pytest.approx(1.0)

tests/collections/asr/decoding/test_ctc_decoding.py

@@ -0,0 +1,487 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import os
+from functools import cached_property, lru_cache
+from pathlib import Path
+
+import pytest
+import torch
+from kaldialign import edit_distance
+from omegaconf import DictConfig, open_dict
+
+from nemo.collections.asr.models import ASRModel
+from nemo.collections.asr.parts.mixins import mixins
+from nemo.collections.asr.parts.submodules.ctc_decoding import (
+    CTCBPEDecoding,
+    CTCBPEDecodingConfig,
+    CTCDecoding,
+    CTCDecodingConfig,
+)
+from nemo.collections.asr.parts.submodules.ngram_lm.ngram_lm_batched import NGramGPULanguageModel
+from nemo.collections.asr.parts.utils.asr_confidence_utils import ConfidenceConfig
+from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
+from nemo.core.utils.cuda_python_utils import skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported
+from tests.collections.asr.decoding.test_timestamps import BaseTimestampsTest
+
+
+@pytest.fixture(scope="module")
+def audio_file(test_data_dir):
+    return os.path.join(test_data_dir, "asr/test/an4/wav/cen3-mjwl-b.wav")
+
+
+CTC_MODEL = "nvidia/stt_en_conformer_ctc_small"
+
+
+@pytest.fixture(scope="module")
+def kenlm_model_path(tmp_path_factory, test_data_dir):
+    lm_path = Path(test_data_dir) / "asr/kenlm_ngram_lm/parakeet-tdt_ctc-110m-libri-1024.kenlm.tmp.arpa"
+    assert os.path.exists(lm_path), f"LM file not found: {lm_path}"
+    lm_nemo_path = tmp_path_factory.mktemp("lm") / f"{lm_path.name}.nemo"
+    NGramGPULanguageModel.from_file(lm_path, vocab_size=1024).save_to(f"{lm_nemo_path}")
+    return f"{lm_nemo_path}"
+
+
+@pytest.fixture(scope="module")
+def ctc_model():
+    model = ASRModel.from_pretrained(model_name=CTC_MODEL, map_location="cpu")
+    model.eval()
+    return model
+
+
+def char_vocabulary():
+    return [' ', 'a', 'b', 'c', 'd', 'e', 'f', '.']
+
+
+@pytest.fixture()
+@lru_cache(maxsize=8)
+def tmp_tokenizer(test_data_dir):
+    cfg = DictConfig({'dir': os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128"), 'type': 'wpe'})
+
+    class _TmpASRBPE(mixins.ASRBPEMixin):
+        def register_artifact(self, _, vocab_path):
+            return vocab_path
+
+    asrbpe = _TmpASRBPE()
+    asrbpe._setup_tokenizer(cfg)
+    return asrbpe.tokenizer
+
+
+class TestCTCDecoding:
+    @pytest.mark.unit
+    def test_constructor(self):
+        cfg = CTCDecodingConfig()
+        vocab = char_vocabulary()
+        decoding = CTCDecoding(decoding_cfg=cfg, vocabulary=vocab)
+        assert decoding is not None
+
+    @pytest.mark.unit
+    def test_constructor_subword(self, tmp_tokenizer):
+        cfg = CTCBPEDecodingConfig()
+        decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=tmp_tokenizer)
+        assert decoding is not None
+
+    @pytest.mark.unit
+    def test_char_decoding_greedy_forward(
+        self,
+    ):
+        cfg = CTCDecodingConfig(strategy='greedy')
+        vocab = char_vocabulary()
+        decoding = CTCDecoding(decoding_cfg=cfg, vocabulary=vocab)
+
+        B, T = 4, 20
+        V = len(char_vocabulary()) + 1
+        input_signal = torch.randn(size=(B, T, V))
+        length = torch.randint(low=1, high=T, size=[B])
+
+        with torch.no_grad():
+            hypotheses = decoding.ctc_decoder_predictions_tensor(
+                input_signal, length, fold_consecutive=True, return_hypotheses=False
+            )
+            texts = [hyp.text for hyp in hypotheses]
+
+            for text in texts:
+                assert isinstance(text, str)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('alignments', [False, True])
+    @pytest.mark.parametrize('timestamps', [False, True])
+    def test_char_decoding_greedy_forward_hypotheses(self, alignments, timestamps):
+        cfg = CTCDecodingConfig(strategy='greedy', preserve_alignments=alignments, compute_timestamps=timestamps)
+        vocab = char_vocabulary()
+        decoding = CTCDecoding(decoding_cfg=cfg, vocabulary=vocab)
+
+        B, T = 4, 20
+        V = len(char_vocabulary()) + 1
+        input_signal = torch.randn(size=(B, T, V))
+        length = torch.randint(low=1, high=T, size=[B])
+
+        with torch.no_grad():
+            hyps = decoding.ctc_decoder_predictions_tensor(
+                input_signal, length, fold_consecutive=True, return_hypotheses=True
+            )
+
+            for idx, hyp in enumerate(hyps):
+                assert isinstance(hyp, Hypothesis)
+                assert torch.is_tensor(hyp.y_sequence)
+                assert isinstance(hyp.text, str)
+
+                # alignments check
+                if alignments:
+                    assert hyp.alignments is not None
+                    assert isinstance(hyp.alignments, tuple)
+                    assert len(hyp.alignments[0]) == length[idx]
+                    assert len(hyp.alignments[1]) == length[idx]
+
+                # timestamps check
+                if timestamps:
+                    BaseTimestampsTest.check_char_timestamps(hyp, decoding)
+
+    @pytest.mark.unit
+    def test_subword_decoding_greedy_forward(self, tmp_tokenizer):
+        cfg = CTCBPEDecodingConfig(strategy='greedy')
+        decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=tmp_tokenizer)
+
+        B, T = 4, 20
+        V = decoding.tokenizer.tokenizer.vocab_size + 1
+        input_signal = torch.randn(size=(B, T, V))
+        length = torch.randint(low=1, high=T, size=[B])
+
+        with torch.no_grad():
+            hypotheses = decoding.ctc_decoder_predictions_tensor(
+                input_signal, length, fold_consecutive=True, return_hypotheses=False
+            )
+            texts = [hyp.text for hyp in hypotheses]
+
+            for text in texts:
+                assert isinstance(text, str)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('alignments', [False, True])
+    @pytest.mark.parametrize('timestamps', [False, True])
+    @pytest.mark.pleasefixme
+    def test_subword_decoding_greedy_forward_hypotheses(self, tmp_tokenizer, alignments, timestamps):
+        cfg = CTCBPEDecodingConfig(strategy='greedy', preserve_alignments=alignments, compute_timestamps=timestamps)
+        decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=tmp_tokenizer)
+
+        B, T = 4, 20
+        V = decoding.tokenizer.tokenizer.vocab_size + 1
+        input_signal = torch.randn(size=(B, T, V))
+        length = torch.randint(low=1, high=T, size=[B])
+
+        with torch.no_grad():
+            hyps = decoding.ctc_decoder_predictions_tensor(
+                input_signal, length, fold_consecutive=True, return_hypotheses=True
+            )
+
+            for idx, hyp in enumerate(hyps):
+                assert isinstance(hyp, Hypothesis)
+                assert torch.is_tensor(hyp.y_sequence)
+                assert isinstance(hyp.text, str)
+
+                # alignments check
+                if alignments:
+                    assert hyp.alignments is not None
+                    assert isinstance(hyp.alignments, tuple)
+                    assert len(hyp.alignments[0]) == length[idx]
+                    assert len(hyp.alignments[1]) == length[idx]
+
+                # timestamps check
+                if timestamps:
+                    BaseTimestampsTest.check_subword_timestamps(hyp, decoding)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('alignments', [False, True])
+    @pytest.mark.parametrize('timestamps', [False, True])
+    @pytest.mark.parametrize('preserve_frame_confidence', [False, True])
+    @pytest.mark.parametrize('length_is_none', [False, True])
+    @pytest.mark.parametrize(
+        "logprobs_device",
+        [
+            torch.device("cpu"),
+            pytest.param(
+                torch.device("cuda"),
+                marks=pytest.mark.skipif(
+                    not torch.cuda.is_available(),
+                    reason='CUDA required for test.',
+                ),
+            ),
+        ],
+    )
+    @pytest.mark.parametrize(
+        "length_device",
+        [
+            torch.device("cpu"),
+            pytest.param(
+                torch.device("cuda"),
+                marks=pytest.mark.skipif(
+                    not torch.cuda.is_available(),
+                    reason='CUDA required for test.',
+                ),
+            ),
+        ],
+    )
+    def test_batched_decoding_logprobs(
+        self,
+        tmp_tokenizer,
+        alignments,
+        timestamps,
+        preserve_frame_confidence,
+        length_is_none,
+        logprobs_device,
+        length_device,
+    ):
+        cfg = CTCBPEDecodingConfig(
+            strategy='greedy',
+            preserve_alignments=alignments,
+            compute_timestamps=timestamps,
+            confidence_cfg=ConfidenceConfig(preserve_frame_confidence=preserve_frame_confidence),
+        )
+        unbatched_decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=tmp_tokenizer)
+
+        cfg.strategy = 'greedy_batch'
+        batched_decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=tmp_tokenizer)
+
+        torch.manual_seed(1)
+        B, T = 4, 20
+        V = unbatched_decoding.tokenizer.tokenizer.vocab_size + 1
+        input_signal = torch.randn(size=(B, T, V), device=logprobs_device)
+        # Set the blank index to a very high probability to make sure
+        # that we always handle at least a few blanks.
+        input_signal[:, 0, unbatched_decoding.tokenizer.tokenizer.vocab_size] = 1000
+        input_signal[:, 1, unbatched_decoding.tokenizer.tokenizer.vocab_size] = 1000
+        if length_is_none:
+            length = None
+        else:
+            length = torch.randint(low=1, high=T, size=[B], device=length_device)
+
+        with torch.inference_mode():
+            hyps = unbatched_decoding.ctc_decoder_predictions_tensor(
+                input_signal, length, fold_consecutive=True, return_hypotheses=True
+            )
+
+            batched_hyps = batched_decoding.ctc_decoder_predictions_tensor(
+                input_signal, length, fold_consecutive=True, return_hypotheses=True
+            )
+
+            assert len(hyps) == len(batched_hyps) == B
+            for hyp, batched_hyp in zip(hyps, batched_hyps):
+                assert torch.abs(hyp.score - batched_hyp.score) <= 1e-5
+                assert torch.all(hyp.y_sequence == batched_hyp.y_sequence)
+                if timestamps:
+                    assert hyp.timestamp == batched_hyp.timestamp
+                if alignments:
+                    assert torch.all(hyp.alignments[0] == batched_hyp.alignments[0])
+                    assert torch.all(hyp.alignments[1] == batched_hyp.alignments[1])
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize('timestamps', [False, True])
+    @pytest.mark.parametrize('length_is_none', [False, True])
+    @pytest.mark.parametrize(
+        "labels_device",
+        [
+            torch.device("cpu"),
+            pytest.param(
+                torch.device("cuda"),
+                marks=pytest.mark.skipif(
+                    not torch.cuda.is_available(),
+                    reason='CUDA required for test.',
+                ),
+            ),
+        ],
+    )
+    @pytest.mark.parametrize(
+        "length_device",
+        [
+            torch.device("cpu"),
+            pytest.param(
+                torch.device("cuda"),
+                marks=pytest.mark.skipif(
+                    not torch.cuda.is_available(),
+                    reason='CUDA required for test.',
+                ),
+            ),
+        ],
+    )
+    def test_batched_decoding_labels(self, tmp_tokenizer, timestamps, length_is_none, labels_device, length_device):
+        cfg = CTCBPEDecodingConfig(strategy='greedy', compute_timestamps=timestamps)
+        unbatched_decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=tmp_tokenizer)
+        cfg.strategy = 'greedy_batch'
+        batched_decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=tmp_tokenizer)
+
+        torch.manual_seed(1)
+        B, T = 4, 20
+        V = unbatched_decoding.tokenizer.tokenizer.vocab_size + 1
+        input_labels = torch.randint(V, size=(B, T), device=labels_device)
+        # Set some indices to blank to make sure that we always handle
+        # at least a few blanks.
+        input_labels[:, 0] = unbatched_decoding.tokenizer.tokenizer.vocab_size
+        input_labels[:, 1] = unbatched_decoding.tokenizer.tokenizer.vocab_size
+        if length_is_none:
+            length = None
+        else:
+            length = torch.randint(low=1, high=T, size=[B], device=length_device)
+
+        with torch.inference_mode():
+            hyps = unbatched_decoding.ctc_decoder_predictions_tensor(
+                input_labels, length, fold_consecutive=True, return_hypotheses=True
+            )
+
+            batched_hyps = batched_decoding.ctc_decoder_predictions_tensor(
+                input_labels, length, fold_consecutive=True, return_hypotheses=True
+            )
+
+            assert len(hyps) == len(batched_hyps) == B
+            for hyp, batched_hyp in zip(hyps, batched_hyps):
+                assert abs(hyp.score - batched_hyp.score) <= 1e-5
+                assert torch.all(hyp.y_sequence == batched_hyp.y_sequence)
+                if timestamps:
+                    assert hyp.timestamp == batched_hyp.timestamp
+
+
+class TestCTCTimestamps(BaseTimestampsTest):
+    """CTC-specific timestamp tests that inherit from BaseTimestampsTest"""
+
+    @cached_property
+    def decoding_char(self):
+        cfg = CTCDecodingConfig()
+        vocab = char_vocabulary()
+        decoding = CTCDecoding(decoding_cfg=cfg, vocabulary=vocab)
+        return decoding
+
+    @cached_property
+    def decoding_subword_wpe(self):
+        cfg = CTCBPEDecodingConfig(compute_timestamps=True)
+        decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=self.tmp_tokenizer)
+        return decoding
+
+    @cached_property
+    def decoding_subword_bpe(self):
+        cfg = CTCBPEDecodingConfig(compute_timestamps=True)
+        decoding = CTCBPEDecoding(decoding_cfg=cfg, tokenizer=self.bpe_tokenizer)
+        return decoding
+
+    @pytest.mark.unit
+    def test_word_offsets_subword_wpe(self, tmp_tokenizer):
+        self.tmp_tokenizer = tmp_tokenizer
+        super().test_word_offsets_subword_wpe()
+
+    @pytest.mark.unit
+    def test_word_offsets_subword_wpe_other_delimiter(self, tmp_tokenizer):
+        self.tmp_tokenizer = tmp_tokenizer
+        super().test_word_offsets_subword_wpe_other_delimiter()
+
+
+class TestCTCGreedyDecodingWithNGPU_LM:
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Test is only GPU-based decoding")
+    def test_ctc_decoding_gpulm(
+        self,
+        audio_file,
+        kenlm_model_path,
+        ctc_model,
+    ):
+        device = torch.device("cuda")
+        model = ctc_model.to(device)
+
+        gt_hyp = model.transcribe([audio_file], num_workers=None)
+
+        decoding_config = copy.deepcopy(model.cfg.decoding)
+        with open_dict(model.decoding.cfg) as cfg:
+            cfg.greedy["ngram_lm_model"] = kenlm_model_path
+            cfg.greedy["ngram_lm_alpha"] = 0.0
+            model.change_decoding_strategy(cfg)
+        lm_hyp = model.transcribe([audio_file], num_workers=None)
+
+        assert gt_hyp[0].text == lm_hyp[0].text
+        assert abs(gt_hyp[0].score - lm_hyp[0].score) <= 1e-3
+
+        with open_dict(model.decoding.cfg) as cfg:
+            cfg.greedy["ngram_lm_model"] = kenlm_model_path
+            cfg.greedy["ngram_lm_alpha"] = 10.0
+        model.change_decoding_strategy(cfg)
+        lm_hyp = model.transcribe([audio_file], num_workers=None)
+        assert gt_hyp[0].text != lm_hyp[0].text
+        assert abs(gt_hyp[0].score - lm_hyp[0].score) > 1e-3
+
+        model.change_decoding_strategy(decoding_config)
+
+
+class TestCTCGreedyDecodingCudaGrpahs:
+    """
+    Tests CudaGraphs implementations from CTC models greedy decoding
+    """
+
+    @pytest.mark.with_downloads
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA decoder can run only on CUDA")
+    @pytest.mark.parametrize("force_mode", ["no_graphs", "no_while_loops", "full_graph"])
+    def test_stated_stateless(self, audio_file, kenlm_model_path, ctc_model, force_mode: str):
+        """
+        Compares pure Pytorch and with three modes of statefull implementations for double floating point precision.
+            1. Pure pytorch, but statefull implementation: no_graphs
+            2. With CudaGrpahs: no_while_loops and full_graph.
+        """
+        if force_mode == "full_graph":
+            skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported()
+
+        device = torch.device("cuda")
+        model = ctc_model.to(device)
+        decoding_config = copy.deepcopy(model.cfg.decoding)
+
+        with open_dict(model.decoding.cfg) as cfg:
+            cfg.greedy["ngram_lm_model"] = kenlm_model_path
+            cfg.greedy["ngram_lm_alpha"] = 0.2
+            cfg.greedy["allow_cuda_graphs"] = False
+
+            model.change_decoding_strategy(cfg)
+
+        actual_hypotheses = model.transcribe([audio_file], num_workers=None)
+        actual_transcripts = [hyp.text for hyp in actual_hypotheses]
+        actual_scores = [hyp.score for hyp in actual_hypotheses]
+        actual_timestamps = [hyp.timestamp for hyp in actual_hypotheses]
+
+        # transcribe with use implementation with cuda graphs
+        model.decoding.cfg["greedy"]["allow_cuda_graphs"] = True
+        model.change_decoding_strategy(model.decoding.cfg)
+        model.decoding.decoding.force_cuda_graphs_mode(mode=force_mode)
+
+        cudagraph_hypotheses = model.transcribe([audio_file], num_workers=None)
+        cudagraph_transcripts = [hyp.text for hyp in cudagraph_hypotheses]
+        cudagraph_scores = [hyp.score for hyp in cudagraph_hypotheses]
+        cudagraph_timestamps = [hyp.timestamp for hyp in cudagraph_hypotheses]
+
+        for batch_idx in range(len(actual_transcripts)):
+            assert len(actual_transcripts[batch_idx]) == len(cudagraph_transcripts[batch_idx])
+            assert cudagraph_scores[batch_idx] == pytest.approx(
+                actual_scores[batch_idx], abs=1e-2
+            ), f"Scores mismatch for batch_idx {batch_idx}"
+            assert (
+                cudagraph_timestamps[batch_idx] == actual_timestamps[batch_idx]
+            ), f"Timestamps mismatch for batch_idx {batch_idx}"
+
+            ref_words = actual_transcripts[batch_idx].split()
+            hyp_words = cudagraph_transcripts[batch_idx].split()
+            wer = edit_distance(ref_words, hyp_words)['total'] / max(len(ref_words), 1)
+
+            assert wer <= 1e-3, "Cuda graph greedy decoder should match original decoder implementation."
+
+            for actual, fast in zip(actual_transcripts[batch_idx], cudagraph_transcripts[batch_idx]):
+                if actual != fast:
+                    print("Erroneous samples in batch:", batch_idx)
+                    print("Original transcript:", actual)
+                    print("New transcript:", fast)
+
+        model.change_decoding_strategy(decoding_config)

tests/collections/asr/decoding/test_cuda_graph_rnnt_greedy_decoding.py

@@ -0,0 +1,368 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import copy
+import glob
+import types
+
+import lightning.pytorch as ptl
+import pytest
+import torch
+from kaldialign import edit_distance
+from omegaconf import DictConfig, open_dict
+
+from nemo.core.config.pytorch_lightning import TrainerConfig
+from nemo.core.utils.cuda_python_utils import skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported
+
+
+# These tests move the model to CUDA before calling transcribe(), so avoid forking DataLoader workers afterwards.
+CUDA_GRAPH_TRANSCRIBE_NUM_WORKERS = 0
+
+
+def test_forced_full_graph_compile_does_not_fallback():
+    from nemo.collections.asr.parts.submodules.transducer_decoding.rnnt_label_looping import (
+        GreedyBatchedRNNTLabelLoopingComputer,
+    )
+
+    accelerator_error = getattr(torch, "AcceleratorError", RuntimeError)
+    computer = GreedyBatchedRNNTLabelLoopingComputer.__new__(GreedyBatchedRNNTLabelLoopingComputer)
+    computer.cuda_graphs_allow_fallback = False
+
+    with pytest.raises(RuntimeError, match="Full CUDA graph decoding failed"):
+        computer._raise_or_warn_no_while_loop_cuda_graphs(accelerator_error("CUDA error: invalid argument"))
+
+
+def test_conditional_node_restores_previous_stream_on_body_error(monkeypatch):
+    from nemo.core.utils import cuda_python_utils
+
+    if not cuda_python_utils.CUDA_PYTHON_AVAILABLE:
+        pytest.skip("cuda-python is required to test with_conditional_node")
+
+    class FakeStream:
+        def __init__(self, name):
+            self.name = name
+            self.cuda_stream = name
+
+    class FakeTorchCuda:
+        def __init__(self):
+            self.parent_stream = FakeStream("parent")
+            self.body_stream = FakeStream("body")
+            self.current = self.parent_stream
+            self.set_calls = []
+
+        def current_stream(self, device=None):
+            return self.current
+
+        def Stream(self, device=None):
+            return self.body_stream
+
+        def set_stream(self, stream):
+            self.current = stream
+            self.set_calls.append(stream)
+
+    class FakeCudart:
+        cudaStreamCaptureStatus = types.SimpleNamespace(cudaStreamCaptureStatusActive="active")
+        cudaStreamUpdateCaptureDependenciesFlags = types.SimpleNamespace(cudaStreamSetCaptureDependencies="set")
+        cudaStreamCaptureMode = types.SimpleNamespace(cudaStreamCaptureModeThreadLocal="thread_local")
+
+        def __init__(self):
+            self.ended_streams = []
+
+        def cudaStreamGetCaptureInfo(self, stream):
+            return ("active", None, "graph", ["dependency"])
+
+        def cudaStreamUpdateCaptureDependencies(self, *args):
+            return ()
+
+        def cudaStreamBeginCaptureToGraph(self, *args):
+            return ()
+
+        def cudaStreamEndCapture(self, stream):
+            self.ended_streams.append(stream)
+            return ()
+
+    class FakeCuda:
+        CUgraphNodeType = types.SimpleNamespace(CU_GRAPH_NODE_TYPE_CONDITIONAL="conditional")
+        CUgraphConditionalNodeType = types.SimpleNamespace(CU_GRAPH_COND_TYPE_WHILE="while")
+
+        class CUgraphNodeParams:
+            def __init__(self):
+                self.conditional = types.SimpleNamespace(phGraph_out=["body_graph"])
+
+        def cuGraphAddNode(self, *args):
+            return ("node",)
+
+        def cuCtxGetCurrent(self):
+            return ("ctx",)
+
+        def cuLaunchKernel(self, *args):
+            return ()
+
+    fake_torch_cuda = FakeTorchCuda()
+    fake_cudart = FakeCudart()
+    fake_args = types.SimpleNamespace(ctypes=types.SimpleNamespace(data=1234))
+    fake_handle = types.SimpleNamespace(getPtr=lambda: 5678)
+
+    monkeypatch.setattr(cuda_python_utils, "cu_call", lambda result: result)
+    monkeypatch.setattr(cuda_python_utils, "cuda", FakeCuda())
+    monkeypatch.setattr(cuda_python_utils, "cudart", fake_cudart)
+    monkeypatch.setattr(cuda_python_utils, "cuda_python_version", "13.0.0")
+    monkeypatch.setattr(cuda_python_utils.torch, "cuda", fake_torch_cuda)
+
+    with pytest.raises(RuntimeError, match="body failed"):
+        with cuda_python_utils.with_conditional_node("kernel", fake_args, fake_handle, device="cuda"):
+            assert fake_torch_cuda.current_stream(device="cuda") is fake_torch_cuda.body_stream
+            raise RuntimeError("body failed")
+
+    assert fake_torch_cuda.current_stream(device="cuda") is fake_torch_cuda.parent_stream
+    assert fake_torch_cuda.set_calls == [fake_torch_cuda.body_stream, fake_torch_cuda.parent_stream]
+    assert fake_cudart.ended_streams == ["body"]
+
+
+@pytest.mark.with_downloads
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA decoder can run only on CUDA")
+@pytest.mark.parametrize(
+    ("model_name", "batch_size", "enable_bfloat16"),
+    [
+        ("stt_en_fastconformer_transducer_large", 8, False),
+        ("stt_en_fastconformer_transducer_large", 8, True),
+    ],
+)
+@pytest.mark.parametrize("loop_labels", [False, True])
+def test_cuda_graph_rnnt_greedy_decoder(model_name, batch_size, enable_bfloat16, loop_labels: bool, request):
+    if not loop_labels:
+        skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported()
+    if enable_bfloat16 and not torch.cuda.is_bf16_supported():
+        pytest.skip("bfloat16 is not supported")
+
+    device = torch.device("cuda")
+    nemo_model = request.getfixturevalue(model_name).to(device)
+    decoding_config = copy.deepcopy(nemo_model.cfg.decoding)
+
+    with open_dict(decoding_config):
+        decoding_config["greedy"]["max_symbols"] = 5
+        decoding_config["greedy"]["loop_labels"] = loop_labels
+        decoding_config["greedy"]["use_cuda_graph_decoder"] = False
+
+    nemo_model.change_decoding_strategy(decoding_config)
+    audio_filepaths = glob.glob("tests/.data/asr/test/an4/wav/*.wav")
+
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=enable_bfloat16):
+        actual_hypotheses = nemo_model.transcribe(
+            audio_filepaths, batch_size=batch_size, num_workers=CUDA_GRAPH_TRANSCRIBE_NUM_WORKERS
+        )
+
+    actual_transcripts = [hyp.text for hyp in actual_hypotheses]
+    actual_y_sequences = [hyp.y_sequence for hyp in actual_hypotheses]
+
+    decoding_config["greedy"]["use_cuda_graph_decoder"] = True
+
+    nemo_model.change_decoding_strategy(decoding_config)
+
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=enable_bfloat16):
+        fast_hypotheses = nemo_model.transcribe(
+            audio_filepaths, batch_size=batch_size, num_workers=CUDA_GRAPH_TRANSCRIBE_NUM_WORKERS
+        )
+
+    fast_transcripts = [hyp.text for hyp in fast_hypotheses]
+    fast_y_sequences = [hyp.y_sequence for hyp in fast_hypotheses]
+
+    total_dist = sum(
+        edit_distance(r.split(), h.split())['total'] for r, h in zip(actual_transcripts, fast_transcripts)
+    )
+    total_words = sum(len(r.split()) for r in actual_transcripts)
+    wer = total_dist / total_words if total_words > 0 else 0.0
+    y_sequence_eq = [torch.equal(act_y, fast_y) for (act_y, fast_y) in zip(actual_y_sequences, fast_y_sequences)]
+
+    assert wer <= 1e-3, "Cuda graph greedy decoder should match original decoder implementation."
+    assert all(y_sequence_eq), "Cuda graph greedy decoder should match original decoder implementation."
+
+    for actual, fast in zip(actual_transcripts, fast_transcripts):
+        if actual != fast:
+            print("erroneous samples:")
+            print("Original transcript:", actual)
+            print("New transcript:", fast)
+
+
+@pytest.mark.with_downloads
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA decoder can run only on CUDA")
+@pytest.mark.parametrize("force_mode", ["no_graphs", "no_while_loops", "full_graph"])
+@pytest.mark.parametrize("enable_bfloat16", [False, True])
+def test_loop_labels_cuda_graph_rnnt_greedy_decoder_forced_mode(
+    stt_en_fastconformer_transducer_large, force_mode: str, enable_bfloat16: bool
+):
+    """
+    Testing Label-Looping algorithm with CUDA graphs in forced mode.
+    This test guarantees that we check that the fallback behavior is working.
+    NB: Since it is impossible to directly debug CUDA graphs, when making changes,
+    start testing and debugging the code with forced "no_graphs" mode.
+    """
+    if enable_bfloat16 and not torch.cuda.is_bf16_supported():
+        pytest.skip("bfloat16 is not supported")
+
+    if force_mode == "full_graph":
+        skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported()
+
+    batch_size = 16
+    device = torch.device("cuda")
+    nemo_model = stt_en_fastconformer_transducer_large.to(device)
+    decoding_config = copy.deepcopy(nemo_model.cfg.decoding)
+
+    with open_dict(decoding_config):
+        decoding_config["greedy"]["max_symbols"] = 5
+        decoding_config["greedy"]["loop_labels"] = True
+        decoding_config["greedy"]["use_cuda_graph_decoder"] = False
+        # test that alignments and confidence do not introduce failures
+        decoding_config["greedy"]["preserve_alignments"] = True
+        decoding_config["greedy"]["preserve_frame_confidence"] = True
+
+    nemo_model.change_decoding_strategy(decoding_config)
+    audio_filepaths = glob.glob("tests/.data/asr/test/an4/wav/*.wav")
+
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=enable_bfloat16):
+        actual_hypotheses = nemo_model.transcribe(
+            audio_filepaths, batch_size=batch_size, num_workers=CUDA_GRAPH_TRANSCRIBE_NUM_WORKERS
+        )
+    actual_transcripts = [hyp.text for hyp in actual_hypotheses]
+
+    # transcribe with use implementation with cuda graphs
+    decoding_config["greedy"]["use_cuda_graph_decoder"] = True
+    nemo_model.change_decoding_strategy(decoding_config)
+    backup_cuda_graph_mode = nemo_model.decoding.decoding.decoding_computer.cuda_graphs_mode
+    try:
+        nemo_model.decoding.decoding.decoding_computer.force_cuda_graphs_mode(mode=force_mode)
+
+        with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=enable_bfloat16):
+            fast_hypotheses = nemo_model.transcribe(
+                audio_filepaths, batch_size=batch_size, num_workers=CUDA_GRAPH_TRANSCRIBE_NUM_WORKERS
+            )
+        fast_transcripts = [hyp.text for hyp in fast_hypotheses]
+
+        total_dist = sum(
+            edit_distance(r.split(), h.split())['total'] for r, h in zip(actual_transcripts, fast_transcripts)
+        )
+        total_words = sum(len(r.split()) for r in actual_transcripts)
+        wer = total_dist / total_words if total_words > 0 else 0.0
+
+        assert wer <= 1e-3, "Cuda graph greedy decoder should match original decoder implementation."
+
+        for actual, fast in zip(actual_transcripts, fast_transcripts):
+            if actual != fast:
+                print("erroneous samples:")
+                print("Original transcript:", actual)
+                print("New transcript:", fast)
+    finally:
+        nemo_model.decoding.decoding.decoding_computer.force_cuda_graphs_mode(mode=backup_cuda_graph_mode)
+
+
+@pytest.mark.with_downloads
+@pytest.mark.skipif(
+    not (torch.cuda.is_available() and torch.cuda.is_bf16_supported()),
+    reason="Test requires CUDA device with bf16 support",
+)
+@pytest.mark.parametrize("is_tdt", [False, True])
+def test_loop_labels_cuda_graph_ddp_mixed_precision(
+    tmp_path_factory,
+    an4_train_manifest_corrected,
+    stt_en_fastconformer_transducer_large,
+    stt_en_fastconformer_tdt_large,
+    is_tdt: bool,
+):
+    """CUDA graphs with DDP and mixed precision have bugs. We need to test that validation works with these settings."""
+    batch_size = 16
+
+    # instantiate trainer with bf16 mixed precision
+    trainer_cfg = TrainerConfig(devices=[0], accelerator="cuda", strategy="ddp", max_epochs=1, precision="bf16-mixed")
+    trainer = ptl.Trainer(**DictConfig(trainer_cfg))
+
+    model = stt_en_fastconformer_tdt_large if is_tdt else stt_en_fastconformer_transducer_large
+
+    # setup validation data
+    val_ds_cfg = model.cfg.validation_ds
+    with open_dict(val_ds_cfg):
+        val_ds_cfg.manifest_filepath = [str(an4_train_manifest_corrected)]
+        val_ds_cfg.batch_size = batch_size
+        val_ds_cfg.is_tarred = False
+        val_ds_cfg.use_lhotse = False
+        if is_tdt:
+            # TDT model has config with missing mandatory values, this results in errors when setting up validation data
+            # we set all the mandatory values to dummy values
+            model.cfg.train_ds.tarred_audio_filepaths = None
+            model.cfg.train_ds.manifest_filepath = None
+            model.cfg.test_ds.manifest_filepath = None
+            model.cfg.tokenizer.dir = None
+
+    model.setup_multiple_validation_data(val_ds_cfg)
+
+    # validate using trainer
+    val_results = trainer.validate(model)
+    wer = val_results[0]["val_wer"]
+
+    # explicitly free resources, then test conditions
+    trainer._teardown()
+    # teardown from the trainer is not enough, and problem with CPU will still exist, related issue:
+    # https://github.com/Lightning-AI/pytorch-lightning/issues/18803)
+    # solution is to destroy torch process group explicitly
+    torch.distributed.destroy_process_group()
+
+    assert wer <= 0.1, f"WER is too high: {wer}"
+
+
+@pytest.mark.with_downloads
+@pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.device_count() < 2, reason="Test requires 2 GPUs")
+@pytest.mark.parametrize("loop_labels", [False, True])
+def test_change_devices(loop_labels: bool, stt_en_fastconformer_transducer_large):
+    if not loop_labels:
+        skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported()
+
+    first_device = torch.device("cuda:0")
+    second_device = torch.device("cuda:1")
+
+    batch_size = 8
+
+    nemo_model = stt_en_fastconformer_transducer_large.to(second_device)
+    decoding_config = copy.deepcopy(nemo_model.cfg.decoding)
+
+    with open_dict(decoding_config):
+        decoding_config["greedy"]["max_symbols"] = 5
+        decoding_config["greedy"]["loop_labels"] = loop_labels
+        decoding_config["greedy"]["use_cuda_graph_decoder"] = True
+
+    nemo_model.change_decoding_strategy(decoding_config)
+
+    # Test that the model can run successfully when it is first
+    # initialized on second_device and then transferred to
+    # true_device
+    nemo_model.to(first_device)
+    audio_filepaths = glob.glob("tests/.data/asr/test/an4/wav/*.wav")
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
+        second_device_hypotheses = nemo_model.transcribe(
+            audio_filepaths, batch_size=batch_size, num_workers=CUDA_GRAPH_TRANSCRIBE_NUM_WORKERS
+        )
+    second_device_transcripts = [hyp.text for hyp in second_device_hypotheses]
+
+    # Test that the model can run successfully back on second_device
+    # after having been first run on first_device. Because the
+    # decoder's data structures are lazily initialized, this activates
+    # slightly different code than the first case (where the decoder
+    # has not run at all), so we want to exercise both cases.
+    nemo_model.to(second_device)
+
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
+        first_device_hypotheses = nemo_model.transcribe(
+            audio_filepaths, batch_size=batch_size, num_workers=CUDA_GRAPH_TRANSCRIBE_NUM_WORKERS
+        )
+    first_device_transcripts = [hyp.text for hyp in first_device_hypotheses]
+    # Sanity check: The device we run on should not change execution
+    # output.
+    assert first_device_transcripts == second_device_transcripts

tests/collections/asr/decoding/test_multi_task_decoding.py

@@ -0,0 +1,342 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from nemo.collections.asr.modules.transformer.transformer import TransformerDecoderNM
+from nemo.collections.asr.modules.transformer.transformer_generators import (
+    BeamSearchSequenceGenerator,
+    BeamSearchSequenceGeneratorWithFusionModels,
+    GreedySequenceGenerator,
+)
+from nemo.collections.asr.parts.context_biasing import GPUBoostingTreeModel
+from nemo.collections.asr.parts.submodules.multitask_beam_decoding import TransformerAEDBeamInfer
+from nemo.collections.asr.parts.submodules.multitask_greedy_decoding import TransformerAEDGreedyInfer
+from nemo.collections.asr.parts.submodules.ngram_lm import NGramGPULanguageModel
+from nemo.collections.asr.parts.submodules.token_classifier import TokenClassifier
+
+
+@pytest.fixture()
+def deterministic_rng():
+    state = torch.get_rng_state()
+    torch.manual_seed(0)
+    yield
+    torch.set_rng_state(state)
+
+
+@pytest.fixture()
+def decoder_nm(deterministic_rng):
+    return TransformerDecoderNM(
+        vocab_size=8,
+        hidden_size=2,
+        num_layers=1,
+        inner_size=4,
+        num_attention_heads=1,
+        max_sequence_length=32,
+    ).eval()
+
+
+@pytest.fixture()
+def nnet(decoder_nm):
+    ans = (
+        decoder_nm.embedding,
+        decoder_nm.decoder,
+        TokenClassifier(hidden_size=2, num_classes=8),
+    )
+    ans = tuple(m.eval() for m in ans)
+    return ans
+
+
+@pytest.fixture()
+def inputs():
+    B, T, C = 1, 5, 2
+    return (
+        torch.tensor([[1]], dtype=torch.long),  # decoder_input_ids
+        torch.ones(B, T, C, dtype=torch.float),  # encoder_hidden_states
+        torch.ones(B, T, dtype=torch.float),  # encoder_input_mask
+    )
+
+
+@pytest.fixture()
+def tokenizer():
+    tok = Mock()
+    tok.pad = 0
+    tok.bos = 1
+    tok.eos = 2
+    return tok
+
+
+@pytest.mark.parametrize('with_confidence', [False, True])
+@pytest.mark.parametrize('return_xattn_scores', [False, True])
+def test_greedy_decoding(inputs, nnet, deterministic_rng, with_confidence, return_xattn_scores):
+    gen = GreedySequenceGenerator(
+        *nnet, return_xattn_scores=return_xattn_scores, preserve_step_confidence=with_confidence
+    )
+    output = gen(*inputs)
+
+    assert len(output) == 4
+    best_path, hypotheses, confidence, xattn_list = output
+
+    assert best_path is not None
+    assert torch.is_tensor(best_path)
+    assert best_path.shape == (1, 25)
+    if return_xattn_scores:
+        assert len(xattn_list) == len(nnet[1].layers)
+        assert xattn_list[0].shape == (1, 1, 24, 5)
+    else:
+        assert xattn_list is None
+
+    assert hypotheses is None
+
+    if with_confidence:
+        assert confidence is not None
+        assert torch.is_tensor(confidence)
+        assert confidence.shape == best_path.shape
+    else:
+        assert confidence is None
+
+
+@pytest.mark.parametrize('return_xattn_scores', [False, True])
+def test_temperature_sampling_decoding(inputs, nnet, return_xattn_scores):
+    gen = GreedySequenceGenerator(*nnet, return_xattn_scores=return_xattn_scores, temperature=10.0, n_samples=2)
+    output = gen(*inputs)
+
+    assert len(output) == 4
+    best_path, hypotheses, _, xatt_list = output
+
+    assert best_path is not None
+    assert torch.is_tensor(best_path)
+    assert best_path.shape[0] == 1
+
+    assert isinstance(hypotheses, list)
+    assert len(hypotheses) == 1
+    (seq0,) = hypotheses
+    assert seq0.shape[0] == 2
+    assert (seq0[0] != seq0[1]).any()
+
+    if return_xattn_scores:
+        assert len(xatt_list) == len(nnet[1].layers)
+        assert xatt_list[0].shape == (2, 1, 24, 5)
+    else:
+        assert xatt_list is None
+
+
+def test_beam_decoding_beam_scores_false(inputs, nnet):
+    gen = BeamSearchSequenceGenerator(*nnet, beam_size=2)
+    output = gen(*inputs, return_beam_scores=False)
+
+    assert len(output) == 1
+    (best_path,) = output
+
+    assert best_path is not None
+    assert torch.is_tensor(best_path)
+    assert best_path.shape == (26,)
+
+
+@pytest.mark.parametrize('return_xattn_scores', [False, True])
+def test_beam_decoding_beam_scores_true(inputs, nnet, return_xattn_scores):
+    gen = BeamSearchSequenceGenerator(*nnet, return_xattn_scores=return_xattn_scores, beam_size=2)
+    output = gen(*inputs, return_beam_scores=True)
+
+    assert len(output) == 4
+    beam_paths, scores, best_path, xatt_scores_list = output
+
+    assert beam_paths is not None
+    assert isinstance(beam_paths, list)
+    assert len(beam_paths) == 1
+    (beam_paths_seq0,) = beam_paths
+    assert torch.is_tensor(beam_paths_seq0)
+    assert beam_paths_seq0.shape == (2, 26)
+
+    assert scores is not None
+    assert isinstance(scores, list)
+    assert len(scores) == 1
+    (scores_seq0,) = scores
+    assert torch.is_tensor(scores_seq0)
+    assert scores_seq0.shape == (2,)
+
+    assert best_path is not None
+    assert torch.is_tensor(best_path)
+    assert best_path.shape == (1, 26)
+
+    if return_xattn_scores:
+        assert xatt_scores_list is not None
+        assert isinstance(xatt_scores_list, list)
+        assert torch.is_tensor(xatt_scores_list[0])
+        assert xatt_scores_list[0].shape == (1, 1, 25, 5)
+    else:
+        assert xatt_scores_list is None
+
+
+def test_beam_decoding_beam_scores_true_with_fusion_models(inputs, nnet):
+    """Test decoding with dummy unigram LM and boosting tree"""
+    # load dummy ngpu-lm
+    lm = NGramGPULanguageModel.dummy_unigram_lm(vocab_size=8)
+
+    # load dummy boosting tree
+    boosting_tree = GPUBoostingTreeModel.dummy_boosting_tree(vocab_size=8)
+
+    fusion_models = [lm, boosting_tree]
+    fusion_models_alpha = [0.2, 0.2]
+
+    gen = BeamSearchSequenceGeneratorWithFusionModels(
+        *nnet,
+        return_xattn_scores=True,
+        fusion_models=fusion_models,
+        fusion_models_alpha=fusion_models_alpha,
+        beam_size=2,
+    )
+    output = gen(*inputs, return_beam_scores=True)
+
+    assert len(output) == 4
+    beam_paths, scores, best_path, xatt_scores_list = output
+
+    assert beam_paths is not None
+    assert isinstance(beam_paths, list)
+    assert len(beam_paths) == 1
+    (beam_paths_seq0,) = beam_paths
+    assert torch.is_tensor(beam_paths_seq0)
+    assert beam_paths_seq0.shape == (2, 26)
+
+    assert scores is not None
+    assert isinstance(scores, list)
+    assert len(scores) == 1
+    (scores_seq0,) = scores
+    assert torch.is_tensor(scores_seq0)
+    assert scores_seq0.shape == (2,)
+
+    assert best_path is not None
+    assert torch.is_tensor(best_path)
+    assert best_path.shape == (1, 26)
+
+    assert xatt_scores_list is not None
+    assert isinstance(xatt_scores_list, list)
+    assert torch.is_tensor(xatt_scores_list[0])
+    assert xatt_scores_list[0].shape == (1, 1, 25, 5)
+
+
+@pytest.fixture()
+def prompted_inputs():
+    B, T, C = 1, 5, 2
+    return (
+        torch.tensor([[1, 0, 2, 3, 4]], dtype=torch.long),  # prompt
+        torch.ones(B, T, C, dtype=torch.float),  # encoder_hidden_states
+        torch.ones(B, T, dtype=torch.float),  # encoder_input_mask
+    )
+
+
+def test_transformer_aed_beam_infer_strips_prompt(prompted_inputs, decoder_nm, nnet, tokenizer):
+    decoder_input_ids, encoder_hidden_states, encoder_input_mask = prompted_inputs
+    *_, classifier = nnet
+
+    # Run the actual top-level module used by MultiTask AED model for decoding.
+    # This module is expected to trim the prompt from the beginning, and eos and pad from the end.
+    gen = TransformerAEDBeamInfer(decoder_nm, classifier, tokenizer)
+    ans = gen(
+        encoder_hidden_states=encoder_hidden_states,
+        encoder_input_mask=encoder_input_mask,
+        decoder_input_ids=decoder_input_ids,
+    )
+    best_path = ans[0][0].y_sequence
+    assert best_path is not None
+    assert torch.is_tensor(best_path)
+
+    # Now run the underlying beam search generator that doesn't trim anything.
+    *_, (untrimmed,), _ = gen.beam_search(*prompted_inputs, return_beam_scores=True)
+    assert untrimmed is not None
+    assert torch.is_tensor(untrimmed)
+
+    # Check that the expected trimming has indeed been done.
+    torch.testing.assert_close(
+        untrimmed[decoder_input_ids.shape[1] :], best_path
+    )  # stripped the prompt from the beggining
+
+
+def test_transformer_aed_greedy_infer_strips_prompt(prompted_inputs, decoder_nm, nnet, tokenizer):
+    decoder_input_ids, encoder_hidden_states, encoder_input_mask = prompted_inputs
+    decoder_input_ids = torch.tensor([[1, 0, 2, 3, 4]], dtype=torch.long)  # prompt
+    *_, classifier = nnet
+
+    # Run the actual top-level module used by MultiTask AED model for decoding.
+    # This module is expected to trim the prompt from the beginning, and eos and pad from the end.
+    gen = TransformerAEDGreedyInfer(decoder_nm, classifier, tokenizer)
+    ans = gen(
+        encoder_hidden_states=encoder_hidden_states,
+        encoder_input_mask=encoder_input_mask,
+        decoder_input_ids=decoder_input_ids,
+    )
+    best_path = ans[0][0].y_sequence
+    assert best_path is not None
+    assert torch.is_tensor(best_path)
+
+    # Now run the underlying beam search generator that doesn't trim anything.
+    (untrimmed,), _, _, _ = gen.greedy_search(*prompted_inputs)
+    assert untrimmed is not None
+    assert torch.is_tensor(untrimmed)
+
+    # Check that the expected trimming has indeed been done.
+    torch.testing.assert_close(
+        untrimmed[decoder_input_ids.shape[1] :], best_path
+    )  # stripped the prompt from the beggining
+
+
+def test_transformer_aed_beam_infer_trims_xatt_scores(prompted_inputs, decoder_nm, nnet, tokenizer):
+    decoder_input_ids, encoder_hidden_states, encoder_input_mask = prompted_inputs
+    *_, classifier = nnet
+
+    # Run the actual top-level module used by MultiTask AED model for decoding.
+    # This module is expected to trim eos and pads in xatt from the end.
+    gen = TransformerAEDBeamInfer(decoder_nm, classifier, tokenizer, return_xattn_scores=True)
+    ans = gen(
+        encoder_hidden_states=encoder_hidden_states,
+        encoder_input_mask=encoder_input_mask,
+        decoder_input_ids=decoder_input_ids,
+    )
+    hyp = ans[0][0]
+
+    assert hyp.xatt_scores is not None
+    seq_len = hyp.y_sequence.shape[0]
+    decoder_input_ids_len = decoder_input_ids.shape[1]
+    total_expected_len = seq_len + decoder_input_ids_len - 1
+
+    # Check that the expected trimming has indeed been done.
+    for layer_scores in hyp.xatt_scores:
+        assert layer_scores.shape[1] == total_expected_len
+
+
+def test_transformer_aed_greedy_infer_trims_xatt_scores(prompted_inputs, decoder_nm, nnet, tokenizer):
+    decoder_input_ids, encoder_hidden_states, encoder_input_mask = prompted_inputs
+    *_, classifier = nnet
+
+    # Run the actual top-level module used by MultiTask AED model for decoding.
+    # This module is expected to trim eos and pads in xatt from the end.
+    gen = TransformerAEDGreedyInfer(decoder_nm, classifier, tokenizer, return_xattn_scores=True)
+    ans = gen(
+        encoder_hidden_states=encoder_hidden_states,
+        encoder_input_mask=encoder_input_mask,
+        decoder_input_ids=decoder_input_ids,
+    )
+    hyp = ans[0][0]
+
+    assert hyp.xatt_scores is not None
+    seq_len = hyp.y_sequence.shape[0]
+    decoder_input_ids_len = decoder_input_ids.shape[1]
+    total_expected_len = seq_len + decoder_input_ids_len - 1
+
+    # Check that the expected trimming has indeed been done.
+    for layer_scores in hyp.xatt_scores:
+        assert layer_scores.shape[1] == total_expected_len

tests/collections/asr/decoding/test_multi_task_streaming_decoding.py

@@ -0,0 +1,192 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+from kaldialign import edit_distance
+from omegaconf import OmegaConf
+from tqdm.auto import tqdm
+
+from nemo.collections.asr.models.aed_multitask_models import lens_to_mask
+from nemo.collections.asr.parts.submodules.aed_decoding import (
+    GreedyBatchedStreamingAEDComputer,
+    return_decoder_input_ids,
+)
+from nemo.collections.asr.parts.submodules.multitask_decoding import (
+    AEDStreamingDecodingConfig,
+    MultiTaskDecodingConfig,
+)
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest
+from nemo.collections.asr.parts.utils.streaming_utils import ContextSize
+from tests.collections.asr.decoding.utils import load_audio, make_preprocessor_deterministic
+
+DEVICES = [torch.device("cpu")]
+if torch.cuda.is_available():
+    DEVICES.append(torch.device("cuda:0"))
+
+if torch.mps.is_available():
+    DEVICES.append(torch.device("mps"))
+
+
+def get_batch_encoder_outputs_from_records(records, model, device):
+    """Helper function to get encoder outputs for a batch of manifest records"""
+    local_batch_size = len(records)
+    filenames = [record["audio_filepath"] for record in records]
+    audio_filepaths = filenames[:local_batch_size]
+
+    with torch.no_grad():
+        make_preprocessor_deterministic(model)
+        model.eval()
+
+        all_inputs, all_lengths = [], []
+        for audio_file in tqdm(audio_filepaths, desc="Loading audio files"):
+            audio_tensor, _ = load_audio(audio_file)
+            all_inputs.append(audio_tensor)
+            all_lengths.append(torch.tensor(audio_tensor.shape[0], dtype=torch.int64))
+
+        input_batch = torch.nn.utils.rnn.pad_sequence(all_inputs, batch_first=True).to(
+            device=device, dtype=torch.float32
+        )
+        length_batch = torch.tensor(all_lengths, dtype=torch.int64).to(device)
+        # get encoder output using full audio signal
+        _, encoded_length, encoded_output, _ = model(input_signal=input_batch, input_signal_length=length_batch)
+
+    return encoded_output, encoded_length
+
+
+@pytest.mark.with_downloads
+@pytest.mark.parametrize(
+    "device,use_cuda_graph_decoder",
+    [(device, False) for device in DEVICES] + [(device, True) for device in DEVICES if device.type == "cuda"],
+)
+@pytest.mark.parametrize("decoding_policy", ["waitk", "alignatt"])
+@pytest.mark.parametrize("chunk_size", [3, 4])
+@pytest.mark.parametrize("batch_size", [4])
+def test_multi_task_streaming_decoding(
+    tmp_path_factory,
+    an4_val_manifest_corrected,
+    canary_180m_flash,
+    device: torch.device,
+    use_cuda_graph_decoder: bool,
+    decoding_policy: str,
+    chunk_size: int,
+    batch_size: int,
+):
+    """Test streaming decoding with multi-task model for different decoding policies"""
+    model = canary_180m_flash
+    model.eval()
+    model.to(device=device)
+
+    # setup streaming decoding config
+    streaming_decoding_cfg = AEDStreamingDecodingConfig()
+    streaming_decoding_cfg.streaming_policy = decoding_policy
+    streaming_decoding_cfg.chunk_secs = 1
+    streaming_decoding_cfg.right_context_secs = 0.0
+    streaming_decoding_cfg.batch_size = batch_size
+    streaming_decoding_cfg.prompt = OmegaConf.create({'pnc': 'no', 'task': 'asr'})
+
+    context_encoder_frames = ContextSize(
+        left=100,
+        chunk=chunk_size,
+        right=0.0,
+    )
+
+    # setup decoding strategy
+    if hasattr(model, 'change_decoding_strategy'):
+        multitask_decoding = MultiTaskDecodingConfig()
+        multitask_decoding.strategy = "greedy"
+        model.change_decoding_strategy(multitask_decoding)
+
+    manifest = read_manifest(an4_val_manifest_corrected)
+
+    all_hyps = []
+    tokens_frame_alignment = []
+    predicted_token_ids = []
+    decoding_computer = GreedyBatchedStreamingAEDComputer(
+        model,
+        frame_chunk_size=chunk_size,
+        decoding_cfg=streaming_decoding_cfg,
+    )
+
+    with torch.no_grad(), torch.inference_mode():
+        for i in range(0, len(manifest), batch_size):
+            encoder_output, encoder_output_len = get_batch_encoder_outputs_from_records(
+                manifest[i : i + batch_size], model=model, device=device
+            )
+            local_batch_size = encoder_output_len.shape[0]
+            decoder_input_ids = return_decoder_input_ids(streaming_decoding_cfg, model)
+            model_state = GreedyBatchedStreamingAEDComputer.initialize_aed_model_state(
+                asr_model=model,
+                decoder_input_ids=decoder_input_ids,
+                batch_size=local_batch_size,
+                context_encoder_frames=context_encoder_frames,
+                chunk_secs=streaming_decoding_cfg.chunk_secs,
+                right_context_secs=streaming_decoding_cfg.right_context_secs,
+            )
+
+            # decode encoder output by chunks, passing state between decoder invocations
+            for t in range(0, encoder_output.shape[1], chunk_size):
+                current_len = torch.full_like(encoder_output_len, fill_value=t + chunk_size)
+                current_len = torch.minimum(current_len, encoder_output_len)
+                model_state.is_last_chunk_batch = current_len >= encoder_output_len
+                encoder_input_mask = lens_to_mask(current_len, encoder_output[:, : t + chunk_size].shape[1]).to(
+                    encoder_output.dtype
+                )
+
+                model_state = decoding_computer(
+                    encoder_output=encoder_output[:, : t + chunk_size],
+                    encoder_output_len=current_len,
+                    encoder_input_mask=encoder_input_mask,
+                    prev_batched_state=model_state,
+                )
+            # get final results for each sample in the batch
+            for j in range(local_batch_size):
+                transcription_idx = model_state.pred_tokens_ids[
+                    j, model_state.decoder_input_ids.size(-1) : model_state.current_context_lengths[j]
+                ]
+                transcription = model.tokenizer.ids_to_text(transcription_idx.tolist()).strip()
+                all_hyps.append(transcription)
+                tokens_frame_alignment.append(model_state.tokens_frame_alignment[j])
+                predicted_token_ids.append(
+                    model_state.pred_tokens_ids[
+                        j, model_state.decoder_input_ids.size(-1) : model_state.current_context_lengths[j]
+                    ]
+                )
+
+    # compare decoding results with reference transcripts
+    ref_transcripts = [item['text'] for item in manifest]
+    assert (
+        edit_distance(ref_transcripts, all_hyps)['total'] <= len(ref_transcripts) * 0.1
+    )  # Expected WER is less than 10%
+
+    # compute latency
+    audio_encoder_fs = 80  # in ms
+    laal_list = None
+    if decoding_policy == "waitk":
+        laal_list = decoding_computer.compute_waitk_lagging(
+            manifest, predicted_token_ids, context_encoder_frames, audio_encoder_fs, BOW_PREFIX="\u2581"
+        )
+    elif decoding_policy == "alignatt":
+        laal_list = decoding_computer.compute_alignatt_lagging(
+            manifest,
+            predicted_token_ids,
+            tokens_frame_alignment,
+            context_encoder_frames,
+            audio_encoder_fs,
+            BOW_PREFIX="\u2581",
+        )
+    else:
+        raise ValueError(f"Decoding policy {decoding_policy} is not supported")
+    laal = sum(laal_list) / len(laal_list)
+    assert 300 <= laal <= 900  # Expected LAAL is between 300ms and 900ms depending on the decoding policy

tests/collections/asr/decoding/test_rnnt_alignments.py

@@ -0,0 +1,134 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from pathlib import Path
+from typing import Union
+
+import pytest
+import torch.cuda
+from examples.asr.transcribe_speech import TranscriptionConfig
+from omegaconf import OmegaConf
+
+from nemo.collections.asr.models import EncDecRNNTBPEModel
+from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
+from nemo.collections.asr.parts.utils.transcribe_utils import prepare_audio_data
+
+DEVICES = []
+
+if torch.cuda.is_available():
+    DEVICES.append('cuda')
+
+if torch.mps.is_available():
+    DEVICES.append('mps')
+
+
+@pytest.fixture(scope="module")
+def stt_en_conformer_transducer_small_model():
+    model = EncDecRNNTBPEModel.from_pretrained(model_name="stt_en_conformer_transducer_small", map_location="cpu")
+    return model
+
+
+def get_rnnt_alignments(
+    strategy: str,
+    manifest_path: Union[Path, str],
+    model: EncDecRNNTBPEModel,
+    loop_labels: bool = True,
+    use_cuda_graph_decoder=False,
+    device="cuda",
+) -> list[Hypothesis]:
+    cfg = OmegaConf.structured(TranscriptionConfig())
+    cfg.rnnt_decoding.confidence_cfg.preserve_frame_confidence = True
+    cfg.rnnt_decoding.confidence_cfg.exclude_blank = False
+    cfg.rnnt_decoding.preserve_alignments = True
+    cfg.rnnt_decoding.strategy = strategy
+    if cfg.rnnt_decoding.strategy == "greedy_batch":
+        cfg.rnnt_decoding.greedy.loop_labels = loop_labels
+        cfg.rnnt_decoding.greedy.use_cuda_graph_decoder = use_cuda_graph_decoder
+    cfg.dataset_manifest = str(manifest_path)
+    filepaths = prepare_audio_data(cfg)[0][:8]  # selecting 8 files only
+    # NB: 9th file has the same transcription but a bit different alignment for batched/non-batched decoding
+
+    model = model.to(device)
+    model.change_decoding_strategy(cfg.rnnt_decoding)
+
+    transcriptions: list[Hypothesis] = model.transcribe(
+        audio=filepaths,
+        batch_size=cfg.batch_size,
+        num_workers=cfg.num_workers,
+        return_hypotheses=True,
+        channel_selector=cfg.channel_selector,
+    )
+
+    for transcription in transcriptions:
+        for align_elem, frame_confidence in zip(transcription.alignments, transcription.frame_confidence):
+            assert len(align_elem) == len(frame_confidence)  # frame confidences have to match alignments
+            assert len(align_elem) > 0  # no empty alignments
+            for idx, pred in enumerate(align_elem):
+                if idx < len(align_elem) - 1:
+                    assert pred[1].item() != model.decoder.blank_idx  # all except last have to be non-blank
+                else:
+                    assert pred[1].item() == model.decoder.blank_idx  # last one has to be blank
+    return transcriptions
+
+
+@pytest.fixture(autouse=True)
+def cleanup_local_folder():
+    """Overriding global fixture to make sure it's not applied for this test.
+
+    Otherwise, there will be errors in the CI in github.
+    """
+    return
+
+
+# TODO: add the same tests for multi-blank RNNT decoding
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("loop_labels", [True, False])
+@pytest.mark.parametrize("use_cuda_graph_decoder", [True, False])
+@pytest.mark.with_downloads
+def test_rnnt_alignments(
+    loop_labels: bool,
+    use_cuda_graph_decoder: bool,
+    device: str,
+    an4_val_manifest_corrected,
+    stt_en_conformer_transducer_small_model,
+):
+    if use_cuda_graph_decoder and device != "cuda":
+        pytest.skip("CUDA decoder works only with CUDA")
+    if not loop_labels and use_cuda_graph_decoder:
+        pytest.skip("Frame-Looping algorithm with CUDA graphs does not yet support alignments")
+    # using greedy as baseline and comparing all other configurations to it
+    ref_transcriptions = get_rnnt_alignments(
+        "greedy",
+        manifest_path=an4_val_manifest_corrected,
+        model=stt_en_conformer_transducer_small_model,
+        device=device,
+    )
+    transcriptions = get_rnnt_alignments(
+        "greedy_batch",
+        loop_labels=loop_labels,
+        use_cuda_graph_decoder=use_cuda_graph_decoder,
+        manifest_path=an4_val_manifest_corrected,
+        model=stt_en_conformer_transducer_small_model,
+        device=device,
+    )
+    # comparing that label sequence in alignments is exactly the same
+    # we can't compare logits as well, because they are expected to be
+    # slightly different in batched and single-sample mode
+    assert len(ref_transcriptions) == len(transcriptions)
+    for ref_transcription, transcription in zip(ref_transcriptions, transcriptions):
+        for ref_align_elem, align_elem in zip(ref_transcription.alignments, transcription.alignments):
+            assert len(ref_align_elem) == len(align_elem)
+            for ref_pred, pred in zip(ref_align_elem, align_elem):
+                assert ref_pred[1].item() == pred[1].item()

tests/collections/asr/decoding/test_rnnt_decoding.py

@@ -0,0 +1,661 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import os
+from functools import cached_property, lru_cache
+from pathlib import Path
+from typing import Optional
+
+import pytest
+import torch
+from omegaconf import DictConfig
+
+from nemo.collections.asr.models import ASRModel
+from nemo.collections.asr.modules import RNNTDecoder, RNNTJoint
+from nemo.collections.asr.parts.context_biasing import BoostingTreeModelConfig, GPUBoostingTreeModel
+from nemo.collections.asr.parts.mixins import mixins
+from nemo.collections.asr.parts.submodules import rnnt_beam_decoding
+from nemo.collections.asr.parts.submodules import rnnt_greedy_decoding as greedy_decode
+from nemo.collections.asr.parts.submodules import tdt_beam_decoding
+from nemo.collections.asr.parts.submodules.ngram_lm import NGramGPULanguageModel
+from nemo.collections.asr.parts.submodules.rnnt_decoding import RNNTBPEDecoding, RNNTDecoding, RNNTDecodingConfig
+from nemo.collections.asr.parts.utils import rnnt_utils
+from nemo.core.utils import numba_utils
+from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__
+from tests.collections.asr.decoding.test_timestamps import BaseTimestampsTest
+
+NUMBA_RNNT_LOSS_AVAILABLE = numba_utils.numba_cpu_is_supported(
+    __NUMBA_MINIMUM_VERSION__
+) or numba_utils.numba_cuda_is_supported(__NUMBA_MINIMUM_VERSION__)
+
+
+def char_vocabulary():
+    return [' ', 'a', 'b', 'c', 'd', 'e', 'f', '.']
+
+
+@pytest.fixture()
+@lru_cache(maxsize=8)
+def tmp_tokenizer(test_data_dir):
+    cfg = DictConfig({'dir': os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128"), 'type': 'wpe'})
+
+    class _TmpASRBPE(mixins.ASRBPEMixin):
+        def register_artifact(self, _, vocab_path):
+            return vocab_path
+
+    asrbpe = _TmpASRBPE()
+    asrbpe._setup_tokenizer(cfg)
+    return asrbpe.tokenizer
+
+
+@lru_cache(maxsize=2)
+def get_rnnt_decoder(vocab_size, decoder_output_size=4):
+    prednet_cfg = {'pred_hidden': decoder_output_size, 'pred_rnn_layers': 1}
+    torch.manual_seed(0)
+    decoder = RNNTDecoder(prednet=prednet_cfg, vocab_size=vocab_size)
+    decoder.freeze()
+    return decoder
+
+
+@lru_cache(maxsize=2)
+def get_rnnt_joint(vocab_size, vocabulary=None, encoder_output_size=4, decoder_output_size=4, joint_output_shape=4):
+    jointnet_cfg = {
+        'encoder_hidden': encoder_output_size,
+        'pred_hidden': decoder_output_size,
+        'joint_hidden': joint_output_shape,
+        'activation': 'relu',
+    }
+    torch.manual_seed(0)
+    joint = RNNTJoint(jointnet_cfg, vocab_size, vocabulary=vocabulary)
+    joint.freeze()
+    return joint
+
+
+@lru_cache(maxsize=1)
+def get_model_encoder_output(data_dir, model_name):
+    # Import inside function to avoid issues with dependencies
+    import librosa
+
+    audio_filepath = os.path.join(data_dir, 'asr', 'test', 'an4', 'wav', 'cen3-fjlp-b.wav')
+
+    with torch.no_grad():
+        model = ASRModel.from_pretrained(model_name, map_location='cpu')  # type: ASRModel
+        model.preprocessor.featurizer.dither = 0.0
+        model.preprocessor.featurizer.pad_to = 0
+        model.eval()
+
+        audio, sr = librosa.load(path=audio_filepath, sr=16000, mono=True)
+
+        input_signal = torch.tensor(audio, dtype=torch.float32).unsqueeze(0)
+        input_signal_length = torch.tensor([len(audio)], dtype=torch.int32)
+
+        encoded, encoded_len = model(input_signal=input_signal, input_signal_length=input_signal_length)
+
+    return model, encoded, encoded_len
+
+
+def decode_text_from_greedy_hypotheses(hyps, decoding):
+    decoded_hyps = decoding.decode_hypothesis(hyps)  # type: List[str]
+
+    return decoded_hyps
+
+
+def decode_text_from_nbest_hypotheses(hyps, decoding):
+    hypotheses = []
+    all_hypotheses = []
+
+    for nbest_hyp in hyps:  # type: rnnt_utils.NBestHypotheses
+        n_hyps = nbest_hyp.n_best_hypotheses  # Extract all hypotheses for this sample
+        decoded_hyps = decoding.decode_hypothesis(n_hyps)  # type: List[str]
+
+        hypotheses.append(decoded_hyps[0])  # best hypothesis
+        all_hypotheses.append(decoded_hyps)
+
+    return hypotheses, all_hypotheses
+
+
+def check_beam_decoding(test_data_dir, beam_config):
+    beam_size = beam_config.pop("beam_size", 1)
+    model, encoded, encoded_len = get_model_encoder_output(test_data_dir, 'nvidia/parakeet-tdt_ctc-110m')
+
+    model_config = model.to_config_dict()
+    durations = list(model_config["model_defaults"]["tdt_durations"])
+
+    beam = tdt_beam_decoding.BeamTDTInfer(
+        model.decoder,
+        model.joint,
+        beam_size=beam_size,
+        return_best_hypothesis=False,
+        durations=durations,
+        **beam_config,
+    )
+
+    enc_out = encoded
+    enc_len = encoded_len
+
+    with torch.no_grad():
+        hyps: rnnt_utils.Hypothesis = beam(encoder_output=enc_out, encoded_lengths=enc_len)[0]
+        _, all_hyps = decode_text_from_nbest_hypotheses(hyps, model.decoding)
+        all_hyps = all_hyps[0]
+
+        print("Beam search algorithm :", beam_config['search_type'])
+        for idx, hyp_ in enumerate(all_hyps):
+            print("Hyp index", idx + 1, "text :", hyp_.text)
+
+            assert len(hyp_.timestamp) > 0
+            print("Timesteps", hyp_.timestamp)
+            print()
+
+
+def check_tdt_greedy_decoding(
+    test_data_dir,
+    use_cuda_graph_decoder: bool,
+    lm_path: Optional[str | Path] = None,
+    boosting_tree: Optional[BoostingTreeModelConfig] = None,
+    enable_per_stream_biasing: bool = False,
+):
+    model, encoded, encoded_len = get_model_encoder_output(test_data_dir, 'nvidia/parakeet-tdt_ctc-110m')
+
+    model_config = model.to_config_dict()
+
+    fusion_models, fusion_models_alpha = None, None
+    if lm_path or boosting_tree:
+        fusion_models = []
+        fusion_models_alpha = []
+
+    if lm_path:
+        fusion_models.append(NGramGPULanguageModel.from_file(lm_path=lm_path, vocab_size=model.decoder.blank_idx))
+        fusion_models_alpha.append(0.5)
+    if boosting_tree:
+        fusion_models.append(GPUBoostingTreeModel.from_config(boosting_tree, tokenizer=model.tokenizer))
+        fusion_models_alpha.append(0.5)
+
+    decoding_algo = greedy_decode.GreedyBatchedTDTInfer(
+        model.decoder,
+        model.joint,
+        blank_index=model.decoder.blank_idx,
+        durations=list(model_config["model_defaults"]["tdt_durations"]),
+        max_symbols_per_step=10,
+        preserve_alignments=False,
+        preserve_frame_confidence=False,
+        use_cuda_graph_decoder=use_cuda_graph_decoder,
+        fusion_models=fusion_models,
+        fusion_models_alpha=fusion_models_alpha,
+        enable_per_stream_biasing=enable_per_stream_biasing,
+    )
+
+    enc_out = encoded
+    enc_len = encoded_len
+
+    with torch.no_grad():
+        hyps: rnnt_utils.Hypothesis = decoding_algo(encoder_output=enc_out, encoded_lengths=enc_len)[0]
+        all_hyps = decode_text_from_greedy_hypotheses(hyps, model.decoding)
+
+        print("Decoding result")
+        for idx, hyp_ in enumerate(all_hyps):
+            print(f"Hyp index {idx + 1} | text : {hyp_.text}")
+            assert len(hyp_.timestamp) > 0
+            print("Timesteps", hyp_.timestamp)
+            print()
+
+
+class TestRNNTDecoding:
+    @pytest.mark.unit
+    def test_constructor(self):
+        cfg = RNNTDecodingConfig()
+        vocab = char_vocabulary()
+        decoder = get_rnnt_decoder(vocab_size=len(vocab))
+        joint = get_rnnt_joint(vocab_size=len(vocab))
+        decoding = RNNTDecoding(decoding_cfg=cfg, decoder=decoder, joint=joint, vocabulary=vocab)
+        assert decoding is not None
+
+    @pytest.mark.unit
+    def test_constructor_subword(self, tmp_tokenizer):
+        cfg = RNNTDecodingConfig()
+        vocab = tmp_tokenizer.vocab
+        decoder = get_rnnt_decoder(vocab_size=len(vocab))
+        joint = get_rnnt_joint(vocab_size=len(vocab))
+        decoding = RNNTBPEDecoding(decoding_cfg=cfg, decoder=decoder, joint=joint, tokenizer=tmp_tokenizer)
+        assert decoding is not None
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    def test_greedy_decoding_preserve_alignments(self, test_data_dir):
+        model, encoded, encoded_len = get_model_encoder_output(test_data_dir, 'stt_en_conformer_transducer_small')
+
+        beam = greedy_decode.GreedyRNNTInfer(
+            model.decoder,
+            model.joint,
+            blank_index=model.joint.num_classes_with_blank - 1,
+            max_symbols_per_step=5,
+            preserve_alignments=True,
+        )
+
+        enc_out = encoded
+        enc_len = encoded_len
+
+        with torch.no_grad():
+            hyps = beam(encoder_output=enc_out, encoded_lengths=enc_len)[0]  # type: rnnt_utils.Hypothesis
+            hyp = decode_text_from_greedy_hypotheses(hyps, model.decoding)
+            hyp = hyp[0]
+
+            assert hyp.alignments is not None
+
+            # Use the following commented print statements to check
+            # the alignment of other algorithms compared to the default
+            print("Text", hyp.text)
+            for t in range(len(hyp.alignments)):
+                t_u = []
+                for u in range(len(hyp.alignments[t])):
+                    logp, label = hyp.alignments[t][u]
+                    assert torch.is_tensor(logp)
+                    assert torch.is_tensor(label)
+
+                    t_u.append(int(label))
+
+                print(f"Tokens at timestamp {t} = {t_u}")
+            print()
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.parametrize("loop_labels", [True, False])
+    def test_batched_greedy_decoding_preserve_alignments(self, test_data_dir, loop_labels: bool):
+        """Test batched greedy decoding using non-batched decoding as a reference"""
+        model, encoded, encoded_len = get_model_encoder_output(test_data_dir, 'stt_en_conformer_transducer_small')
+
+        search_algo = greedy_decode.GreedyBatchedRNNTInfer(
+            model.decoder,
+            model.joint,
+            blank_index=model.joint.num_classes_with_blank - 1,
+            max_symbols_per_step=5,
+            preserve_alignments=True,
+            loop_labels=loop_labels,
+        )
+
+        etalon_search_algo = greedy_decode.GreedyRNNTInfer(
+            model.decoder,
+            model.joint,
+            blank_index=model.joint.num_classes_with_blank - 1,
+            max_symbols_per_step=5,
+            preserve_alignments=True,
+        )
+
+        enc_out = encoded
+        enc_len = encoded_len
+
+        with torch.no_grad():
+            hyps: list[rnnt_utils.Hypothesis] = search_algo(encoder_output=enc_out, encoded_lengths=enc_len)[0]
+            hyp = decode_text_from_greedy_hypotheses(hyps, model.decoding)[0]
+            etalon_hyps: list[rnnt_utils.Hypothesis] = etalon_search_algo(
+                encoder_output=enc_out, encoded_lengths=enc_len
+            )[0]
+            etalon_hyp = decode_text_from_greedy_hypotheses(etalon_hyps, model.decoding)[0]
+
+            assert hyp.alignments is not None
+            assert etalon_hyp.alignments is not None
+
+            assert hyp.text == etalon_hyp.text
+            assert len(hyp.alignments) == len(etalon_hyp.alignments)
+
+            for t in range(len(hyp.alignments)):
+                t_u = []
+                for u in range(len(hyp.alignments[t])):
+                    logp, label = hyp.alignments[t][u]
+                    assert torch.is_tensor(logp)
+                    assert torch.is_tensor(label)
+                    etalon_logp, etalon_label = etalon_hyp.alignments[t][u]
+                    assert label == etalon_label
+                    assert torch.allclose(logp, etalon_logp, atol=1e-4, rtol=1e-4)
+
+                    t_u.append(int(label))
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {"search_type": "greedy"},
+            {
+                "search_type": "default",
+                "beam_size": 2,
+            },
+            {
+                "search_type": "alsd",
+                "alsd_max_target_len": 0.5,
+                "beam_size": 2,
+            },
+            {
+                "search_type": "tsd",
+                "tsd_max_sym_exp_per_step": 3,
+                "beam_size": 2,
+            },
+            {"search_type": "maes", "maes_num_steps": 2, "maes_expansion_beta": 2, "beam_size": 2},
+            {"search_type": "maes", "maes_num_steps": 3, "maes_expansion_beta": 1, "beam_size": 2},
+        ],
+    )
+    def test_rnnt_beam_decoding_preserve_alignments(self, test_data_dir, beam_config):
+        beam_size = beam_config.pop("beam_size", 1)
+        model, encoded, encoded_len = get_model_encoder_output(test_data_dir, 'stt_en_conformer_transducer_small')
+        beam = rnnt_beam_decoding.BeamRNNTInfer(
+            model.decoder,
+            model.joint,
+            beam_size=beam_size,
+            return_best_hypothesis=False,
+            preserve_alignments=True,
+            **beam_config,
+        )
+
+        enc_out = encoded
+        enc_len = encoded_len
+        blank_id = torch.tensor(model.joint.num_classes_with_blank - 1, dtype=torch.int32)
+
+        with torch.no_grad():
+            hyps = beam(encoder_output=enc_out, encoded_lengths=enc_len)[0]  # type: rnnt_utils.Hypothesis
+            hyp, all_hyps = decode_text_from_nbest_hypotheses(hyps, model.decoding)
+            hyp = hyp[0]  # best hypothesis
+            all_hyps = all_hyps[0]
+
+            assert hyp.alignments is not None
+
+            if beam_config['search_type'] == 'alsd':
+                assert len(all_hyps) <= int(beam_config['alsd_max_target_len'] * float(enc_len[0]))
+
+            print("Beam search algorithm :", beam_config['search_type'])
+            # Use the following commented print statements to check
+            # the alignment of other algorithms compared to the default
+            for idx, hyp_ in enumerate(all_hyps):  # type: (int, rnnt_utils.Hypothesis)
+                print("Hyp index", idx + 1, "text :", hyp_.text)
+
+                # Alignment length (T) must match audio length (T)
+                # NOTE: increase length threshold to two to prevent intermittent failures when a word is split into subwords
+                assert abs(len(hyp_.alignments) - enc_len[0]) <= 2  # 1
+
+                for t in range(len(hyp_.alignments)):
+                    t_u = []
+                    for u in range(len(hyp_.alignments[t])):
+                        logp, label = hyp_.alignments[t][u]
+                        assert torch.is_tensor(logp)
+                        assert torch.is_tensor(label)
+
+                        t_u.append(int(label))
+
+                    # Blank token must be the last token in the current
+                    if len(t_u) > 1:
+                        assert t_u[-1] == blank_id
+
+                        # No blank token should be present in the current timestamp other than at the end
+                        for token in t_u[:-1]:
+                            assert token != blank_id
+
+                    print(f"Tokens at timestamp {t} = {t_u}")
+                print()
+
+                assert len(hyp_.timestamp) > 0
+                print("Timesteps", hyp_.timestamp)
+                print()
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "model_name, decoding_strategy",
+        [
+            ("stt_en_conformer_transducer_small", "greedy"),
+            ("stt_en_conformer_transducer_small", "greedy_batch"),
+            ("stt_en_conformer_transducer_small", "beam"),
+            # ("stt_en_conformer_transducer_small", "tsd"),
+            ("stt_en_conformer_transducer_small", "alsd"),
+            ("nvidia/parakeet-tdt_ctc-110m", "greedy"),
+            ("nvidia/parakeet-tdt_ctc-110m", "greedy_batch"),
+        ],
+    )
+    def test_subword_decoding_compute_timestamps(self, test_data_dir, decoding_strategy, model_name):
+
+        model, encoded, encoded_len = get_model_encoder_output(test_data_dir, model_name)
+
+        cfg = DictConfig(model.cfg.decoding)
+        cfg['strategy'] = decoding_strategy
+        cfg['preserve_alignments'] = True
+        cfg['compute_timestamps'] = True
+
+        decoding = RNNTBPEDecoding(
+            decoding_cfg=cfg, decoder=model.decoder, joint=model.joint, tokenizer=model.tokenizer
+        )
+
+        hyps = decoding.rnnt_decoder_predictions_tensor(encoded, encoded_len, return_hypotheses=True)
+        if isinstance(hyps[0], list):
+            BaseTimestampsTest.check_subword_timestamps(hyps[0][0], decoding)
+        else:
+            BaseTimestampsTest.check_subword_timestamps(hyps[0], decoding)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "model_name, decoding_strategy",
+        [
+            ("stt_en_conformer_transducer_small", "greedy"),
+            ("stt_en_conformer_transducer_small", "greedy_batch"),
+            ("stt_en_conformer_transducer_small", "beam"),
+            # ("stt_en_conformer_transducer_small", "tsd"),
+            ("stt_en_conformer_transducer_small", "alsd"),
+            ("nvidia/parakeet-tdt_ctc-110m", "greedy"),
+            ("nvidia/parakeet-tdt_ctc-110m", "greedy_batch"),
+        ],
+    )
+    def test_char_decoding_compute_timestamps(self, test_data_dir, decoding_strategy, model_name):
+
+        model, encoded, encoded_len = get_model_encoder_output(test_data_dir, model_name)
+
+        cfg = DictConfig(model.cfg.decoding)
+        cfg['strategy'] = decoding_strategy
+        cfg['preserve_alignments'] = True
+        cfg['compute_timestamps'] = True
+
+        vocab = [t[0] for t in model.tokenizer.vocab]
+
+        decoding = RNNTDecoding(decoding_cfg=cfg, decoder=model.decoder, joint=model.joint, vocabulary=vocab)
+
+        hyps = decoding.rnnt_decoder_predictions_tensor(encoded, encoded_len, return_hypotheses=True)
+
+        if isinstance(hyps[0], list):
+            BaseTimestampsTest.check_char_timestamps(hyps[0][0], decoding)
+        else:
+            BaseTimestampsTest.check_char_timestamps(hyps[0], decoding)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.parametrize("use_cuda_graph_decoder", [True, False])
+    @pytest.mark.parametrize("use_lm", [True, False])
+    @pytest.mark.parametrize("use_boosting_tree", [True, False])
+    @pytest.mark.parametrize("enable_per_stream_biasing", [True, False])
+    def test_tdt_greedy_decoding(
+        self,
+        test_data_dir,
+        use_cuda_graph_decoder: bool,
+        use_lm: bool,
+        use_boosting_tree: bool,
+        enable_per_stream_biasing: bool,
+    ):
+        kenlm_model_path = Path(test_data_dir) / "asr/kenlm_ngram_lm/parakeet-tdt_ctc-110m-libri-1024.kenlm.tmp.arpa"
+        boosting_tree = BoostingTreeModelConfig(key_phrases_list=["hello", "nvidia"]) if use_boosting_tree else None
+        check_tdt_greedy_decoding(
+            test_data_dir,
+            use_cuda_graph_decoder=use_cuda_graph_decoder,
+            lm_path=kenlm_model_path if use_lm else None,
+            boosting_tree=boosting_tree,
+            enable_per_stream_biasing=enable_per_stream_biasing,
+        )
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {
+                "search_type": "default",
+                "beam_size": 2,
+            },
+            {"search_type": "maes", "maes_num_steps": 2, "maes_expansion_beta": 2, "beam_size": 2},
+            {"search_type": "maes", "maes_num_steps": 2, "maes_expansion_beta": 1, "beam_size": 4},
+        ],
+    )
+    def test_tdt_beam_decoding(self, test_data_dir, beam_config):
+        check_beam_decoding(test_data_dir, beam_config)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.with_downloads
+    @pytest.mark.unit
+    @pytest.mark.parametrize(
+        "beam_config",
+        [
+            {
+                "search_type": "maes",
+                "maes_num_steps": 2,
+                "maes_expansion_beta": 1,
+                "beam_size": 4,
+                "ngram_lm_alpha": 0.3,
+            },
+        ],
+    )
+    def test_tdt_beam_decoding_with_kenlm(self, test_data_dir, beam_config):
+        # skipping if kenlm is not installed
+        pytest.importorskip("kenlm", reason="Skipping test because 'kenlm' is not installed.")
+
+        kenlm_model_path = os.path.join(
+            test_data_dir, "asr", "kenlm_ngram_lm", "parakeet-tdt_ctc-110m-libri-1024.kenlm.tmp.arpa"
+        )
+        beam_config["ngram_lm_model"] = kenlm_model_path
+        check_beam_decoding(test_data_dir, beam_config)
+
+
+class TestRNNTTimestamps(BaseTimestampsTest):
+    """RNNT-specific timestamp tests that inherit from BaseTimestampsTest"""
+
+    def _convert_offsets(self, offsets):
+        result = copy.deepcopy(offsets)
+        for offset in result:
+            offset['char'] = [offset['char']]
+        return result
+
+    @property
+    def char_offsets_chars(self):
+        return self._convert_offsets(super().char_offsets_chars)
+
+    @property
+    def char_offsets_wpe(self):
+        return self._convert_offsets(super().char_offsets_wpe)
+
+    @property
+    def char_offsets_bpe(self):
+        return self._convert_offsets(super().char_offsets_bpe)
+
+    @property
+    def encoded_char_offsets_bpe(self):
+        return self._convert_offsets(super().encoded_char_offsets_bpe)
+
+    @cached_property
+    def decoding_char(self):
+        cfg = RNNTDecodingConfig()
+        vocab = char_vocabulary()
+        decoder = get_rnnt_decoder(vocab_size=len(vocab))
+        joint = get_rnnt_joint(vocab_size=len(vocab))
+        decoding = RNNTDecoding(decoding_cfg=cfg, decoder=decoder, joint=joint, vocabulary=vocab)
+        return decoding
+
+    @cached_property
+    def decoding_subword_wpe(self):
+        cfg = RNNTDecodingConfig()
+        vocab = self.tmp_tokenizer.vocab
+        decoder = get_rnnt_decoder(vocab_size=len(vocab))
+        joint = get_rnnt_joint(vocab_size=len(vocab))
+        decoding = RNNTBPEDecoding(decoding_cfg=cfg, decoder=decoder, joint=joint, tokenizer=self.tmp_tokenizer)
+        return decoding
+
+    @cached_property
+    def decoding_subword_bpe(self):
+        vocab = self.bpe_tokenizer.vocab
+        cfg = RNNTDecodingConfig()
+        decoder = get_rnnt_decoder(vocab_size=len(vocab))
+        joint = get_rnnt_joint(vocab_size=len(vocab))
+        decoding = RNNTBPEDecoding(decoding_cfg=cfg, decoder=decoder, joint=joint, tokenizer=self.bpe_tokenizer)
+        return decoding
+
+    @pytest.mark.unit
+    def test_word_offsets_subword_wpe(self, tmp_tokenizer):
+        self.tmp_tokenizer = tmp_tokenizer
+        super().test_word_offsets_subword_wpe()
+
+    @pytest.mark.unit
+    def test_word_offsets_subword_wpe_other_delimiter(self, tmp_tokenizer):
+        self.tmp_tokenizer = tmp_tokenizer
+        super().test_word_offsets_subword_wpe_other_delimiter()
+
+
+@pytest.mark.unit
+@pytest.mark.with_downloads
+def test_transcribe_timestamps_no_decoder_reinstantiation(stt_en_fastconformer_transducer_large, test_data_dir):
+    """
+    Test that calling transcribe with timestamps=True multiple times
+    does not reinstantiate the decoder.
+
+    Regression test for the fix that avoids calling change_decoding_strategy()
+    when compute_timestamps is already set to the desired value.
+    """
+    model = stt_en_fastconformer_transducer_large
+    audio_file = os.path.join(test_data_dir, "asr/test/an4/wav/cen3-mjwl-b.wav")
+
+    # First call - may change decoding strategy
+    _ = model.transcribe(audio_file, timestamps=True)
+
+    # Get reference to decoding algorithm after first call
+    decoding_after_first_call = model.decoding.decoding
+
+    # Second call - should NOT reinstantiate decoder
+    _ = model.transcribe(audio_file, timestamps=True)
+
+    # Verify decoder is the same object (not reinstantiated)
+    assert model.decoding.decoding is decoding_after_first_call, (
+        "Decoder was reinstantiated on second transcribe call with timestamps=True. "
+        "This indicates change_decoding_strategy() was called unnecessarily."
+    )

tests/collections/asr/decoding/test_streaming_beam_decoding.py

@@ -0,0 +1,600 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Streaming beam-search (MALSD + MAES) decoding tests.
+
+Beam-search analogue of ``test_streaming_decoding.py``. Exercises the streaming
+path of the batched beam-search computers by manually feeding the encoder
+output to ``model.decoding.decoding.decoding_computer`` in chunks and threading
+``prev_batched_state`` (a ``BatchedBeamState``):
+
+  - :func:`test_malsd_streaming_batched_state` -- covers RNNT and TDT MALSD
+    (:class:`ModifiedALSDBatchedRNNTComputer`, :class:`ModifiedALSDBatchedTDTComputer`)
+    across the eager path and both captured-graph variants (``full_graph`` and
+    ``no_while_loops``).
+  - :func:`test_malsd_streaming_batched_state_with_word_boosting` -- same MALSD matrix
+    but with a ``GPUBoostingTreeModel`` fusion model plugged in (``boosting_tree.
+    key_phrases_list``); exercises cross-chunk restoration of per-beam fusion states
+    in :meth:`_init_decoding_state`.
+  - :func:`test_maes_streaming_batched_state` -- covers RNNT MAES
+    (:class:`ModifiedAESBatchedRNNTComputer`); MAES is RNNT-only and pure-PyTorch,
+    so there is no ``is_tdt`` / ``cuda_graphs_mode`` axis.
+
+Per-chunk results are merged into a single ``BatchedBeamHyps`` via
+``flatten_()`` + ``merge_(..., is_chunk_continuation=True,
+boundary_prev_ptr=...)`` -- the same accumulation pattern used by the
+cache-aware / chunked streaming inference scripts.
+
+Streamed transcripts are asserted to be identical to the non-streaming
+reference produced by ``model.transcribe`` with the same beam settings:
+beam search with ``prev_batched_state`` is chunk-invariant because all
+cross-chunk per-beam state (scores, ``last_label``, decoded lengths,
+decoder + fusion states, ``last_timestamp_lasts``, ...) is preserved across
+boundaries.
+"""
+
+import copy
+from typing import Optional
+
+import pytest
+import torch
+from omegaconf import open_dict
+from tqdm.auto import tqdm
+
+from nemo.collections.asr.models import ASRModel
+from nemo.collections.asr.parts.context_biasing.biasing_multi_model import BiasingRequestItemConfig
+from nemo.collections.asr.parts.context_biasing.boosting_graph_batched import BoostingTreeModelConfig
+from nemo.collections.asr.parts.submodules.transducer_decoding.label_looping_base import BatchedBeamState
+from nemo.collections.asr.parts.utils.batched_beam_decoding_utils import BatchedBeamHyps
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest
+from tests.collections.asr.decoding.utils import load_audio, make_preprocessor_deterministic
+
+
+def get_devices_for_testing(use_cpu_always: bool = False) -> list[torch.device]:
+    devices = [torch.device("cpu")] if use_cpu_always else []
+    if torch.cuda.is_available():
+        devices.append(torch.device("cuda:0"))
+
+    if torch.mps.is_available():
+        devices.append(torch.device("mps"))
+
+    if len(devices) == 0:
+        # no fast device for testing, add CPU
+        devices.append(torch.device("cpu"))
+    return devices
+
+
+DEVICES = get_devices_for_testing(use_cpu_always=True)
+
+
+def _make_device_param_matrix() -> list:
+    """
+    Build the ``(device, cuda_graphs_mode)`` parametrize entries with explicit, readable
+    pytest IDs (``cpu-no-graphs``, ``cuda-full-graph``, ``cuda-no-while-loops``, ...) so the
+    test matrix shows which device + graph-mode pair is exercised instead of opaque
+    ``device0`` / ``device1`` ids.
+
+    ``cuda_graphs_mode`` is ``None`` for the eager path or one of the
+    :class:`ModifiedALSDBatched{RNNT,TDT}Computer.CudaGraphsMode` string values
+    (``"full_graph"`` / ``"no_while_loops"``) for the two captured-graph variants. The test
+    uses ``force_cuda_graphs_mode`` to pin the variant explicitly so each captured path is
+    actually exercised (otherwise ``maybe_enable_cuda_graphs`` would always pick
+    ``full_graph`` whenever conditional nodes are supported).
+
+    Coverage:
+      - every device in ``DEVICES`` with ``cuda_graphs_mode=None`` (eager path).
+      - every CUDA device additionally with both ``"full_graph"`` and ``"no_while_loops"``.
+    """
+    entries: list = []
+    for device in DEVICES:
+        entries.append(pytest.param(device, None, id=f"{device.type}-no-graphs"))
+    for device in DEVICES:
+        if device.type == "cuda":
+            entries.append(pytest.param(device, "full_graph", id=f"{device.type}-full-graph"))
+            entries.append(pytest.param(device, "no_while_loops", id=f"{device.type}-no-while-loops"))
+    return entries
+
+
+DEVICE_PARAM_MATRIX = _make_device_param_matrix()
+
+
+def _make_maes_device_param_matrix() -> list:
+    """Build readable ``device`` parametrize entries for MAES tests.
+
+    MAES has no CUDA-graphs path (it's pure-PyTorch), so the matrix is just one entry per
+    available device with explicit IDs (``cpu``, ``cuda``, ...) instead of opaque
+    ``device0`` / ``device1``.
+    """
+    return [pytest.param(device, id=device.type) for device in DEVICES]
+
+
+MAES_DEVICE_PARAM_MATRIX = _make_maes_device_param_matrix()
+
+
+def get_model_encoder_output(
+    test_audio_filenames,
+    num_samples: int,
+    model: ASRModel,
+    device: torch.device = torch.device("cpu"),
+    dtype: torch.dtype = torch.float32,
+):
+    audio_filepaths = test_audio_filenames[:num_samples]
+
+    with torch.no_grad():
+        make_preprocessor_deterministic(model)
+        model.eval()
+
+        all_inputs, all_lengths = [], []
+        for audio_file in tqdm(audio_filepaths, desc="Loading audio files"):
+            audio_tensor, _ = load_audio(audio_file)
+            all_inputs.append(audio_tensor)
+            all_lengths.append(torch.tensor(audio_tensor.shape[0], dtype=torch.int64))
+
+        input_batch = torch.nn.utils.rnn.pad_sequence(all_inputs, batch_first=True).to(device=device, dtype=dtype)
+        length_batch = torch.tensor(all_lengths, dtype=torch.int64).to(device)
+
+        encoded_outputs, encoded_length = model(input_signal=input_batch, input_signal_length=length_batch)
+
+    return encoded_outputs, encoded_length
+
+
+def get_batch_encoder_outputs_from_records(records, model, device):
+    """Helper function to get encoder outputs for a batch of manifest records"""
+    filenames = [record["audio_filepath"] for record in records]
+    local_batch_size = len(filenames)
+    encoder_output, encoder_output_len = get_model_encoder_output(
+        test_audio_filenames=filenames, model=model, num_samples=local_batch_size, device=device
+    )
+    return encoder_output, encoder_output_len
+
+
+def _configure_malsd_decoding(
+    model: ASRModel,
+    cuda_graphs_mode: Optional[str],
+    beam_size: int,
+    max_symbols: int,
+    key_phrases_list: Optional[list[str]] = None,
+    boosting_tree_alpha: float = 1.0,
+    enable_per_stream_biasing: bool = False,
+) -> None:
+    """Switch ``model`` to the ``malsd_batch`` beam-search strategy used by the streaming tests.
+
+    ``cuda_graphs_mode`` is the CudaGraphsMode string (``"full_graph"`` or
+    ``"no_while_loops"``) when CUDA graphs should be used, or ``None`` for the eager path.
+    When non-None we also call ``force_cuda_graphs_mode`` to pin the variant after the
+    decoding strategy is swapped in -- ``maybe_enable_cuda_graphs`` would otherwise auto-pick
+    ``full_graph`` whenever conditional nodes are supported, leaving the ``no_while_loops``
+    branch effectively untested.
+
+    When ``key_phrases_list`` is given, a ``boosting_tree`` fusion model is plugged in
+    (``BoostingTreeModelConfig.key_phrases_list``) so the streaming path exercises
+    cross-chunk fusion-state restoration in :meth:`_init_decoding_state`.
+    """
+    decoding_cfg = copy.deepcopy(model.cfg.decoding)
+    decoding_cfg.strategy = "malsd_batch"
+    with open_dict(decoding_cfg):
+        decoding_cfg.beam.beam_size = beam_size
+        decoding_cfg.beam.max_symbols = max_symbols
+        decoding_cfg.beam.allow_cuda_graphs = cuda_graphs_mode is not None
+        decoding_cfg.beam.return_best_hypothesis = True
+        decoding_cfg.beam.score_norm = True
+        if key_phrases_list is not None:
+            decoding_cfg.beam.boosting_tree = {"key_phrases_list": list(key_phrases_list)}
+            decoding_cfg.beam.boosting_tree_alpha = boosting_tree_alpha
+        if enable_per_stream_biasing:
+            decoding_cfg.beam.enable_per_stream_biasing = True
+    model.change_decoding_strategy(decoding_cfg)
+    if cuda_graphs_mode is not None:
+        model.decoding.decoding.decoding_computer.force_cuda_graphs_mode(cuda_graphs_mode)
+
+
+def _configure_maes_decoding(
+    model: ASRModel,
+    beam_size: int,
+    maes_num_steps: int,
+    maes_expansion_beta: int,
+    maes_expansion_gamma: float,
+) -> None:
+    """Switch ``model`` to the ``maes_batch`` beam-search strategy used by the streaming tests.
+
+    MAES is RNNT-only and currently pure-PyTorch (CUDA graphs are not implemented;
+    ``allow_cuda_graphs`` is accepted only for API parity with MALSD and is ignored by
+    :class:`ModifiedAESBatchedRNNTComputer`).
+    """
+    decoding_cfg = copy.deepcopy(model.cfg.decoding)
+    decoding_cfg.strategy = "maes_batch"
+    with open_dict(decoding_cfg):
+        decoding_cfg.beam.beam_size = beam_size
+        decoding_cfg.beam.maes_num_steps = maes_num_steps
+        decoding_cfg.beam.maes_expansion_beta = maes_expansion_beta
+        decoding_cfg.beam.maes_expansion_gamma = maes_expansion_gamma
+        decoding_cfg.beam.allow_cuda_graphs = False
+        decoding_cfg.beam.return_best_hypothesis = True
+        decoding_cfg.beam.score_norm = True
+    model.change_decoding_strategy(decoding_cfg)
+
+
+def _reset_decoding_computer_state(model: ASRModel) -> None:
+    decoding_computer = model.decoding.decoding.decoding_computer
+    if hasattr(decoding_computer, "reset_cuda_graphs_state"):
+        decoding_computer.reset_cuda_graphs_state()
+
+
+def _decode_malsd_encoder_in_chunks(
+    decoding_computer,
+    encoder_output: torch.Tensor,
+    encoder_output_len: torch.Tensor,
+    chunk_size: int,
+    multi_biasing_ids: Optional[torch.Tensor] = None,
+) -> BatchedBeamHyps:
+    encoder_output = encoder_output.transpose(1, 2)
+    state: Optional[BatchedBeamState] = None
+    current_batched_hyps: BatchedBeamHyps | None = None
+    decode_kwargs = {}
+    if multi_biasing_ids is not None:
+        decode_kwargs["multi_biasing_ids"] = multi_biasing_ids
+
+    for t in range(0, encoder_output.shape[1], chunk_size):
+        rest_len = encoder_output_len - t
+        current_len = torch.full_like(encoder_output_len, fill_value=chunk_size)
+        current_len = torch.minimum(current_len, rest_len)
+        current_len = torch.maximum(current_len, torch.zeros_like(current_len))
+        chunk_batched_hyps, state = decoding_computer(
+            x=encoder_output[:, t : t + chunk_size],
+            out_len=current_len,
+            prev_batched_state=state,
+            **decode_kwargs,
+        )
+        chunk_root_ptrs = chunk_batched_hyps.flatten_()
+        if current_batched_hyps is None:
+            current_batched_hyps = chunk_batched_hyps
+        else:
+            current_batched_hyps.merge_(
+                chunk_batched_hyps,
+                is_chunk_continuation=True,
+                boundary_prev_ptr=chunk_root_ptrs,
+            )
+
+    assert current_batched_hyps is not None
+    return current_batched_hyps
+
+
+def _register_per_stream_biasing(
+    decoding_computer,
+    tokenizer,
+    boost_texts: list[str],
+    device: torch.device,
+    boosting_model_alpha: float = 10.0,
+) -> tuple[torch.Tensor, list[BiasingRequestItemConfig | None]]:
+    batch_size = len(boost_texts)
+    multi_biasing_ids = torch.full([batch_size], fill_value=-1, dtype=torch.long, device=device)
+    biasing_requests: list[BiasingRequestItemConfig | None] = []
+
+    for batch_idx, boost_text in enumerate(boost_texts):
+        if not boost_text:
+            biasing_requests.append(None)
+            continue
+        request = BiasingRequestItemConfig(
+            boosting_model_cfg=BoostingTreeModelConfig(key_phrases_list=[boost_text], unk_score=-100),
+            boosting_model_alpha=boosting_model_alpha,
+        )
+        request.add_to_multi_model(
+            tokenizer=tokenizer,
+            biasing_multi_model=decoding_computer.biasing_multi_model,
+        )
+        if request.multi_model_id is not None:
+            multi_biasing_ids[batch_idx] = request.multi_model_id
+        biasing_requests.append(request)
+
+    return multi_biasing_ids, biasing_requests
+
+
+def _unregister_per_stream_biasing(decoding_computer, biasing_requests: list[BiasingRequestItemConfig | None]) -> None:
+    for request in biasing_requests:
+        if request is not None and request.multi_model_id is not None:
+            decoding_computer.biasing_multi_model.remove_model(request.multi_model_id)
+            request.multi_model_id = None
+
+
+def _run_malsd_streaming_manifest(
+    model: ASRModel,
+    manifest_path,
+    device: torch.device,
+    chunk_size: int,
+    batch_size: int,
+    boost_texts: Optional[list[str]] = None,
+    boosting_model_alpha: float = 10.0,
+) -> list[str]:
+    manifest = read_manifest(manifest_path)
+    decoding_computer = model.decoding.decoding.decoding_computer
+    all_transcripts: list[str] = []
+
+    with torch.no_grad(), torch.inference_mode():
+        for i in range(0, len(manifest), batch_size):
+            batch_records = manifest[i : i + batch_size]
+            encoder_output, encoder_output_len = get_batch_encoder_outputs_from_records(
+                batch_records, model=model, device=device
+            )
+
+            multi_biasing_ids = None
+            biasing_requests: list[BiasingRequestItemConfig | None] = []
+            if boost_texts is not None:
+                assert decoding_computer.biasing_multi_model is not None
+                batch_boost_texts = boost_texts[i : i + batch_size]
+                multi_biasing_ids, biasing_requests = _register_per_stream_biasing(
+                    decoding_computer,
+                    model.tokenizer,
+                    batch_boost_texts,
+                    device,
+                    boosting_model_alpha=boosting_model_alpha,
+                )
+
+            batched_hyps = _decode_malsd_encoder_in_chunks(
+                decoding_computer,
+                encoder_output,
+                encoder_output_len,
+                chunk_size,
+                multi_biasing_ids=multi_biasing_ids,
+            )
+            if boost_texts is not None:
+                _unregister_per_stream_biasing(decoding_computer, biasing_requests)
+
+            all_transcripts.extend(
+                model.tokenizer.ids_to_text(hyp.y_sequence.tolist())
+                for hyp in batched_hyps.to_hyps_list(score_norm=True)
+            )
+
+    return all_transcripts
+
+
+def _run_streaming_batched_state(
+    model: ASRModel,
+    manifest_path,
+    device: torch.device,
+    chunk_size: int,
+    batch_size: int,
+) -> tuple[list[str], list[str]]:
+    """Drive the model's beam-search ``decoding_computer`` chunk-by-chunk and return
+    ``(ref_transcripts, streaming_transcripts)``.
+
+    Shared between the MALSD and MAES streaming tests: both decoders return a
+    ``(BatchedBeamHyps, None, BatchedBeamState)`` triple and accept
+    ``prev_batched_state`` for cross-chunk state threading. The per-chunk results are
+    flattened and merged into a single accumulator via ``flatten_()`` +
+    ``merge_(..., is_chunk_continuation=True, boundary_prev_ptr=...)`` -- the same
+    accumulation pattern used by the cache-aware / chunked streaming inference scripts.
+    """
+    manifest = read_manifest(manifest_path)
+    transcriptions = model.transcribe(audio=str(manifest_path.absolute()), batch_size=batch_size)
+    ref_transcripts = [hyp.text for hyp in transcriptions]
+
+    all_hyps = []
+    decoding_computer = model.decoding.decoding.decoding_computer
+    with torch.no_grad(), torch.inference_mode():
+        for i in range(0, len(manifest), batch_size):
+            encoder_output, encoder_output_len = get_batch_encoder_outputs_from_records(
+                manifest[i : i + batch_size], model=model, device=device
+            )
+            state: Optional[BatchedBeamState] = None
+            current_batched_hyps: BatchedBeamHyps | None = None
+            encoder_output = encoder_output.transpose(1, 2)  # (B, T, D)
+            for t in range(0, encoder_output.shape[1], chunk_size):
+                rest_len = encoder_output_len - t
+                current_len = torch.full_like(encoder_output_len, fill_value=chunk_size)
+                current_len = torch.minimum(current_len, rest_len)
+                current_len = torch.maximum(current_len, torch.zeros_like(current_len))
+                chunk_batched_hyps, state = decoding_computer(
+                    x=encoder_output[:, t : t + chunk_size],
+                    out_len=current_len,
+                    prev_batched_state=state,
+                )
+                # Flatten this chunk's prefix tree and thread the cross-chunk beam
+                # permutation (``root_ptrs``) into the accumulator so the final
+                # ``flatten_sort_`` walks back through the right beam history.
+                chunk_root_ptrs = chunk_batched_hyps.flatten_()
+                if current_batched_hyps is None:
+                    current_batched_hyps = chunk_batched_hyps
+                else:
+                    current_batched_hyps.merge_(
+                        chunk_batched_hyps,
+                        is_chunk_continuation=True,
+                        boundary_prev_ptr=chunk_root_ptrs,
+                    )
+            assert current_batched_hyps is not None
+            # ``to_hyps_list`` mutates the prefix tree via ``flatten_sort_``, but we're done
+            # with ``current_batched_hyps`` here so an in-place call is fine.
+            all_hyps.extend(current_batched_hyps.to_hyps_list(score_norm=True))
+
+    streaming_transcripts = [model.tokenizer.ids_to_text(hyp.y_sequence.tolist()) for hyp in all_hyps]
+    return ref_transcripts, streaming_transcripts
+
+
+@pytest.mark.with_downloads
+@pytest.mark.parametrize("device,cuda_graphs_mode", DEVICE_PARAM_MATRIX)
+@pytest.mark.parametrize("is_tdt", [False, True])
+@pytest.mark.parametrize("chunk_size", [1, 3])
+@pytest.mark.parametrize("batch_size", [4])
+@pytest.mark.parametrize("beam_size", [4])
+@pytest.mark.parametrize("max_symbols", [10])
+def test_malsd_streaming_batched_state(
+    an4_val_manifest_corrected,
+    stt_en_fastconformer_transducer_large,
+    stt_en_fastconformer_tdt_large,
+    device: torch.device,
+    cuda_graphs_mode: Optional[str],
+    is_tdt: bool,
+    chunk_size: int,
+    batch_size: int,
+    beam_size: int,
+    max_symbols: int,
+):
+    """Streaming MALSD decoding with batched beam state passed across chunks."""
+    model = stt_en_fastconformer_tdt_large if is_tdt else stt_en_fastconformer_transducer_large
+    model.eval()
+    model.to(device=device)
+
+    _configure_malsd_decoding(model, cuda_graphs_mode, beam_size=beam_size, max_symbols=max_symbols)
+
+    ref_transcripts, streaming_transcripts = _run_streaming_batched_state(
+        model=model,
+        manifest_path=an4_val_manifest_corrected,
+        device=device,
+        chunk_size=chunk_size,
+        batch_size=batch_size,
+    )
+    assert ref_transcripts == streaming_transcripts
+
+
+@pytest.mark.with_downloads
+@pytest.mark.parametrize("device", MAES_DEVICE_PARAM_MATRIX)
+@pytest.mark.parametrize("chunk_size", [1, 3])
+@pytest.mark.parametrize("batch_size", [4])
+@pytest.mark.parametrize("beam_size", [4])
+@pytest.mark.parametrize("maes_num_steps", [2])
+@pytest.mark.parametrize("maes_expansion_beta", [2])
+@pytest.mark.parametrize("maes_expansion_gamma", [2.3])
+def test_maes_streaming_batched_state(
+    an4_val_manifest_corrected,
+    stt_en_fastconformer_transducer_large,
+    device: torch.device,
+    chunk_size: int,
+    batch_size: int,
+    beam_size: int,
+    maes_num_steps: int,
+    maes_expansion_beta: int,
+    maes_expansion_gamma: float,
+):
+    """Streaming MAES decoding with batched beam state passed across chunks.
+
+    MAES is RNNT-only and pure-PyTorch (no CUDA graphs path), so the device matrix is
+    just the set of available devices and there is no ``cuda_graphs_mode`` / ``is_tdt``
+    parameter.
+    """
+    model = stt_en_fastconformer_transducer_large
+    model.eval()
+    model.to(device=device)
+
+    _configure_maes_decoding(
+        model,
+        beam_size=beam_size,
+        maes_num_steps=maes_num_steps,
+        maes_expansion_beta=maes_expansion_beta,
+        maes_expansion_gamma=maes_expansion_gamma,
+    )
+
+    ref_transcripts, streaming_transcripts = _run_streaming_batched_state(
+        model=model,
+        manifest_path=an4_val_manifest_corrected,
+        device=device,
+        chunk_size=chunk_size,
+        batch_size=batch_size,
+    )
+    assert ref_transcripts == streaming_transcripts
+
+
+# Phrases chosen from the AN4 val transcripts so the boosting tree is actually exercised
+# (an empty / unseen list collapses to the no-boosting path and would not test fusion-state
+# restoration across chunks).
+_WB_KEY_PHRASES: list[str] = ["nineteen", "forty", "fifty", "repeat", "stop", "yes"]
+
+
+@pytest.mark.with_downloads
+@pytest.mark.parametrize("device,cuda_graphs_mode", DEVICE_PARAM_MATRIX)
+@pytest.mark.parametrize("is_tdt", [False, True])
+@pytest.mark.parametrize("chunk_size", [1, 3])
+@pytest.mark.parametrize("batch_size", [4])
+@pytest.mark.parametrize("beam_size", [4])
+@pytest.mark.parametrize("max_symbols", [10])
+def test_malsd_streaming_batched_state_with_word_boosting(
+    an4_val_manifest_corrected,
+    stt_en_fastconformer_transducer_large,
+    stt_en_fastconformer_tdt_large,
+    device: torch.device,
+    cuda_graphs_mode: Optional[str],
+    is_tdt: bool,
+    chunk_size: int,
+    batch_size: int,
+    beam_size: int,
+    max_symbols: int,
+):
+    """Streaming MALSD with word-boosting (``boosting_tree``) is chunk-invariant.
+
+    Adds a ``GPUBoostingTreeModel`` fusion model on top of the standard streaming MALSD
+    test. The reference (``model.transcribe``) and the streaming path are configured
+    identically, so the boosting tree's per-beam fusion states must be restored across
+    chunks via ``_init_decoding_state`` for the two transcripts to match.
+    """
+    model = stt_en_fastconformer_tdt_large if is_tdt else stt_en_fastconformer_transducer_large
+    model.eval()
+    model.to(device=device)
+
+    _configure_malsd_decoding(
+        model,
+        cuda_graphs_mode,
+        beam_size=beam_size,
+        max_symbols=max_symbols,
+        key_phrases_list=_WB_KEY_PHRASES,
+    )
+
+    ref_transcripts, streaming_transcripts = _run_streaming_batched_state(
+        model=model,
+        manifest_path=an4_val_manifest_corrected,
+        device=device,
+        chunk_size=chunk_size,
+        batch_size=batch_size,
+    )
+    assert ref_transcripts == streaming_transcripts
+
+
+@pytest.mark.with_downloads
+@pytest.mark.parametrize("device,cuda_graphs_mode", DEVICE_PARAM_MATRIX)
+@pytest.mark.parametrize("is_tdt", [False, True])
+@pytest.mark.parametrize("chunk_size", [1])
+@pytest.mark.parametrize("batch_size", [4])
+@pytest.mark.parametrize("beam_size", [4])
+@pytest.mark.parametrize("max_symbols", [10])
+def test_malsd_streaming_boosting_with_ref_transcripts(
+    an4_val_manifest_corrected,
+    stt_en_fastconformer_transducer_large,
+    stt_en_fastconformer_tdt_large,
+    device: torch.device,
+    cuda_graphs_mode: Optional[str],
+    is_tdt: bool,
+    chunk_size: int,
+    batch_size: int,
+    beam_size: int,
+    max_symbols: int,
+):
+    """Metamorphic test analogous to ``test_label_looping_streaming_boosting_with_ref_transcripts``."""
+    model = stt_en_fastconformer_tdt_large if is_tdt else stt_en_fastconformer_transducer_large
+    model.eval()
+    model.to(device=device)
+
+    _configure_malsd_decoding(model, cuda_graphs_mode, beam_size, max_symbols)
+    ref_transcripts = [
+        hyp.text for hyp in model.transcribe(audio=str(an4_val_manifest_corrected.absolute()), batch_size=batch_size)
+    ]
+
+    _configure_malsd_decoding(model, cuda_graphs_mode, beam_size, max_symbols, enable_per_stream_biasing=True)
+    _reset_decoding_computer_state(model)
+
+    streaming_transcripts = _run_malsd_streaming_manifest(
+        model,
+        an4_val_manifest_corrected,
+        device,
+        chunk_size,
+        batch_size,
+        boost_texts=ref_transcripts,
+    )
+    assert ref_transcripts == streaming_transcripts

tests/collections/asr/decoding/test_streaming_buffer.py

@@ -0,0 +1,469 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for `StreamingBatchedAudioBuffer` and accompanying helper
+classes defined in
+`nemo.collections.asr.parts.utils.streaming_utils`.
+"""
+
+from __future__ import annotations
+
+import math
+
+import pytest
+import torch
+
+from nemo.collections.asr.parts.utils.streaming_utils import (
+    ContextSize,
+    ContextSizeBatch,
+    DynamicLengthTensor,
+    StreamingBatchedAudioBuffer,
+)
+
+# -----------------------------------------------------------------------------
+# Helper constants / fixtures
+# -----------------------------------------------------------------------------
+
+DEVICES: list[torch.device] = [torch.device("cpu")]
+if torch.cuda.is_available():
+    DEVICES.append(torch.device("cuda:0"))
+
+
+def _create_audio_batch(batch_size: int, length: int, device: torch.device, dtype: torch.dtype = torch.float32):
+    """Create a dummy audio batch of shape (batch_size, length)."""
+    # Use a simple ramp signal to ease debugging.
+    vals = torch.arange(batch_size * length, device=device, dtype=dtype)
+    return vals.view(batch_size, length)
+
+
+def _make_chunk(batch_size: int, length: int, channels: int, start: float, device: torch.device) -> torch.Tensor:
+    """Create a deterministic chunk of shape (batch_size, length, channels)."""
+    n = batch_size * length * channels
+    return (start + torch.arange(n, device=device, dtype=torch.float32)).view(batch_size, length, channels)
+
+
+def _make_ndim_chunk(
+    batch_size: int, length: int, dim_shape: list[int], start: float, device: torch.device
+) -> torch.Tensor:
+    """Create a deterministic chunk of shape (batch_size, length, *dim_shape)."""
+    n = batch_size * length * math.prod(dim_shape)
+    return (start + torch.arange(n, device=device, dtype=torch.float32)).view(batch_size, length, *dim_shape)
+
+
+# -----------------------------------------------------------------------------
+# Tests for ContextSize and ContextSizeBatch
+# -----------------------------------------------------------------------------
+
+
+class TestContextSize:
+    @pytest.mark.unit
+    def test_context_size_total_and_subsample(self):
+        ctx = ContextSize(left=4, chunk=2, right=1)
+        assert ctx.total() == 7
+
+        half_ctx = ctx.subsample(factor=2)
+        assert isinstance(half_ctx, ContextSize)
+        assert half_ctx.left == 2 and half_ctx.chunk == 1 and half_ctx.right == 0
+        assert half_ctx.total() == math.floor(7 / 2)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_context_size_batch_total_and_subsample(self, device: torch.device):
+        left = torch.tensor([4, 4], dtype=torch.long, device=device)
+        chunk = torch.tensor([2, 2], dtype=torch.long, device=device)
+        right = torch.tensor([2, 2], dtype=torch.long, device=device)
+        batch_ctx = ContextSizeBatch(left=left, chunk=chunk, right=right)
+
+        # total() should equal element-wise sum
+        expected_total = left + chunk + right
+        assert torch.equal(batch_ctx.total(), expected_total)
+
+        # After subsampling by 2 each component should be halved (floor division)
+        half_ctx = batch_ctx.subsample(2)
+        assert torch.equal(half_ctx.left, left // 2)
+        assert torch.equal(half_ctx.chunk, chunk // 2)
+        assert torch.equal(half_ctx.right, right // 2)
+
+
+# -----------------------------------------------------------------------------
+# Tests for StreamingBatchedAudioBuffer
+# -----------------------------------------------------------------------------
+
+
+class TestStreamingBatchedAudioBuffer:
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_streaming_batched_audio_buffer(self, device: torch.device):
+        batch_size = 2
+        expected_ctx = ContextSize(left=4, chunk=2, right=1)  # total = 7
+        buffer = StreamingBatchedAudioBuffer(
+            batch_size=batch_size,
+            context_samples=expected_ctx,
+            dtype=torch.float32,
+            device=device,
+        )
+
+        # ------------------------------------------------------------------
+        # First add : chunk + right (filling initial buffer)
+        # ------------------------------------------------------------------
+        first_len = expected_ctx.chunk + expected_ctx.right  # 3
+        audio_batch = _create_audio_batch(batch_size, first_len, device)
+        audio_lens = torch.full(
+            [
+                batch_size,
+            ],
+            first_len,
+            dtype=torch.long,
+            device=device,
+        )
+        buffer.add_audio_batch_(
+            audio_batch=audio_batch,
+            audio_lengths=audio_lens,
+            is_last_chunk=False,
+            is_last_chunk_batch=torch.zeros(batch_size, dtype=torch.bool, device=device),
+        )
+
+        # Validate context sizes
+        assert buffer.context_size.left == 0
+        assert buffer.context_size.chunk == expected_ctx.chunk
+        assert buffer.context_size.right == expected_ctx.right
+        assert buffer.samples.shape[1] == first_len  # No truncation yet
+
+        # ------------------------------------------------------------------
+        # Second add : only chunk length
+        # ------------------------------------------------------------------
+        chunk_len = expected_ctx.chunk  # 2
+        audio_batch = _create_audio_batch(batch_size, chunk_len, device)
+        audio_lens.fill_(chunk_len)
+        buffer.add_audio_batch_(
+            audio_batch=audio_batch,
+            audio_lengths=audio_lens,
+            is_last_chunk=False,
+            is_last_chunk_batch=torch.zeros(batch_size, dtype=torch.bool, device=device),
+        )
+
+        # After second add, left should have grown by previous chunk (2)
+        assert buffer.context_size.left == 2
+        assert buffer.context_size.chunk == expected_ctx.chunk
+        assert buffer.context_size.right == expected_ctx.right
+        assert buffer.samples.shape[1] == 5  # 2 (left) + 2 (chunk) + 1 (right)
+
+        # ------------------------------------------------------------------
+        # Third add : another chunk, buffer should now reach full capacity (7)
+        # ------------------------------------------------------------------
+        buffer.add_audio_batch_(
+            audio_batch=audio_batch,
+            audio_lengths=audio_lens,
+            is_last_chunk=False,
+            is_last_chunk_batch=torch.zeros(batch_size, dtype=torch.bool, device=device),
+        )
+
+        assert buffer.samples.shape[1] == expected_ctx.total()
+        assert buffer.context_size.total() == expected_ctx.total()
+
+        # ------------------------------------------------------------------
+        # Fourth add : buffer overflows by 2 samples; implementation should
+        # drop the excess from the left context.
+        # ------------------------------------------------------------------
+        buffer.add_audio_batch_(
+            audio_batch=audio_batch,
+            audio_lengths=audio_lens,
+            is_last_chunk=False,
+            is_last_chunk_batch=torch.zeros(batch_size, dtype=torch.bool, device=device),
+        )
+
+        # Buffer length remains constant (total context size)
+        assert buffer.samples.shape[1] == expected_ctx.total()
+        assert buffer.context_size.total() == expected_ctx.total()
+
+        # Left context should have been clipped by 2 samples (from 6 to 4)
+        assert buffer.context_size.left == expected_ctx.left  # 4
+
+        # ------------------------------------------------------------------
+        # Final add : mark last chunk with shorter length; right context
+        # should go to 0 afterwards.
+        # ------------------------------------------------------------------
+        last_len = 1
+        audio_batch = _create_audio_batch(batch_size, last_len, device)
+        audio_lens.fill_(last_len)
+        buffer.add_audio_batch_(
+            audio_batch=audio_batch,
+            audio_lengths=audio_lens,
+            is_last_chunk=True,
+            is_last_chunk_batch=torch.ones(batch_size, dtype=torch.bool, device=device),
+        )
+
+        # After last chunk, right context must be zero and total size preserved
+        assert buffer.context_size.right == 0
+        assert buffer.context_size.total() == expected_ctx.total()
+        assert buffer.samples.shape[1] == expected_ctx.total()
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_streaming_batched_audio_buffer_raises_on_too_long_chunk(self, device: torch.device):
+        """`add_audio_batch_` should raise if provided chunk is larger than chunk + right."""
+
+        expected_ctx = ContextSize(left=0, chunk=2, right=1)
+        buffer = StreamingBatchedAudioBuffer(
+            batch_size=1,
+            context_samples=expected_ctx,
+            dtype=torch.float32,
+            device=device,
+        )
+
+        # Attempt to add a chunk that is too long (4 > 3)
+        too_long_chunk_size = expected_ctx.chunk + expected_ctx.right + 1
+        audio = _create_audio_batch(1, too_long_chunk_size, device)
+        audio_lens = torch.tensor([too_long_chunk_size], dtype=torch.long, device=device)
+
+        with pytest.raises(ValueError):
+            buffer.add_audio_batch_(
+                audio_batch=audio,
+                audio_lengths=audio_lens,
+                is_last_chunk=False,
+                is_last_chunk_batch=torch.tensor([False], dtype=torch.bool, device=device),
+            )
+
+
+# -----------------------------------------------------------------------------
+# Tests for DynamicLengthTensor
+# -----------------------------------------------------------------------------
+
+
+class TestDynamicLengthTensor:
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    @pytest.mark.parametrize(
+        "dim_shape, expected_dim_shape",
+        [
+            (None, []),
+            (3, [3]),
+            ([4, 5], [4, 5]),
+        ],
+    )
+    def test_init(self, device, dim_shape, expected_dim_shape):
+        batch_size, init_length = 2, 5
+        t = DynamicLengthTensor(
+            batch_size=batch_size,
+            init_length=init_length,
+            dim_shape=dim_shape,
+            device=device,
+            dtype=torch.float32,
+        )
+
+        assert t.dim_shape == expected_dim_shape
+        assert list(t.data.shape) == [batch_size, init_length, *expected_dim_shape]
+        assert list(t.lengths.shape) == [batch_size]
+        assert t.lengths.dtype == torch.long
+        assert t.data.dtype == torch.float32
+        # Freshly created storage is zeroed and reports no content.
+        assert torch.count_nonzero(t.lengths) == 0
+        assert torch.count_nonzero(t.data) == 0
+
+    @pytest.mark.unit
+    def test_init_minimum_length(self):
+        """`init_length` is clamped to at least 1 so doubling-based growth works."""
+        t = DynamicLengthTensor(batch_size=2, init_length=0, dim_shape=1)
+        assert t._max_length == 1
+        assert t.data.shape[1] == 1
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_append_with_lengths(self, device):
+        """Per-batch `lengths` control how many frames become valid for each item."""
+        t = DynamicLengthTensor(batch_size=2, init_length=4, dim_shape=1, device=device, dtype=torch.float32)
+
+        # First chunk: batch item 0 keeps 2 frames, item 1 keeps 1 frame.
+        chunk1 = _make_chunk(batch_size=2, length=2, channels=1, start=10.0, device=device)
+        # chunk1 == [[[10], [11]], [[12], [13]]]
+        t.append_(data=chunk1, lengths=torch.tensor([2, 1], device=device))
+        assert t.lengths.tolist() == [2, 1]
+
+        # Second chunk: item 0 keeps 1 frame, item 1 keeps 2 frames. Item 1 should
+        # overwrite the previously written "garbage" frame at position 1.
+        chunk2 = _make_chunk(batch_size=2, length=2, channels=1, start=30.0, device=device)
+        # chunk2 == [[[30], [31]], [[32], [33]]]
+        t.append_(data=chunk2, lengths=torch.tensor([1, 2], device=device))
+        assert t.lengths.tolist() == [3, 3]
+
+        # Valid frames are everything up to the per-item length.
+        item0 = t.data[0, : t.lengths[0], 0].tolist()
+        item1 = t.data[1, : t.lengths[1], 0].tolist()
+        assert item0 == [10.0, 11.0, 30.0]
+        assert item1 == [12.0, 32.0, 33.0]
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_append_without_lengths(self, device):
+        """Without `lengths`, every frame in the chunk is appended for all items."""
+        t = DynamicLengthTensor(batch_size=2, init_length=2, dim_shape=1, device=device, dtype=torch.float32)
+
+        chunk = _make_chunk(batch_size=2, length=3, channels=1, start=0.0, device=device)
+        t.append_(data=chunk)
+
+        assert t.lengths.tolist() == [3, 3]
+        assert torch.equal(t.data[:, :3], chunk)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    @pytest.mark.parametrize("dim_shape", [[], [3], [2, 3], [2, 3, 4]])
+    def test_append_without_lengths_multidim(self, device, dim_shape):
+        """Append must scatter the whole feature vector for arbitrary trailing `dim_shape`."""
+        batch_size = 2
+        # init_length < appended length so this also exercises the growth path with multi-dim shapes.
+        t = DynamicLengthTensor(
+            batch_size=batch_size, init_length=2, dim_shape=dim_shape, device=device, dtype=torch.float32
+        )
+
+        chunk = _make_ndim_chunk(batch_size, length=3, dim_shape=dim_shape, start=0.0, device=device)
+        t.append_(data=chunk)
+
+        assert t.lengths.tolist() == [3, 3]
+        assert torch.equal(t.data[:, :3], chunk)
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_append_with_lengths_multidim(self, device):
+        """Per-batch `lengths` must place the full multi-dim feature vectors at the right offsets."""
+        dim_shape = [2, 3]
+        t = DynamicLengthTensor(batch_size=2, init_length=4, dim_shape=dim_shape, device=device, dtype=torch.float32)
+
+        # First chunk: item 0 keeps 2 frames, item 1 keeps 1 frame.
+        chunk1 = _make_ndim_chunk(2, length=2, dim_shape=dim_shape, start=10.0, device=device)
+        t.append_(data=chunk1, lengths=torch.tensor([2, 1], device=device))
+        assert t.lengths.tolist() == [2, 1]
+
+        # Second chunk: item 0 keeps 1 frame, item 1 keeps 2 frames. Item 1's second frame
+        # must overwrite the previously written "garbage" frame at position 1.
+        chunk2 = _make_ndim_chunk(2, length=2, dim_shape=dim_shape, start=100.0, device=device)
+        t.append_(data=chunk2, lengths=torch.tensor([1, 2], device=device))
+        assert t.lengths.tolist() == [3, 3]
+
+        # Item 0: chunk1[0, 0], chunk1[0, 1], chunk2[0, 0]; each is a full (2, 3) feature vector.
+        assert torch.equal(t.data[0, 0], chunk1[0, 0])
+        assert torch.equal(t.data[0, 1], chunk1[0, 1])
+        assert torch.equal(t.data[0, 2], chunk2[0, 0])
+        # Item 1: chunk1[1, 0], chunk2[1, 0] (overwrites garbage at pos 1), chunk2[1, 1].
+        assert torch.equal(t.data[1, 0], chunk1[1, 0])
+        assert torch.equal(t.data[1, 1], chunk2[1, 0])
+        assert torch.equal(t.data[1, 2], chunk2[1, 1])
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_growth_preserves_data(self, device):
+        """Appending more than the initial capacity reallocates and keeps content."""
+        t = DynamicLengthTensor(batch_size=1, init_length=2, dim_shape=1, device=device, dtype=torch.float32)
+        initial_capacity = t._max_length
+
+        big_len = 10
+        chunk = _make_chunk(batch_size=1, length=big_len, channels=1, start=0.0, device=device)
+        t.append_(data=chunk, lengths=torch.tensor([big_len], device=device))
+
+        assert t._max_length > initial_capacity
+        assert t._max_length >= big_len
+        assert t.lengths.tolist() == [big_len]
+        assert t.data[0, :big_len, 0].tolist() == list(range(big_len))
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_incremental_appends_double_capacity(self, device):
+        """Repeated single-frame appends grow capacity geometrically (amortized O(1))."""
+        n_appends = 9
+        t = DynamicLengthTensor(batch_size=1, init_length=1, dim_shape=1, device=device, dtype=torch.float32)
+
+        capacities = []
+        for i in range(n_appends):
+            frame = torch.full((1, 1, 1), float(i), device=device)
+            t.append_(data=frame)
+            capacities.append(t._max_length)
+
+        # Everything that was appended is retained, in order.
+        assert t.lengths.tolist() == [n_appends]
+        assert t.data[0, :n_appends, 0].tolist() == [float(i) for i in range(n_appends)]
+        # Capacity is always at least what is stored, and grew far less than linearly.
+        assert t._max_length >= n_appends
+        assert len(set(capacities)) < n_appends  # capacity reused across appends, not bumped every time
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_clear(self, device):
+        """`clear_` resets both lengths and storage to zero while keeping capacity."""
+        t = DynamicLengthTensor(batch_size=2, init_length=4, dim_shape=1, device=device, dtype=torch.float32)
+        t.append_(data=_make_chunk(2, 3, 1, start=1.0, device=device), lengths=torch.tensor([3, 3], device=device))
+
+        capacity_before = t._max_length
+        t.clear_()
+
+        assert t.lengths.tolist() == [0, 0]
+        assert torch.count_nonzero(t.data) == 0
+        assert t._max_length == capacity_before
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_merge(self, device):
+        """`merge_` concatenates another tensor's content along the length dim."""
+        a = DynamicLengthTensor(batch_size=2, init_length=2, dim_shape=1, device=device, dtype=torch.float32)
+        a.append_(data=_make_chunk(2, 2, 1, start=1.0, device=device), lengths=torch.tensor([2, 2], device=device))
+
+        b = DynamicLengthTensor(batch_size=2, init_length=2, dim_shape=1, device=device, dtype=torch.float32)
+        b.append_(data=_make_chunk(2, 2, 1, start=100.0, device=device), lengths=torch.tensor([1, 2], device=device))
+
+        a_item0_before = a.data[0, : a.lengths[0], 0].tolist()
+        a_item1_before = a.data[1, : a.lengths[1], 0].tolist()
+        b_item0 = b.data[0, : b.lengths[0], 0].tolist()
+        b_item1 = b.data[1, : b.lengths[1], 0].tolist()
+
+        ret = a.merge_(b)
+        assert ret is a  # in-place, returns self
+        assert a.lengths.tolist() == [3, 4]
+        assert a.data[0, : a.lengths[0], 0].tolist() == a_item0_before + b_item0
+        assert a.data[1, : a.lengths[1], 0].tolist() == a_item1_before + b_item1
+
+    @pytest.mark.unit
+    @pytest.mark.parametrize("device", DEVICES)
+    def test_clone_is_independent(self, device):
+        """`clone` returns a deep copy: same data/lengths, but independent storage."""
+        t = DynamicLengthTensor(batch_size=2, init_length=4, dim_shape=3, device=device, dtype=torch.float32)
+        t.append_(data=_make_chunk(2, 2, 3, start=1.0, device=device), lengths=torch.tensor([2, 1], device=device))
+
+        clone = t.clone()
+        assert clone is not t
+        assert clone.dim_shape == t.dim_shape
+        assert clone.data.shape == t.data.shape
+        assert torch.equal(clone.lengths, t.lengths)
+        assert torch.equal(clone.data, t.data)
+
+        # Mutating the clone must not affect the original.
+        clone.append_(
+            data=_make_chunk(2, 1, 3, start=50.0, device=device), lengths=torch.tensor([1, 1], device=device)
+        )
+        assert clone.lengths.tolist() == [3, 2]
+        assert t.lengths.tolist() == [2, 1]
+
+    @pytest.mark.unit
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA to verify cross-device move")
+    def test_to_device(self):
+        """`to_device` moves the underlying storage (not just the bookkeeping attr)."""
+        t = DynamicLengthTensor(
+            batch_size=2, init_length=4, dim_shape=1, device=torch.device("cpu"), dtype=torch.float32
+        )
+        t.append_(data=_make_chunk(2, 2, 1, start=1.0, device=torch.device("cpu")), lengths=torch.tensor([2, 2]))
+
+        ret = t.to_device("cuda:0")
+        assert ret is t
+        assert t.device == "cuda:0"
+        assert t.data.device.type == "cuda"
+        assert t.lengths.device.type == "cuda"
+        # Content survives the move.
+        assert t.data[0, :2, 0].tolist() == [1.0, 2.0]

tests/collections/asr/decoding/test_streaming_decoding.py

@@ -0,0 +1,608 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+from typing import Optional
+
+import pytest
+import torch
+import torch.nn.functional as F
+from omegaconf import open_dict
+from tqdm.auto import tqdm
+
+from nemo.collections.asr.models import ASRModel
+from nemo.collections.asr.parts.context_biasing.biasing_multi_model import BiasingRequestItemConfig
+from nemo.collections.asr.parts.context_biasing.boosting_graph_batched import BoostingTreeModelConfig
+from nemo.collections.asr.parts.submodules.transducer_decoding.label_looping_base import (
+    BatchedLabelLoopingState,
+    GreedyBatchedLabelLoopingComputerBase,
+)
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest
+from nemo.collections.asr.parts.utils.rnnt_utils import BatchedHyps, Hypothesis, batched_hyps_to_hypotheses
+from tests.collections.asr.decoding.utils import load_audio, make_preprocessor_deterministic
+
+
+def get_devices_for_testing(use_cpu_always: bool = False) -> list[torch.device]:
+    devices = [torch.device("cpu")] if use_cpu_always else []
+    if torch.cuda.is_available():
+        devices.append(torch.device("cuda:0"))
+
+    if torch.mps.is_available():
+        devices.append(torch.device("mps"))
+
+    if len(devices) == 0:
+        # no fast device for testing, add CPU
+        devices.append(torch.device("cpu"))
+    return devices
+
+
+DEVICES = get_devices_for_testing(use_cpu_always=False)
+
+
+def get_model_encoder_output(
+    test_audio_filenames,
+    num_samples: int,
+    model: ASRModel,
+    device: torch.device = torch.device("cpu"),
+    dtype: torch.dtype = torch.float32,
+):
+    audio_filepaths = test_audio_filenames[:num_samples]
+
+    with torch.no_grad():
+        make_preprocessor_deterministic(model)
+        model.eval()
+
+        all_inputs, all_lengths = [], []
+        for audio_file in tqdm(audio_filepaths, desc="Loading audio files"):
+            audio_tensor, _ = load_audio(audio_file)
+            all_inputs.append(audio_tensor)
+            all_lengths.append(torch.tensor(audio_tensor.shape[0], dtype=torch.int64))
+
+        input_batch = torch.nn.utils.rnn.pad_sequence(all_inputs, batch_first=True).to(device=device, dtype=dtype)
+        length_batch = torch.tensor(all_lengths, dtype=torch.int64).to(device)
+
+        encoded_outputs, encoded_length = model(input_signal=input_batch, input_signal_length=length_batch)
+
+    return encoded_outputs, encoded_length
+
+
+def get_batch_encoder_outputs_from_records(records, model, device):
+    """Helper function to get encoder outputs for a batch of manifest records"""
+    filenames = [record["audio_filepath"] for record in records]
+    local_batch_size = len(filenames)
+    encoder_output, encoder_output_len = get_model_encoder_output(
+        test_audio_filenames=filenames, model=model, num_samples=local_batch_size, device=device
+    )
+    return encoder_output, encoder_output_len
+
+
+@pytest.mark.with_downloads
+@pytest.mark.parametrize(
+    "device,use_cuda_graph_decoder",
+    [(device, False) for device in DEVICES] + [(device, True) for device in DEVICES if device.type == "cuda"],
+)
+@pytest.mark.parametrize("is_tdt", [False, True])
+@pytest.mark.parametrize("chunk_size", [1, 3])
+@pytest.mark.parametrize("batch_size", [4])
+@pytest.mark.parametrize("max_symbols", [10])
+def test_label_looping_streaming_batched_state(
+    tmp_path_factory,
+    an4_val_manifest_corrected,
+    stt_en_fastconformer_transducer_large,
+    stt_en_fastconformer_tdt_large,
+    device: torch.device,
+    use_cuda_graph_decoder: bool,
+    is_tdt: bool,
+    chunk_size: int,
+    batch_size: int,
+    max_symbols: int,
+):
+    """Test streaming decoding with batched state"""
+    model = stt_en_fastconformer_tdt_large if is_tdt else stt_en_fastconformer_transducer_large
+    model.eval()
+    model.to(device=device)
+
+    decoding_cfg = copy.deepcopy(model.cfg.decoding)
+    decoding_cfg.strategy = "greedy_batch"
+    with open_dict(decoding_cfg):
+        decoding_cfg.greedy.use_cuda_graph_decoder = use_cuda_graph_decoder
+        decoding_cfg.greedy.max_symbols = max_symbols
+    model.change_decoding_strategy(decoding_cfg)
+
+    manifest = read_manifest(an4_val_manifest_corrected)
+    transcriptions = model.transcribe(audio=str(an4_val_manifest_corrected.absolute()), batch_size=batch_size)
+    ref_transcripts = [hyp.text for hyp in transcriptions]
+
+    all_hyps = []
+    decoding_computer: GreedyBatchedLabelLoopingComputerBase = model.decoding.decoding.decoding_computer
+    with torch.no_grad(), torch.inference_mode():
+        for i in range(0, len(manifest), batch_size):
+            encoder_output, encoder_output_len = get_batch_encoder_outputs_from_records(
+                manifest[i : i + batch_size], model=model, device=device
+            )
+            local_batch_size = encoder_output_len.shape[0]
+            # decode encoder output by chunks, passing state between decoder invocations
+            state: Optional[BatchedLabelLoopingState] = None
+            batched_hyps: BatchedHyps | None = None
+            encoder_output = encoder_output.transpose(1, 2)
+            for t in range(0, encoder_output.shape[1], chunk_size):
+                rest_len = encoder_output_len - t
+                current_len = torch.full_like(encoder_output_len, fill_value=chunk_size)
+                current_len = torch.minimum(current_len, rest_len)
+                current_len = torch.maximum(current_len, torch.zeros_like(current_len))
+                batched_hyps_chunk, state = decoding_computer(
+                    x=encoder_output[:, t : t + chunk_size],
+                    out_len=current_len,
+                    prev_batched_state=state,
+                )
+                if batched_hyps is None:
+                    batched_hyps = batched_hyps_chunk
+                else:
+                    batched_hyps.merge_(batched_hyps_chunk)
+            assert batched_hyps is not None
+            all_hyps.extend(batched_hyps_to_hypotheses(batched_hyps, batch_size=local_batch_size))
+
+    streaming_transcripts = []
+    for hyp in all_hyps:
+        streaming_transcripts.append(model.tokenizer.ids_to_text(hyp.y_sequence.tolist()))
+    assert ref_transcripts == streaming_transcripts
+
+
+@pytest.mark.with_downloads
+@pytest.mark.parametrize(
+    "device,use_cuda_graph_decoder",
+    [(device, False) for device in DEVICES] + [(device, True) for device in DEVICES if device.type == "cuda"],
+)
+@pytest.mark.parametrize("is_tdt", [False, True])
+@pytest.mark.parametrize("chunk_size", [1, 3])
+@pytest.mark.parametrize("batch_size", [4])
+@pytest.mark.parametrize("max_symbols", [10])
+def test_label_looping_streaming_partial_hypotheses(
+    tmp_path_factory,
+    an4_val_manifest_corrected,
+    stt_en_fastconformer_transducer_large,
+    stt_en_fastconformer_tdt_large,
+    device: torch.device,
+    use_cuda_graph_decoder: bool,
+    is_tdt: bool,
+    chunk_size: int,
+    batch_size: int,
+    max_symbols: int,
+):
+    """Test streaming decoding with partial hypotheses"""
+    model = stt_en_fastconformer_tdt_large if is_tdt else stt_en_fastconformer_transducer_large
+    model.eval()
+    model.to(device=device)
+
+    decoding_cfg = copy.deepcopy(model.cfg.decoding)
+    decoding_cfg.strategy = "greedy_batch"
+    with open_dict(decoding_cfg):
+        decoding_cfg.greedy.use_cuda_graph_decoder = use_cuda_graph_decoder
+        decoding_cfg.greedy.max_symbols = max_symbols
+    model.change_decoding_strategy(decoding_cfg)
+
+    manifest = read_manifest(an4_val_manifest_corrected)
+    transcriptions = model.transcribe(audio=str(an4_val_manifest_corrected.absolute()), batch_size=batch_size)
+    ref_transcripts = [hyp.text for hyp in transcriptions]
+
+    all_hyps = []
+    rnnt_infer = model.decoding.decoding
+    with torch.no_grad(), torch.inference_mode():
+        for i in range(0, len(manifest), batch_size):
+            encoder_output, encoder_output_len = get_batch_encoder_outputs_from_records(
+                manifest[i : i + batch_size], model=model, device=device
+            )
+            # decode encoder output by chunks, passing state between decoder invocations
+            hyps: list[Hypothesis] | None = None
+            for t in range(0, encoder_output.shape[2], chunk_size):
+                rest_len = encoder_output_len - t
+                current_len = torch.full_like(encoder_output_len, fill_value=chunk_size)
+                current_len = torch.minimum(current_len, rest_len)
+                current_len = torch.maximum(current_len, torch.zeros_like(current_len))
+                (hyps,) = rnnt_infer(
+                    encoder_output=encoder_output[:, :, t : t + chunk_size],
+                    encoded_lengths=current_len,
+                    partial_hypotheses=hyps,
+                )
+            # free up memory by resetting decoding state
+            for hyp in hyps:
+                hyp.clean_decoding_state_()
+            all_hyps.extend(hyps)
+    streaming_transcripts = []
+    for hyp in all_hyps:
+        streaming_transcripts.append(model.tokenizer.ids_to_text(hyp.y_sequence.tolist()))
+    assert ref_transcripts == streaming_transcripts
+
+
+@pytest.mark.with_downloads
+@pytest.mark.parametrize(
+    "device,use_cuda_graph_decoder",
+    [(device, False) for device in DEVICES] + [(device, True) for device in DEVICES if device.type == "cuda"],
+)
+@pytest.mark.parametrize("is_tdt", [False, True])
+@pytest.mark.parametrize("chunk_size", [1, 3])
+@pytest.mark.parametrize("batch_size", [4])
+@pytest.mark.parametrize("max_symbols", [10])
+def test_label_looping_continuous_streaming_batched_state(
+    tmp_path_factory,
+    an4_val_manifest_corrected,
+    stt_en_fastconformer_transducer_large,
+    stt_en_fastconformer_tdt_large,
+    device: torch.device,
+    use_cuda_graph_decoder: bool,
+    is_tdt: bool,
+    chunk_size: int,
+    batch_size: int,
+    max_symbols: int,
+):
+    """Test streaming continuos decoding with partial hypotheses"""
+    model = stt_en_fastconformer_tdt_large if is_tdt else stt_en_fastconformer_transducer_large
+    model.eval()
+    model.to(device=device)
+
+    decoding_cfg = copy.deepcopy(model.cfg.decoding)
+    decoding_cfg.strategy = "greedy_batch"
+    with open_dict(decoding_cfg):
+        decoding_cfg.greedy.use_cuda_graph_decoder = use_cuda_graph_decoder
+        decoding_cfg.greedy.max_symbols = max_symbols
+    model.change_decoding_strategy(decoding_cfg)
+
+    manifest = read_manifest(an4_val_manifest_corrected)
+    transcriptions = model.transcribe(audio=str(an4_val_manifest_corrected.absolute()), batch_size=batch_size)
+    ref_transcripts = [hyp.text for hyp in transcriptions]
+
+    all_hyps = [None for _ in range(len(manifest))]
+    decoding_computer: GreedyBatchedLabelLoopingComputerBase = model.decoding.decoding.decoding_computer
+    assert batch_size < len(
+        manifest
+    ), "Batch size should be less than the number of records in the manifest for continuous streaming test."
+    with torch.no_grad(), torch.inference_mode():
+        # get first 2 batches
+        encoder_output, encoder_output_len = get_batch_encoder_outputs_from_records(
+            manifest[:batch_size], model=model, device=device
+        )
+        encoder_output_next, encoder_output_len_next = get_batch_encoder_outputs_from_records(
+            manifest[batch_size : batch_size + batch_size], model=model, device=device
+        )
+        # we always work with encoder_output, getting next utterances from encoder_output_next
+        # so we need to pad encoder_output if it is shorter than encoder_output_next
+        if encoder_output.shape[2] < encoder_output_next.shape[2]:
+            encoder_output = F.pad(encoder_output, (0, encoder_output_next.shape[2] - encoder_output.shape[2]))
+        expanded_batch_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand(-1, chunk_size)
+        next_batch_i = 0
+        next_batch_global_i = batch_size
+        next_query_utterance_i = batch_size + batch_size
+        has_next = True  # if we have anything in next batch to decode
+        hyps: list[Hypothesis | None] = [None for _ in range(batch_size)]
+        hyps_global_indices = list(range(batch_size))
+        encoder_output_t = torch.zeros_like(encoder_output_len)
+        state = None  # decoding state
+        # while there is something to decode
+        while ((rest_len := encoder_output_len - encoder_output_t) > 0).any():
+            frame_indices = encoder_output_t[:, None] + torch.arange(chunk_size, device=device)[None, :]
+            frame_indices = torch.minimum(frame_indices, encoder_output_len[:, None] - 1)
+            current_len = torch.full_like(encoder_output_len, fill_value=chunk_size)
+            current_len = torch.minimum(current_len, rest_len)
+            encoder_frames = encoder_output[expanded_batch_indices, :, frame_indices]
+            batched_hyps, state = decoding_computer(
+                x=encoder_frames,
+                out_len=current_len,
+                prev_batched_state=state,
+            )
+            hyps_continuations = batched_hyps_to_hypotheses(batched_hyps, batch_size=batch_size)
+            for i, (hyp, hyp_continuation) in enumerate(zip(hyps, hyps_continuations)):
+                if hyp is None:
+                    hyps[i] = hyp_continuation
+                else:
+                    hyp.merge_(hyp_continuation)
+            encoder_output_t += current_len
+            rest_len -= current_len
+
+            decoding_computer.reset_state_by_mask(state, rest_len == 0)
+            finished_decoding_indices = torch.nonzero(rest_len == 0, as_tuple=True)[0].cpu().tolist()
+            for idx in finished_decoding_indices:
+                hyp = hyps[idx]
+                if all_hyps[hyps_global_indices[idx]] is None:
+                    all_hyps[hyps_global_indices[idx]] = hyp
+                hyps[idx] = None  # reset to None
+                if has_next:
+                    # get next utterance to decode for finished hypothesis
+                    encoder_output[idx] = encoder_output_next[next_batch_i]
+                    encoder_output_len[idx] = encoder_output_len_next[next_batch_i]
+                    hyps_global_indices[idx] = next_batch_global_i
+                    encoder_output_t[idx] = 0
+                    next_batch_i += 1
+                    next_batch_global_i += 1
+                    # if next_batch_i is out of bounds, get next batch of encoder outputs
+                    if next_batch_i >= encoder_output_len_next.shape[0]:
+                        if next_query_utterance_i < len(manifest):
+                            encoder_output_next, encoder_output_len_next = get_batch_encoder_outputs_from_records(
+                                manifest[next_query_utterance_i : next_query_utterance_i + batch_size],
+                                model=model,
+                                device=device,
+                            )
+                            # pad if needed to allow futher assignment of encoder_output_next to encoder_output
+                            if encoder_output.shape[2] < encoder_output_next.shape[2]:
+                                encoder_output = F.pad(
+                                    encoder_output, (0, encoder_output_next.shape[2] - encoder_output.shape[2])
+                                )
+                            next_batch_i = 0
+                            next_query_utterance_i += batch_size
+                        else:
+                            has_next = False
+
+    streaming_transcripts = []
+    for hyp in all_hyps:
+        streaming_transcripts.append(model.tokenizer.ids_to_text(hyp.y_sequence.tolist()))
+    assert ref_transcripts == streaming_transcripts
+
+
+@pytest.mark.with_downloads
+@pytest.mark.parametrize(
+    "device,use_cuda_graph_decoder",
+    [(device, False) for device in DEVICES] + [(device, True) for device in DEVICES if device.type == "cuda"],
+)
+@pytest.mark.parametrize("is_tdt", [False, True])
+@pytest.mark.parametrize("chunk_size", [1, 3])
+@pytest.mark.parametrize("batch_size", [4])
+@pytest.mark.parametrize("max_symbols", [10])
+def test_label_looping_continuous_streaming_partial_hypotheses(
+    tmp_path_factory,
+    an4_val_manifest_corrected,
+    stt_en_fastconformer_transducer_large,
+    stt_en_fastconformer_tdt_large,
+    device: torch.device,
+    use_cuda_graph_decoder: bool,
+    is_tdt: bool,
+    chunk_size: int,
+    batch_size: int,
+    max_symbols: int,
+):
+    """Test streaming continuos decoding with partial hypotheses"""
+    model = stt_en_fastconformer_tdt_large if is_tdt else stt_en_fastconformer_transducer_large
+    model.to(device=device)
+
+    decoding_cfg = copy.deepcopy(model.cfg.decoding)
+    decoding_cfg.strategy = "greedy_batch"
+    with open_dict(decoding_cfg):
+        decoding_cfg.greedy.use_cuda_graph_decoder = use_cuda_graph_decoder
+        decoding_cfg.greedy.max_symbols = max_symbols
+    model.change_decoding_strategy(decoding_cfg)
+
+    manifest = read_manifest(an4_val_manifest_corrected)
+    transcriptions = model.transcribe(audio=str(an4_val_manifest_corrected.absolute()), batch_size=batch_size)
+    ref_transcripts = [hyp.text for hyp in transcriptions]
+
+    all_hyps = [None for _ in range(len(manifest))]
+    rnnt_infer = model.decoding.decoding
+    assert batch_size < len(
+        manifest
+    ), "Batch size should be less than the number of records in the manifest for continuous streaming test."
+    with torch.no_grad(), torch.inference_mode():
+        # get first 2 batches
+        encoder_output, encoder_output_len = get_batch_encoder_outputs_from_records(
+            manifest[:batch_size], model=model, device=device
+        )
+        encoder_output_next, encoder_output_len_next = get_batch_encoder_outputs_from_records(
+            manifest[batch_size : batch_size + batch_size], model=model, device=device
+        )
+        # we always work with encoder_output, getting next utterances from encoder_output_next
+        # so we need to pad encoder_output if it is shorter than encoder_output_next
+        if encoder_output.shape[2] < encoder_output_next.shape[2]:
+            encoder_output = F.pad(encoder_output, (0, encoder_output_next.shape[2] - encoder_output.shape[2]))
+        expanded_batch_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand(-1, chunk_size)
+        # NB: we assume that encoder_output_len and encoder_output_len_next
+        # have no zero elements (no empty utterances), and we do not check this condition further
+        next_batch_i = 0
+        next_batch_global_i = batch_size
+        next_query_utterance_i = batch_size + batch_size
+        has_next = True  # if we have anything in next batch to decode
+        hyps: list[Hypothesis | None] = [None for _ in range(batch_size)]
+        hyps_global_indices = list(range(batch_size))
+        encoder_output_t = torch.zeros_like(encoder_output_len)
+        # while there is something to decode
+        while ((rest_len := encoder_output_len - encoder_output_t) > 0).any():
+            frame_indices = encoder_output_t[:, None] + torch.arange(chunk_size, device=device)[None, :]
+            frame_indices = torch.minimum(frame_indices, encoder_output_len[:, None] - 1)
+            current_len = torch.full_like(encoder_output_len, fill_value=chunk_size)
+            current_len = torch.minimum(current_len, rest_len)
+            encoder_frames = encoder_output[expanded_batch_indices, :, frame_indices].transpose(1, 2)
+            (hyps,) = rnnt_infer(
+                encoder_output=encoder_frames,
+                encoded_lengths=current_len,
+                partial_hypotheses=hyps,
+            )
+            encoder_output_t += current_len
+            rest_len -= current_len
+            finished_decoding_indices = torch.nonzero(rest_len == 0, as_tuple=True)[0].cpu().tolist()
+            for idx in finished_decoding_indices:
+                hyp = hyps[idx]
+                all_hyps[hyps_global_indices[idx]] = hyp
+                # NB: we clean decoding state and set hyp to None only if we have next utterances to decode
+                # otherwise for each decoder invocation with 0 length it will recreate the hypothesis object,
+                # which is computationally expensive
+                # decoding current hyp with 0 length will not change the hypothesis
+                if has_next:
+                    hyp.clean_decoding_state_()
+                    hyps[idx] = None  # reset to None
+                    # get next utterance to decode for finished hypothesis
+                    encoder_output[idx] = encoder_output_next[next_batch_i]
+                    encoder_output_len[idx] = encoder_output_len_next[next_batch_i]
+                    hyps_global_indices[idx] = next_batch_global_i
+                    encoder_output_t[idx] = 0
+                    next_batch_i += 1
+                    next_batch_global_i += 1
+                    # if next_batch_i is out of bounds, get next batch of encoder outputs
+                    if next_batch_i >= encoder_output_len_next.shape[0]:
+                        if next_query_utterance_i < len(manifest):
+                            encoder_output_next, encoder_output_len_next = get_batch_encoder_outputs_from_records(
+                                manifest[next_query_utterance_i : next_query_utterance_i + batch_size],
+                                model=model,
+                                device=device,
+                            )
+                            # pad if needed to allow futher assignment of encoder_output_next to encoder_output
+                            if encoder_output.shape[2] < encoder_output_next.shape[2]:
+                                encoder_output = F.pad(
+                                    encoder_output, (0, encoder_output_next.shape[2] - encoder_output.shape[2])
+                                )
+                            next_batch_i = 0
+                            next_query_utterance_i += batch_size
+                        else:
+                            has_next = False
+    for hyp in hyps:
+        if hyp is not None:
+            hyp.clean_decoding_state_()
+
+    streaming_transcripts = []
+    for hyp in all_hyps:
+        streaming_transcripts.append(model.tokenizer.ids_to_text(hyp.y_sequence.tolist()))
+    assert ref_transcripts == streaming_transcripts
+
+
+@pytest.mark.with_downloads
+@pytest.mark.parametrize(
+    "device,use_cuda_graph_decoder",
+    [(device, False) for device in DEVICES] + [(device, True) for device in DEVICES if device.type == "cuda"],
+)
+@pytest.mark.parametrize("is_tdt", [False, True])
+@pytest.mark.parametrize("chunk_size", [1])  # Small chunk size to trigger more state updates
+@pytest.mark.parametrize("batch_size", [4])
+@pytest.mark.parametrize("max_symbols", [10])
+def test_label_looping_streaming_boosting_with_ref_transcripts(
+    tmp_path_factory,
+    an4_val_manifest_corrected,
+    stt_en_fastconformer_transducer_large,
+    stt_en_fastconformer_tdt_large,
+    device: torch.device,
+    use_cuda_graph_decoder: bool,
+    is_tdt: bool,
+    chunk_size: int,
+    batch_size: int,
+    max_symbols: int,
+):
+    """
+    Metamorphic test: boosting with reference transcripts should yield identical results.
+
+    This test validates that when we boost with the exact transcripts that the model
+    would produce without boosting, the results remain the same. This is a metamorphic
+    property that should hold for correct implementations.
+
+    This test specifically validates the fix for TDT streaming boosting where
+    fusion states were incorrectly updated using `active_mask` instead of `found_labels_mask`.
+    """
+    model = stt_en_fastconformer_tdt_large if is_tdt else stt_en_fastconformer_transducer_large
+    model.eval()
+    model.to(device=device)
+
+    # First, get reference transcripts without boosting
+    decoding_cfg = copy.deepcopy(model.cfg.decoding)
+    decoding_cfg.strategy = "greedy_batch"
+    with open_dict(decoding_cfg):
+        decoding_cfg.greedy.use_cuda_graph_decoder = use_cuda_graph_decoder
+        decoding_cfg.greedy.max_symbols = max_symbols
+    model.change_decoding_strategy(decoding_cfg)
+
+    manifest = read_manifest(an4_val_manifest_corrected)
+    transcriptions = model.transcribe(audio=str(an4_val_manifest_corrected.absolute()), batch_size=batch_size)
+    ref_transcripts = [hyp.text for hyp in transcriptions]
+
+    # Now set up per-stream boosting with reference transcripts
+    decoding_cfg_boosted = copy.deepcopy(model.cfg.decoding)
+    decoding_cfg_boosted.strategy = "greedy_batch"
+    with open_dict(decoding_cfg_boosted):
+        decoding_cfg_boosted.greedy.use_cuda_graph_decoder = use_cuda_graph_decoder
+        decoding_cfg_boosted.greedy.max_symbols = max_symbols
+        decoding_cfg_boosted.greedy.enable_per_stream_biasing = True
+    model.change_decoding_strategy(decoding_cfg_boosted)
+
+    all_hyps = []
+    decoding_computer: GreedyBatchedLabelLoopingComputerBase = model.decoding.decoding.decoding_computer
+
+    with torch.no_grad(), torch.inference_mode():
+        for i in range(0, len(manifest), batch_size):
+            batch_records = manifest[i : i + batch_size]
+            batch_ref_transcripts = ref_transcripts[i : i + batch_size]
+
+            encoder_output, encoder_output_len = get_batch_encoder_outputs_from_records(
+                batch_records, model=model, device=device
+            )
+            local_batch_size = encoder_output_len.shape[0]
+
+            # Create biasing requests for each sample in the batch
+            biasing_requests = []
+            multi_biasing_ids = torch.full([local_batch_size], fill_value=-1, dtype=torch.long, device=device)
+
+            for batch_idx, ref_text in enumerate(batch_ref_transcripts):
+                if ref_text:  # Only boost non-empty transcripts
+                    request = BiasingRequestItemConfig(
+                        boosting_model_cfg=BoostingTreeModelConfig(key_phrases_list=[ref_text], unk_score=-100),
+                        boosting_model_alpha=10.0,
+                    )
+                    request.add_to_multi_model(
+                        tokenizer=model.tokenizer,
+                        biasing_multi_model=decoding_computer.biasing_multi_model,
+                    )
+                    if request.multi_model_id is not None:
+                        multi_biasing_ids[batch_idx] = request.multi_model_id
+                    biasing_requests.append(request)
+                else:
+                    biasing_requests.append(None)
+
+            # Decode encoder output by chunks, passing state between decoder invocations
+            state: Optional[BatchedLabelLoopingState] = None
+            batched_hyps: BatchedHyps | None = None
+            encoder_output = encoder_output.transpose(1, 2)
+
+            for t in range(0, encoder_output.shape[1], chunk_size):
+                rest_len = encoder_output_len - t
+                current_len = torch.full_like(encoder_output_len, fill_value=chunk_size)
+                current_len = torch.minimum(current_len, rest_len)
+                current_len = torch.maximum(current_len, torch.zeros_like(current_len))
+
+                batched_hyps_chunk, state = decoding_computer(
+                    x=encoder_output[:, t : t + chunk_size],
+                    out_len=current_len,
+                    prev_batched_state=state,
+                    multi_biasing_ids=multi_biasing_ids,
+                )
+
+                if batched_hyps is None:
+                    batched_hyps = batched_hyps_chunk
+                else:
+                    batched_hyps.merge_(batched_hyps_chunk)
+
+            # Clean up biasing models
+            for request in biasing_requests:
+                if request is not None and request.multi_model_id is not None:
+                    decoding_computer.biasing_multi_model.remove_model(request.multi_model_id)
+                    request.multi_model_id = None
+
+            assert batched_hyps is not None
+            all_hyps.extend(batched_hyps_to_hypotheses(batched_hyps, batch_size=local_batch_size))
+
+    streaming_transcripts = []
+    for hyp in all_hyps:
+        streaming_transcripts.append(model.tokenizer.ids_to_text(hyp.y_sequence.tolist()))
+
+    # The key assertion: boosting with ref transcripts should yield same results
+    assert ref_transcripts == streaming_transcripts, (
+        f"Boosting with reference transcripts should yield identical results.\n"
+        f"This failure indicates a bug in fusion state handling during streaming decoding.\n"
+        f"Differences found:\n"
+        + "\n".join(
+            f"  [{i}] ref: {ref!r} != boosted: {boosted!r}"
+            for i, (ref, boosted) in enumerate(zip(ref_transcripts, streaming_transcripts))
+            if ref != boosted
+        )
+    )

tests/collections/asr/decoding/test_timestamps.py

@@ -0,0 +1,317 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os.path
+import re
+from functools import cached_property
+from typing import Any
+
+from nemo.collections.asr.models import ASRModel
+from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
+from nemo.collections.asr.parts.utils.timestamp_utils import get_segment_offsets, get_words_offsets
+
+
+class BaseTimestampsTest:
+    """
+    Base class for testing timestamps in decoders (CTC and RNNT).
+    This class defines common test methods that can be inherited by both
+    test_ctc_decoding.py and test_rnnt_decoding.py.
+    """
+
+    @cached_property
+    def bpe_tokenizer(self):
+        model_path = "/home/TestData/asr/stt_en_conformer_transducer_small.nemo"
+        if os.path.exists(model_path):
+            model = ASRModel.restore_from(model_path, map_location="cpu")
+        else:
+            model = ASRModel.from_pretrained("stt_en_conformer_transducer_small", map_location="cpu")
+        return model.tokenizer
+
+    @property
+    def char_offsets_chars(self):
+        char_offsets = [
+            {"char": "e", "start_offset": 0, "end_offset": 1},
+            {"char": " ", "start_offset": 2, "end_offset": 2},
+            {"char": "e", "start_offset": 3, "end_offset": 4},
+            {"char": " ", "start_offset": 5, "end_offset": 5},
+            {"char": ".", "start_offset": 6, "end_offset": 7},
+            {"char": " ", "start_offset": 8, "end_offset": 9},
+            {"char": "e", "start_offset": 10, "end_offset": 11},
+            {"char": " ", "start_offset": 12, "end_offset": 13},
+            {"char": "?", "start_offset": 14, "end_offset": 15},
+            {"char": " ", "start_offset": 16, "end_offset": 17},
+        ]
+        return char_offsets
+
+    @property
+    def word_offsets_chars_expected_output(self):
+        return [
+            {'word': 'e', 'start_offset': 0, 'end_offset': 1},
+            {'word': 'e.', 'start_offset': 3, 'end_offset': 7},
+            {'word': 'e?', 'start_offset': 10, 'end_offset': 15},
+        ]
+
+    @property
+    def word_offsets_chars_expected_output_other_delimiter(self):
+        return [
+            {'word': 'e e ', 'start_offset': 0, 'end_offset': 5},
+            {'word': ' e? ', 'start_offset': 8, 'end_offset': 17},
+        ]
+
+    @property
+    def segment_offsets_expected_output(self):
+        return [
+            {'segment': 'e e.', 'start_offset': 0, 'end_offset': 7},
+            {'segment': 'e?', 'start_offset': 10, 'end_offset': 15},
+        ]
+
+    @property
+    def segment_offsets_expected_output_gap(self):
+        return [
+            {'segment': 'e e. e?', 'start_offset': 0, 'end_offset': 15},
+        ]
+
+    @property
+    def char_offsets_wpe(self):
+        char_offsets = [
+            {"char": "nineteen", "start_offset": 0, "end_offset": 1},
+            {"char": "##th", "start_offset": 2, "end_offset": 2},
+            {"char": "re", "start_offset": 3, "end_offset": 4},
+            {"char": "seven", "start_offset": 5, "end_offset": 6},
+            {"char": "##ty", "start_offset": 6, "end_offset": 7},
+            {"char": "eighty", "start_offset": 8, "end_offset": 9},
+        ]
+
+        return char_offsets
+
+    @property
+    def word_offsets_wpe_expected_output(self):
+        return [
+            {'word': 'nineteenth', 'start_offset': 0, 'end_offset': 2},
+            {'word': 're', 'start_offset': 3, 'end_offset': 4},
+            {'word': 'seventy', 'start_offset': 5, 'end_offset': 7},
+            {'word': 'eighty', 'start_offset': 8, 'end_offset': 9},
+        ]
+
+    @property
+    def word_offsets_wpe_expected_output_other_delimiter(self):
+        return [
+            {'word': 'nineteenth', 'start_offset': 0, 'end_offset': 2},
+            {'word': 'seventy eighty', 'start_offset': 5, 'end_offset': 9},
+        ]
+
+    @property
+    def char_offsets_bpe(self):
+        char_offsets = [
+            {"char": "discuss", "start_offset": 0, "end_offset": 2},
+            {"char": "absolute", "start_offset": 2, "end_offset": 4},
+            {"char": "'", "start_offset": 5, "end_offset": 5},
+            {"char": "really", "start_offset": 5, "end_offset": 6},
+            {"char": "friend", "start_offset": 6, "end_offset": 7},
+            {"char": "ship", "start_offset": 8, "end_offset": 9},
+        ]
+
+        return char_offsets
+
+    @property
+    def encoded_char_offsets_bpe(self):
+        char_offsets = [
+            {"char": "▁discuss", "start_offset": 0, "end_offset": 2},
+            {"char": "▁absolute", "start_offset": 2, "end_offset": 4},
+            {"char": "'", "start_offset": 5, "end_offset": 5},
+            {"char": "▁really", "start_offset": 5, "end_offset": 6},
+            {"char": "▁friend", "start_offset": 6, "end_offset": 7},
+            {"char": "ship", "start_offset": 8, "end_offset": 9},
+        ]
+
+        return char_offsets
+
+    @property
+    def word_offsets_bpe_expected_output(self):
+        return [
+            {'word': "discuss", 'start_offset': 0, 'end_offset': 2},
+            {'word': "absolute'", 'start_offset': 2, 'end_offset': 5},
+            {'word': "really", 'start_offset': 5, 'end_offset': 6},
+            {'word': "friendship", 'start_offset': 6, 'end_offset': 9},
+        ]
+
+    @property
+    def word_offsets_bpe_expected_output_other_delimiter(self):
+        return [
+            {'word': "discuss absolute'", 'start_offset': 0, 'end_offset': 5},
+            {'word': "friendship", 'start_offset': 6, 'end_offset': 9},
+        ]
+
+    @staticmethod
+    def check_char_timestamps(hyp: Hypothesis, decoding: Any):
+        """Test character-level timestamps for both CTC and RNNT"""
+        assert hyp.timestamp is not None
+        assert isinstance(hyp.timestamp, dict)
+        assert 'timestep' in hyp.timestamp
+        assert 'char' in hyp.timestamp
+        assert 'word' in hyp.timestamp
+        assert 'segment' in hyp.timestamp
+
+        hypothesis_text = re.sub(r'\s+', ' ', hyp.text.strip())
+
+        words = hyp.text.split(decoding.word_seperator)
+        words = list(filter(lambda x: x != '', words))
+        assert len(hyp.timestamp['word']) == len(words)
+
+        words_from_timestamps = [ts['word'] for ts in hyp.timestamp['word']]
+        assert hypothesis_text == decoding.word_seperator.join(words_from_timestamps)
+
+        segments = []
+        segment = []
+
+        for word in words:
+            segment.append(word)
+            if word[-1] in decoding.segment_seperators:
+                segments.append(' '.join(segment))
+                segment = []
+
+        if segment:
+            segments.append(' '.join(segment))
+
+        assert len(hyp.timestamp['segment']) == len(segments)
+
+        segments_from_timestamps = [ts['segment'] for ts in hyp.timestamp['segment']]
+        assert hypothesis_text == decoding.word_seperator.join(segments_from_timestamps)
+
+    @staticmethod
+    def check_subword_timestamps(hyp: Hypothesis, decoding: Any):
+        """Test subword-level timestamps for both CTC and RNNT"""
+        assert hyp.timestamp is not None
+        assert isinstance(hyp.timestamp, dict)
+        assert 'timestep' in hyp.timestamp
+        assert 'char' in hyp.timestamp
+        assert 'word' in hyp.timestamp
+        assert 'segment' in hyp.timestamp
+
+        chars = list(hyp.text)
+        chars = list(filter(lambda x: x not in ['', ' ', '#'], chars))
+        all_chars = [list(decoding.tokenizer.tokens_to_text(data['char'])) for data in hyp.timestamp['char']]
+        all_chars = [char for subword in all_chars for char in subword]
+        all_chars = list(filter(lambda x: x not in ['', ' ', '#'], all_chars))
+        assert len(chars) == len(all_chars)
+
+        hypothesis_text = re.sub(r'\s+', ' ', hyp.text.strip())
+
+        words_from_timestamps = [ts['word'] for ts in hyp.timestamp['word']]
+        assert hypothesis_text == decoding.word_seperator.join(words_from_timestamps)
+
+        segments_count = sum([hyp.text.count(seperator) for seperator in decoding.segment_seperators])
+        if hyp.text[-1] not in decoding.segment_seperators:
+            segments_count += 1
+
+        if hyp.text in decoding.segment_seperators:
+            segments_count = 0
+
+        assert len(hyp.timestamp['segment']) == segments_count
+
+        segments_from_timestamps = [ts['segment'] for ts in hyp.timestamp['segment']]
+        assert hypothesis_text == decoding.word_seperator.join(segments_from_timestamps)
+
+    def test_word_offsets_chars(self):
+        word_offsets = get_words_offsets(
+            char_offsets=self.char_offsets_chars,
+            encoded_char_offsets=None,
+            word_delimiter_char=" ",
+            tokenizer_type="char",
+            supported_punctuation={'.', '!', '?'},
+            decode_tokens_to_str=self.decoding_char.decode_tokens_to_str,
+        )
+
+        assert word_offsets == self.word_offsets_chars_expected_output
+
+    def test_word_offsets_char_other_delimiter(self):
+        word_offsets = get_words_offsets(
+            char_offsets=self.char_offsets_chars,
+            encoded_char_offsets=None,
+            tokenizer_type="char",
+            word_delimiter_char=".",
+            supported_punctuation={'.', '!', '?'},
+            decode_tokens_to_str=self.decoding_char.decode_tokens_to_str,
+        )
+
+        assert word_offsets == self.word_offsets_chars_expected_output_other_delimiter
+
+    def test_word_offsets_subword_wpe(self):
+
+        word_offsets = get_words_offsets(
+            char_offsets=self.char_offsets_wpe,
+            encoded_char_offsets=None,
+            word_delimiter_char=" ",
+            tokenizer_type="wpe",
+            supported_punctuation={'.', '!', '?'},
+            decode_tokens_to_str=self.decoding_subword_wpe.decode_tokens_to_str,
+        )
+
+        assert word_offsets == self.word_offsets_wpe_expected_output
+
+    def test_word_offsets_subword_wpe_other_delimiter(self):
+
+        word_offsets = get_words_offsets(
+            char_offsets=self.char_offsets_wpe,
+            encoded_char_offsets=None,
+            word_delimiter_char="re",
+            tokenizer_type="wpe",
+            supported_punctuation={'.', '!', '?'},
+            decode_tokens_to_str=self.decoding_subword_wpe.decode_tokens_to_str,
+        )
+
+        assert word_offsets == self.word_offsets_wpe_expected_output_other_delimiter
+
+    def test_word_offsets_subword_bpe(self):
+
+        word_offsets = get_words_offsets(
+            char_offsets=self.char_offsets_bpe,
+            encoded_char_offsets=self.encoded_char_offsets_bpe,
+            word_delimiter_char=" ",
+            tokenizer_type="bpe",
+            supported_punctuation={'.', '!', '?'},
+            decode_tokens_to_str=self.decoding_subword_bpe.decode_tokens_to_str,
+        )
+
+        assert word_offsets == self.word_offsets_bpe_expected_output
+
+    def test_word_offsets_subword_bpe_other_delimiter(self):
+        word_offsets = get_words_offsets(
+            char_offsets=self.char_offsets_bpe,
+            encoded_char_offsets=self.encoded_char_offsets_bpe,
+            word_delimiter_char="really",
+            tokenizer_type="bpe",
+            supported_punctuation={'.', '!', '?'},
+            decode_tokens_to_str=self.decoding_subword_bpe.decode_tokens_to_str,
+        )
+
+        assert word_offsets == self.word_offsets_bpe_expected_output_other_delimiter
+
+    def test_segment_offsets_delimiter(self):
+        segment_offsets = get_segment_offsets(
+            word_offsets=self.word_offsets_chars_expected_output,
+            segment_delimiter_tokens=['.', '!', '?'],
+            supported_punctuation={'.', '!', '?'},
+        )
+
+        assert segment_offsets == self.segment_offsets_expected_output
+
+    def test_segment_offsets_gap(self):
+        segment_offsets = get_segment_offsets(
+            word_offsets=self.word_offsets_chars_expected_output,
+            segment_delimiter_tokens=[],
+            supported_punctuation={},
+            segment_gap_threshold=10,
+        )
+
+        assert segment_offsets == self.segment_offsets_expected_output_gap

tests/collections/asr/decoding/utils.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+from contextlib import contextmanager
+
+import librosa
+import torch
+
+from nemo.collections.asr.models import ASRModel
+
+
+@contextmanager
+def preserve_decoding_cfg_and_cpu_device(model: ASRModel):
+    """
+    Context manager to preserve the decoding strategy and device of the model.
+    This is useful for tests that modify the model's decoding strategy or device
+    to avoid side effects or costly model reloading.
+    """
+    backup_decoding_cfg = copy.deepcopy(model.cfg.decoding)
+
+    try:
+        yield
+    finally:
+        model.to(device="cpu")
+        if model.cfg.decoding != backup_decoding_cfg:
+            model.change_decoding_strategy(backup_decoding_cfg)
+
+
+def load_audio(file_path, target_sr=16000) -> tuple[torch.Tensor, int]:
+    audio, sr = librosa.load(file_path, sr=target_sr)
+    return torch.tensor(audio, dtype=torch.float32), sr
+
+
+@contextmanager
+def avoid_sync_operations(device: torch.device):
+    try:
+        if device.type == "cuda":
+            torch.cuda.set_sync_debug_mode(2)  # fail if a blocking operation
+        yield
+    finally:
+        if device.type == "cuda":
+            torch.cuda.set_sync_debug_mode(0)  # default, blocking operations are allowed
+
+
+def make_preprocessor_deterministic(model: ASRModel):
+    model.preprocessor.featurizer.dither = 0.0
+    model.preprocessor.featurizer.pad_to = 0