From eec70286ec56f2172605184f9b8ed373e97483f0 Mon Sep 17 00:00:00 2001 From: yunkchen Date: Wed, 1 Apr 2026 16:17:14 +0800 Subject: [PATCH] [Bugfix] Fix effective_mamba_size over-allocation (#20858) Co-authored-by: Shangming Cai --- python/sglang/srt/disaggregation/decode.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index e75582736..d99061fd3 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -195,9 +195,18 @@ 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 - effective_mamba_size = ( - mamba_size if mamba_size is not None else size - ) + pre_alloc_size + if mamba_size is not None: + effective_mamba_size = min(mamba_size, size + pre_alloc_size) + if mamba_size > size + pre_alloc_size: + logger.warning( + "mamba_size (%d) exceeds size + pre_alloc_size (%d), " + "capping effective_mamba_size to %d", + mamba_size, + size + pre_alloc_size, + effective_mamba_size, + ) + else: + effective_mamba_size = size + pre_alloc_size self.start_layer = start_layer if start_layer is not None else 0 self.layer_transfer_counter = None self._init_mamba_pool(