chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[BUGFIX] Skip Mamba Cache Slot 0 to Avoid Using Dummy Cache (#17404)

Browse Source
This commit is contained in:
Jincong Chen
2026-01-22 23:08:38 +08:00
committed by GitHub
parent 42523d0364
commit 72f790bf6f
1 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/memory_pool.py
+5 -2
View File
@@ -288,8 +288,9 @@ class MambaPool:
f"conv_state size: {get_tensor_size_bytes(conv_state) / GB:.2f}GB, " f"conv_state size: {get_tensor_size_bytes(conv_state) / GB:.2f}GB, "
f"ssm_state size: {get_tensor_size_bytes(temporal_state) / GB:.2f}GB " f"ssm_state size: {get_tensor_size_bytes(temporal_state) / GB:.2f}GB "
) )
# The padded slot 0 is used for writing dummy outputs from padded tokens.
self.free_slots = torch.arange( self.free_slots = torch.arange(
self.size, dtype=torch.int64, device=self.device 1, self.size + 1, dtype=torch.int64, device=self.device
) )
self.mem_usage = self.mamba_cache.mem_usage_bytes() / GB self.mem_usage = self.mamba_cache.mem_usage_bytes() / GB
self.num_mamba_layers = num_mamba_layers self.num_mamba_layers = num_mamba_layers
@@ -323,7 +324,9 @@ class MambaPool:
self.free_slots = torch.cat((self.free_slots, free_index)) self.free_slots = torch.cat((self.free_slots, free_index))
def clear(self): def clear(self):
self.free_slots = torch.arange(self.size, dtype=torch.int64, device=self.device) self.free_slots = torch.arange(
1, self.size + 1, dtype=torch.int64, device=self.device
)
def copy_from(self, src_index: torch.Tensor, dst_index: torch.Tensor): def copy_from(self, src_index: torch.Tensor, dst_index: torch.Tensor):
for i in range(len(self.mamba_cache.conv)): for i in range(len(self.mamba_cache.conv)):