diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 293e69c42..fa7431048 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/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: diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py index bd08fce53..90fc130f3 100644 --- a/python/sglang/srt/managers/detokenizer_manager.py +++ b/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, diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index c5910434c..879e1bfa6 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/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): diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py index 978be0cd0..0e9d314e3 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/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), diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 8b807613e..b9dd921bf 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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 diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index 451ea80f7..e40586c24 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/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, diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index c6f6cc22c..f4fc29e29 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/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] diff --git a/python/sglang/srt/sampling/sampling_batch_info.py b/python/sglang/srt/sampling/sampling_batch_info.py index 20f020baa..a8f17c754 100644 --- a/python/sglang/srt/sampling/sampling_batch_info.py +++ b/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()