kvcache-ai/sglang
1
0
Fork 0
mirror of https://github.com/kvcache-ai/sglang.git synced 2026-07-07 07:47:25 +00:00
Code Issues 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

7
python/sglang/srt/mem_cache/memory_pool.py
View File

@@ -288,8 +288,9 @@ class MambaPool:
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 "
)
# The padded slot 0 is used for writing dummy outputs from padded tokens.
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.num_mamba_layers = num_mamba_layers
@@ -323,7 +324,9 @@ class MambaPool:
self.free_slots = torch.cat((self.free_slots, free_index))
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):
for i in range(len(self.mamba_cache.conv)):