diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py old mode 100644 new mode 100755 index de7405434..ac4695cef --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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): diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 778c7e598..a9719e5a4 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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}" diff --git a/python/sglang/srt/mem_cache/allocator.py b/python/sglang/srt/mem_cache/allocator.py old mode 100644 new mode 100755 index c0ce15e20..eff191b10 --- a/python/sglang/srt/mem_cache/allocator.py +++ b/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): - return self._kvcache.get_cpu_copy(indices) + def get_cpu_copy(self, indices, **kwargs): + return self._kvcache.get_cpu_copy(indices, **kwargs) - def load_cpu_copy(self, kv_cache_cpu, indices): - return self._kvcache.load_cpu_copy(kv_cache_cpu, indices) + def load_cpu_copy(self, kv_cache_cpu, indices, **kwargs): + 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): - return self._kvcache.get_cpu_copy(indices) + def get_cpu_copy(self, indices, **kwargs): + return self._kvcache.get_cpu_copy(indices, **kwargs) - def load_cpu_copy(self, kv_cache_cpu, indices): - return self._kvcache.load_cpu_copy(kv_cache_cpu, indices) + def load_cpu_copy(self, kv_cache_cpu, indices, **kwargs): + return self._kvcache.load_cpu_copy(kv_cache_cpu, indices, **kwargs) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 55605f98f..df4ca52a5 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/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): diff --git a/test/registered/unit/mem_cache/test_mamba_unittest.py b/test/registered/unit/mem_cache/test_mamba_unittest.py old mode 100644 new mode 100755 index b339aee02..bc7c75a9a --- a/test/registered/unit/mem_cache/test_mamba_unittest.py +++ b/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.