[core] Compute dimensions/return_pooled_hidden_states in ForwardBatch.init_new (#26779)

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]

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