diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index efa1dad93..c43295e8f 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -33,7 +33,7 @@ from contextlib import nullcontext from datetime import datetime from enum import Enum from http import HTTPStatus -from typing import Any, Awaitable, Dict, List, Optional, Tuple, Union +from typing import Any, Awaitable, Dict, Iterable, List, Optional, Tuple, Union import fastapi import pybase64 @@ -205,6 +205,9 @@ class ReqState: output_top_logprobs: List[Any] = dataclasses.field(default_factory=list) input_token_ids_logprobs: List[Any] = dataclasses.field(default_factory=list) output_token_ids_logprobs: List[Any] = dataclasses.field(default_factory=list) + customized_info_accumulated: Dict[str, List[Any]] = dataclasses.field( + default_factory=dict + ) # For return_prompt_token_ids: stores prompt token IDs captured after tokenization prompt_token_ids: Optional[List[int]] = None @@ -213,9 +216,13 @@ class ReqState: def _slice_streaming_output_meta_info( meta_info: Dict[Any, Any], last_output_offset: int, + customized_info_keys: Optional[Iterable[str]] = None, ) -> None: """Align output-side metadata with the current incremental streaming chunk.""" - for key in meta_info.keys() & set(_INCREMENTAL_STREAMING_META_INFO_KEYS): + streaming_meta_info_keys = set(_INCREMENTAL_STREAMING_META_INFO_KEYS) + if customized_info_keys is not None: + streaming_meta_info_keys.update(customized_info_keys) + for key in meta_info.keys() & streaming_meta_info_keys: meta_info[key] = meta_info[key][last_output_offset:] @@ -1288,6 +1295,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): self, out_list: list, rid: str, + customized_info_keys: Optional[Iterable[str]] = None, ) -> dict: """Coalesce multiple incremental streaming chunks into one. @@ -1309,7 +1317,10 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): if "meta_info" in out: meta_info_list = [chunk["meta_info"] for chunk in out_list] meta_info = dict(meta_info_list[-1]) - for key in _INCREMENTAL_STREAMING_META_INFO_KEYS: + incremental_streaming_keys = set(_INCREMENTAL_STREAMING_META_INFO_KEYS) + if customized_info_keys is not None: + incremental_streaming_keys.update(customized_info_keys) + for key in incremental_streaming_keys: if any(key in m for m in meta_info_list): meta_info[key] = [ item for m in meta_info_list for item in m.get(key, []) @@ -1401,7 +1412,11 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): is_stream and self.server_args.incremental_streaming_output ) if incremental_stream and len(out_list) > 1: - out = self._coalesce_streaming_chunks(out_list, obj.rid) + out = self._coalesce_streaming_chunks( + out_list, + obj.rid, + state.customized_info_accumulated.keys(), + ) else: out = out_list[-1] @@ -1847,6 +1862,12 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): meta_info["cached_tokens_details"] = recv_obj.cached_tokens_details[ i ] + if recv_obj.customized_info is not None: + for k, v in recv_obj.customized_info.items(): + if k not in state.customized_info_accumulated: + state.customized_info_accumulated[k] = [] + state.customized_info_accumulated[k].extend(v[i]) + meta_info[k] = state.customized_info_accumulated[k] if getattr(recv_obj, "output_hidden_states", None): hidden_states = recv_obj.output_hidden_states[i] @@ -1866,9 +1887,6 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): if isinstance(val, torch.Tensor): val = pybase64.b64encode(val.numpy().tobytes()).decode("utf-8") meta_info["indexer_topk"] = val - if getattr(recv_obj, "customized_info", None): - for k, v in recv_obj.customized_info.items(): - meta_info[k] = v[i] if getattr(recv_obj, "dp_ranks", None): meta_info["dp_rank"] = recv_obj.dp_ranks[i] @@ -1888,7 +1906,11 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): if is_stream: if incremental: output_token_ids = delta_output_ids - _slice_streaming_output_meta_info(meta_info, output_offset) + _slice_streaming_output_meta_info( + meta_info, + output_offset, + state.customized_info_accumulated.keys(), + ) state.last_output_offset = len(state.output_ids) out_dict = { "text": delta_text, @@ -1932,7 +1954,11 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): if is_stream: if incremental: output_token_ids = delta_output_ids - _slice_streaming_output_meta_info(meta_info, output_offset) + _slice_streaming_output_meta_info( + meta_info, + output_offset, + state.customized_info_accumulated.keys(), + ) state.last_output_offset = len(state.output_ids) out_dict = { "output_ids": output_token_ids, diff --git a/test/registered/unit/managers/test_customized_info_streaming.py b/test/registered/unit/managers/test_customized_info_streaming.py new file mode 100644 index 000000000..ebaccf2dc --- /dev/null +++ b/test/registered/unit/managers/test_customized_info_streaming.py @@ -0,0 +1,172 @@ +from __future__ import annotations + +import unittest +from typing import TYPE_CHECKING, List + +import torch + +from sglang.srt.entrypoints.engine import Engine +from sglang.srt.layers.sampler import Sampler, register_sampler_backend +from sglang.srt.managers.scheduler import run_scheduler_process +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.mock_model.utils import MOCK_MODEL_PATH +from sglang.test.test_utils import CustomTestCase + +if TYPE_CHECKING: + from sglang.srt.layers.logits_processor import LogitsProcessorOutput + from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo + +register_cuda_ci(est_time=120, stage="base-b", runner_config="1-gpu-small") + + +CUSTOMIZED_INFO_FIELD = "sampled_token_ids_copy" +CUSTOMIZED_INFO_SAMPLER_BACKEND = "customized_info_probe" +_INPUT_IDS = [464, 9345, 3958, 1752, 13] +_MAX_NEW_TOKENS = 17 + + +class CustomizedInfoSampler(Sampler): + """Sampler probe that mirrors every sampled token into customized_info. + + The scheduler already appends sampled token ids to each request's output_ids. + By copying the same values into customized_info at the sampler boundary, the + test can assert that customized_info is sliced and accumulated exactly like + output_ids throughout the scheduler -> tokenizer manager -> Engine path. + """ + + def forward( + self, + logits_output: "LogitsProcessorOutput", + sampling_info: "SamplingBatchInfo", + return_logprob: bool, + top_logprobs_nums: List[int], + token_ids_logprobs: List[List[int]], + positions: torch.Tensor, + ) -> torch.Tensor: + batch_next_token_ids = super().forward( + logits_output, + sampling_info, + return_logprob, + top_logprobs_nums, + token_ids_logprobs, + positions, + ) + + if logits_output.customized_info is None: + logits_output.customized_info = {} + logits_output.customized_info[CUSTOMIZED_INFO_FIELD] = ( + batch_next_token_ids.detach().cpu().tolist() + ) + return batch_next_token_ids + + +def install_customized_info_sampler() -> None: + # Register before ServerArgs validation in the parent and before sampler + # construction in the scheduler subprocess. + register_sampler_backend( + CUSTOMIZED_INFO_SAMPLER_BACKEND, + CustomizedInfoSampler, + ) + + +def run_scheduler_process_with_customized_info_sampler(*args, **kwargs): + # Engine launches the scheduler in a subprocess. Install the sampler there + # too so create_sampler() can resolve CUSTOMIZED_INFO_SAMPLER_BACKEND. + install_customized_info_sampler() + return run_scheduler_process(*args, **kwargs) + + +class _CustomizedInfoEngine(Engine): + run_scheduler_process_func = staticmethod( + run_scheduler_process_with_customized_info_sampler + ) + + +class TestCustomizedInfoStreaming(CustomTestCase): + @classmethod + def setUpClass(cls): + install_customized_info_sampler() + cls.engine = _CustomizedInfoEngine( + model_path=MOCK_MODEL_PATH, + load_format="dummy", + sampling_backend=CUSTOMIZED_INFO_SAMPLER_BACKEND, + incremental_streaming_output=True, + skip_tokenizer_init=True, + disable_cuda_graph=True, + disable_piecewise_cuda_graph=True, + disable_radix_cache=True, + random_seed=0, + log_level="error", + mem_fraction_static=0.5, + max_total_tokens=1024, + ) + + @classmethod + def tearDownClass(cls): + cls.engine.shutdown() + + def _sampling_params(self, *, stream_interval: int | None = None) -> dict: + sampling_params = { + "temperature": 0.0, + "max_new_tokens": _MAX_NEW_TOKENS, + "ignore_eos": True, + } + if stream_interval is not None: + sampling_params["stream_interval"] = stream_interval + return sampling_params + + def _generate(self, *, stream: bool, stream_interval: int | None = None): + self.engine.flush_cache() + # skip_tokenizer_init keeps this test focused on streaming output + # handling; input_ids bypass tokenizer setup while the real Engine, + # scheduler, and tokenizer-manager response path still run. + return self.engine.generate( + input_ids=_INPUT_IDS, + sampling_params=self._sampling_params(stream_interval=stream_interval), + stream=stream, + ) + + def _assert_customized_info_matches_output_ids(self, output: dict): + # For streaming chunks this should compare per-chunk lists. For the + # non-streaming final response it should compare fully accumulated + # lists. Either failure means customized_info drifted from output_ids. + self.assertIn("output_ids", output) + self.assertIn("meta_info", output) + self.assertIn(CUSTOMIZED_INFO_FIELD, output["meta_info"]) + self.assertEqual( + output["meta_info"][CUSTOMIZED_INFO_FIELD], output["output_ids"] + ) + + def test_non_streaming_returns_accumulated_customized_info(self): + output = self._generate(stream=False) + + self._assert_customized_info_matches_output_ids(output) + self.assertEqual(len(output["output_ids"]), _MAX_NEW_TOKENS) + + def test_incremental_streaming_returns_chunk_customized_info(self): + chunks = list(self._generate(stream=True, stream_interval=1)) + + self.assertEqual(len(chunks), _MAX_NEW_TOKENS) + output_ids = [] + for chunk in chunks: + self._assert_customized_info_matches_output_ids(chunk) + output_ids.extend(chunk["output_ids"]) + self.assertEqual(len(output_ids), _MAX_NEW_TOKENS) + + def test_incremental_streaming_interval_returns_chunk_customized_info(self): + chunks = list(self._generate(stream=True, stream_interval=4)) + + # stream_interval should coalesce multiple scheduler token events into + # at least one multi-token Engine chunk while preserving per-chunk + # customized_info alignment. + self.assertGreater(len(chunks), 1) + self.assertTrue(any(len(chunk["output_ids"]) > 1 for chunk in chunks)) + output_ids = [] + for chunk in chunks: + self._assert_customized_info_matches_output_ids(chunk) + output_ids.extend(chunk["output_ids"]) + self.assertEqual(len(output_ids), _MAX_NEW_TOKENS) + + +if __name__ == "__main__": + unittest.main()