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
+15 -8
View File
@@ -667,18 +667,25 @@ 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."""
need_size = len(indices) if not _is_npu:
for i in range(len(self.mamba_cache.conv)): need_size = len(indices)
t = self.mamba_cache.conv[i] for i in range(len(self.mamba_cache.conv)):
t = self.mamba_cache.conv[i]
z = torch.zeros(1, dtype=t.dtype, device=t.device).expand(
t.shape[0], need_size, *t.shape[2:]
)
t[:, indices] = z
t = self.mamba_cache.temporal
z = torch.zeros(1, dtype=t.dtype, device=t.device).expand( z = torch.zeros(1, dtype=t.dtype, device=t.device).expand(
t.shape[0], need_size, *t.shape[2:] t.shape[0], need_size, *t.shape[2:]
) )
t[:, indices] = z t[:, indices] = z
t = self.mamba_cache.temporal else:
z = torch.zeros(1, dtype=t.dtype, device=t.device).expand( for i in range(len(self.mamba_cache.conv)):
t.shape[0], need_size, *t.shape[2:] t = self.mamba_cache.conv[i]
) t[:, indices] = 0
t[:, indices] = z 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.