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

misc: make prefill-delayer compatible with multiple types of mem pool (#20979)

Browse Source
This commit is contained in:
Yilong Zhao
2026-03-20 00:05:53 -07:00
committed by GitHub
parent 193bbf9b66
commit 26f709e97d
1 changed files with 18 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
python/sglang/srt/managers/scheduler.py
View File

@@ -2135,7 +2135,24 @@ class Scheduler(
def get_new_batch_prefill(self) -> Optional[ScheduleBatch]:
prefill_delayer_single_pass = None
if self.prefill_delayer:
_, token_usage, _, _ = self._get_token_info()
# Get token usage from several pools
token_usage = None
if self.is_hybrid_swa:
_, _, full_token_usage, swa_token_usage, *_ = self._get_swa_token_info()
token_usage = max(full_token_usage, swa_token_usage)
if self.is_hybrid_ssm:
_, _, full_token_usage, mamba_token_usage, *_ = (
self._get_mamba_token_info()
)
token_usage = (
max(token_usage, mamba_token_usage)
if token_usage is not None
else max(full_token_usage, mamba_token_usage)
)
if token_usage is None:
_, token_usage, _, _ = self._get_token_info()
assert token_usage is not None
prefill_delayer_single_pass = PrefillDelayerSinglePassExecutor(
self.prefill_delayer, token_usage=token_usage
)