Inline the single-use split-prefill setup at its caller (#25722)

python/sglang/srt/managers/schedule_batch.py

@@ -2078,11 +2078,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             req.mamba_last_track_seqlen = mamba_track_seqlen_aligned
         mamba_track_seqlens_cpu.append(mamba_track_seqlen)
 
-    def prepare_for_split_prefill(self):
-        self.prepare_for_extend()
-        # For split prefill, we need to set the forward mode to SPLIT_PREFILL
-        self.forward_mode = ForwardMode.SPLIT_PREFILL
-
     def mix_with_running(self, running_batch: "ScheduleBatch"):
         self.forward_mode = ForwardMode.MIXED
         running_bs = running_batch.batch_size()

test/manual/test_forward_split_prefill.py

@@ -15,7 +15,7 @@ import torch
 from sglang.bench_one_batch import TreeCacheNamespace
 from sglang.srt.configs.model_config import ModelConfig
 from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
-from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.sampling.sampling_params import SamplingParams
 from sglang.srt.server_args import PortArgs, ServerArgs
@@ -115,10 +115,10 @@ class TestForwardSplitPrefill(CustomTestCase):
             enable_overlap=False,
             spec_algorithm=SpeculativeAlgorithm.NONE,
         )
-        if is_split_prefill:
-            batch.prepare_for_split_prefill()
-        else:
         batch.prepare_for_extend()
+        if is_split_prefill:
+            # For split prefill, we need to set the forward mode to SPLIT_PREFILL
+            batch.forward_mode = ForwardMode.SPLIT_PREFILL
 
         # Create forward batch
         forward_batch = ForwardBatch.init_new(batch, self.model_runner)