diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index ea516217f..8b7219e80 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/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), ( - "--enable-unified-memory is not yet compatible with hierarchical / " - "host-tiered KV cache (--enable-hierarchical-cache / --enable-lmcache): " - "the unified-memory-pool init wires up no host pools, and its device mamba / " - "full-attention slots are VIRTUAL — the host-offload path does not " - "translate them to physical." + assert not cfg.enable_lmcache, ( + "--enable-unified-memory is not yet compatible with --enable-lmcache: " + "the LMCache offload path indexes the device buffers with the ids it " + "is handed, and under the unified pool those are VIRTUAL." ) if cfg.dcp_size > 1: _validate_unified_memory_dcp(server_args) diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 3d5b5594a..c28184141 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/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, diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 78bac308f..b694ba635 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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() diff --git a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py index 43d0dfb85..c28b28c5f 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py +++ b/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 - mapping, and HiCache for shared SWA is out of scope.""" + """Binding load-back rows already updates the shared SWA v2p mapping.""" 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) ) diff --git a/python/sglang/srt/mem_cache/allocator/unified_mamba.py b/python/sglang/srt/mem_cache/allocator/unified_mamba.py index 75af7b7a5..ce3272e87 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_mamba.py +++ b/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) diff --git a/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py b/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py index 8eb09ec21..3d337b258 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py +++ b/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(): - # Not realizable: a PD transfer blocks the neighbour's compaction, so - # crediting these bytes would admit work no flush can satisfy. + if neighbor.moves_blocked(): + # A blocked neighbor cannot reclaim holes to satisfy an allocation. 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 - `available_size()` instead. Memoized on the chain capacity epoch. + """Tokens allocatable after flushing a neighbor, including reclaimable holes. + + Allocation checks use available_size(). Cache by capacity and gate state. """ - epoch = self._chain_capacity_epoch() - if self._sched_avail_memo_epoch != epoch: + key = self._schedulable_capacity_key() + 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, ( - f"_compact_pending({self.sub_pool_name!r}): eager compaction ran with " - "a PD-disaggregation move gate installed; PD requires lazy_compaction." + 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 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 diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index be69c4803..04e5ab8c8 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/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) diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index 798646d86..f592b6973 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/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: diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index 1e91f0364..60f281ba3 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/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_attn_allocator.alloc if swa_attn_allocator is not None else None - ), - device_free_fn=( - swa_attn_allocator.free if swa_attn_allocator is not None else None + **_swa_allocation_callbacks( + swa_attn_allocator, + swa_indices_from_anchor_fn, + swa_free_from_anchor_fn, ), 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, - device_free_fn=swa_attn_allocator.free, + **_swa_allocation_callbacks( + 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, diff --git a/python/sglang/srt/mem_cache/l2_transfer.py b/python/sglang/srt/mem_cache/l2_transfer.py index abb4fa898..0cd7e5e30 100644 --- a/python/sglang/srt/mem_cache/l2_transfer.py +++ b/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): - if indices.is_cuda: - indices.record_stream(stream) + tensors.extend((transfer.host_indices, transfer.device_indices)) + # Keep temporary translated indices alive until the transfer completes. + tensors.extend(resolved) + for indices in tensors: + if indices is not None and indices.is_cuda: + indices.record_stream(stream) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 1ffc80733..9d138da11 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/pool_host/base.py b/python/sglang/srt/mem_cache/pool_host/base.py index 8c4a6674e..c2d4c239e 100644 --- a/python/sglang/srt/mem_cache/pool_host/base.py +++ b/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. diff --git a/python/sglang/srt/mem_cache/pool_host/group.py b/python/sglang/srt/mem_cache/pool_host/group.py index 12dd6f4eb..48fa88775 100644 --- a/python/sglang/srt/mem_cache/pool_host/group.py +++ b/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, ...] = () diff --git a/python/sglang/srt/mem_cache/pool_host/mamba.py b/python/sglang/srt/mem_cache/pool_host/mamba.py index 328443837..79accc17e 100644 --- a/python/sglang/srt/mem_cache/pool_host/mamba.py +++ b/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( diff --git a/python/sglang/srt/mem_cache/pool_host/mla.py b/python/sglang/srt/mem_cache/pool_host/mla.py index 86abe18c9..b60325f8b 100644 --- a/python/sglang/srt/mem_cache/pool_host/mla.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/unified_cache/components/swa.py b/python/sglang/srt/mem_cache/unified_cache/components/swa.py index f125b6148..0fdac250a 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/swa.py +++ b/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( - [n.component_data[ct].value for n in unbacked_swa_nodes] + device_indices=( + 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], ) diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index 0250b07ab..e3f0554e4 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/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( diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index 69ca18b3a..3cca20a7d 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/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 = UnifiedMambaSlotAllocator( - allocator.mamba_allocator, - max_size=req_to_token_pool._shared_mamba_size, + mamba_slot_allocator = _wire_mamba_slot_allocator( + mamba_end=allocator.mamba_allocator, + req_to_token_pool=req_to_token_pool, device=device, ) - # Inert: this allocator implements neither reader (see HybridLinearKVPool). - req_to_token_pool.mamba_allocator = mamba_slot_allocator + # Only HybridLinearKVPool's retraction CPU-copy path uses this hook. 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 - # radix MambaComponent / model-side sconv reads consume. - mamba_slot_allocator = UnifiedMambaSlotAllocator( - allocator.mamba_allocator, - max_size=req_to_token_pool._shared_mamba_size, + _wire_mamba_slot_allocator( + mamba_end=allocator.mamba_allocator, + req_to_token_pool=req_to_token_pool, device=device, ) - req_to_token_pool.mamba_allocator = mamba_slot_allocator logger.info( "[unified-memory-pool] ============================================================" diff --git a/scripts/ci/rerun_test_groups.json b/scripts/ci/rerun_test_groups.json index f5b13d1f2..ad2a0346d 100644 --- a/scripts/ci/rerun_test_groups.json +++ b/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", diff --git a/test/registered/disaggregation/test_disaggregation_basic.py b/test/registered/disaggregation/test_disaggregation_basic.py index 706225025..1764b7f03 100644 --- a/test/registered/disaggregation/test_disaggregation_basic.py +++ b/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( - session, + async with session.post( 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) - await asyncio.sleep(0.1) + # /v1/loads can still report a previous batch. Wait for every + # 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 requests finished before decode retract pause was issued.", + all(not task.done() for task in tasks), + "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( diff --git a/test/registered/e2e/hicache/test_hicache_unified_memory.py b/test/registered/e2e/hicache/test_hicache_unified_memory.py new file mode 100644 index 000000000..ffccb2476 --- /dev/null +++ b/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() diff --git a/test/registered/unit/disaggregation/test_unified_memory_move_gate.py b/test/registered/unit/disaggregation/test_unified_memory_move_gate.py index 296168966..e50cd4ff8 100644 --- a/test/registered/unit/disaggregation/test_unified_memory_move_gate.py +++ b/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): - peer = self._Peer(gate) + def _credit(self, gate, host_gate=None): + 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 - # that matters: it inherits the setter, so an enumeration written inside - # that setter would silently leave the third member ungated. + # Inherited gate setters must cover every member, including tri-pool Mamba. _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): - """A gate that reaches only some members is not a weaker gate, it is no - gate: the ungated end relocates its own pages under the very transfer - the gate was installed for. - """ + def test_every_gate_reaches_every_member(self): + """Both transfer gates must protect every sub-pool from relocation.""" for name, expected in self._EXPECTED_COVERAGE.items(): - with self.subTest(composite=name): - self.assertEqual( - self._members_reached(name, "disagg_move_gate"), - expected, - f"{name}.disagg_move_gate does not cover every member", - ) + for slot in ("disagg_move_gate", "host_transfer_move_gate"): + with self.subTest(composite=name, slot=slot): + self.assertEqual( + self._members_reached(name, slot), + 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): - body = inspect.getsource(cls.set_disagg_move_gate) - self.assertIn("install_move_gate", body) - self.assertNotIn("_move_gate = ", body) + for setter in ("set_disagg_move_gate", "set_host_transfer_move_gate"): + if setter not in vars(cls): + continue # inherited, and the inherited one is checked above + 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 diff --git a/test/registered/unit/mem_cache/test_hicache_dcp_host_pool.py b/test/registered/unit/mem_cache/test_hicache_dcp_host_pool.py index 1fb9254ce..af0d8db9e 100644 --- a/test/registered/unit/mem_cache/test_hicache_dcp_host_pool.py +++ b/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, diff --git a/test/registered/unit/mem_cache/test_kv_index_translator.py b/test/registered/unit/mem_cache/test_kv_index_translator.py index 8f9a92ffe..45768e04a 100644 --- a/test/registered/unit/mem_cache/test_kv_index_translator.py +++ b/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, ) diff --git a/test/registered/unit/mem_cache/test_unified_hicache_regressions.py b/test/registered/unit/mem_cache/test_unified_hicache_regressions.py new file mode 100644 index 000000000..4bd2da3a9 --- /dev/null +++ b/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() diff --git a/test/registered/unit/mem_cache/test_unified_hicache_strided_state.py b/test/registered/unit/mem_cache/test_unified_hicache_strided_state.py new file mode 100644 index 000000000..476f65801 --- /dev/null +++ b/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}")