Add MambaPool kvcache offloading during retraction (#22493)

python/sglang/srt/managers/schedule_batch.py

@@ -1250,13 +1250,19 @@ class Req(ReqDllmMixin):
         token_indices = req_to_token_pool.req_to_token[
             self.req_pool_idx, : self.seqlen - 1
         ]
-        self.kv_cache_cpu = token_to_kv_pool_allocator.get_cpu_copy(token_indices)
+        # Copies over both the kv cache and mamba state if available
+        self.kv_cache_cpu = token_to_kv_pool_allocator.get_cpu_copy(
+            token_indices, mamba_indices=self.mamba_pool_idx
+        )
 
     def load_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator):
         token_indices = req_to_token_pool.req_to_token[
             self.req_pool_idx, : self.seqlen - 1
         ]
-        token_to_kv_pool_allocator.load_cpu_copy(self.kv_cache_cpu, token_indices)
+        # Loads both the kv cache and mamba state if exists
+        token_to_kv_pool_allocator.load_cpu_copy(
+            self.kv_cache_cpu, token_indices, mamba_indices=self.mamba_pool_idx
+        )
         del self.kv_cache_cpu
 
     def log_time_stats(self):

python/sglang/srt/managers/scheduler.py

@@ -2686,11 +2686,20 @@ class Scheduler(
         ):
             old_available_tokens = self.token_to_kv_pool_allocator.available_size()
             old_ratio = self.new_token_ratio
+            mamba_pool = getattr(self.tree_cache.req_to_token_pool, "mamba_pool", None)
+            old_mamba_available = (
+                mamba_pool.available_size() if mamba_pool is not None else None
+            )
             retracted_reqs, new_token_ratio, reqs_to_abort = batch.retract_decode(
                 self.server_args
             )
             new_available_tokens = self.token_to_kv_pool_allocator.available_size()
             new_token_gained = new_available_tokens - old_available_tokens
+            mamba_num_gained = (
+                mamba_pool.available_size() - old_mamba_available
+                if mamba_pool is not None
+                else None
+            )
 
             self.num_retracted_reqs = len(retracted_reqs)
             if self.enable_metrics and len(retracted_reqs) > 0:
@@ -2720,6 +2729,8 @@ class Scheduler(
                 else "Testing retraction. "
             )
             msg_details = f"#retracted_reqs: {len(retracted_reqs)}, #new_tokens_gained: {new_token_gained}"
