minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Bug] Size mamba mappings from req pool, not mamba pool (#24244)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-05-01 15:45:20 -07:00
committed by GitHub
parent 79bc2505a5
commit 8a530468fd
2 changed files with 9 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode.py
+1 -1
View File
@@ -228,7 +228,7 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool):
self.start_layer = start_layer if start_layer is not None else 0 self.start_layer = start_layer if start_layer is not None else 0
self.layer_transfer_counter = None self.layer_transfer_counter = None
self._init_mamba_pool( self._init_mamba_pool(
size=effective_mamba_size, mamba_size=effective_mamba_size,
mamba_spec_state_size=size + pre_alloc_size, mamba_spec_state_size=size + pre_alloc_size,
cache_params=cache_params, cache_params=cache_params,
mamba_layer_ids=mamba_layer_ids, mamba_layer_ids=mamba_layer_ids,
python/sglang/srt/mem_cache/memory_pool.py
+8 -5
View File
@@ -509,7 +509,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
self.start_layer = start_layer if start_layer is not None else 0 self.start_layer = start_layer if start_layer is not None else 0
self.layer_transfer_counter = None self.layer_transfer_counter = None
self._init_mamba_pool( self._init_mamba_pool(
size=mamba_size, mamba_size=mamba_size,
mamba_spec_state_size=mamba_spec_state_size, mamba_spec_state_size=mamba_spec_state_size,
cache_params=cache_params, cache_params=cache_params,
mamba_layer_ids=mamba_layer_ids, mamba_layer_ids=mamba_layer_ids,
@@ -520,7 +520,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
def _init_mamba_pool( def _init_mamba_pool(
self, self,
size: int, mamba_size: int,
mamba_spec_state_size: int, mamba_spec_state_size: int,
cache_params: BaseLinearStateParams, cache_params: BaseLinearStateParams,
mamba_layer_ids: List[int], mamba_layer_ids: List[int],
@@ -529,7 +529,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
speculative_num_draft_tokens: int = None, speculative_num_draft_tokens: int = None,
): ):
self.mamba_pool = MambaPool( self.mamba_pool = MambaPool(
size=size, size=mamba_size,
spec_state_size=mamba_spec_state_size, spec_state_size=mamba_spec_state_size,
cache_params=cache_params, cache_params=cache_params,
mamba_layer_ids=mamba_layer_ids, 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.mamba_map = {layer_id: i for i, layer_id in enumerate(mamba_layer_ids)}
self.device = device 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( 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: if enable_mamba_extra_buffer:
self.req_index_to_mamba_ping_pong_track_buffer_mapping: torch.Tensor = ( self.req_index_to_mamba_ping_pong_track_buffer_mapping: torch.Tensor = (
torch.zeros( torch.zeros(
(size, self.mamba_ping_pong_track_buffer_size), (req_pool_size, self.mamba_ping_pong_track_buffer_size),
dtype=torch.int32, dtype=torch.int32,
device=self.device, device=self.device,
) )