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

python/sglang/srt/mem_cache/memory_pool.py

@@ -667,18 +667,25 @@ class MambaPool:
 
     def clear_slots(self, indices: torch.Tensor):
         """Zero out mamba state at the given pool indices. Must run on forward stream."""
-        need_size = len(indices)
-        for i in range(len(self.mamba_cache.conv)):
-            t = self.mamba_cache.conv[i]
+        if not _is_npu:
+            need_size = len(indices)
+            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(
                 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(
-            t.shape[0], need_size, *t.shape[2:]
-        )
-        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):
         """Clone mamba state (conv + temporal) from src slots into dst slots.