fix(asr): stop TDT greedy decoding skipping an extra encoder frame (#16096)

`GreedyTDTInfer._greedy_decode` advanced `time_idx` by one extra frame whenever the
per-frame symbol budget was exhausted, not only when the budget was exhausted on a
zero-duration prediction. With `max_symbols_per_step=1` the budget is exhausted on
every frame, so the decoder skipped `duration + 1` frames every time: tokens were
dropped and the surviving timestamps were late, while `token_duration` still reported
the predicted duration, making the hypothesis self-contradictory.

The inner loop can only exit with `skip == 0` because the budget ran out, so the
forced advance belongs in the existing zero-duration branch and the trailing
conditional can go. This also keeps the `preserve_alignments` /
`preserve_frame_confidence` buffers, which are extended by `skip` entries, in step
with `time_idx`.

With random weights over 468 single-vs-batched utterance comparisons,
`GreedyTDTInfer` disagreed with `GreedyBatchedTDTInfer` on 135 of them before the
change and on none after it.

Signed-off-by: Udaya Tejas <udayatejas2004@gmail.com>
Co-authored-by: Vladimir Bataev <vbataev@nvidia.com>

Author: udsy19

nemo/collections/asr/parts/submodules/rnnt_greedy_decoding.py

@@ -2728,11 +2728,12 @@ def _greedy_decode(
                 time_idx += skip
                 need_loop = skip == 0
 
-            # this rarely happens, but we manually increment the `skip` number
-            # if blank is emitted and duration=0 is predicted. This prevents possible
-            # infinite loops.
+            # The inner loop exits either because a non-zero duration was predicted (`time_idx` has already
+            # been advanced by it) or because the symbol budget ran out on a zero-duration prediction. Only
+            # the latter needs a manual advance, which also prevents possible infinite loops.
             if skip == 0:
                 skip = 1
+                time_idx += 1
 
             if self.preserve_alignments:
                 # convert Ti-th logits into a torch array
@@ -2743,9 +2744,6 @@ def _greedy_decode(
                 for i in range(skip):
                     hypothesis.frame_confidence.append([])  # blank buffer for next timestep
 
-            if symbols_added == self.max_symbols:
-                time_idx += 1
-
         # Remove trailing empty list of Alignments
         if self.preserve_alignments:
             if len(hypothesis.alignments[-1]) == 0:

tests/collections/asr/test_asr_rnnt_encdec_model.py

@@ -670,6 +670,153 @@ def test_multiblank_rnnt_greedy_decoding(self, greedy_class):
             with torch.no_grad():
                 _ = greedy(encoder_output=enc_out, encoded_lengths=enc_len)
 
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    @pytest.mark.parametrize("max_symbols_per_step", [1, 2, 5])
+    def test_tdt_greedy_decoding_advances_by_predicted_duration(self, max_symbols_per_step: int):
+        """TDT greedy decoding must advance the time index by exactly the predicted duration.
+
+        A joint that always predicts token 1 with duration 2 makes the answer computable by hand:
+        the tokens have to land on t = 0, 2, 4, ... for every value of `max_symbols_per_step`.
+        """
+        token_list = [" ", "a", "b", "c"]
+        vocab_size = len(token_list)
+        durations = [0, 1, 2, 4]
+        duration_index = 2
+        num_frames = 10
+
+        encoder_output_size = 4
+        decoder_output_size = 4
+        joint_output_shape = 4
+
+        class ConstantTDTJoint(RNNTJoint):
+            """Joint that always predicts token 1 with duration `durations[duration_index]`."""
+
+            def joint_after_projection(self, f: torch.Tensor, g: torch.Tensor) -> torch.Tensor:
+                batch_size = f.shape[0]
+                time_steps = f.shape[1] if f.dim() == 3 else 1
+                target_steps = g.shape[1] if g.dim() == 3 else 1
+                logits = torch.full(
+                    (batch_size, time_steps, target_steps, vocab_size + 1 + len(durations)), -10.0, device=f.device
+                )
+                logits[..., 1] = 5.0
+                logits[..., vocab_size + 1 + duration_index] = 5.0
+                return logits
+
+        prednet_cfg = {'pred_hidden': decoder_output_size, 'pred_rnn_layers': 1}
+        jointnet_cfg = {
+            'encoder_hidden': encoder_output_size,
+            'pred_hidden': decoder_output_size,
+            'joint_hidden': joint_output_shape,
+            'activation': 'relu',
+        }
+
+        decoder = RNNTDecoder(prednet_cfg, vocab_size)
+        joint_net = ConstantTDTJoint(jointnet_cfg, vocab_size, vocabulary=token_list, num_extra_outputs=len(durations))
+
+        duration = durations[duration_index]
+        expected_timestamp = list(range(0, num_frames, duration))
+
+        # (B, D, T)
+        enc_out = torch.zeros(1, encoder_output_size, num_frames)
+        enc_len = torch.tensor([num_frames], dtype=torch.int32)
+
+        for greedy_class, additional_decoding_kwargs in [
+            (greedy_decode.GreedyTDTInfer, {}),
+            (greedy_decode.GreedyBatchedTDTInfer, {"use_cuda_graph_decoder": False}),
+        ]:
+            greedy = greedy_class(
+                decoder,
+                joint_net,
+                blank_index=vocab_size,
+                durations=durations,
+                max_symbols_per_step=max_symbols_per_step,
+                include_duration=True,
+                **additional_decoding_kwargs,
+            )
+
+            with torch.no_grad():
+                hyp = greedy(encoder_output=enc_out, encoded_lengths=enc_len)[0][0]
+
+            assert [int(t) for t in hyp.timestamp] == expected_timestamp, greedy_class.__name__
+            assert [int(d) for d in hyp.token_duration] == [duration] * len(expected_timestamp)
+
+    @pytest.mark.skipif(
+        not NUMBA_RNNT_LOSS_AVAILABLE,
+        reason='RNNTLoss has not been compiled with appropriate numba version.',
+    )
+    @pytest.mark.unit
+    def test_tdt_greedy_decoding_exhausted_symbol_budget(self):
+        """Exhausting the symbol budget on a non-zero duration must not advance the time index further.
+
+        The joint alternates duration 0 and duration 2 on successive calls, so with
+        `max_symbols_per_step=2` every time step emits one token with duration 0 and one with duration 2,
+        exhausting the budget exactly when the time index has already advanced by 2.
+        """
+        token_list = [" ", "a", "b", "c"]
+        vocab_size = len(token_list)
+        durations = [0, 1, 2, 4]
+        num_frames = 10
+
+        encoder_output_size = 4
+        decoder_output_size = 4
+        joint_output_shape = 4
+
+        class AlternatingTDTJoint(RNNTJoint):
+            """Joint that always predicts token 1, with duration 0 and 2 on alternating calls."""
+
+            def __init__(self, *args, **kwargs):
+                super().__init__(*args, **kwargs)
+                self.num_calls = 0
+
+            def joint_after_projection(self, f: torch.Tensor, g: torch.Tensor) -> torch.Tensor:
+                batch_size = f.shape[0]
+                time_steps = f.shape[1] if f.dim() == 3 else 1
+                target_steps = g.shape[1] if g.dim() == 3 else 1
+                logits = torch.full(
+                    (batch_size, time_steps, target_steps, vocab_size + 1 + len(durations)), -10.0, device=f.device
+                )
+                logits[..., 1] = 5.0
+                logits[..., vocab_size + 1 + (0 if self.num_calls % 2 == 0 else 2)] = 5.0
+                self.num_calls += 1
+                return logits
+
+        prednet_cfg = {'pred_hidden': decoder_output_size, 'pred_rnn_layers': 1}
+        jointnet_cfg = {
+            'encoder_hidden': encoder_output_size,
+            'pred_hidden': decoder_output_size,
+            'joint_hidden': joint_output_shape,
+            'activation': 'relu',
+        }
+
+        decoder = RNNTDecoder(prednet_cfg, vocab_size)
+        joint_net = AlternatingTDTJoint(
+            jointnet_cfg, vocab_size, vocabulary=token_list, num_extra_outputs=len(durations)
+        )
+
+        greedy = greedy_decode.GreedyTDTInfer(
+            decoder,
+            joint_net,
+            blank_index=vocab_size,
+            durations=durations,
+            max_symbols_per_step=2,
+            include_duration=True,
+        )
+
+        # (B, D, T)
+        enc_out = torch.zeros(1, encoder_output_size, num_frames)
+        enc_len = torch.tensor([num_frames], dtype=torch.int32)
+
+        with torch.no_grad():
+            hyp = greedy(encoder_output=enc_out, encoded_lengths=enc_len)[0][0]
+
+        # two tokens per time step, and the time step advances by the duration of the second one
+        assert [int(t) for t in hyp.timestamp] == [t for t in range(0, num_frames, 2) for _ in range(2)]
+        assert [int(d) for d in hyp.token_duration] == [0, 2] * (num_frames // 2)
+
     @pytest.mark.skipif(
         not NUMBA_RNNT_LOSS_AVAILABLE,
         reason='RNNTLoss has not been compiled with appropriate numba version.',