[unified-memory] Hierarchical cache for every unified pool shape (#37507)

python/sglang/srt/arg_groups/kv_cache_hook.py

@@ -507,12 +507,10 @@ def handle_unified_memory_pool(server_args: Any) -> None:
         "write loc, so a captured decode replay raises. "
         "TODO(ch-wan): carry out_cache_loc_virtual into the child view."
     )
-    assert not (cfg.enable_hierarchical_cache or cfg.enable_lmcache), (
+    assert not cfg.enable_lmcache, (
-        "--enable-unified-memory is not yet compatible with hierarchical / "
+        "--enable-unified-memory is not yet compatible with --enable-lmcache: "
-        "host-tiered KV cache (--enable-hierarchical-cache / --enable-lmcache): "
+        "the LMCache offload path indexes the device buffers with the ids it "
-        "the unified-memory-pool init wires up no host pools, and its device mamba / "
+        "is handed, and under the unified pool those are VIRTUAL."
-        "full-attention slots are VIRTUAL — the host-offload path does not "
-        "translate them to physical."
     )
     if cfg.dcp_size > 1:
         _validate_unified_memory_dcp(server_args)

python/sglang/srt/managers/cache_controller.py

@@ -790,6 +790,15 @@ class HiCacheController:
             self.prefetch_sync_thread.start()
             self.backup_thread.start()
 
+    def has_inflight_device_transfers(self) -> bool:
+        """Whether queued or unacknowledged L2 transfers still use device rows."""
+        return bool(
+            self.write_queue
+            or self.load_queue
+            or self.ack_write_queue
+            or self.ack_load_queue
+        )
+
     def write(
         self,
         device_indices: torch.Tensor,

python/sglang/srt/managers/scheduler.py

@@ -595,6 +595,12 @@ class Scheduler(
                 cache_controller.load_fence_stream = (
                     self.tp_worker.model_runner.forward_stream
                 )
+                if self.enable_unified_memory:
+                    # Keep device rows stable until host transfers are acknowledged.
+                    # Queue reads and relocation both run on the scheduler thread.
+                    self.token_to_kv_pool_allocator.set_host_transfer_move_gate(
+                        lambda c=cache_controller: not c.has_inflight_device_transfers()
+                    )
         self.emit_metrics_constants()
         self.maybe_init_hccl_dp_prewarm()
 

python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py

@@ -152,6 +152,21 @@ class UnifiedSWAAllocatorBase(SWATokenToKVPoolAllocator):
             full_allocator=self.full_attn_allocator,
             swa_allocator=self.swa_attn_allocator,
         )
+        # Size host pools in tokens; sub-pool `size` counts kernel-facing rows.
+        kvcache.full_kv_pool.host_capacity_tokens = self._size_full
+        kvcache.swa_kv_pool.host_capacity_tokens = self._size_swa
+        for name, pool in (
+            ("full", kvcache.full_kv_pool),
+            ("swa", kvcache.swa_kv_pool),
+        ):
+            pool.host_capacity_bytes = (
+                pool.host_capacity_tokens * unified_buffer.spec(name).entry_bytes()
+            )
+        # Full-attention transfers use virtual IDs. SWA transfers already use
+        # kernel-facing IDs from translate_loc_from_full_to_swa.
+        kvcache.full_kv_pool.host_transfer_translate = (
+            self.full_attn_allocator.translate_kv_loc_for_kernel
+        )
 
         self.free_group = None
         self.free_page_reps_group: Optional[List[torch.Tensor]] = None
@@ -329,6 +344,43 @@ class UnifiedSWAAllocatorBase(SWATokenToKVPoolAllocator):
             lazy_compaction=self.lazy_compaction,
         )
 
+    def bind_swa_for_loaded_rows(
+        self, full_token_ids: torch.Tensor
+    ) -> Optional[torch.Tensor]:
+        """Bind SWA pages to resident or newly loaded full-attention virtual IDs.
+
+        Bind before translating: unbound pages translate to the padding sink.
+        Return kernel-facing IDs, or None if capacity cannot be reclaimed.
+        """
+        ids = full_token_ids.to(torch.int64)
+        if ids.numel() == 0:
+            return ids
+        ps = self.page_size
+        pages = torch.unique(ids // ps)
+        # Tombstones (-1) and the padding sink (0) both need a physical binding.
+        unbound = pages[self.swa_attn_allocator.virtual_to_physical[pages] <= 0]
+        need = int(unbound.numel()) * ps
+        if need:
+            if need > self.swa_available_size():
+                return None
+            if (
+                need > self.swa_attn_allocator.available_size()
+                and not _relieve_for_alloc(self.swa_attn_allocator, need)
+            ):
+                return None
+            self.swa_attn_allocator.alloc_with_virtual(unbound)
+        return self.translate_loc_from_full_to_swa(ids)
+
+    def set_host_transfer_move_gate(self, gate: Callable[[], bool]) -> None:
+        """Block page relocation while host transfers use resolved device indices."""
+        install_move_gate(
+            self._move_gate_targets(),
+            slot="host_transfer_move_gate",
+            gate=gate,
+            feature="HiCache",
+            lazy_compaction=self.lazy_compaction,
+        )
+
     def translate_kv_loc_for_kernel(
         self,
         loc: torch.Tensor,
@@ -648,8 +700,7 @@ class UnifiedSWAAllocatorBase(SWATokenToKVPoolAllocator):
     def set_full_to_swa_mapping(
         self, full_indices: torch.Tensor, swa_indices: torch.Tensor
     ) -> None:
-        """No-op stub for HiCache load-back: in shared mode the swa v2p IS the
+        """Binding load-back rows already updates the shared SWA v2p mapping."""
-        mapping, and HiCache for shared SWA is out of scope."""
         return
 
     def clear_full_to_swa_mapping(self, full_indices: torch.Tensor) -> None:
@@ -999,7 +1050,7 @@ class UnifiedSWATokenToKVPoolAllocator(UnifiedSWAAllocatorBase):
 
     def _compaction_allowed(self) -> bool:
         return all(
-            allocator.disagg_move_gate is None or allocator.disagg_move_gate()
+            not allocator.moves_blocked()
             for allocator in (self.full_attn_allocator, self.swa_attn_allocator)
         )
 

python/sglang/srt/mem_cache/allocator/unified_mamba.py

@@ -127,6 +127,11 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
             self.mamba_allocator.available_size(),
         )
 
+        # HiCache indexes the full sub-pool's per-layer views with kernel-facing IDs.
+        kvcache.full_kv_pool.host_transfer_translate = (
+            self.full_attn_allocator.translate_kv_loc_for_kernel
+        )
+
     # -- size: dynamic --
     @property
     def size(self) -> int:
@@ -346,6 +351,16 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
             lazy_compaction=self.lazy_compaction,
         )
 
+    def set_host_transfer_move_gate(self, gate: Callable[[], bool]) -> None:
+        """Block page relocation while host transfers use resolved device indices."""
+        install_move_gate(
+            self._move_gate_targets(),
+            slot="host_transfer_move_gate",
+            gate=gate,
+            feature="HiCache",
+            lazy_compaction=self.lazy_compaction,
+        )
+
     def is_slot_allocated(self, slot: int) -> bool:
         return self.full_attn_allocator.is_slot_allocated(slot)
 

python/sglang/srt/mem_cache/allocator/unified_sub_pool.py

@@ -418,8 +418,10 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
             _STATS_INSTANCES.add(self)
             _install_signal_handlers_once()
         self.live_page_count = 0
-        # While this returns False, `_flush` must not relocate any page.
+        # RDMA and HiCache install independent gates to protect published device
+        # addresses. Either returning False blocks page relocation in `_flush`.
         self.disagg_move_gate: Optional[Callable[[], bool]] = None
+        self.host_transfer_move_gate: Optional[Callable[[], bool]] = None
         self._latest_forward_done_event: Optional[torch.cuda.Event] = None
         # Most-recent forward's (done_event, out_cache_loc_virtual) for `_flush`'s
         # write-race check. Single slot: at most ONE forward in flight per call
@@ -436,7 +438,7 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
         # schedulers read them O(queue) times per step.
         self._avail_memo_epoch: Optional[int] = None
         self._avail_memo_tokens: int = 0
-        self._sched_avail_memo_epoch: Optional[int] = None
+        self._sched_avail_memo_key: Optional[tuple] = None
         self._sched_avail_memo_tokens: int = 0
 
         self.clear()
@@ -610,7 +612,7 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
                     f"[{self.sub_pool_name}] stale available_size memo: "
                     f"cached={self._avail_memo_tokens}, actual={actual}"
                 )
-        if self._sched_avail_memo_epoch == epoch:
+        if self._sched_avail_memo_key == self._schedulable_capacity_key():
             actual = self._available_tokens(
                 extra_gap_bytes=self._peer_drainable_hole_bytes()
             )
@@ -725,22 +727,38 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
         neighbor = self._growth_side_neighbor()
         if neighbor is None or not neighbor.lazy_compaction:
             return 0
-        if neighbor.disagg_move_gate is not None and not neighbor.disagg_move_gate():
+        if neighbor.moves_blocked():
-            # Not realizable: a PD transfer blocks the neighbour's compaction, so
+            # A blocked neighbor cannot reclaim holes to satisfy an allocation.
-            # crediting these bytes would admit work no flush can satisfy.
             return 0
         return len(neighbor._free_phys_pages) * neighbor.entry_bytes_per_page
 
+    def moves_blocked(self) -> bool:
+        """Whether any installed gate currently forbids relocating pages."""
+        for gate in (self.disagg_move_gate, self.host_transfer_move_gate):
+            if gate is not None and not gate():
+                return True
+        return False
+
+    def _schedulable_capacity_key(self) -> tuple:
+        gates = [self.moves_blocked()]
+        for direction in ("low_peer", "high_peer"):
+            neighbor = getattr(self, direction)
+            while neighbor is not None:
+                gates.append(neighbor.moves_blocked())
+                neighbor = getattr(neighbor, direction)
+        return self._chain_capacity_epoch(), tuple(gates)
+
     def schedulable_available_size(self) -> int:
