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(