Support returning customized info (#16144)

python/sglang/srt/layers/logits_processor.py

@@ -15,7 +15,7 @@
 
 import dataclasses
 import logging
-from typing import List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
 import triton
@@ -101,6 +101,9 @@ class LogitsProcessorOutput:
     ## Part 4: Diffusion LLM only.
     full_logits: Optional[torch.Tensor] = None
 
+    ## Part 5: Customized Info
+    customized_info: Optional[Dict[str, List[Any]]] = None
+
 
 @dataclasses.dataclass
 class LogitsMetadata:

python/sglang/srt/managers/detokenizer_manager.py

@@ -341,6 +341,7 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
             output_token_entropy_val=recv_obj.output_token_entropy_val,
             output_hidden_states=recv_obj.output_hidden_states,
             output_routed_experts=output_routed_experts,
+            customized_info=recv_obj.customized_info,
             placeholder_tokens_idx=None,
             placeholder_tokens_val=None,
             retraction_counts=recv_obj.retraction_counts,

python/sglang/srt/managers/io_struct.py

@@ -984,6 +984,8 @@ class BatchTokenIDOutput(
 
     # Load for DP balance
     load: GetLoadReqOutput = None
+    # Customized info
+    customized_info: Optional[Dict[str, List[Any]]] = None
 
 
 @dataclass
@@ -1070,6 +1072,9 @@ class BatchStrOutput(
     # Load for DP balance
     load: GetLoadReqOutput = None
 
+    # Customized info
+    customized_info: Optional[Dict[str, List[Any]]] = None
+
 
 @dataclass
 class BatchMultimodalOutput(BaseBatchReq):

python/sglang/srt/managers/multi_tokenizer_mixin.py

@@ -104,6 +104,14 @@ def _extract_field_by_index(
     if field is None:
         return None
 
+    if isinstance(field, dict):
+        new_field = {}
+        for k, v in field.items():
+            if len(v) <= index:
+                new_field[k] = None
+            new_field[k] = v[index]
+        return new_field
+
     if check_length:
         if len(field) <= index:
             return None
@@ -273,6 +281,9 @@ def _handle_output_by_index(output, i):
             output_routed_experts=_extract_field_by_index(
                 output, "output_routed_experts", i, check_length=False
             ),
+            customized_info=_extract_field_by_index(
+                output, "customized_info", i, check_length=False
+            ),
             placeholder_tokens_idx=None,
             placeholder_tokens_val=None,
             retraction_counts=_extract_field_by_index(output, "retraction_counts", i),

python/sglang/srt/managers/schedule_batch.py

@@ -91,6 +91,8 @@ from sglang.srt.utils import flatten_nested_list
 from sglang.srt.utils.cuda_ipc_transport_utils import CudaIpcTensorTransportProxy
 
 if TYPE_CHECKING:
+    from typing import Any, Dict
+
     from sglang.srt.configs.model_config import ModelConfig
     from sglang.srt.speculative.eagle_info import EagleDraftInput
     from sglang.srt.speculative.spec_info import SpecInput, SpeculativeAlgorithm
@@ -699,6 +701,8 @@ class Req:
         self.routed_experts: Optional[torch.Tensor] = (
             None  # cpu tensor: shape (seqlen, topk)
         )
+        # Customized info
+        self.customized_info: Optional[Dict[str, List[Any]]] = None
 
         # Embedding (return values)
         self.embedding = None
@@ -1068,6 +1072,7 @@ class Req:
         self.last_node = None
         self.swa_uuid_for_lock = None
         self.extend_input_len = 0
+        self.customized_info = None
         self.is_retracted = True
         self.retracted_stain = True
         self.input_token_logprobs = None

python/sglang/srt/managers/scheduler_output_processor_mixin.py

@@ -71,6 +71,17 @@ class SchedulerOutputProcessorMixin:
             req_to_token_pool=self.req_to_token_pool,
         )
 
+    def maybe_collect_customized_info(
+        self: Scheduler, i: int, req: Req, logits_output: LogitsProcessorOutput
+    ):
+        if logits_output is not None and logits_output.customized_info is not None:
+            if req.customized_info is None:
+                req.customized_info = {}
+            for k, v in logits_output.customized_info.items():
+                if k not in req.customized_info:
+                    req.customized_info[k] = []
+                req.customized_info[k].append(v[i])
+
     def process_batch_result_prefill(
         self: Scheduler,
         batch: ScheduleBatch,
@@ -132,6 +143,8 @@ class SchedulerOutputProcessorMixin:
                         # This updates radix so others can match
                         self.tree_cache.cache_unfinished_req(req)
 
+                    self.maybe_collect_customized_info(i, req, logits_output)
+
                     if batch.return_logprob:
                         assert extend_logprob_start_len_per_req is not None
                         assert extend_input_len_per_req is not None
@@ -408,6 +421,8 @@ class SchedulerOutputProcessorMixin:
 
                 req.time_stats.completion_time = time.perf_counter()
 
+            self.maybe_collect_customized_info(i, req, logits_output)
+
             if req.return_logprob and batch.spec_algorithm.is_none():
                 # speculative worker handles logprob in speculative decoding
                 req.output_token_logprobs_val.append(next_token_logprobs[i])
@@ -841,6 +856,7 @@ class SchedulerOutputProcessorMixin:
         output_hidden_states = None
         load = self.get_load()
         output_routed_experts = None
+        customized_info = {}
 
         queue_times = []
         forward_entry_times = []
@@ -1040,6 +1056,12 @@ class SchedulerOutputProcessorMixin:
                         output_routed_experts = []
                     output_routed_experts.append(req.routed_experts)
 
+                if req.customized_info is not None:
+                    for k, v in req.customized_info.items():
+                        if k not in customized_info:
+                            customized_info[k] = []
+                        customized_info[k].append(v)
+
             if (
                 req.finished()
                 and self.attn_tp_rank == 0
@@ -1089,6 +1111,7 @@ class SchedulerOutputProcessorMixin:
                     output_token_entropy_val=None,
                     output_hidden_states=output_hidden_states,
                     output_routed_experts=output_routed_experts,
+                    customized_info=customized_info,
                     placeholder_tokens_idx=None,
                     placeholder_tokens_val=None,
                     retraction_counts=retraction_counts,

python/sglang/srt/managers/tokenizer_manager.py

@@ -1506,9 +1506,11 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
 
             if getattr(recv_obj, "output_hidden_states", None):
                 meta_info["hidden_states"] = recv_obj.output_hidden_states[i]
-
             if getattr(recv_obj, "output_routed_experts", None):
                 meta_info["routed_experts"] = recv_obj.output_routed_experts[i]
+            if getattr(recv_obj, "customized_info", None):
+                for k, v in recv_obj.customized_info.items():
+                    meta_info[k] = v[i]
 
             if isinstance(recv_obj, BatchStrOutput):
                 state.text += recv_obj.output_strs[i]

python/sglang/srt/sampling/sampling_batch_info.py

@@ -187,6 +187,10 @@ class SamplingBatchInfo:
     def adjusted_from_schedule_batch(self, batch: ScheduleBatch, vocab_size: int):
         pass
 
+    # placeholder for override
+    def adjusted_merge_batch(self, other: "SamplingBatchInfo"):
+        pass
+
     def __len__(self):
         return len(self.temperatures)
 
@@ -372,6 +376,8 @@ class SamplingBatchInfo:
         self.need_top_k_sampling |= other.need_top_k_sampling
         self.need_min_p_sampling |= other.need_min_p_sampling
 
+        self.adjusted_merge_batch(other)
+
     def copy_for_forward(self):
         # Accumulate the penalty into a pre-allocated buffer to get rid of the dependency of `penalizer_orchestrator` later
         self.update_penalties()