-        """Tokens allocatable AFTER a neighbor urgent-flush; alloc gates use
+        """Tokens allocatable after flushing a neighbor, including reclaimable holes.
-        `available_size()` instead. Memoized on the chain capacity epoch.
+
+        Allocation checks use available_size(). Cache by capacity and gate state.
         """
-        epoch = self._chain_capacity_epoch()
+        key = self._schedulable_capacity_key()
-        if self._sched_avail_memo_epoch != epoch:
+        if self._sched_avail_memo_key != key:
             self._sched_avail_memo_tokens = self._available_tokens(
                 extra_gap_bytes=self._peer_drainable_hole_bytes()
             )
-            self._sched_avail_memo_epoch = epoch
+            self._sched_avail_memo_key = key
         return self._sched_avail_memo_tokens
 
     def _flush_targets(self):
@@ -1424,9 +1442,10 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
             self._compact_pending_impl(freed_physical_pages)
 
     def _compact_pending_impl(self, freed_physical_pages: torch.Tensor) -> None:
-        assert self.disagg_move_gate is None, (
+        assert self.disagg_move_gate is None and self.host_transfer_move_gate is None, (
-            f"_compact_pending({self.sub_pool_name!r}): eager compaction ran with "
+            f"_compact_pending({self.sub_pool_name!r}): eager compaction ran "
-            "a PD-disaggregation move gate installed; PD requires lazy_compaction."
+            "with a move gate installed; PD disaggregation and HiCache both "
+            "require lazy_compaction."
         )
         freed_set = set(int(x) for x in freed_physical_pages.tolist())
         if not freed_set:
@@ -1780,7 +1799,7 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
         """
         if not self.lazy_compaction:
             return 0
-        if self.disagg_move_gate is not None and not self.disagg_move_gate():
+        if self.moves_blocked():
             # Holes stay in the free list; the next flush picks them up.
             return 0
         self._stats_n_flush_calls += 1
@@ -2170,7 +2189,7 @@ class FloatMultiEndedAllocator(MultiEndedAllocator):
             p = p.low_peer if side == "low" else p.high_peer
         if p is None or not p.lazy_compaction:
             return 0
-        if p.disagg_move_gate is not None and not p.disagg_move_gate():
+        if p.moves_blocked():
             return 0
         return len(p._free_phys_pages) * p.entry_bytes_per_page
 

python/sglang/srt/mem_cache/hicache_storage.py

@@ -122,6 +122,9 @@ class PoolTransfer:
     hit_policy: PoolHitPolicy = PoolHitPolicy.ALL_PAGES
     nodes_to_load: Optional[List[Any]] = None
     indices_from_pool: Optional[PoolName] = None
+    # Full IDs backing a dependent device allocation: resident tensors or
+    # slices of the full rows allocated by this load, in transfer order.
+    anchor_index_parts: Optional[List[torch.Tensor | slice]] = None
 
 
 @dataclass(frozen=True)

python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py

@@ -1336,6 +1336,7 @@ class HybridCacheController(BaseHiCacheController):
             return None
         newly_allocated: list[tuple[PoolTransfer, Callable, torch.Tensor]] = []
         derived_transfers: list[PoolTransfer] = []
+        anchor_transfers = []
 
         def rollback_allocated() -> None:
             for prev_pool, prev_free_fn, prev_indices in newly_allocated:
@@ -1351,6 +1352,11 @@ class HybridCacheController(BaseHiCacheController):
                 continue
             if pool.device_indices is not None or pool.host_indices is None:
                 continue
+            if entry.device_indices_from_anchor_fn is not None:
+                # Allocate independent pools first: their allocation/eviction
+                # can compact SWA before its kernel-facing IDs are captured.
+                anchor_transfers.append((pool, entry))
+                continue
             # device_alloc_fn / device_free_fn override entry.device_pool's
             # methods for pools whose device_pool is a raw KV pool (layout)
             # rather than an allocator (e.g. SWA).
@@ -1369,6 +1375,28 @@ class HybridCacheController(BaseHiCacheController):
             pool.device_indices = indices
             newly_allocated.append((pool, free_fn, indices))
 
+        for pool, entry in anchor_transfers:
+            if kv_device_indices is None or not pool.anchor_index_parts:
+                rollback_allocated()
+                return None
+            anchor_indices = torch.cat(
+                [
+                    kv_device_indices[part] if isinstance(part, slice) else part
+                    for part in pool.anchor_index_parts
+                ]
+            )
+            assert len(anchor_indices) == len(pool.host_indices)
+            bind = entry.device_indices_from_anchor_fn
+            indices = bind(anchor_indices)
+            if indices is None and entry.device_evict_fn:
+                entry.device_evict_fn(len(anchor_indices))
+                indices = bind(anchor_indices)
+            if indices is None:
+                rollback_allocated()
+                return None
+            pool.device_indices = indices
+            newly_allocated.append((pool, entry.device_free_fn, anchor_indices))
+
         # Assign indices to deferred pools from their source.
         for pool in derived_transfers:
             if pool.indices_from_pool == PoolName.KV:

python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py

@@ -181,7 +181,9 @@ def _split_hicache_size(
 ) -> tuple[float, ...]:
     device_pool_sizes = []
     for kv_pool in kv_pools:
-        size_bytes = kv_pool.get_kv_size_bytes()
+        size_bytes = getattr(kv_pool, "host_capacity_bytes", None)
+        if size_bytes is None:
+            size_bytes = kv_pool.get_kv_size_bytes()
         device_pool_sizes.append(
             sum(size_bytes) if isinstance(size_bytes, tuple) else size_bytes
         )