+            if mamba_num_gained is not None:
+                msg_details += f", #mamba_num_gained: {mamba_num_gained}"
             if kv_full_retract_flag:
                 msg_details += (
                     f", #new_token_ratio: {old_ratio:.4f} -> {new_token_ratio:.4f}"

python/sglang/srt/mem_cache/allocator.py

@@ -164,11 +164,11 @@ class TokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         else:
             self.free_group.append(free_index)
 
-    def get_cpu_copy(self, indices):
+    def get_cpu_copy(self, indices, **kwargs):
-        return self._kvcache.get_cpu_copy(indices)
+        return self._kvcache.get_cpu_copy(indices, **kwargs)
 
-    def load_cpu_copy(self, kv_cache_cpu, indices):
+    def load_cpu_copy(self, kv_cache_cpu, indices, **kwargs):
-        return self._kvcache.load_cpu_copy(kv_cache_cpu, indices)
+        return self._kvcache.load_cpu_copy(kv_cache_cpu, indices, **kwargs)
 
 
 def alloc_extend_naive(
@@ -512,8 +512,8 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         self.free_group = []
         self.release_pages = torch.empty((0,), dtype=torch.int64, device=self.device)
 
-    def get_cpu_copy(self, indices):
+    def get_cpu_copy(self, indices, **kwargs):
-        return self._kvcache.get_cpu_copy(indices)
+        return self._kvcache.get_cpu_copy(indices, **kwargs)
 
-    def load_cpu_copy(self, kv_cache_cpu, indices):
+    def load_cpu_copy(self, kv_cache_cpu, indices, **kwargs):
-        return self._kvcache.load_cpu_copy(kv_cache_cpu, indices)
+        return self._kvcache.load_cpu_copy(kv_cache_cpu, indices, **kwargs)

python/sglang/srt/mem_cache/memory_pool.py

@@ -389,6 +389,28 @@ class MambaPool:
         self.copy_from(src_index, dst_index)
         return dst_index
 
+    def get_cpu_copy(self, indices, **kwargs):
+        torch.cuda.synchronize()
+        conv_cpu = [
+            conv[:, indices].to("cpu", non_blocking=True)
+            for conv in self.mamba_cache.conv
+        ]
+        temporal_cpu = self.mamba_cache.temporal[:, indices].to(
+            "cpu", non_blocking=True
+        )
+        torch.cuda.synchronize()
+        return conv_cpu, temporal_cpu
+
+    def load_cpu_copy(self, mamba_cache_cpu, indices, **kwargs):
+        conv_cpu, temporal_cpu = mamba_cache_cpu
+        torch.cuda.synchronize()
+        for i, conv in enumerate(self.mamba_cache.conv):
+            conv[:, indices] = conv_cpu[i].to(conv.device, non_blocking=True)
+        self.mamba_cache.temporal[:, indices] = temporal_cpu.to(
+            self.mamba_cache.temporal.device, non_blocking=True
+        )
+        torch.cuda.synchronize()
+
     def get_contiguous_buf_infos(self):
         """
         Get buffer info for RDMA registration.
@@ -729,10 +751,10 @@ class KVCache(abc.ABC):
     def register_layer_transfer_counter(self, layer_transfer_counter: LayerDoneCounter):
         self.layer_transfer_counter = layer_transfer_counter
 
-    def get_cpu_copy(self, indices):
+    def get_cpu_copy(self, indices, **kwargs):
         raise NotImplementedError()
 
-    def load_cpu_copy(self, kv_cache_cpu, indices):
+    def load_cpu_copy(self, kv_cache_cpu, indices, **kwargs):
         raise NotImplementedError()
 
     def maybe_get_custom_mem_pool(self):
@@ -934,7 +956,7 @@ class MHATokenToKVPool(KVCache):
         ]
         return kv_data_ptrs, kv_data_lens, kv_item_lens
 
-    def get_cpu_copy(self, indices):
+    def get_cpu_copy(self, indices, **kwargs):
         torch.cuda.synchronize()
         kv_cache_cpu = []
         chunk_size = self.cpu_offloading_chunk_size
@@ -952,7 +974,7 @@ class MHATokenToKVPool(KVCache):
         torch.cuda.synchronize()
         return kv_cache_cpu
 
-    def load_cpu_copy(self, kv_cache_cpu, indices):
+    def load_cpu_copy(self, kv_cache_cpu, indices, **kwargs):
         torch.cuda.synchronize()
         chunk_size = self.cpu_offloading_chunk_size
         for layer_id in range(self.layer_num):
@@ -1418,6 +1440,21 @@ class HybridLinearKVPool(KVCache):
     def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor):
         self.full_kv_pool.move_kv_cache(tgt_loc, src_loc)
 
+    def get_cpu_copy(self, indices, mamba_indices=None, **kwargs):
+        kv_cpu = self.full_kv_pool.get_cpu_copy(indices)
+        mamba_cpu = (
+            self.mamba_pool.get_cpu_copy(mamba_indices)
+            if mamba_indices is not None
+            else None
+        )
+        return kv_cpu, mamba_cpu
+
+    def load_cpu_copy(self, cache_cpu, indices, mamba_indices=None, **kwargs):
+        kv_cpu, mamba_cpu = cache_cpu
+        self.full_kv_pool.load_cpu_copy(kv_cpu, indices)
+        if mamba_cpu is not None and mamba_indices is not None:
+            self.mamba_pool.load_cpu_copy(mamba_cpu, mamba_indices)
+
     def get_v_head_dim(self):
         return self.full_kv_pool.get_value_buffer(0).shape[-1]
 
@@ -1649,7 +1686,7 @@ class MLATokenToKVPool(KVCache):
         get_mla_kv_buffer_triton(kv_buffer, loc, cache_k_nope, cache_k_rope)
         return cache_k_nope, cache_k_rope
 
-    def get_cpu_copy(self, indices):
+    def get_cpu_copy(self, indices, **kwargs):
         torch.cuda.synchronize()
         kv_cache_cpu = []
         chunk_size = self.cpu_offloading_chunk_size
@@ -1664,7 +1701,7 @@ class MLATokenToKVPool(KVCache):
         torch.cuda.synchronize()
         return kv_cache_cpu
 
-    def load_cpu_copy(self, kv_cache_cpu, indices):
+    def load_cpu_copy(self, kv_cache_cpu, indices, **kwargs):
         torch.cuda.synchronize()
         chunk_size = self.cpu_offloading_chunk_size
         for layer_id in range(self.layer_num):

test/registered/unit/mem_cache/test_mamba_unittest.py

@@ -392,6 +392,129 @@ class TestMamba(unittest.TestCase):
 
         return tree, allocator, req_to_token_pool, make_dummy_req
 
+    def test_mamba_pool_cpu_offload(self):
+        """MambaPool.get_cpu_copy / load_cpu_copy round-trips conv and temporal state."""
+        _, _, req_to_token_pool, _ = self._setup_tree_and_allocator()
+        mamba_pool = req_to_token_pool.mamba_pool
+        n = 3
+        indices = mamba_pool.alloc(n)
+        self.assertIsNotNone(indices)
+
+        # Write known sentinel values at the allocated slots.
+        for conv in mamba_pool.mamba_cache.conv:
+            conv[:, indices] = 1.0
+        mamba_pool.mamba_cache.temporal[:, indices] = 2.0
+
+        # Save to CPU.
+        conv_cpu, temporal_cpu = mamba_pool.get_cpu_copy(indices)
+
+        # Verify CPU tensors match what was written.
+        for i, conv in enumerate(mamba_pool.mamba_cache.conv):
+            expected = conv[:, indices].cpu()
+            self.assertTrue(
+                torch.allclose(conv_cpu[i].float(), expected.float()),
+                f"conv[{i}] CPU copy mismatch",
+            )
+        expected_t = mamba_pool.mamba_cache.temporal[:, indices].cpu()
+        self.assertTrue(
+            torch.allclose(temporal_cpu.float(), expected_t.float()),
+            "temporal CPU copy mismatch",
+        )
+
+        # Zero out GPU slots and restore from CPU copy.
+        for conv in mamba_pool.mamba_cache.conv:
+            conv[:, indices] = 0.0
+        mamba_pool.mamba_cache.temporal[:, indices] = 0.0
+
+        mamba_pool.load_cpu_copy((conv_cpu, temporal_cpu), indices)
+
+        # Verify restored values match the sentinels.
+        for conv in mamba_pool.mamba_cache.conv:
+            restored = conv[:, indices]
+            self.assertTrue(
+                torch.all(restored == 1.0),
+                "conv not restored after load_cpu_copy",
+            )
+        self.assertTrue(
+            torch.all(mamba_pool.mamba_cache.temporal[:, indices] == 2.0),
+            "temporal not restored after load_cpu_copy",
+        )
+
+    def test_hybrid_kv_pool_cpu_offload(self):
+        """HybridLinearKVPool.get_cpu_copy / load_cpu_copy saves and restores both
+        the full-attention KV cache and Mamba state in a single round-trip."""
+        _, allocator, req_to_token_pool, _ = self._setup_tree_and_allocator()
+        mamba_pool = req_to_token_pool.mamba_pool
+        hybrid_pool = allocator._kvcache  # HybridLinearKVPool
+
+        self.assertIsInstance(hybrid_pool, HybridLinearKVPool)
+
+        n_tokens = 4
+        kv_indices = allocator.alloc(n_tokens)
+        self.assertIsNotNone(kv_indices)
+        mamba_indices = mamba_pool.alloc(1)
+        self.assertIsNotNone(mamba_indices)
+
+        # Write sentinel values into KV buffers (all full-attention layers).
+        for layer_id in range(hybrid_pool.full_kv_pool.layer_num):
+            hybrid_pool.full_kv_pool.k_buffer[layer_id][kv_indices] = 3.0
+            hybrid_pool.full_kv_pool.v_buffer[layer_id][kv_indices] = 4.0
+
+        # Write sentinel values into Mamba state.
+        for conv in mamba_pool.mamba_cache.conv:
+            conv[:, mamba_indices] = 5.0
+        mamba_pool.mamba_cache.temporal[:, mamba_indices] = 6.0
+
+        # --- Round-trip with Mamba indices provided ---
+        cpu_copy = allocator.get_cpu_copy(kv_indices, mamba_indices=mamba_indices)
+        kv_cpu, mamba_cpu = cpu_copy
+        self.assertIsNotNone(
+            mamba_cpu, "mamba_cpu should be saved when mamba_indices given"
+        )
+
+        # Zero out GPU.
+        for layer_id in range(hybrid_pool.full_kv_pool.layer_num):
+            hybrid_pool.full_kv_pool.k_buffer[layer_id][kv_indices] = 0.0
+            hybrid_pool.full_kv_pool.v_buffer[layer_id][kv_indices] = 0.0
+        for conv in mamba_pool.mamba_cache.conv:
+            conv[:, mamba_indices] = 0.0
+        mamba_pool.mamba_cache.temporal[:, mamba_indices] = 0.0
+
+        allocator.load_cpu_copy(cpu_copy, kv_indices, mamba_indices=mamba_indices)
+
+        # Verify KV restored.
+        for layer_id in range(hybrid_pool.full_kv_pool.layer_num):
+            self.assertTrue(
+                torch.all(
+                    hybrid_pool.full_kv_pool.k_buffer[layer_id][kv_indices] == 3.0
+                ),
+                f"k_buffer layer {layer_id} not restored",
+            )
+            self.assertTrue(
+                torch.all(
+                    hybrid_pool.full_kv_pool.v_buffer[layer_id][kv_indices] == 4.0
+                ),
+                f"v_buffer layer {layer_id} not restored",
+            )
+
+        # Verify Mamba restored.
+        for conv in mamba_pool.mamba_cache.conv:
+            self.assertTrue(
+                torch.all(conv[:, mamba_indices] == 5.0),
+                "conv not restored after load_cpu_copy",
+            )
+        self.assertTrue(
+            torch.all(mamba_pool.mamba_cache.temporal[:, mamba_indices] == 6.0),
+            "temporal not restored after load_cpu_copy",
+        )
+
+        # --- Without mamba_indices: mamba_cpu must be None ---
+        cpu_copy_no_mamba = allocator.get_cpu_copy(kv_indices, mamba_indices=None)
+        _, mamba_cpu_none = cpu_copy_no_mamba
+        self.assertIsNone(
+            mamba_cpu_none, "mamba_cpu should be None when mamba_indices=None"
+        )
+
     def test_insert_prev_prefix_len(self):
         """Test that prev_prefix_len correctly controls which KV indices are freed
         during insert, covering: full free, partial free across multi-node, and no free.