[mem_cache][12/N] refactor: stop allocators round-tripping through the KV pool

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,

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")
 

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)
         )

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
-        )

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."""

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
-        )

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"

test/registered/unit/mem_cache/test_retraction_mamba_backup.py

@@ -22,21 +22,26 @@ 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"
 
     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)