@@ -204,6 +206,7 @@ def build_pool_entry(
     device_evict_fn: Optional[Callable[[int], Any]] = None,
     device_alloc_fn: Optional[Callable[[int], Any]] = None,
     device_free_fn: Optional[Callable[[Any], Any]] = None,
+    device_indices_from_anchor_fn: Optional[Callable[[Any], Any]] = None,
     packed_draft_device_pools: tuple[Any, ...] = (),
 ) -> PoolEntry:
     return PoolEntry(
@@ -216,6 +219,7 @@ def build_pool_entry(
         device_evict_fn=device_evict_fn,
         device_alloc_fn=device_alloc_fn,
         device_free_fn=device_free_fn,
+        device_indices_from_anchor_fn=device_indices_from_anchor_fn,
         packed_draft_device_pools=packed_draft_device_pools,
     )
 
@@ -263,6 +267,19 @@ def build_kv_only_group(
     )
 
 
+def _swa_allocation_callbacks(allocator, bind=None, free_bound=None) -> dict:
+    """Keep allocation and rollback in the same ID space for every SWA stack."""
+    if bind is not None:
+        assert free_bound is not None
+        return dict(
+            device_indices_from_anchor_fn=bind,
+            device_free_fn=free_bound,
+        )
+    if allocator is None:
+        return {}
+    return dict(device_alloc_fn=allocator.alloc, device_free_fn=allocator.free)
+
+
 def build_hybrid_swa_group(
     *,
     page_size: int,
@@ -276,6 +293,8 @@ def build_hybrid_swa_group(
     host_swa_evict_fn: Optional[Callable[[int], Any]] = None,
     device_swa_evict_fn: Optional[Callable[[int], Any]] = None,
     swa_attn_allocator: Any = None,
+    swa_indices_from_anchor_fn: Optional[Callable[[Any], Any]] = None,
+    swa_free_from_anchor_fn: Optional[Callable[[Any], Any]] = None,
     mtp_swa_device_pools: tuple[Any, ...] = (),
 ) -> HostPoolGroup:
     """Anchor (full) + SWA host pool group for a hybrid-SWA device pool."""
@@ -322,11 +341,10 @@ def build_hybrid_swa_group(
                 transfer_layer_id_max=transfer_layer_id_max + len(mtp_swa_device_pools),
                 host_evict_fn=host_swa_evict_fn,
                 device_evict_fn=device_swa_evict_fn,
-                device_alloc_fn=(
+                **_swa_allocation_callbacks(
-                    swa_attn_allocator.alloc if swa_attn_allocator is not None else None
+                    swa_attn_allocator,
-                ),
+                    swa_indices_from_anchor_fn,
-                device_free_fn=(
+                    swa_free_from_anchor_fn,
-                    swa_attn_allocator.free if swa_attn_allocator is not None else None
                 ),
                 packed_draft_device_pools=mtp_swa_device_pools,
             ),
@@ -423,6 +441,17 @@ def build_hybrid_swa_stack(
         device_swa_evict_fn=device_swa_evict_fn,
         # For SWA hybrid, device allocation goes through the inner allocator.
         swa_attn_allocator=params.token_to_kv_pool_allocator.swa_attn_allocator,
+        # Unified SWA binds pages to the full pool's virtual IDs instead.
+        swa_indices_from_anchor_fn=(
+            params.token_to_kv_pool_allocator.bind_swa_for_loaded_rows
+            if get_memory().enable_unified_memory
+            else None
+        ),
+        swa_free_from_anchor_fn=(
+            params.token_to_kv_pool_allocator.free_swa
+            if get_memory().enable_unified_memory
+            else None
+        ),
         mtp_swa_device_pools=mtp_swa_device_pools,
     )
     cache_controller = HybridCacheController(
@@ -1226,8 +1255,15 @@ def build_hybrid_mamba_swa_stack(
             transfer_layer_id_max=transfer_layer_id_max,
             host_evict_fn=host_swa_evict_fn,
             device_evict_fn=device_swa_evict_fn,
-            device_alloc_fn=swa_attn_allocator.alloc,
+            **_swa_allocation_callbacks(
-            device_free_fn=swa_attn_allocator.free,
+                swa_attn_allocator,
+                params.token_to_kv_pool_allocator.bind_swa_for_loaded_rows
+                if get_memory().enable_unified_memory
+                else None,
+                params.token_to_kv_pool_allocator.free_swa
+                if get_memory().enable_unified_memory
+                else None,
+            ),
         ),
         build_pool_entry(
             name=PoolName.MAMBA,

python/sglang/srt/mem_cache/l2_transfer.py

@@ -54,21 +54,49 @@ class L2TransferEngine:
         self.device_to_host_stream = device_module.Stream()
         self.host_to_device_stream = device_module.Stream()
 
+    @staticmethod
+    def _resolve_device_indices(transfer: L2Transfer) -> torch.Tensor:
+        """Resolve controller IDs into this pool's device-buffer indices.
+
+        Call on the transfer stream after the producer event. The host-transfer
+        move gate prevents relocation until the transfer completes.
+        """
+        # Mamba state pools do not inherit KVCache's default attributes.
+        translate = getattr(transfer.device_pool, "host_transfer_translate", None)
+        if translate is None:
+            return transfer.device_indices
+        original_device = transfer.device_indices.device
+        # The direct backend supplies CPU indices even for a CUDA pool.
+        # Translate alongside the v2p table, then restore the backend's device.
+        indices = transfer.device_indices.to(
+            getattr(transfer.device_pool, "device", original_device)
+        ).contiguous()
+        dcp_size = getattr(transfer.host_pool, "dcp_size", 1)
+        if dcp_size > 1:
+            # The MLA host pool selects this rank and collapses logical IDs.
+            # Translate in local virtual space, then preserve that widened
+            # interface (including token order) for the host pool.
+            resolved = translate(indices // dcp_size) * dcp_size + indices % dcp_size
+        else:
+            resolved = translate(indices)
+        return resolved.to(original_device)
+
     def submit_device_to_host(self, transfers: list[L2Transfer]) -> TransferCompletion:
         start_event = self._start_event(None)
         ack_start, ack_finish, timing_enabled = make_timing_event_pair()
         with device_module.stream(self.device_to_host_stream):
             start_event.wait(self.device_to_host_stream)
+            device_indices = [self._resolve_device_indices(t) for t in transfers]
             ack_start.record()
-            for transfer in transfers:
+            for transfer, dev_idx in zip(transfers, device_indices):
                 transfer.host_pool.backup_from_device_all_layer(
                     transfer.device_pool,
                     transfer.host_indices,
-                    transfer.device_indices,
+                    dev_idx,
                     self.io_backend,
                 )
             ack_finish.record()
-            self._record_stream(transfers, self.device_to_host_stream)
+            self._record_stream(transfers, self.device_to_host_stream, device_indices)
         return TransferCompletion(ack_start, ack_finish, timing_enabled)
 
     def submit_host_to_device(
@@ -84,9 +112,10 @@ class L2TransferEngine:
         primary = transfers[0] if transfers else None
         with device_module.stream(self.host_to_device_stream):
             start_event.wait(self.host_to_device_stream)
+            device_indices = [self._resolve_device_indices(t) for t in transfers]
             ack_start.record()
             for layer_id in range(transfer_layer_id_max):
-                for transfer in transfers:
+                for transfer, dev_idx in zip(transfers, device_indices):
                     local_layer_id = (
                         transfer.layer_mapper(layer_id)
                         if transfer.layer_mapper is not None
@@ -101,7 +130,7 @@ class L2TransferEngine:
                     transfer.host_pool.load_to_device_per_layer(
                         transfer.device_pool,
                         transfer.host_indices,
-                        transfer.device_indices,
+                        dev_idx,
                         local_layer_id,
                         self.io_backend,
                         is_draft=transfer.is_draft,
@@ -109,7 +138,7 @@ class L2TransferEngine:
                 if on_layer_done is not None:
                     on_layer_done(layer_id)
             ack_finish.record()
-            self._record_stream(transfers, self.host_to_device_stream)
+            self._record_stream(transfers, self.host_to_device_stream, device_indices)
         return TransferCompletion(ack_start, ack_finish, timing_enabled)
 
     @staticmethod
@@ -120,8 +149,12 @@ class L2TransferEngine:
         return start_event
 
     @staticmethod
-    def _record_stream(transfers: list[L2Transfer], stream) -> None:
+    def _record_stream(transfers: list[L2Transfer], stream, resolved=()) -> None:
+        tensors = []
         for transfer in transfers:
-            for indices in (transfer.host_indices, transfer.device_indices):
+            tensors.extend((transfer.host_indices, transfer.device_indices))
-                if indices.is_cuda:
+        # Keep temporary translated indices alive until the transfer completes.
-                    indices.record_stream(stream)
+        tensors.extend(resolved)
+        for indices in tensors:
+            if indices is not None and indices.is_cuda:
+                indices.record_stream(stream)

python/sglang/srt/mem_cache/memory_pool.py

@@ -1924,6 +1924,15 @@ class KVCache(abc.ABC):
     ) -> None:
         raise NotImplementedError()
 
+    # Optional translation from controller IDs to this pool's buffer indices.
+    # L2TransferEngine resolves it on the transfer stream; move gates prevent
+    # relocation until the transfer is acknowledged.
+    host_transfer_translate: Optional[Callable[[torch.Tensor], torch.Tensor]] = None
+    # Token capacity for host sizing; `size` may count rows in per-layer views.
+    host_capacity_tokens: Optional[int] = None
+    # Host-budget weight; get_kv_size_bytes may be zero for shared-buffer views.
+    host_capacity_bytes: Optional[int] = None
+
     def register_layer_transfer_counter(self, layer_transfer_counter: LayerDoneCounter):
         self.layer_transfer_counter = layer_transfer_counter
 

python/sglang/srt/mem_cache/pool_host/base.py

@@ -170,26 +170,31 @@ class HostKVCache(abc.ABC):
 
         self.dtype = device_pool.store_dtype
         self.size_per_token = self.get_size_per_token()
+        # Unified pools report token capacity separately from their buffer-row count.
+        device_capacity = getattr(device_pool, "host_capacity_tokens", None)
+        if device_capacity is None:
+            device_capacity = device_pool.size
+        self.device_capacity_tokens = device_capacity
         if host_size > 0:
             self.size = sync_fixed_hicache_size(
                 int(host_size * 1e9 // self.size_per_token), host_size
             )
         else:
-            self.size = int(device_pool.size * host_to_device_ratio)
+            self.size = int(device_capacity * host_to_device_ratio)
         # Align up the host memory pool size to the page size
         self.page_num = self.size // self.page_size + 1
         self.size = self.page_num * self.page_size
         self.start_layer = device_pool.start_layer
         self.end_layer = device_pool.end_layer
 
-        if self.size <= device_pool.size:
+        if self.size <= device_capacity:
             logger.warning(
                 "HiCache %s host pool (%d tokens) is smaller than the device pool (%d tokens);"
                 "L2 cache effectiveness is reduced."
                 "Consider increasing --hicache-ratio (or --hicache-size) for higher L2 cache hit rate.",
                 pool_label,
                 self.size,
-                device_pool.size,
+                device_capacity,
             )
 
         # Verify there is enough available host memory.

python/sglang/srt/mem_cache/pool_host/group.py

@@ -22,6 +22,10 @@ class PoolEntry:
     device_evict_fn: Callable[[int], Any] | None = None
     device_alloc_fn: Callable[[int], Any] | None = None
     device_free_fn: Callable[[Any], Any] | None = None
+    # Bind rows to the anchor's virtual IDs when pools share an ID space.
+    # Return buffer indices, or None if allocation fails. Rollback through
+    # device_free_fn takes the anchor's virtual IDs, not the returned indices.
+    device_indices_from_anchor_fn: Callable[[Any], Any] | None = None
     packed_draft_device_pools: tuple[Any, ...] = ()
 
 

python/sglang/srt/mem_cache/pool_host/mamba.py

@@ -109,23 +109,26 @@ class MambaPoolHost(HostKVCache):
         self.dtype = self.conv_dtype
         self.size_per_token = self.get_size_per_token()
 
+        device_capacity = getattr(device_pool, "host_capacity_tokens", None)
+        if device_capacity is None:
+            device_capacity = device_pool.size
         if host_size > 0:
             self.size = sync_fixed_hicache_size(
                 int(host_size * 1e9 // self.size_per_token), host_size
             )
         else:
-            self.size = int(device_pool.size * host_to_device_ratio)
+            self.size = int(device_capacity * host_to_device_ratio)
 
         self.page_num = self.size // self.page_size + 1
         self.size = self.page_num * self.page_size
 
-        if self.size <= device_pool.size:
+        if self.size <= device_capacity:
             logger.warning(
                 "HiCache host KV pool (%d tokens) is smaller than the device pool (%d tokens);"
                 "L2 cache effectiveness is reduced."
                 "Consider increasing --hicache-ratio (or --hicache-size) for higher L2 cache hit rate.",
                 self.size,
-                device_pool.size,
+                device_capacity,
             )
 
         requested_bytes = self.size * self.size_per_token
@@ -351,6 +354,15 @@ class MambaPoolHost(HostKVCache):
             return 0
         return int(tensor[0].numel() * tensor.element_size())
 
+    @staticmethod
+    def _slots_are_strided(tensor: torch.Tensor) -> bool:
+        """Whether slot stride differs from the slot size transfer kernels expect."""
+        return (
+            tensor.dim() >= 1
+            and tensor.shape[0] > 0
+            and tensor.stride(0) != tensor[0].numel()
+        )
+
     @staticmethod
     def _copy_tensor(
         src: torch.Tensor,
@@ -361,6 +373,34 @@ class MambaPoolHost(HostKVCache):
     ) -> None:
         if src_indices.numel() == 0:
             return
+        # Unified conv/SSM views span a whole state envelope per slot. Stage
+        # contiguously for transfer kernels; torch indexing respects the strides
+        # and runs on the caller's transfer stream.
+        if MambaPoolHost._slots_are_strided(src):
+            staged = src.index_select(0, src_indices.to(src.device))
+            MambaPoolHost._copy_tensor(
+                staged,
+                dst,
+                torch.arange(staged.shape[0], device=staged.device),
+                dst_indices,
+                io_backend,
+            )
+            return
+        if MambaPoolHost._slots_are_strided(dst):
+            staged = torch.empty(
+                (dst_indices.numel(), *dst.shape[1:]),
+                dtype=dst.dtype,
+                device=dst.device,
+            )
+            MambaPoolHost._copy_tensor(
+                src,
+                staged,
+                src_indices,
+                torch.arange(staged.shape[0], device=staged.device),
+                io_backend,
+            )
+            dst.index_copy_(0, dst_indices.to(dst.device), staged)
+            return
         if io_backend == "kernel":
             # TODO: Rename the interface for clarity.
             # Here, transfer_kv_per_layer_mla is reused to transfer the Mamba state.
@@ -402,6 +442,24 @@ class MambaPoolHost(HostKVCache):
     ) -> None:
         if src_indices.numel() == 0:
             return
+        if MambaPoolHost._slots_are_strided(dst):
+            # Transfer into contiguous staging, then scatter into the strided view.
+            staged = torch.empty(
+                (dst_indices.numel(), *dst.shape[1:]),
+                dtype=dst.dtype,
+                device=dst.device,
+            )
+            MambaPoolHost._copy_tensor_pf_lf(
+                src,
+                staged,
+                src_indices,
+                torch.arange(staged.shape[0], device=staged.device),
+                layer_id,
+                num_layers,
+                io_backend,
+            )
+            dst.index_copy_(0, dst_indices.to(dst.device), staged)
+            return
         if io_backend == "kernel":
             item_size = MambaPoolHost._item_size_per_index(dst)
             # Mamba JIT kernel expects all index tensors on CUDA.
@@ -460,6 +518,31 @@ class MambaPoolHost(HostKVCache):
     ) -> None:
         if src_indices.numel() == 0:
             return
+        if MambaPoolHost._slots_are_strided(src_layers[0]):
+            # Stage contiguous slots per layer and pass the staging buffer's pointers.
+            staged = torch.stack(
+                [
+                    src_layers[i].index_select(0, src_indices.to(src_layers.device))
+                    for i in range(num_layers)
+                ]
+            )
+            staged_ptrs = torch.tensor(
+                [staged[i].data_ptr() for i in range(num_layers)],
+                dtype=torch.uint64,
+                device=staged.device,
+            )
+            MambaPoolHost._copy_tensor_all_layers_lf_pf(
+                staged,
+                dst,
+                torch.arange(staged.shape[1], device=staged.device),
+                dst_indices,
+                num_layers,
+                io_backend,
+                staged_ptrs,
+                staging=staging,
+                can_use_jit=can_use_jit,
+            )
+            return
         if io_backend == "kernel":
             item_size = MambaPoolHost._item_size_per_index(src_layers[0])
             transfer_kv_mamba_lf_pf(

python/sglang/srt/mem_cache/pool_host/mla.py

@@ -185,7 +185,10 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
                 int(host_size * 1e9 // self.size_per_token), host_size
             )
         else:
-            self.size = int(device_pool.size * host_to_device_ratio)
+            self.size = int(
+                (getattr(device_pool, "host_capacity_tokens", None) or device_pool.size)
+                * host_to_device_ratio
+            )
         self.page_num = self.size // self.page_size + 1
         self.size = self.page_num * self.page_size
         self.start_layer = device_pool.start_layer

python/sglang/srt/mem_cache/unified_cache/components/swa.py

@@ -1051,8 +1051,19 @@ class SWAComponent(TreeComponent):
             return [
                 PoolTransfer(
                     name=PoolName.SWA,
-                    device_indices=torch.cat(
+                    device_indices=(
-                        [n.component_data[ct].value for n in unbacked_swa_nodes]
+                        self._translate_full_to_swa(
+                            torch.cat(
+                                [
+                                    n.component_data[BASE_COMPONENT_TYPE].value
+                                    for n in unbacked_swa_nodes
+                                ]
+                            )
+                        )
+                        if self._unified_allocator() is not None
+                        else torch.cat(
+                            [n.component_data[ct].value for n in unbacked_swa_nodes]
+                        )
                     ).to(torch.int64),
                     nodes_to_load=[n.id for n in unbacked_swa_nodes],
                 )

python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py

@@ -2321,6 +2321,21 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
                 nodes_to_load=[],
             )
             return empty_kv, {}
+        # SWA can be evicted independently of FULL, including holes between
+        # resident SWA nodes. Describe precisely which full rows back it.
+        full_load_slices = {}
+        offset = 0
+        for nid in kv_xfer.nodes_to_load or ():
+            count = len(self.node_by_id(nid).key)
+            full_load_slices[nid] = slice(offset, offset + count)
+            offset += count
+        for xfer in comp_xfers.get(ComponentType.SWA, ()):
+            xfer.anchor_index_parts = [
+                full_load_slices[nid]
+                if nid in full_load_slices
+                else self.node_by_id(nid).component_data[BASE_COMPONENT_TYPE].value
+                for nid in xfer.nodes_to_load or ()
+            ]
         return kv_xfer, comp_xfers
 
     def prefetch_anchor_info(

python/sglang/srt/mem_cache/unified_memory_pool.py

@@ -596,6 +596,8 @@ class UnifiedMHATokenToKVPool(MHATokenToKVPool):
     def _create_buffers(self):
         self.k_buffer = self._k_views
         self.v_buffer = self._v_views
+        # This override must initialize the pointers and row strides used by HiCache.
+        self._init_data_ptrs_and_strides()
 
     def _clear_buffers(self):
         # Lifetime owned by UnifiedKVPool; do not delete the views.
@@ -1218,6 +1220,28 @@ def _check_bs1_feasibility_floor(
     )
 
 
+def _wire_mamba_slot_allocator(
+    *,
+    mamba_end,
+    req_to_token_pool,
+    device,
+) -> UnifiedMambaSlotAllocator:
+    """Install Mamba slot allocation, host capacities, and transfer translation."""
+    slot_allocator = UnifiedMambaSlotAllocator(
+        mamba_end,
+        max_size=req_to_token_pool._shared_mamba_size,
+        device=device,
+    )
+    req_to_token_pool.mamba_allocator = slot_allocator
+    state_pool = req_to_token_pool.mamba_pool
+    state_pool.host_transfer_translate = slot_allocator.translate
+    state_pool.host_capacity_tokens = req_to_token_pool._shared_mamba_size
+    state_pool.host_capacity_bytes = (
+        state_pool.host_capacity_tokens * mamba_end.entry_bytes
+    )
+    return slot_allocator
+
+
 def init_unified_mamba_pools(
     *,
     device: str,
@@ -1387,15 +1411,19 @@ def init_unified_mamba_pools(
         forward_stream=forward_stream,
         lazy_compaction=lazy_compaction,
     )
+    # Size host storage from the configured token cap, not the dynamic buffer view.
+    full_pool = token_to_kv_pool.full_kv_pool
+    full_pool.host_capacity_tokens = max_total_num_tokens
+    full_pool.host_capacity_bytes = (
+        max_total_num_tokens * allocator.full_attn_allocator.entry_bytes
+    )
 
-    # Wrap the composite's mamba MultiEndedAllocator in a slot allocator (PHYSICAL view).
+    mamba_slot_allocator = _wire_mamba_slot_allocator(
-    mamba_slot_allocator = UnifiedMambaSlotAllocator(
+        mamba_end=allocator.mamba_allocator,
-        allocator.mamba_allocator,
+        req_to_token_pool=req_to_token_pool,
-        max_size=req_to_token_pool._shared_mamba_size,
         device=device,
     )
-    # Inert: this allocator implements neither reader (see HybridLinearKVPool).
+    # Only HybridLinearKVPool's retraction CPU-copy path uses this hook.
-    req_to_token_pool.mamba_allocator = mamba_slot_allocator
     token_to_kv_pool._mamba_translate = mamba_slot_allocator.translate
     # No full-KV translate hook is wired: both MLA doors now receive
     # KERNEL-FACING ids -- writes from the ForwardBatch rebind, reads
@@ -2060,14 +2088,11 @@ def init_unified_mamba_swa_pools(
         forward_stream=forward_stream,
         lazy_compaction=lazy_compaction,
     )
-    # Wrap the composite's mamba end in the slot allocator (PHYSICAL view) the
+    _wire_mamba_slot_allocator(
-    # radix MambaComponent / model-side sconv reads consume.
+        mamba_end=allocator.mamba_allocator,
-    mamba_slot_allocator = UnifiedMambaSlotAllocator(
+        req_to_token_pool=req_to_token_pool,
-        allocator.mamba_allocator,
-        max_size=req_to_token_pool._shared_mamba_size,
         device=device,
     )
-    req_to_token_pool.mamba_allocator = mamba_slot_allocator
 
     logger.info(
         "[unified-memory-pool] ============================================================"

scripts/ci/rerun_test_groups.json

@@ -11,6 +11,7 @@
     "registered/disaggregation/test_disaggregation_unified_memory.py",
     "registered/e2e/disaggregation/test_disaggregation_unified_memory_swa.py",
     "registered/e2e/disaggregation/test_disaggregation_unified_memory_tri.py",
+    "registered/e2e/hicache/test_hicache_unified_memory.py",
     "registered/e2e/models/test_inkling_unified.py",
     "registered/e2e/models/test_kimi_linear_models.py",
     "registered/e2e/models/test_kimi_linear_unified_memory.py",
@@ -37,6 +38,8 @@
     "registered/unit/mem_cache/test_unified_capacity_memo.py",
     "registered/unit/mem_cache/test_unified_free_no_host_sync.py",
     "registered/unit/mem_cache/test_unified_handout_zeroing.py",
+    "registered/unit/mem_cache/test_unified_hicache_regressions.py",
+    "registered/unit/mem_cache/test_unified_hicache_strided_state.py",
     "registered/unit/mem_cache/test_unified_mamba_views.py",
     "registered/unit/mem_cache/test_unified_mha_views.py",
     "registered/unit/mem_cache/test_unified_mla_gpu_parity.py",

test/registered/disaggregation/test_disaggregation_basic.py

@@ -459,29 +459,11 @@ class TestDisaggregationPauseResumeDecodeRetract(PDDisaggregationServerBase):
             self._run_pause_on_decode_running_batch("retract", weight_update=True)
         )
 
-    async def _get_decode_num_running_reqs(self, session):
-        """Query current decode running_batch size from /v1/loads."""
-        async with session.get(
-            self.decode_url + "/v1/loads?include=core",
-            timeout=aiohttp.ClientTimeout(total=5),
-        ) as resp:
-            resp.raise_for_status()
-            body = await resp.json()
-            return sum(load["num_running_reqs"] for load in body["loads"])
-
-    async def _wait_for_decode_running_batch(self, session, timeout):
-        deadline = asyncio.get_running_loop().time() + timeout
-        while asyncio.get_running_loop().time() < deadline:
-            if await self._get_decode_num_running_reqs(session) > 0:
-                return
-            await asyncio.sleep(0.2)
-
-        self.fail("Timed out waiting for decode running_batch to become non-empty")
-
     async def _run_pause_on_decode_running_batch(self, mode, weight_update=False):
         num_requests = 2
         max_new_tokens = 512
         prompt = "Write a detailed numbered explanation of distributed inference. " * 12
+        decode_started = [asyncio.Event() for _ in range(num_requests)]
 
         async def _post(session, url, json_data, timeout=30):
             async with session.post(
@@ -493,20 +475,37 @@ class TestDisaggregationPauseResumeDecodeRetract(PDDisaggregationServerBase):
                 return await resp.json()
 
         async def _generate(session, request_id):
-            return await _post(
+            async with session.post(
-                session,
                 self.lb_url + "/generate",
-                {
+                json={
                     "text": f"Request {request_id}: {prompt}",
                     "background": True,
+                    "stream": True,
                     "sampling_params": {
                         "temperature": 0,
                         "ignore_eos": True,
                         "max_new_tokens": max_new_tokens,
                     },
                 },
-                timeout=180,
+                timeout=aiohttp.ClientTimeout(total=180),
-            )
+            ) as resp:
+                resp.raise_for_status()
+                response = None
+                async for line in resp.content:
+                    line = line.strip()
+                    if not line.startswith(b"data: "):
+                        continue
+                    data = line[len(b"data: ") :]
+                    if data == b"[DONE]":
+                        break
+                    response = json.loads(data)
+                    self.assertNotIn("error", response)
+                    # Prefill produces the first token. A later token proves this
+                    # request has reached running_batch on the decode worker.
+                    if response["meta_info"]["completion_tokens"] > 1:
+                        decode_started[request_id].set()
+                self.assertIsNotNone(response, "Generation stream returned no output")
+                return response
 
         async with aiohttp.ClientSession() as session:
             tasks = [
@@ -515,12 +514,17 @@ class TestDisaggregationPauseResumeDecodeRetract(PDDisaggregationServerBase):
             decode_paused = False
 
             try:
-                await self._wait_for_decode_running_batch(session, timeout=30)
+                # /v1/loads can still report a previous batch. Wait for every
-                await asyncio.sleep(0.1)
+                # current request to decode so none can arrive in the prealloc
+                # queue after the pause and prevent the weight-update flush.
+                await asyncio.wait_for(
+                    asyncio.gather(*(event.wait() for event in decode_started)),
+                    timeout=30,
+                )
 
                 self.assertTrue(
-                    any(not task.done() for task in tasks),
+                    all(not task.done() for task in tasks),
-                    "All requests finished before decode retract pause was issued.",
+                    "A request finished before decode retract pause was issued.",
                 )
 
                 await _post(
@@ -580,6 +584,9 @@ class TestDisaggregationPauseResumeDecodeRetract(PDDisaggregationServerBase):
             for response in responses:
                 self.assertIn("text", response)
                 self.assertGreater(len(response["text"]), 0)
+                self.assertEqual(
+                    response["meta_info"]["completion_tokens"], max_new_tokens
+                )
 
             self.assertGreater(
                 sum(

test/registered/e2e/hicache/test_hicache_unified_memory.py

@@ -0,0 +1,240 @@
+"""Compare unified-memory HiCache reloads against a resident-cache reference.
+
+Evict a target prefix with distinct filler requests, require a host hit on
+reload, and compare generated text and output logprobs. Both servers use the
+same unified-memory configuration to keep attention reduction order comparable.
+Covers GDN, SWA, tri-pool, and MLA layouts.
+"""
+
+import os
+import time
+import unittest
+
+import requests
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+register_cuda_ci(est_time=600, stage="extra-a", runner_config="2-gpu-large")
+
+_COMMON_ARGS = [
+    "--trust-remote-code",
+    "--enable-unified-memory",
+    "--enable-cache-report",
+    "--max-running-requests",
+    "1",
+    "--context-length",
+    "4096",
+]
+
+# Distinct filler prefixes must evict the target from device memory.
+_SMALL_POOL = ["--max-total-tokens", "8192"]
+
+_PREFIX = (
+    "The following is a detailed technical description of a distributed inference "
+    "system with paged attention, radix prefix caching and hierarchical offload. "
+) * 90
+_TARGET = _PREFIX + " Question one:"
+_CONTINUATION = _TARGET + " Explain how it works."
+
+
+def _generate(base_url, text, max_new_tokens=32, logprobs=True):
+    payload = {
+        "text": text,
+        "sampling_params": {"temperature": 0.0, "max_new_tokens": max_new_tokens},
+    }
+    if logprobs:
+        payload["return_logprob"] = True
+        # Output logprobs suffice; asking for prompt logprobs from zero
+        # caps the reusable prefix at zero and bypasses HiCache entirely.
+        payload["logprob_start_len"] = -1
+    resp = requests.post(f"{base_url}/generate", json=payload, timeout=600)
+    assert resp.status_code == 200, resp.text
+    data = resp.json()
+    lp = (
+        [t[0] for t in data["meta_info"]["output_token_logprobs"]] if logprobs else None
+    )
+    return data["text"], lp, data["meta_info"]
+
+
+class UnifiedMemoryHiCacheBase(CustomTestCase):
+    """Compare identical unified-memory configurations with and without HiCache."""
+
+    model: str = ""
+    extra_args: list = []
+    server_env: dict = {}
+
+    @classmethod
+    def setUpClass(cls):
+        if cls is UnifiedMemoryHiCacheBase:
+            raise unittest.SkipTest("base class")
+        base_args = _COMMON_ARGS + cls.extra_args
+        cls.hicache_url = "http://127.0.0.1:8157"
+        cls.reference_url = "http://127.0.0.1:8158"
+        env = {**os.environ, **cls.server_env} if cls.server_env else None
+        hicache_args = ["--enable-hierarchical-cache"]
+        if "--hicache-size" not in base_args:
+            hicache_args += ["--hicache-ratio", "4"]
+        cls.process_hicache = popen_launch_server(
+            cls.model,
+            cls.hicache_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=base_args + hicache_args,
+            env=env,
+        )
+        cls.addClassCleanup(kill_process_tree, cls.process_hicache.pid)
+        cls.process_reference = popen_launch_server(
+            cls.model,
+            cls.reference_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=base_args + ["--base-gpu-id", "1"],
+            env=env,
+        )
+        cls.addClassCleanup(kill_process_tree, cls.process_reference.pid)
+
+    def _force_host_round_trip(self):
+        """Evict the target off the device so the next hit must come from L2."""
+        for i in range(8):
+            _generate(
+                self.hicache_url,
+                f"Document {i}. "
+                + (f"Unique filler {i} about an unrelated subject. " * 300),
+                max_new_tokens=8,
+                logprobs=False,
+            )
+
+    def _flush_both(self):
+        """Reset cache state to match prefill boundaries and reduction order."""
+        for url in (self.hicache_url, self.reference_url):
+            requests.post(f"{url}/flush_cache", timeout=180)
+        time.sleep(3)
+
+    def test_load_back_matches_no_hicache(self):
+        """Host reloads preserve generated text and logprobs within tolerance."""
+        self._flush_both()
+        cold_text, cold_lp, _ = _generate(self.hicache_url, _TARGET)
+        ref_cold_text, ref_cold_lp, _ = _generate(self.reference_url, _TARGET)
+        self._force_host_round_trip()
+        # Extend the prefix so both servers compute new KV rows. Repeating it
+        # would let only the resident reference reuse its original final-token KV.
+        warm_text, warm_lp, warm_meta = _generate(self.hicache_url, _CONTINUATION)
+        ref_text, ref_lp, _ = _generate(self.reference_url, _CONTINUATION)
+
+        self.assertGreater(
+            (warm_meta.get("cached_tokens_details") or {}).get("host", 0),
+            0,
+            msg=f"Target did not reload from host: {warm_meta}",
+        )
+        self.assertEqual(cold_text, ref_cold_text)
+        self.assertEqual(warm_text, ref_text)
+        # Match the reference's prefill boundary in each comparison: cold
+        # against cold, and an L2 prefix hit against a resident prefix hit.
+        for label, lp, reference in (
+            ("cold", cold_lp, ref_cold_lp),
+            ("after-L2-reload", warm_lp, ref_lp),
+        ):
+            self.assertEqual(len(lp), len(reference))
+            delta = max(abs(a - b) for a, b in zip(lp, reference))
+            self.assertAlmostEqual(
+                delta,
+                0.0,
+                places=5,
+                msg=f"{label} diverged from the no-HiCache reference by {delta}",
+            )
+
+    def test_server_survives_the_round_trip(self):
+        """Cache churn must leave both schedulers healthy."""
+        self._force_host_round_trip()
+        for url in (self.hicache_url, self.reference_url):
+            resp = requests.get(f"{url}/health", timeout=30)
+            self.assertEqual(resp.status_code, 200)
+
+
+class TestUnifiedMemoryHiCacheGDN(UnifiedMemoryHiCacheBase):
+    """MHA full attention with envelope-strided gated-delta-net state."""
+
+    model = "yujiepan/qwen3.5-tiny-random"
+    extra_args = _SMALL_POOL + [
+        "--linear-attn-backend",
+        "triton",
+        "--mamba-backend",
+        "triton",
+        "--max-mamba-cache-size",
+        "8",
+        "--mem-fraction-static",
+        "0.6",
+    ]
+
+
+class TestUnifiedMemoryHiCacheSWA(UnifiedMemoryHiCacheBase):
+    """Hybrid SWA reloads bind pages to the full-attention pool's virtual IDs."""
+
+    model = "yujiepan/gemma-4e-tiny-random"
+    extra_args = _SMALL_POOL + [
+        "--attention-backend",
+        "triton",
+        "--mem-fraction-static",
+        "0.7",
+    ]
+
+
+class TestUnifiedMemoryHiCacheTriPool(UnifiedMemoryHiCacheBase):
+    """Full attention, sliding-window attention, and ShortConv state together."""
+
+    # The test revision is the reduced checkpoint used by Inkling CI.
+    model = "thinkingmachines/Inkling"
+    server_env = {"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}
+    extra_args = _SMALL_POOL + [
+        "--revision",
+        "test",
+        "--attention-backend",
+        "triton",
+        "--page-size",
+        "128",
+        "--mamba-radix-cache-strategy",
+        "extra_buffer",
+        "--swa-full-tokens-ratio",
+        "0.8",
+        "--max-mamba-cache-size",
+        "8",
+        "--mamba-full-memory-ratio",
+        "0.1",
+        "--mem-fraction-static",
+        "0.5",
+        "--cuda-graph-backend-prefill",
+        "disabled",
+        # Bound total host memory across all three component pools.
+        "--hicache-size",
+        "8",
+    ]
+
+
+class TestUnifiedMemoryHiCacheMLA(UnifiedMemoryHiCacheBase):
+    """MLA full attention with KDA state and MLA-specific transfer pointers."""
+
+    model = "yujiepan/kimi-linear-tiny-random"
+    extra_args = _SMALL_POOL + [
+        "--max-mamba-cache-size",
+        "8",
+        "--mem-fraction-static",
+        "0.5",
+        "--linear-attn-backend",
+        "triton",
+        "--mamba-backend",
+        "triton",
+        "--attention-backend",
+        "triton",
+        "--cuda-graph-backend-decode",
+        "disabled",
+        "--cuda-graph-backend-prefill",
+        "disabled",
+    ]
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/disaggregation/test_unified_memory_move_gate.py

@@ -146,15 +146,19 @@ class TestGatedPeerHolesAreNotSchedulable(CustomTestCase):
     """
 
     class _Peer:
-        def __init__(self, gate):
+        def __init__(self, gate, host_gate=None):
             self.lazy_compaction = True
             self._free_phys_pages = [0, 1, 2, 3]  # only len() is read
             self.entry_bytes_per_page = 512
             self.disagg_move_gate = gate
+            self.host_transfer_move_gate = host_gate
 
         def _is_frontier_transparent(self):
             return False
 
+        # Exercise the production predicate when checking each gate.
+        moves_blocked = MultiEndedAllocator.moves_blocked
+
     class _Owner:
         """Stands in for a grow-up END pool: the credit walks the chain from
         `_growth_side_neighbor()`, so the stub must expose what that walk reads,
@@ -167,8 +171,8 @@ class TestGatedPeerHolesAreNotSchedulable(CustomTestCase):
 
         _growth_side_neighbor = MultiEndedAllocator._growth_side_neighbor
 
-    def _credit(self, gate):
+    def _credit(self, gate, host_gate=None):
-        peer = self._Peer(gate)
+        peer = self._Peer(gate, host_gate)
         owner = self._Owner(peer)
         return MultiEndedAllocator._peer_drainable_hole_bytes(owner)
 
@@ -179,6 +183,10 @@ class TestGatedPeerHolesAreNotSchedulable(CustomTestCase):
         self.assertEqual(self._credit(gate=lambda: True), 4 * 512)
         # Gate closed: an urgent flush would move nothing, so credit nothing.
         self.assertEqual(self._credit(gate=lambda: False), 0)
+        # Either the RDMA gate or the HiCache gate can block compaction.
+        self.assertEqual(self._credit(gate=None, host_gate=lambda: True), 4 * 512)
+        self.assertEqual(self._credit(gate=None, host_gate=lambda: False), 0)
+        self.assertEqual(self._credit(gate=lambda: True, host_gate=lambda: False), 0)
 
 
 class TestMoveGateRejectsNonPdNode(CustomTestCase):
@@ -232,9 +240,7 @@ class TestUnifiedAllocatorsPublishTheTransferContract(CustomTestCase):
     # REACHES rather than on what the stub was given.
     _MEMBER_ATTRS = ("full_attn_allocator", "swa_attn_allocator", "mamba_allocator")
 
-    # The members each composite's gate must reach. The tri-pool row is the one
+    # Inherited gate setters must cover every member, including tri-pool Mamba.
-    # that matters: it inherits the setter, so an enumeration written inside
-    # that setter would silently leave the third member ungated.
     _EXPECTED_COVERAGE = {
         "UnifiedMambaTokenToKVPoolAllocator": {
             "full_attn_allocator",
@@ -269,25 +275,26 @@ class TestUnifiedAllocatorsPublishTheTransferContract(CustomTestCase):
         def gate() -> bool:
             return True
 
-        alloc.set_disagg_move_gate(gate)
+        if slot == "disagg_move_gate":
+            alloc.set_disagg_move_gate(gate)
+        else:
+            alloc.set_host_transfer_move_gate(gate)
         return {
             attr
             for attr in self._MEMBER_ATTRS
             if getattr(getattr(alloc, attr), slot) is gate
         }
 
-    def test_the_gate_reaches_every_member(self):
+    def test_every_gate_reaches_every_member(self):
-        """A gate that reaches only some members is not a weaker gate, it is no
+        """Both transfer gates must protect every sub-pool from relocation."""
-        gate: the ungated end relocates its own pages under the very transfer
-        the gate was installed for.
-        """
         for name, expected in self._EXPECTED_COVERAGE.items():
-            with self.subTest(composite=name):
+            for slot in ("disagg_move_gate", "host_transfer_move_gate"):
-                self.assertEqual(
+                with self.subTest(composite=name, slot=slot):
-                    self._members_reached(name, "disagg_move_gate"),
+                    self.assertEqual(
-                    expected,
+                        self._members_reached(name, slot),
-                    f"{name}.disagg_move_gate does not cover every member",
+                        expected,
-                )
+                        f"{name}.{slot} does not cover every member",
+                    )
 
     def test_gate_setters_do_not_enumerate_members_themselves(self):
         """The structural half of the rule above: a setter that names its
@@ -298,10 +305,13 @@ class TestUnifiedAllocatorsPublishTheTransferContract(CustomTestCase):
 
         for name in self._EXPECTED_COVERAGE:
             cls = self._allocator_class(name)
-            with self.subTest(composite=name):
+            for setter in ("set_disagg_move_gate", "set_host_transfer_move_gate"):
-                body = inspect.getsource(cls.set_disagg_move_gate)
+                if setter not in vars(cls):
-                self.assertIn("install_move_gate", body)
+                    continue  # inherited, and the inherited one is checked above
-                self.assertNotIn("_move_gate = ", body)
+                with self.subTest(composite=name, setter=setter):
+                    body = inspect.getsource(getattr(cls, setter))
+                    self.assertIn("install_move_gate", body)
+                    self.assertNotIn("_move_gate = ", body)
 
     def test_swa_composite_translates_the_swa_side_separately(self):
         """The SWA sub-pool runs its OWN compaction, so a full-side physical id

test/registered/unit/mem_cache/test_hicache_dcp_host_pool.py

@@ -31,6 +31,8 @@ WIDENED_PAGE = PHYSICAL_PAGE * DCP_SIZE
 def _fake_mla_device_pool(size: int = 1024) -> SimpleNamespace:
     return SimpleNamespace(
         size=size,
+        # Match KVCache's default: this static pool's size already counts tokens.
+        host_capacity_tokens=None,
         store_dtype=torch.float16,
         kv_lora_rank=8,
         qk_rope_head_dim=4,

test/registered/unit/mem_cache/test_kv_index_translator.py

@@ -27,7 +27,7 @@ import unittest
 from types import SimpleNamespace
 
 import torch
-from test_multi_ended_allocator import _FakeUnifiedSWAKVPool
+from test_multi_ended_allocator import _FakeKVCache, _FakeUnifiedSWAKVPool
 
 from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
     UnifiedSWATokenToKVPoolAllocator,
@@ -673,7 +673,10 @@ class TestWriteLoc(CustomTestCase):
                 )
                 allocator = UnifiedMambaTokenToKVPoolAllocator(
                     unified_buffer=pool,
-                    kvcache=SimpleNamespace(full_kv_pool=None, mamba_pool=None),
+                    kvcache=SimpleNamespace(
+                        full_kv_pool=_FakeKVCache(pool.max_slots("full")),
+                        mamba_pool=_FakeKVCache(pool.max_slots("mamba")),
+                    ),
                     device="cpu",
                     page_size=4,
                 )

test/registered/unit/mem_cache/test_unified_hicache_regressions.py

@@ -0,0 +1,345 @@
+"""Host reloads must preserve ID domains, allocation ownership, and stream order."""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import MagicMock, Mock, patch
+
+import torch
+
+from sglang.srt.layers.dcp.layout import maybe_dcp_kernel_indices
+from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer
+from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
+    HybridCacheController,
+)
+from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import _split_hicache_size
+from sglang.srt.mem_cache.l2_transfer import L2Transfer, L2TransferEngine
+from sglang.srt.mem_cache.pool_host.group import PoolEntry
+from sglang.srt.mem_cache.unified_cache.component_type import ComponentType
+from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=15, stage="extra-a", runner_config="1-gpu-small")
+
+
+class TestHiCacheIndexDomains(unittest.TestCase):
+    def test_dcp_translation_uses_local_virtual_ids(self):
+        # Two non-adjacent pages, relocated to different physical pages.
+        page_size, dcp_size = 4, 2
+        logical = torch.cat((torch.arange(8, 16), torch.arange(24, 32)))
+        v2p = torch.tensor([0, 5, 4, 2])
+
+        def translate(ids):
+            return v2p[ids // page_size] * (page_size * 3) + ids % page_size
+
+        transfer = L2Transfer(
+            SimpleNamespace(dcp_size=dcp_size),
+            SimpleNamespace(host_transfer_translate=translate),
+            logical.clone(),
+            logical,
+        )
+        resolved = L2TransferEngine._resolve_device_indices(transfer)
+        for rank in range(dcp_size):
+            expected = translate(maybe_dcp_kernel_indices(logical, dcp_size, rank))
+            torch.testing.assert_close(
+                maybe_dcp_kernel_indices(resolved, dcp_size, rank), expected
+            )
+
+    def test_fixed_size_uses_host_capacity_for_shared_buffers(self):
+        pools = [
+            SimpleNamespace(host_capacity_bytes=n, get_kv_size_bytes=lambda: (0, 0))
+            for n in (600, 300, 100)
+        ]
+        self.assertEqual(_split_hicache_size(10, tuple(pools)), (6, 3, 1))
+
+
+class TestHostGateCapacity(unittest.TestCase):
+    def test_schedulable_memo_tracks_gate_without_allocator_mutation(self):
+        from test_unified_capacity_memo import _build
+
+        inst, allocator, kvcache = _build(lazy=True)
+        ids = inst._alloc(allocator, kvcache, 8)
+        allocator.free_swa(ids[2:6])
+        full = allocator.full_attn_allocator
+        state = {"open": True}
+        allocator.set_host_transfer_move_gate(lambda: state["open"])
+        epoch = full._chain_capacity_epoch()
+        before = full.schedulable_available_size()
+        state["open"] = False
+        self.assertEqual(full._chain_capacity_epoch(), epoch)
+        self.assertEqual(full.schedulable_available_size(), full.available_size())
+        self.assertGreater(before, full.schedulable_available_size())
+        self.assertFalse(allocator._compaction_allowed())
+        self.assertEqual(allocator.verify_byte_accounting(), [])
+        state["open"] = True
+        self.assertEqual(full.schedulable_available_size(), before)
+
+
+class TestSwaLoadAllocation(unittest.TestCase):
+    def _controller(self, bind, free=None, evict=None):
+        controller = object.__new__(HybridCacheController)
+        entry = PoolEntry(
+            name=PoolName.SWA,
+            host_pool=SimpleNamespace(),
+            device_pool=SimpleNamespace(),
+            layer_mapper=lambda i: i,
+            device_indices_from_anchor_fn=bind,
+            device_free_fn=free or Mock(),
+            device_evict_fn=evict,
+        )
+        controller.mem_pool_host = SimpleNamespace(entry_map={PoolName.SWA: entry})
+        return controller
+
+    def _transfer(self, parts, count):
+        transfer = PoolTransfer(name=PoolName.SWA, host_indices=torch.arange(count))
+        transfer.anchor_index_parts = parts
+        return transfer
+
+    def test_swa_only_load_uses_resident_full_ids(self):
+        bind = Mock(side_effect=lambda x: x + 100)
+        controller = self._controller(bind)
+        transfer = self._transfer([torch.tensor([13, 14])], 2)
+        result = controller._resolve_device_transfers(
+            [transfer], torch.empty(0, dtype=torch.int64)
+        )
+        self.assertIsNotNone(result)
+        torch.testing.assert_close(transfer.device_indices, torch.tensor([113, 114]))
+
+    def test_mixed_load_skips_resident_swa_nodes(self):
+        bind = Mock(side_effect=lambda x: x + 100)
+        controller = self._controller(bind)
+        transfer = self._transfer([torch.tensor([13, 14]), slice(2, 4)], 4)
+        controller._resolve_device_transfers(
+            [transfer], torch.tensor([20, 21, 22, 23, 24, 25])
+        )
+        torch.testing.assert_close(
+            transfer.device_indices, torch.tensor([113, 114, 122, 123])
+        )
+
+    def test_binding_retries_after_eviction(self):
+        bind = Mock(side_effect=[None, torch.tensor([41, 42])])
+        evict = Mock()
+        controller = self._controller(bind, evict=evict)
+        transfer = self._transfer([slice(0, 2)], 2)
+        self.assertIsNotNone(
+            controller._resolve_device_transfers([transfer], torch.tensor([11, 12]))
+        )
+        evict.assert_called_once_with(2)
+        self.assertEqual(bind.call_count, 2)
+
+    def test_rollback_releases_swa_binding_by_virtual_ids(self):
+        free = Mock()
+        controller = self._controller(lambda x: x + 100, free=free)
+        transfer = self._transfer([torch.tensor([13, 14])], 2)
+        # A missing sidecar source fails after SWA has bound its pages.
+        sidecar = PoolTransfer(name=PoolName.MAMBA, indices_from_pool=PoolName.INDEXER)
+        result = controller._resolve_device_transfers(
+            [transfer, sidecar], torch.empty(0, dtype=torch.int64)
+        )
+        self.assertIsNone(result)
+        torch.testing.assert_close(free.call_args.args[0], torch.tensor([13, 14]))
+        self.assertIsNone(transfer.device_indices)
+
+    def test_independent_allocations_precede_kernel_id_resolution(self):
+        events = []
+        controller = self._controller(lambda x: events.append("bind") or x + 100)
+        controller.mem_pool_host.entry_map[PoolName.MAMBA] = PoolEntry(
+            name=PoolName.MAMBA,
+            host_pool=SimpleNamespace(),
+            device_pool=SimpleNamespace(),
+            layer_mapper=lambda i: i,
+            device_alloc_fn=lambda n: events.append("mamba") or torch.arange(n),
+            device_free_fn=Mock(),
+        )
+        swa = self._transfer([slice(0, 2)], 2)
+        mamba = PoolTransfer(name=PoolName.MAMBA, host_indices=torch.arange(1))
+        self.assertIsNotNone(
+            controller._resolve_device_transfers([swa, mamba], torch.tensor([10, 11]))
+        )
+        self.assertEqual(events, ["mamba", "bind"])
+
+    def test_tree_spec_preserves_node_correspondence(self):
+        kv = PoolTransfer(
+            name=PoolName.KV, host_indices=torch.arange(4), nodes_to_load=[2, 3]
+        )
+        swa = PoolTransfer(
+            name=PoolName.SWA, host_indices=torch.arange(4), nodes_to_load=[1, 3]
+        )
+        nodes = {
+            i: SimpleNamespace(
+                id=i,
+                key=[i, i],
+                load_back_pending_id=None,
+                component_data={
+                    ComponentType.FULL: SimpleNamespace(
+                        value=torch.tensor([10, 11]) if i == 1 else None
+                    )
+                },
+            )
+            for i in (1, 2, 3)
+        }
+        full_component = SimpleNamespace(
+            component_type=ComponentType.FULL,
+            build_hicache_transfers=lambda *a, **k: [kv],
+        )
+        swa_component = SimpleNamespace(
+            component_type=ComponentType.SWA,
+            build_hicache_transfers=lambda *a, **k: [swa],
+        )
+        core = SimpleNamespace(
+            node_by_id=nodes.__getitem__,
+            components=[full_component, swa_component],
+            components_by_type={
+                ComponentType.FULL: full_component,
+                ComponentType.SWA: swa_component,
+            },
+        )
+        UnifiedTreeCore.build_load_back_spec(core, 3)
+        parts = swa.anchor_index_parts
+        torch.testing.assert_close(parts[0], torch.tensor([10, 11]))
+        self.assertEqual(parts[1], slice(2, 4))
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "CUDA required")
+class TestTransferStreamOrdering(unittest.TestCase):
+    def test_load_translation_follows_supplied_start_event(self):
+        # The start event precedes translation; transfer must wait for both.
+        engine = L2TransferEngine("kernel")
+        ids = torch.tensor([2, 4, 7], device="cuda")
+        output = torch.full_like(ids, -1)
+        resolved = torch.full_like(ids, -2)
+        torch.cuda.synchronize()
+        start = torch.cuda.Event()
+        start.record()
+
+        def translate(indices):
+            self.assertEqual(torch.cuda.current_stream(), engine.host_to_device_stream)
+            torch.cuda._sleep(20_000_000)
+            resolved.copy_(indices + 100)
+            return resolved
+
+        host = SimpleNamespace(
+            layer_num=1,
+            load_to_device_per_layer=lambda pool, h, d, layer, backend, **kw: (
+                output.copy_(d)
+            ),
+        )
+        transfer = L2Transfer(
+            host,
+            SimpleNamespace(host_transfer_translate=translate),
+            torch.arange(3),
+            ids,
+        )
+        completion = engine.submit_host_to_device(
+            [transfer], transfer_layer_id_max=1, start_event=start
+        )
+        completion.finish_event.synchronize()
+        torch.testing.assert_close(output.cpu(), torch.tensor([102, 104, 107]))
+
+
+class TestSwaBackupAfterCompaction(unittest.TestCase):
+    def test_backup_resolves_current_binding(self):
+        from sglang.srt.mem_cache.unified_cache.components.base import (
+            CacheTransferPhase,
+        )
+        from sglang.srt.mem_cache.unified_cache.components.swa import SWAComponent
+
+        node = SimpleNamespace(
+            component_data={
+                ComponentType.FULL: SimpleNamespace(value=torch.tensor([3, 7])),
+                ComponentType.SWA: SimpleNamespace(value=torch.tensor([103, 107])),
+            },
+            id=1,
+        )
+        component = SimpleNamespace(
+            component_type=ComponentType.SWA,
+            tree_core=SimpleNamespace(has_swa_host_pool=True),
+            _collect_unbacked_swa_nodes=lambda n: [n],
+            _unified_allocator=lambda: object(),
+            _translate_full_to_swa=lambda x: x + 200,
+        )
+        transfers = SWAComponent.build_hicache_transfers(
+            component, node, CacheTransferPhase.BACKUP_HOST
+        )
+        torch.testing.assert_close(
+            transfers[0].device_indices, torch.tensor([203, 207])
+        )
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "CUDA required")
+class TestDirectBackendTranslation(unittest.TestCase):
+    def test_cpu_indices_translate_on_device_and_return_to_cpu(self):
+        mapping = torch.tensor([0, 9, 6], device="cuda")
+
+        def translate(ids):
+            self.assertTrue(ids.is_cuda)
+            return mapping[ids]
+
+        transfer = L2Transfer(
+            SimpleNamespace(),
+            SimpleNamespace(device="cuda", host_transfer_translate=translate),
+            torch.tensor([0, 1]),
+            torch.tensor([2, 1]),
+        )
+        result = L2TransferEngine._resolve_device_indices(transfer)
+        self.assertEqual(result.device.type, "cpu")
+        torch.testing.assert_close(result, torch.tensor([6, 9]))
+
+
+class TestTriPoolAssembly(unittest.TestCase):
+    def test_swa_allocation_and_rollback_match_pool_id_ownership(self):
+        from sglang.srt.mem_cache.hybrid_cache import hybrid_pool_assembler as assembler
+
+        for unified in (False, True):
+            with self.subTest(unified=unified):
+                swa_allocator = SimpleNamespace(alloc=Mock(), free=Mock())
+                composite = SimpleNamespace(
+                    swa_attn_allocator=swa_allocator,
+                    bind_swa_for_loaded_rows=Mock(),
+                    free_swa=Mock(),
+                )
+                params = SimpleNamespace(
+                    token_to_kv_pool_allocator=composite,
+                    req_to_token_pool=SimpleNamespace(
+                        mamba_allocator=SimpleNamespace(alloc=Mock(), free=Mock())
+                    ),
+                )
+                memory = MagicMock(enable_unified_memory=unified, hicache_size=0)
+                host = MagicMock()
+                with (
+                    patch.object(assembler, "get_memory", return_value=memory),
+                    patch.object(
+                        assembler, "_get_allocator_type", return_value="default"
+                    ),
+                    patch.object(assembler, "build_kv_host_pool", return_value=host),
+                    patch.object(assembler, "MambaPoolHost", return_value=host),
+                    patch.object(assembler, "HybridCacheController"),
+                ):
+                    group, _ = assembler.build_hybrid_mamba_swa_stack(
+                        params=params,
+                        full_kv_pool=object(),
+                        swa_kv_pool=object(),
+                        mamba_pool=object(),
+                        full_layer_mapping={0: 0},
+                        swa_layer_mapping={1: 0},
+                        mamba_layer_mapping={2: 0},
+                        page_size=1,
+                        tp_group=None,
+                        load_cache_event=None,
+                        storage_backend=None,
+                    )
+                entry = group.entry_map[PoolName.SWA]
+                if unified:
+                    self.assertIs(
+                        entry.device_indices_from_anchor_fn,
+                        composite.bind_swa_for_loaded_rows,
+                    )
+                    self.assertIs(entry.device_free_fn, composite.free_swa)
+                    self.assertIsNone(entry.device_alloc_fn)
+                else:
+                    self.assertIs(entry.device_alloc_fn, swa_allocator.alloc)
+                    self.assertIs(entry.device_free_fn, swa_allocator.free)
+                    self.assertIsNone(entry.device_indices_from_anchor_fn)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/test_unified_hicache_strided_state.py

@@ -0,0 +1,115 @@
+"""Cover strided Mamba state staging and shared allocator wiring for HiCache."""
+
+import unittest
+
+import torch
+
+from sglang.srt.mem_cache.pool_host.mamba import MambaPoolHost
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=20, suite="base-a-test-cpu")
+
+
+class TestStridedStateDetection(CustomTestCase):
+    def test_contiguous_slots_are_not_strided(self):
+        for shape in ((8, 4), (8, 4, 3), (1, 5)):
+            with self.subTest(shape=shape):
+                self.assertFalse(
+                    MambaPoolHost._slots_are_strided(torch.zeros(shape)),
+                    "a contiguous slot array must take the direct path",
+                )
+
+    def test_envelope_strided_slots_are_detected(self):
+        num_slots, per_slot, envelope = 6, 4, 10
+        raw = torch.zeros(num_slots * envelope)
+        view = torch.as_strided(raw, size=(num_slots, per_slot), stride=(envelope, 1))
+        self.assertTrue(MambaPoolHost._slots_are_strided(view))
+
+    def test_empty_tensor_is_not_strided(self):
+        self.assertFalse(MambaPoolHost._slots_are_strided(torch.zeros((0, 4))))
+
+    def test_staging_round_trip_preserves_slot_contents(self):
+        """Gather and scatter preserve selected slots without changing their neighbors."""
+        num_slots, per_slot, envelope = 6, 4, 10
+        raw = torch.arange(num_slots * envelope, dtype=torch.float32)
+        view = torch.as_strided(raw, size=(num_slots, per_slot), stride=(envelope, 1))
+        indices = torch.tensor([4, 1, 3])
+
+        staged = view.index_select(0, indices)
+        self.assertTrue(staged.is_contiguous())
+        for row, slot in enumerate(indices.tolist()):
+            self.assertTrue(torch.equal(staged[row], view[slot]))
+
+        dst = torch.zeros_like(raw)
+        dst_view = torch.as_strided(
+            dst, size=(num_slots, per_slot), stride=(envelope, 1)
+        )
+        dst_view.index_copy_(0, indices, staged)
+        for slot in indices.tolist():
+            self.assertTrue(torch.equal(dst_view[slot], view[slot]))
+        # A partial transfer must leave unselected slots untouched.
+        for slot in set(range(num_slots)) - set(indices.tolist()):
+            self.assertTrue(torch.all(dst_view[slot] == 0))
+
+
+if __name__ == "__main__":
+    unittest.main()
+
+
+class TestMambaSlotWiringIsShared(CustomTestCase):
+    """Both Mamba factories must install slot allocation and transfer translation.
+
+    Inspect the AST so this check does not need GPU-backed pools.
+    """
+
+    @staticmethod
+    def _assignments_to(attr: str):
+        """Find functions assigning the attribute, excluding None initialization."""
+        import ast
+        import inspect
+
+        from sglang.srt.mem_cache import unified_memory_pool
+
+        tree = ast.parse(inspect.getsource(unified_memory_pool))
+        found = []
+        for fn in ast.walk(tree):
+            if not isinstance(fn, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                continue
+            for node in ast.walk(fn):
+                if not isinstance(node, ast.Assign):
+                    continue
+                if isinstance(node.value, ast.Constant) and node.value.value is None:
+                    continue
+                for tgt in node.targets:
+                    if isinstance(tgt, ast.Attribute) and tgt.attr == attr:
+                        found.append(fn.name)
+        return found
+
+    def test_only_the_shared_hook_wraps_the_slot_allocator(self):
+        self.assertEqual(
+            sorted(set(self._assignments_to("mamba_allocator"))),
+            ["_wire_mamba_slot_allocator"],
+            "a factory wraps the mamba end itself; it will miss "
+            "host_transfer_translate exactly as the tri-pool factory did",
+        )
+
+    def test_every_unified_mamba_factory_calls_the_shared_hook(self):
+        import ast
+        import inspect
+
+        from sglang.srt.mem_cache import unified_memory_pool
+
+        tree = ast.parse(inspect.getsource(unified_memory_pool))
+        # Factories whose composite holds a mamba end.
+        expected = {"init_unified_mamba_pools", "init_unified_mamba_swa_pools"}
+        callers = {
+            fn.name
+            for fn in ast.walk(tree)
+            if isinstance(fn, ast.FunctionDef)
+            for node in ast.walk(fn)
+            if isinstance(node, ast.Call)
+            and isinstance(node.func, ast.Name)
+            and node.func.id == "_wire_mamba_slot_allocator"
+        }
+        self.assertEqual(expected, expected & callers, f"missing: {expected - callers}")