minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
fzyzcjy
2026-05-19 09:21:29 +08:00
committed by GitHub
parent 1f3e5aa1e0
commit 3fd6a58e6c
2 changed files with 4 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/schedule_batch.py
-5
View File
@@ -2078,11 +2078,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
req.mamba_last_track_seqlen = mamba_track_seqlen_aligned req.mamba_last_track_seqlen = mamba_track_seqlen_aligned
mamba_track_seqlens_cpu.append(mamba_track_seqlen) 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"): def mix_with_running(self, running_batch: "ScheduleBatch"):
self.forward_mode = ForwardMode.MIXED self.forward_mode = ForwardMode.MIXED
running_bs = running_batch.batch_size() running_bs = running_batch.batch_size()
test/manual/test_forward_split_prefill.py
+4 -4
View File
@@ -15,7 +15,7 @@ import torch
from sglang.bench_one_batch import TreeCacheNamespace from sglang.bench_one_batch import TreeCacheNamespace
from sglang.srt.configs.model_config import ModelConfig from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.managers.schedule_batch import Req, ScheduleBatch 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.model_executor.model_runner import ModelRunner
from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.sampling.sampling_params import SamplingParams
from sglang.srt.server_args import PortArgs, ServerArgs from sglang.srt.server_args import PortArgs, ServerArgs
@@ -115,10 +115,10 @@ class TestForwardSplitPrefill(CustomTestCase):
enable_overlap=False, enable_overlap=False,
spec_algorithm=SpeculativeAlgorithm.NONE, spec_algorithm=SpeculativeAlgorithm.NONE,
) )
batch.prepare_for_extend()
if is_split_prefill: if is_split_prefill:
batch.prepare_for_split_prefill() # For split prefill, we need to set the forward mode to SPLIT_PREFILL
else: batch.forward_mode = ForwardMode.SPLIT_PREFILL
batch.prepare_for_extend()
# Create forward batch # Create forward batch
forward_batch = ForwardBatch.init_new(batch, self.model_runner) forward_batch = ForwardBatch.init_new(batch, self.model_runner)