From 2d2be5d7b247626c3259fe33c145d0281db6ac4f Mon Sep 17 00:00:00 2001 From: yudian0504 <138860534+yudian0504@users.noreply.github.com> Date: Thu, 30 Apr 2026 10:57:55 +0800 Subject: [PATCH] [PD][Bugfix] fix mamba cache capping (#22462) Co-authored-by: hzh0425 Co-authored-by: yizhang2077 <1109276519@qq.com> --- python/sglang/srt/disaggregation/decode.py | 18 +++++++++++++----- python/sglang/srt/server_args.py | 6 ++++++ 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index f5be51e83..a0aa735a4 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -197,18 +197,26 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool): self.mamba_ping_pong_track_buffer_size = 2 if enable_overlap_schedule else 1 self.enable_mamba_extra_buffer = enable_mamba_extra_buffer self.enable_memory_saver = enable_memory_saver + # Each request needs 1 main mamba slot + ping-pong slots when extra_buffer is enabled. + # Cap the pool at max concurrent requests * slots_per_req to avoid allocating failed. + slots_per_req = 1 + ( + self.mamba_ping_pong_track_buffer_size if enable_mamba_extra_buffer else 0 + ) + max_slots_needed = (size + pre_alloc_size) * slots_per_req if mamba_size is not None: - effective_mamba_size = min(mamba_size, size + pre_alloc_size) - if mamba_size > size + pre_alloc_size: + effective_mamba_size = max(mamba_size, max_slots_needed) + if mamba_size < max_slots_needed: logger.warning( - "mamba_size (%d) exceeds size + pre_alloc_size (%d), " - "capping effective_mamba_size to %d", + "mamba_size (%d) is less than decode side's max_slots_needed (%d = %d reqs * %d slots/req), " + "raising effective_mamba_size to %d", mamba_size, + max_slots_needed, size + pre_alloc_size, + slots_per_req, effective_mamba_size, ) else: - effective_mamba_size = size + pre_alloc_size + effective_mamba_size = max_slots_needed self.start_layer = start_layer if start_layer is not None else 0 self.layer_transfer_counter = None self._init_mamba_pool( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 25242d665..c4cba5346 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -3683,6 +3683,12 @@ class ServerArgs: if self.disaggregation_mode == "decode": self.disable_radix_cache = True logger.warning("KV cache is forced as chunk cache for decode server") + if self.enable_mamba_extra_buffer(): + logger.warning( + "Mamba extra_buffer is disabled because decode disaggregation " + "currently forces chunk cache. Falling back to no_buffer." + ) + self.mamba_scheduler_strategy = "no_buffer" elif self.disaggregation_mode == "prefill": assert (