From 8541b1118d6cb396fe9c8f580ea597996635f8fb Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Thu, 19 Feb 2026 22:31:26 -0800 Subject: [PATCH] [Fix][Qwen3.5] Pass max_mamba_cache_size to mamba pool in disaggregation decode path (#19002) --- python/sglang/srt/disaggregation/decode.py | 6 +++++- .../srt/model_executor/model_runner_kv_cache_mixin.py | 1 + 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index fa653448b..1d8baf002 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -158,6 +158,7 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool): speculative_num_draft_tokens: int, enable_mamba_extra_buffer: bool, pre_alloc_size: int, + mamba_size: int = None, ): DecodeReqToTokenPool.__init__( self, @@ -172,8 +173,11 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool): ) 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 self._init_mamba_pool( - size=size + pre_alloc_size, + size=effective_mamba_size, mamba_spec_state_size=size + pre_alloc_size, cache_params=cache_params, device=device, diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index 9852aaff3..0fccfd027 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -445,6 +445,7 @@ class ModelRunnerKVCacheMixin: speculative_num_draft_tokens=self.server_args.speculative_num_draft_tokens, enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(), pre_alloc_size=pre_alloc_size, + mamba_size=self.server_args.max_mamba_cache_size, ) else: self.req_to_token_pool = DecodeReqToTokenPool(