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

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,
         )
+        batch.prepare_for_extend()
         if is_split_prefill:
-            batch.prepare_for_split_prefill()
-        else:
-            batch.prepare_for_extend()
+            # 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)