Batch detokenization across positions in detokenize_top_logprobs_tokens

python/sglang/srt/managers/tokenizer_manager.py

@@ -2009,33 +2009,62 @@ def detokenize_logprob_tokens(
                 (logprob, token_id, None)
                 for logprob, token_id in zip(token_logprobs_val, token_logprobs_idx)
             ]
-        else:
-            assert self.tokenizer is not None
-            # In transformers v5, batch_decode([1, 2, 3]) concatenates all tokens
-            # into one string. Wrap each ID in its own list so they decode separately.
-            token_texts = self.tokenizer.batch_decode(
-                [[idx] for idx in token_logprobs_idx]
+
+        assert self.tokenizer is not None
+        token_texts = self._batch_decode_token_ids(token_logprobs_idx)
+        return list(zip(token_logprobs_val, token_logprobs_idx, token_texts))
+
+    def _batch_decode_token_ids(self, token_ids: List[int]) -> List[str]:
+        assert self.tokenizer is not None
+        # In transformers v5, batch_decode([1, 2, 3]) concatenates all tokens
+        # into one string. Wrap each ID in its own list so they decode separately.
+        token_id_seqs = [[idx] for idx in token_ids]
+        backend_tokenizer = getattr(self.tokenizer, "backend_tokenizer", None)
+        if backend_tokenizer is not None and hasattr(backend_tokenizer, "decode_batch"):
+            return backend_tokenizer.decode_batch(
+                token_id_seqs, skip_special_tokens=False
             )
-            return list(zip(token_logprobs_val, token_logprobs_idx, token_texts))
+        return self.tokenizer.batch_decode(token_id_seqs)
 
     def detokenize_top_logprobs_tokens(
         self,
-        token_logprobs_val: List[float],
-        token_logprobs_idx: List[int],
+        token_logprobs_val: List[Optional[List[float]]],
+        token_logprobs_idx: List[Optional[List[int]]],
         decode_to_text: bool,
     ):
-        # TODO: The current implementation only batches the detokenization for top-k tokens per single position.
-        # We should batch all top-k tokens in all positions.
-        ret = []
-        for i in range(len(token_logprobs_val)):
-            if token_logprobs_val[i]:
-                ret.append(
-                    self.detokenize_logprob_tokens(
-                        token_logprobs_val[i], token_logprobs_idx[i], decode_to_text
-                    )
-                )
-            else:
-                ret.append(None)
+        ret: List[Optional[List]] = [None] * len(token_logprobs_val)
+
+        if not decode_to_text:
+            for i, vals in enumerate(token_logprobs_val):
+                if vals:
+                    idxs = token_logprobs_idx[i]
+                    assert idxs is not None
+                    ret[i] = self.detokenize_logprob_tokens(vals, idxs, decode_to_text)
+            return ret
+
+        assert self.tokenizer is not None
+        flat_ids: List[int] = []
+        nonempty_positions: List[int] = []
+        lengths: List[int] = []
+        for i, vals in enumerate(token_logprobs_val):
+            if vals:
+                idxs = token_logprobs_idx[i]
+                assert idxs is not None
+                flat_ids.extend(idxs)
+                lengths.append(len(idxs))
+                nonempty_positions.append(i)
+
+        if not flat_ids:
+            return ret
+
+        decoded = self._batch_decode_token_ids(flat_ids)
+        offset = 0
+        for pos, length in zip(nonempty_positions, lengths):
+            chunk_texts = decoded[offset : offset + length]
+            offset += length
+            ret[pos] = list(
+                zip(token_logprobs_val[pos], token_logprobs_idx[pos], chunk_texts)
+            )
         return ret
 
     def _calculate_spec_decoding_metrics(

test/manual/test_tokenizer_manager.py

@@ -413,6 +413,143 @@ def test_full_workflow_batch_strings(self):
         self.assertIsNone(result_token_type_ids)
 
 
+class TestDetokenizeTopLogprobsTokens(unittest.TestCase):
+    """Test cases for detokenize_top_logprobs_tokens batched decoding.
+
+    We avoid constructing a real TokenizerManager (which requires GPU-dependent
+    setup) and instead invoke the method as an unbound function against a bare
+    stand-in object that only exposes the attributes the method actually uses.
+    """
+
+    def setUp(self):
+        self.fn = TokenizerManager.detokenize_top_logprobs_tokens
+        self.stub = Mock(
+            spec=["tokenizer", "detokenize_logprob_tokens", "_batch_decode_token_ids"]
+        )
+        self.stub.tokenizer = Mock(spec=["batch_decode", "backend_tokenizer"])
+        self.stub.tokenizer.backend_tokenizer = None
+        self.stub.tokenizer.batch_decode = Mock(
+            side_effect=lambda ids: [f"tok_{i[0]}" for i in ids]
+        )
+
+        def batch_decode_token_ids(token_ids):
+            return TokenizerManager._batch_decode_token_ids(self.stub, token_ids)
+
+        def detokenize_logprob_tokens(vals, idxs, decode_to_text):
+            return TokenizerManager.detokenize_logprob_tokens(
+                self.stub, vals, idxs, decode_to_text
+            )
+
+        self.stub._batch_decode_token_ids = batch_decode_token_ids
+        # Delegate to the real helper so we exercise the production path.
+        self.stub.detokenize_logprob_tokens = detokenize_logprob_tokens
+
+    def _call(self, vals, idxs, decode_to_text):
+        return self.fn(self.stub, vals, idxs, decode_to_text)
+
+    def _reference_impl(self, vals, idxs, decode_to_text):
+        """Per-position reference for what the old implementation produced."""
+        ret = []
+        for i in range(len(vals)):
+            if vals[i]:
+                if not decode_to_text:
+                    ret.append([(lp, tid, None) for lp, tid in zip(vals[i], idxs[i])])
+                else:
+                    texts = [f"tok_{tid}" for tid in idxs[i]]
+                    ret.append(list(zip(vals[i], idxs[i], texts)))
+            else:
+                ret.append(None)
+        return ret
+
+    def test_decode_to_text_false_skips_tokenizer(self):
+        """When decode_to_text=False, batch_decode should not be called."""
+        vals = [[-0.1, -0.2], [-0.3]]
+        idxs = [[10, 20], [30]]
+
+        result = self._call(vals, idxs, decode_to_text=False)
+
+        self.stub.tokenizer.batch_decode.assert_not_called()
+        self.assertEqual(
+            result,
+            [[(-0.1, 10, None), (-0.2, 20, None)], [(-0.3, 30, None)]],
+        )
+
+    def test_all_empty_positions_returns_nones(self):
+        """All-empty input returns [None, ...] and never calls batch_decode."""
+        vals = [[], [], None]
+        idxs = [[], [], None]
+
+        result = self._call(vals, idxs, decode_to_text=True)
+
+        self.assertEqual(result, [None, None, None])
+        self.stub.tokenizer.batch_decode.assert_not_called()
+
+    def test_empty_input_list(self):
+        """Zero-length input returns an empty list."""
+        result = self._call([], [], decode_to_text=True)
+        self.assertEqual(result, [])
+        self.stub.tokenizer.batch_decode.assert_not_called()
+
+    def test_mixed_empty_and_nonempty_positions(self):
+        """Texts are sliced back to the correct positions; empties stay None."""
+        vals = [[-0.1, -0.2], [], [-0.5, -0.6, -0.7], None, [-0.9]]
+        idxs = [[10, 20], [], [30, 40, 50], None, [60]]
+
+        result = self._call(vals, idxs, decode_to_text=True)
+
+        expected = [
+            [(-0.1, 10, "tok_10"), (-0.2, 20, "tok_20")],
+            None,
+            [(-0.5, 30, "tok_30"), (-0.6, 40, "tok_40"), (-0.7, 50, "tok_50")],
+            None,
+            [(-0.9, 60, "tok_60")],
+        ]
+        self.assertEqual(result, expected)
+
+    def test_single_batch_decode_call_with_flattened_ids(self):
+        """Efficiency guarantee: batch_decode is called once with flattened token ids."""
+        vals = [[-0.1, -0.2], [], [-0.5, -0.6, -0.7], [-0.9]]
+        idxs = [[10, 20], [], [30, 40, 50], [60]]
+
+        self._call(vals, idxs, decode_to_text=True)
+
+        self.assertEqual(self.stub.tokenizer.batch_decode.call_count, 1)
+        (called_ids,), _ = self.stub.tokenizer.batch_decode.call_args
+        self.assertEqual(list(called_ids), [[10], [20], [30], [40], [50], [60]])
+
+    def test_prefers_backend_decode_batch_when_available(self):
+        """Use the tokenizer backend to avoid per-sequence Python decode loops."""
+        vals = [[-0.1, -0.2], [], [-0.5, -0.6, -0.7], [-0.9]]
+        idxs = [[10, 20], [], [30, 40, 50], [60]]
+        backend_tokenizer = Mock(spec=["decode_batch"])
+        backend_tokenizer.decode_batch = Mock(
+            side_effect=lambda ids, skip_special_tokens: [
+                f"tok_{i[0]}_{skip_special_tokens}" for i in ids
+            ]
+        )
+        self.stub.tokenizer.backend_tokenizer = backend_tokenizer
+
+        result = self._call(vals, idxs, decode_to_text=True)
+
+        self.stub.tokenizer.batch_decode.assert_not_called()
+        backend_tokenizer.decode_batch.assert_called_once_with(
+            [[10], [20], [30], [40], [50], [60]], skip_special_tokens=False
+        )
+        self.assertEqual(result[0][0], (-0.1, 10, "tok_10_False"))
+        self.assertEqual(result[2][2], (-0.7, 50, "tok_50_False"))
+
+    def test_matches_reference_per_position_implementation(self):
+        """Batched result is equivalent to per-position decoding."""
+        vals = [[-0.1, -0.2, -0.3], [], [-0.5], None, [-0.7, -0.8]]
+        idxs = [[10, 20, 30], [], [40], None, [50, 60]]
+
+        for decode_to_text in (True, False):
+            self.stub.tokenizer.batch_decode.reset_mock()
+            batched = self._call(vals, idxs, decode_to_text=decode_to_text)
+            expected = self._reference_impl(vals, idxs, decode_to_text)
+            self.assertEqual(batched, expected)
+
+
 def _make_state() -> ReqState:
     """Create a minimal ReqState for testing."""
     obj = Mock(spec=GenerateReqInput)