diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 350c84cb4..17aa028e0 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -123,8 +123,10 @@ class PrefillBootstrapQueue: self.bootstrap_port = bootstrap_port self.queue: List[Req] = [] self.gloo_group = gloo_group - self.max_total_num_tokens = max_total_num_tokens self.scheduler = scheduler + self.max_total_num_tokens = ( + self.scheduler.tp_worker.model_runner.max_token_pool_size + ) self.transfer_backend = transfer_backend if envs.SGLANG_DISAGG_STAGING_BUFFER.get() and self.is_mla_backend: raise RuntimeError( @@ -133,13 +135,6 @@ class PrefillBootstrapQueue: ) self.kv_manager = self._init_kv_manager() - if self.scheduler.tp_worker.is_hybrid_swa: - # FIXME: current SWA allocation allocate full kv cache size in prefill - self.max_total_num_tokens = min( - self.max_total_num_tokens, - self.scheduler.tp_worker.model_runner.swa_max_total_num_tokens, - ) - def _init_kv_manager(self) -> CommonKVManager: kv_args_class = get_kv_class(self.transfer_backend, KVClassType.KVARGS) kv_args = kv_args_class()