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

Fix MambaPool.clear_slots OOM by replacing expand-based tensor allocation with scalar zeroing (#27923)

Browse Source
This commit is contained in:
iridiumine
2026-07-02 09:26:55 +08:00
committed by GitHub
parent 0c1a0be3b2
commit 70df09b833
1 changed files with 15 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/memory_pool.py
+7
View File
@@ -667,6 +667,7 @@ class MambaPool:
def clear_slots(self, indices: torch.Tensor): def clear_slots(self, indices: torch.Tensor):
"""Zero out mamba state at the given pool indices. Must run on forward stream.""" """Zero out mamba state at the given pool indices. Must run on forward stream."""
if not _is_npu:
need_size = len(indices) need_size = len(indices)
for i in range(len(self.mamba_cache.conv)): for i in range(len(self.mamba_cache.conv)):
t = self.mamba_cache.conv[i] t = self.mamba_cache.conv[i]
@@ -679,6 +680,12 @@ class MambaPool:
t.shape[0], need_size, *t.shape[2:] t.shape[0], need_size, *t.shape[2:]
) )
t[:, indices] = z t[:, indices] = z
else:
for i in range(len(self.mamba_cache.conv)):
t = self.mamba_cache.conv[i]
t[:, indices] = 0
t = self.mamba_cache.temporal
t[:, indices] = 0
def copy_from(self, src_indices: torch.Tensor, dst_indices: torch.Tensor): def copy_from(self, src_indices: torch.Tensor, dst_indices: torch.Tensor):
"""Clone mamba state (conv + temporal) from src slots into dst slots. """Clone mamba state (conv + temporal) from src slots into dst slots.