chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Support returning customized info (#16144)

Browse Source
This commit is contained in:
Qiaolin Yu
2026-01-02 02:32:18 -05:00
committed by GitHub
parent 6c8587b5db
commit dd93e4451c
8 changed files with 58 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
python/sglang/srt/layers/logits_processor.py
View File

@@ -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:

1
python/sglang/srt/managers/detokenizer_manager.py
View File

@@ -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,

5
python/sglang/srt/managers/io_struct.py
View File

@@ -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):

11
python/sglang/srt/managers/multi_tokenizer_mixin.py
View File

@@ -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),

5
python/sglang/srt/managers/schedule_batch.py
View File

@@ -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

23
python/sglang/srt/managers/scheduler_output_processor_mixin.py
View File

@@ -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,

4
python/sglang/srt/managers/tokenizer_manager.py
View File

@@ -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]

6
python/sglang/srt/sampling/sampling_batch_info.py
View File

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