minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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] ============================================================"