From 3fd6a58e6c1e04095e226eda4c7f00c596887ea1 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Tue, 19 May 2026 09:21:29 +0800 Subject: [PATCH] Inline the single-use split-prefill setup at its caller (#25722) --- python/sglang/srt/managers/schedule_batch.py | 5 ----- test/manual/test_forward_split_prefill.py | 8 ++++---- 2 files changed, 4 insertions(+), 9 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index a7a97cdd4..4290a61e4 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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() diff --git a/test/manual/test_forward_split_prefill.py b/test/manual/test_forward_split_prefill.py index ab83965a6..2712bfaaa 100644 --- a/test/manual/test_forward_split_prefill.py +++ b/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)