From 26f709e97df255bd1f7729e2d5a94d7c6b69fed6 Mon Sep 17 00:00:00 2001 From: Yilong Zhao <74357408+happierpig@users.noreply.github.com> Date: Fri, 20 Mar 2026 00:05:53 -0700 Subject: [PATCH] misc: make prefill-delayer compatible with multiple types of mem pool (#20979) --- python/sglang/srt/managers/scheduler.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index da1a9ed58..fce74e237 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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 )