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)