From 3ca8cb470f710c19c36feb6d73ba7db29d94b4d9 Mon Sep 17 00:00:00 2001 From: Cameron Quilici Date: Thu, 28 May 2026 15:31:55 -0500 Subject: [PATCH] [BugFix] preserve cached token details in multi-tokenizer output (#26590) --- .../srt/managers/multi_tokenizer_mixin.py | 3 + .../managers/test_multi_tokenizer_mixin.py | 68 +++++++++++++++++++ 2 files changed, 71 insertions(+) create mode 100644 test/registered/unit/managers/test_multi_tokenizer_mixin.py diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py index 23a63f408..5f6cc6309 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py @@ -246,6 +246,9 @@ def _handle_output_by_index(output, i): completion_tokens=_extract_field_by_index(output, "completion_tokens", i), reasoning_tokens=_extract_field_by_index(output, "reasoning_tokens", i), cached_tokens=_extract_field_by_index(output, "cached_tokens", i), + cached_tokens_details=_extract_field_by_index( + output, "cached_tokens_details", i + ), input_token_logprobs_val=_extract_field_by_index( output, "input_token_logprobs_val", i, check_length=False ), diff --git a/test/registered/unit/managers/test_multi_tokenizer_mixin.py b/test/registered/unit/managers/test_multi_tokenizer_mixin.py new file mode 100644 index 000000000..5a8abfab9 --- /dev/null +++ b/test/registered/unit/managers/test_multi_tokenizer_mixin.py @@ -0,0 +1,68 @@ +import unittest + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.managers.io_struct import BatchStrOutput +from sglang.srt.managers.multi_tokenizer_mixin import _handle_output_by_index + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _make_batch_str_output() -> BatchStrOutput: + return BatchStrOutput( + rids=["rid-0", "rid-1"], + spec_verify_ct=[0, 0], + spec_num_correct_drafts=[0, 0], + spec_correct_drafts_histogram=[[], []], + finished_reasons=[None, {"type": "length"}], + output_strs=["first", "second"], + output_ids=[[1], [2]], + prompt_tokens=[10, 20], + completion_tokens=[1, 2], + reasoning_tokens=[0, 0], + cached_tokens=[3, 4], + cached_tokens_details=[ + {"device": 3, "host": 0}, + {"device": 1, "host": 3}, + ], + input_token_logprobs_val=[[], []], + input_token_logprobs_idx=[[], []], + output_token_logprobs_val=[[], []], + output_token_logprobs_idx=[[], []], + input_top_logprobs_val=[[], []], + input_top_logprobs_idx=[[], []], + output_top_logprobs_val=[[], []], + output_top_logprobs_idx=[[], []], + input_token_ids_logprobs_val=[[], []], + input_token_ids_logprobs_idx=[[], []], + output_token_ids_logprobs_val=[[], []], + output_token_ids_logprobs_idx=[[], []], + output_token_entropy_val=[0.0, 0.0], + output_hidden_states=[None, None], + routed_experts=[None, None], + indexer_topk=[None, None], + placeholder_tokens_idx=[None, None], + placeholder_tokens_val=[None, None], + retraction_counts=[0, 0], + ) + + +class TestMultiTokenizerMixin(unittest.TestCase): + def test_batch_str_output_preserves_cached_tokens_details(self): + output = _make_batch_str_output() + + single_output = _handle_output_by_index(output, 1) + + self.assertEqual(single_output.rids, ["rid-1"]) + self.assertEqual(single_output.cached_tokens, [4]) + self.assertEqual( + single_output.cached_tokens_details, + [{"device": 1, "host": 3}], + ) + + +if __name__ == "__main__": + unittest.main()