From 8a530468fd2f8adc34fdc08212251b65aca8b5ee Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Fri, 1 May 2026 15:45:20 -0700 Subject: [PATCH] [Bug] Size mamba mappings from req pool, not mamba pool (#24244) --- python/sglang/srt/disaggregation/decode.py | 2 +- python/sglang/srt/mem_cache/memory_pool.py | 13 ++++++++----- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index e54e711ce..b33315b8f 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -228,7 +228,7 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool): self.start_layer = start_layer if start_layer is not None else 0 self.layer_transfer_counter = None self._init_mamba_pool( - size=effective_mamba_size, + mamba_size=effective_mamba_size, mamba_spec_state_size=size + pre_alloc_size, cache_params=cache_params, mamba_layer_ids=mamba_layer_ids, diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 5f8ed05b7..bbc25a7e6 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -509,7 +509,7 @@ class HybridReqToTokenPool(ReqToTokenPool): self.start_layer = start_layer if start_layer is not None else 0 self.layer_transfer_counter = None self._init_mamba_pool( - size=mamba_size, + mamba_size=mamba_size, mamba_spec_state_size=mamba_spec_state_size, cache_params=cache_params, mamba_layer_ids=mamba_layer_ids, @@ -520,7 +520,7 @@ class HybridReqToTokenPool(ReqToTokenPool): def _init_mamba_pool( self, - size: int, + mamba_size: int, mamba_spec_state_size: int, cache_params: BaseLinearStateParams, mamba_layer_ids: List[int], @@ -529,7 +529,7 @@ class HybridReqToTokenPool(ReqToTokenPool): speculative_num_draft_tokens: int = None, ): self.mamba_pool = MambaPool( - size=size, + size=mamba_size, spec_state_size=mamba_spec_state_size, cache_params=cache_params, mamba_layer_ids=mamba_layer_ids, @@ -540,13 +540,16 @@ class HybridReqToTokenPool(ReqToTokenPool): self.mamba_map = {layer_id: i for i, layer_id in enumerate(mamba_layer_ids)} self.device = device + # Indexed by req_pool_idx, so size from the req pool buffer + # (self.req_to_token.shape[0]), not from the mamba state pool size. + req_pool_size = self.req_to_token.shape[0] self.req_index_to_mamba_index_mapping: torch.Tensor = torch.zeros( - size, dtype=torch.int32, device=self.device + req_pool_size, dtype=torch.int32, device=self.device ) if enable_mamba_extra_buffer: self.req_index_to_mamba_ping_pong_track_buffer_mapping: torch.Tensor = ( torch.zeros( - (size, self.mamba_ping_pong_track_buffer_size), + (req_pool_size, self.mamba_ping_pong_track_buffer_size), dtype=torch.int32, device=self.device, )