Fix customized_info incremental streaming (#27205)

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,

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()