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 )