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

Do not cap DeepSeek V4 PD prefill by SWA pool size (#26607)

Browse Source
This commit is contained in:
Chi McIsaac
2026-06-01 19:29:45 +08:00
committed by GitHub
parent 1f8d3c7a42
commit 931765e23e
1 changed files with 3 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/prefill.py
+3 -8
View File
@@ -123,8 +123,10 @@ class PrefillBootstrapQueue:
self.bootstrap_port = bootstrap_port self.bootstrap_port = bootstrap_port
self.queue: List[Req] = [] self.queue: List[Req] = []
self.gloo_group = gloo_group self.gloo_group = gloo_group
self.max_total_num_tokens = max_total_num_tokens
self.scheduler = scheduler self.scheduler = scheduler
self.max_total_num_tokens = (
self.scheduler.tp_worker.model_runner.max_token_pool_size
)
self.transfer_backend = transfer_backend self.transfer_backend = transfer_backend
if envs.SGLANG_DISAGG_STAGING_BUFFER.get() and self.is_mla_backend: if envs.SGLANG_DISAGG_STAGING_BUFFER.get() and self.is_mla_backend:
raise RuntimeError( raise RuntimeError(
@@ -133,13 +135,6 @@ class PrefillBootstrapQueue:
) )
self.kv_manager = self._init_kv_manager() 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: def _init_kv_manager(self) -> CommonKVManager:
kv_args_class = get_kv_class(self.transfer_backend, KVClassType.KVARGS) kv_args_class = get_kv_class(self.transfer_backend, KVClassType.KVARGS)
kv_args = kv_args_class() kv_args = kv_args_class()