minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 6 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Cheng Wan
2026-09-21 16:50:37 -07:00
committed by GitHub
parent 22587fb15c
commit 506698761d
26 changed files with 1194 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/kv_cache_hook.py
+4 -6
View File
@@ -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)
python/sglang/srt/managers/cache_controller.py
+9
View File
@@ -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
+6
View File
@@ -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
+54 -3
View File
@@ -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)
)
python/sglang/srt/mem_cache/allocator/unified_mamba.py
+15
View File
@@ -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
+35 -16
View File
@@ -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
python/sglang/srt/mem_cache/hicache_storage.py
+3
View File
@@ -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
+28
View File
@@ -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
+44 -8
View File
@@ -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,
python/sglang/srt/mem_cache/l2_transfer.py
+43 -10
View File
@@ -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)
python/sglang/srt/mem_cache/memory_pool.py
+9
View File
@@ -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
+8 -3
View File
@@ -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
+4
View File
@@ -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
+86 -3
View File
@@ -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
+4 -1
View File
@@ -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
+13 -2
View File
@@ -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],
)
python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py
+15
View File
@@ -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
+37 -12
View File
@@ -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] ============================================================"
scripts/ci/rerun_test_groups.json
+3
View File
@@ -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
+35 -28
View File
@@ -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(
test/registered/e2e/hicache/test_hicache_unified_memory.py
+240
View File
@@ -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
+32 -22
View File
@@ -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
test/registered/unit/mem_cache/test_hicache_dcp_host_pool.py
+2
View File
@@ -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
+5 -2
View File
@@ -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
+345
View File
@@ -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
+115
View File
@@ -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}")