From 9b4be9c574351cdb4fee8f5fa514575edd452776 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sat, 30 May 2026 17:38:01 -0700 Subject: [PATCH] [core] Compute dimensions/return_pooled_hidden_states in ForwardBatch.init_new (#26779) --- python/sglang/srt/managers/schedule_batch.py | 41 ------------------- .../srt/model_executor/forward_batch_info.py | 37 +++++++++++++++-- 2 files changed, 34 insertions(+), 44 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 754793141..cf7e304e7 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1582,12 +1582,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # Speculative decoding spec_algorithm: SpeculativeAlgorithm = None - # For matryoshka embeddings - dimensions: Optional[list[int]] = None - - # Whether to return pooled hidden states (pre-head transformer output) - return_pooled_hidden_states: bool = False - # Whether to return hidden states return_hidden_states: bool = False @@ -1615,9 +1609,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): encoder_cached: Optional[List[bool]] = None encoder_lens_cpu: Optional[List[int]] = None - # Multi-item scoring delimiter indices (set during prepare_for_extend) - multi_item_delimiter_indices: Optional[List[torch.Tensor]] = None - # For extend and mixed chunekd prefill prefix_lens: List[int] = None extend_lens: List[int] = None @@ -1818,21 +1809,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): prefix_lens = [len(r.prefix_indices) for r in reqs] extend_lens = [r.extend_input_len for r in reqs] - # For matryoshka embeddings - if self.model_config.is_matryoshka and any( - r.dimensions is not None for r in reqs - ): - self.dimensions = [ - r.dimensions if r.dimensions else self.model_config.hidden_size - for r in reqs - ] - - # OR across the batch so ForwardBatch matches a single fused forward; requests - # that did not ask for PHS still skip attaching it in the output processor. - self.return_pooled_hidden_states = any( - r.return_pooled_hidden_states for r in reqs - ) - token_type_ids = [ r.token_type_ids for r in reqs if r.token_type_ids is not None ] @@ -2052,23 +2028,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.token_type_ids = token_type_ids_tensor self.seq_lens_sum = sum(seq_lens) - # Pre-compute delimiter indices as CPU tensors for MIS. - # When --enable-mis is on, every request in the batch is expected to - # carry delimiter indices (the score endpoint always produces MIS-structured - # requests). Consumers index this list without None-checking. - if get_global_server_args().enable_mis and any( - r.multi_item_delimiter_indices is not None for r in reqs - ): - assert all( - r.multi_item_delimiter_indices is not None for r in reqs - ), "MIS batch must have delimiter indices on every request" - self.multi_item_delimiter_indices = [ - torch.tensor(r.multi_item_delimiter_indices, dtype=torch.int64) - for r in reqs - ] - else: - self.multi_item_delimiter_indices = None - if self.return_logprob: self.top_logprobs_nums = [r.logprob.top_logprobs_num for r in reqs] self.token_ids_logprobs = [r.logprob.token_ids_logprob for r in reqs] diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index ed18949e1..a978df1fd 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -544,8 +544,6 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): is_prefill_only=batch.is_prefill_only, spec_algorithm=batch.spec_algorithm, capture_hidden_mode=capture_hidden_mode, - dimensions=batch.dimensions, - return_pooled_hidden_states=batch.return_pooled_hidden_states, return_hidden_states_before_norm=return_hidden_states_before_norm, tbo_split_seq_index=batch.tbo_split_seq_index, # Host-side metadata @@ -554,7 +552,6 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): mm_inputs=batch.multimodal_inputs, encoder_cached=batch.encoder_cached, encoder_lens_cpu=batch.encoder_lens_cpu, - multi_item_delimiter_indices=batch.multi_item_delimiter_indices, lora_ids=[req.lora_id for req in batch.reqs], rids=[req.rid for req in batch.reqs], # Compound (carry their own device tensors) @@ -562,6 +559,8 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): spec_info=batch.spec_info, ) + ret._maybe_init_prefill_only(batch) + device = model_runner.device if batch.extend_input_logprob_token_ids is not None: @@ -679,6 +678,38 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): return ret + def _maybe_init_prefill_only(self, batch: ScheduleBatch): + """Derive per-request fields from ``batch.reqs`` for non-generation + (embedding / reward / scoring) forwards; a no-op otherwise.""" + if not self.is_prefill_only: + return + + if batch.model_config.is_matryoshka and any( + r.dimensions is not None for r in batch.reqs + ): + self.dimensions = [ + r.dimensions if r.dimensions else batch.model_config.hidden_size + for r in batch.reqs + ] + + self.return_pooled_hidden_states = any( + r.return_pooled_hidden_states for r in batch.reqs + ) + + # --enable-mis: every request must carry delimiter indices (the score + # endpoint always produces MIS-structured requests; consumers index + # without None-checking). + if get_global_server_args().enable_mis and any( + r.multi_item_delimiter_indices is not None for r in batch.reqs + ): + assert all( + r.multi_item_delimiter_indices is not None for r in batch.reqs + ), "MIS batch must have delimiter indices on every request" + self.multi_item_delimiter_indices = [ + torch.tensor(r.multi_item_delimiter_indices, dtype=torch.int64) + for r in batch.reqs + ] + def adjust_num_token_non_padded_for_attn_tp(self, server_args) -> None: """Make num_token_non_padded local to this attention-TP rank.""" from sglang.srt.utils.common import require_mlp_tp_gather