diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 2323aa7af0b6..4799d3afc2c3 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1824,7 +1824,7 @@ def offload_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator): req_to_token_pool, token_to_kv_pool_allocator ) self.kv.retraction_backup = RetractionBackup( - cpu_tensors=token_to_kv_pool_allocator.get_cpu_copy( + cpu_tensors=token_to_kv_pool_allocator.get_kvcache().get_cpu_copy( token_indices, mamba_indices=self.kv.mamba_pool_idx ), mamba_cpu=( @@ -1845,7 +1845,7 @@ def load_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator): req_to_token_pool.mamba_pool.load_cpu_copy( mamba_cpu, self.kv.mamba_pool_idx.unsqueeze(0) ) - token_to_kv_pool_allocator.load_cpu_copy( + token_to_kv_pool_allocator.get_kvcache().load_cpu_copy( self.kv.retraction_backup.cpu_tensors, token_indices, mamba_indices=self.kv.mamba_pool_idx, diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index 1417c00b5615..31f55d8dd819 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -125,14 +125,6 @@ def translate_kv_indices_for_transfer( """ return kv_indices - def get_cpu_copy(self, indices, mamba_indices=None): - # FIXME: reuse the get_cpu_copy after paged allocator is implemented - raise NotImplementedError() - - def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): - # FIXME: reuse the load_cpu_copy after paged allocator is implemented - raise NotImplementedError() - def alloc_extend(self, *args, **kwargs): raise NotImplementedError("alloc_extend is only for paged allocator") diff --git a/python/sglang/srt/mem_cache/allocator/hisparse.py b/python/sglang/srt/mem_cache/allocator/hisparse.py index 4a10891f8514..908b1cdb1c1c 100644 --- a/python/sglang/srt/mem_cache/allocator/hisparse.py +++ b/python/sglang/srt/mem_cache/allocator/hisparse.py @@ -171,7 +171,9 @@ def get_last_loc_compressed(self, last_locs: torch.Tensor): return last_locs def get_last_loc_hisparse_device(self, last_locs: torch.Tensor): - return self._kvcache._translate_loc_to_hisparse_device(last_locs) + # full_to_hisparse_device_index_mapping is owned and written by this + # allocator; index it directly rather than round-tripping through the pool. + return self.full_to_hisparse_device_index_mapping[last_locs] def alloc_extend( self, @@ -233,7 +235,7 @@ def alloc_decode( ) def free_hisparse(self, free_indices: torch.Tensor): - hisparse_indices = self._kvcache._translate_loc_to_hisparse_device(free_indices) + hisparse_indices = self.full_to_hisparse_device_index_mapping[free_indices] hisparse_indices = hisparse_indices[hisparse_indices > 0] self.free_hisparse_indices(hisparse_indices) self.full_to_hisparse_device_index_mapping[free_indices] = 0 @@ -482,9 +484,11 @@ def get_last_loc_compressed(self, last_locs: torch.Tensor): return (last_locs - 3) // self.compress_ratio def get_last_loc_hisparse_device(self, last_locs: torch.Tensor): - return self.hisparse_kvcache._translate_loc_to_hisparse_device( + # Index the allocator-owned mapping directly (mirrors the pool's + # _translate_loc_to_hisparse_device, which reads this same tensor). + return self.full_to_hisparse_device_index_mapping[ self.get_last_loc_compressed(last_locs) - ) + ] def alloc_extend( self, @@ -526,6 +530,9 @@ def alloc_extend( ) assert logical_indices is not None, "Logical allocation failed in alloc_extend" + # Kept as a pool call on purpose: full->compressed is genuine C4-layout + # arithmetic ((i+1) % compress_ratio), a pool responsibility, not the + # pointless owned-mapping round-trips removed elsewhere in this allocator. compressed_logical_indices = ( self.hisparse_kvcache.translate_loc_from_full_to_compressed(logical_indices) ) @@ -558,14 +565,18 @@ def alloc_decode( ) def free_compressed(self, compressed_indices: torch.Tensor): - hisparse_indices = self.hisparse_kvcache.translate_loc_to_hisparse_device( + # Index the allocator-owned mapping directly; the .to(int32) matches the + # pool's translate_loc_to_hisparse_device, which reads this same tensor. + hisparse_indices = self.full_to_hisparse_device_index_mapping[ compressed_indices - ) + ].to(torch.int32) hisparse_indices = hisparse_indices[hisparse_indices > 0] self.free_hisparse_indices(hisparse_indices) self.full_to_hisparse_device_index_mapping[compressed_indices] = 0 def free_hisparse(self, free_indices: torch.Tensor): + # full->compressed is genuine C4-layout arithmetic; kept as a pool call + # (see alloc_extend) rather than an owned-mapping round-trip. compressed_indices = ( self.hisparse_kvcache.translate_loc_from_full_to_compressed(free_indices) ) diff --git a/python/sglang/srt/mem_cache/allocator/paged.py b/python/sglang/srt/mem_cache/allocator/paged.py index f21c077f623b..346841b70aaa 100755 --- a/python/sglang/srt/mem_cache/allocator/paged.py +++ b/python/sglang/srt/mem_cache/allocator/paged.py @@ -336,11 +336,3 @@ def clear(self): self.free_group = None self.free_page_reps_group = [] self.release_pages = torch.empty((0,), dtype=torch.int64, device=self.device) - - def get_cpu_copy(self, indices, mamba_indices=None): - return self._kvcache.get_cpu_copy(indices, mamba_indices=mamba_indices) - - def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): - return self._kvcache.load_cpu_copy( - kv_cache_cpu, indices, mamba_indices=mamba_indices - ) diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index b7f94ea5545c..4a1fd9ea1d07 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -146,8 +146,10 @@ def get_kvcache(self): return self._kvcache def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor): - assert self._kvcache.full_to_swa_index_mapping is not None - return self._kvcache.translate_loc_from_full_to_swa(kv_indices) + # The allocator owns full_to_swa_index_mapping (it writes it on every + # alloc/free), so index it directly instead of round-tripping through the + # pool, which would only index this very same tensor. + return self.full_to_swa_index_mapping[kv_indices] def alloc(self, need_size: int): assert self.page_size == 1 @@ -450,14 +452,6 @@ def clear(self): self.swa_free_group = [] self.full_free_group = [] - def get_cpu_copy(self, indices, mamba_indices=None): - return self._kvcache.get_cpu_copy(indices, mamba_indices=mamba_indices) - - def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): - return self._kvcache.load_cpu_copy( - kv_cache_cpu, indices, mamba_indices=mamba_indices - ) - class PureSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): """Single-pool allocator for models whose every layer is sliding-window attention.""" diff --git a/python/sglang/srt/mem_cache/allocator/token.py b/python/sglang/srt/mem_cache/allocator/token.py index fd0f326829d5..eaebc5288174 100644 --- a/python/sglang/srt/mem_cache/allocator/token.py +++ b/python/sglang/srt/mem_cache/allocator/token.py @@ -73,11 +73,3 @@ def free(self, free_index: torch.Tensor): self.free_pages = torch.cat((self.free_pages, free_index)) else: self.free_group.append(self._copy_for_free_group(free_index)) - - def get_cpu_copy(self, indices, mamba_indices=None): - return self._kvcache.get_cpu_copy(indices, mamba_indices=mamba_indices) - - def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): - return self._kvcache.load_cpu_copy( - kv_cache_cpu, indices, mamba_indices=mamba_indices - ) diff --git a/test/registered/unit/mem_cache/test_mamba_unittest.py b/test/registered/unit/mem_cache/test_mamba_unittest.py index 313de93ae412..dadab9f0b326 100755 --- a/test/registered/unit/mem_cache/test_mamba_unittest.py +++ b/test/registered/unit/mem_cache/test_mamba_unittest.py @@ -705,7 +705,9 @@ def test_hybrid_kv_pool_cpu_offload(self): 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) + cpu_copy = allocator.get_kvcache().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" @@ -719,7 +721,9 @@ def test_hybrid_kv_pool_cpu_offload(self): 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) + allocator.get_kvcache().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): @@ -748,7 +752,9 @@ def test_hybrid_kv_pool_cpu_offload(self): ) # --- Without mamba_indices: mamba_cpu must be None --- - cpu_copy_no_mamba = allocator.get_cpu_copy(kv_indices, mamba_indices=None) + cpu_copy_no_mamba = allocator.get_kvcache().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" diff --git a/test/registered/unit/mem_cache/test_retraction_mamba_backup.py b/test/registered/unit/mem_cache/test_retraction_mamba_backup.py index 7d90e77ef624..c79563944d32 100644 --- a/test/registered/unit/mem_cache/test_retraction_mamba_backup.py +++ b/test/registered/unit/mem_cache/test_retraction_mamba_backup.py @@ -22,14 +22,11 @@ def load_cpu_copy(self, state, indices): self.loaded = state -class _Allocator: +class _KVPool: def __init__(self, carries_mamba: bool): - self._kv = type("_KV", (), {"cpu_copy_carries_mamba": carries_mamba})() + self.cpu_copy_carries_mamba = carries_mamba self.loaded_kv = None - def get_kvcache(self): - return self._kv - def get_cpu_copy(self, indices, mamba_indices=None): return "kv" @@ -37,6 +34,14 @@ def load_cpu_copy(self, cpu_tensors, indices, mamba_indices=None): self.loaded_kv = cpu_tensors +class _Allocator: + def __init__(self, carries_mamba: bool): + self._kv = _KVPool(carries_mamba) + + def get_kvcache(self): + return self._kv + + def _req_and_pool(): req = object.__new__(Req) req.kv = ReqKvInfo(req_pool_idx=0)