diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index 94fe07af5734..f993d0b51d15 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/python/sglang/srt/arg_groups/kv_cache_hook.py @@ -262,10 +262,6 @@ def handle_unified_memory_pool(server_args: Any) -> None: ) if cfg.disaggregation_decode_retraction_backup == "host_pool": model_config = model_config_of(server_args) - assert not cfg.disaggregation_decode_enable_radix_cache, ( - "--enable-unified-memory host-pool decode retraction does not " - "support decode radix-cache H2D/D2H transfers yet." - ) assert mambaish_config(model_config) is None, ( "--enable-unified-memory host-pool decode retraction does not " "support hybrid-Mamba models." @@ -305,13 +301,49 @@ 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." ) + assert not cfg.enable_unified_cache_external_linker, ( + "--enable-unified-memory does not support " + "--enable-unified-cache-external-linker: direct L3 transfers do not " + "preserve unified page-envelope indices and compaction lifetimes. " + "Use --enable-hierarchical-cache for supported L2/L3 transfers." + ) + if cfg.enable_hierarchical_cache: + model_config = model_config_of(server_args) + assert not use_mla_backend(server_args), ( + "--enable-unified-memory with hierarchical cache does not support " + "MLA models yet." + ) + assert mambaish_config(model_config) is None, ( + "--enable-unified-memory with hierarchical cache does not support " + "recurrent-state models yet." + ) + assert cfg.speculative_algorithm is None, ( + "--enable-unified-memory with hierarchical cache does not support " + "speculative decoding yet." + ) + assert cfg.hicache_io_backend in {"kernel", "direct"}, ( + "--enable-unified-memory with hierarchical cache supports only " + "the kernel and direct I/O backends." + ) + assert cfg.pp_size == 1, ( + "--enable-unified-memory with hierarchical cache does not support " + "pipeline parallelism (--pp-size > 1)." + ) + supported_storage_backends = {None, "file", "sim", "mori", "shm"} + if cfg.hicache_storage_backend not in supported_storage_backends: + raise ValueError( + "--enable-unified-memory with hierarchical cache does not " + "support this storage backend yet; supported backends are " + "file, sim, mori, and shm. Got " + f"{cfg.hicache_storage_backend!r}." + ) + assert not envs.SGLANG_DISABLE_LAZY_COMPACTION.get(), ( + "--enable-unified-memory with hierarchical cache requires lazy " + "compaction so pending H2D physical reservations remain stable." + ) if cfg.dcp_size > 1: _validate_unified_memory_dcp(server_args) # Prefill cuda-graph capture IS wired for the unified pool: the captured diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 3e812f749e2e..30cbbc6ea09e 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -762,12 +762,19 @@ def _match_prefix_and_lock(self, req: Req) -> DecodePrefixMatch: Match a request against the decode-side radix cache, lock the matched node to prevent eviction, and return the matched prefix information. """ + max_prefix_len = None + if self._uses_swa_tail_prealloc(): + fill_len = self._pre_alloc_fill_len(req) + max_prefix_len = fill_len - self._swa_tail_len(fill_len) + # Match and lock only reusable FULL KV. The entire SWA tail must be + # freshly allocated, including when the prefix comes from L2/L3. result = match_prefix_for_req( self.tree_cache, req, req.origin_input_ids, cow_mamba=self.tree_cache.supports_mamba(), include_req=True, + max_prefix_len=max_prefix_len, ) # Keep aggregated scheduling semantics while preserving the SWA lock # boundary needed for the matching dec_lock_ref; the full receipt @@ -775,7 +782,9 @@ def _match_prefix_and_lock(self, req: Req) -> DecodePrefixMatch: req.lock_receipt = self.tree_cache.inc_lock_ref( result.last_device_node ).to_dec_params() - return self._build_decode_prefix_match(req, result) + return self._build_decode_prefix_match( + req, result, max_prefix_len=max_prefix_len + ) def _resolve_prefill_dp_rank(self, req: Req) -> Optional[int]: prefill_info = self.kv_manager.prefill_info_table.get(_bootstrap_addr(req)) @@ -1338,20 +1347,6 @@ def pop_preallocated( fill_len = self._pre_alloc_fill_len(decode_req.req) - # Cap full-attention prefix reuse at the sliding-window start so - # the SWA window lands entirely in the fresh delta, keeping - # alloc_extend_swa_tail's tail->full mapping in range. Costs reuse - # of only the last ~window_size full-attention tokens. - if uses_swa_tail_prealloc and prefix_len > 0: - swa_prefix_cap = fill_len - self._swa_tail_len(fill_len) - if prefix_len > swa_prefix_cap: - prefix_len = swa_prefix_cap - prefix_indices = prefix_indices[:prefix_len] - # Cap the prefill-committed prefix too: tokens past the - # cap are not device-resident, so prefill must transfer - # them. - total_prefix_len = prefix_len - # Decode transfers the SWA tail fresh, so retain only the # full-attention prefix lock needed for reuse. if ( @@ -2127,9 +2122,9 @@ def alloc_for_decode_prealloc( # the live window tail. kv_loc = allocator.alloc_extend_swa_tail( prefix_lens=torch.tensor( - [prefix_len], dtype=torch.int64, device=device + [total_prefix_len], dtype=torch.int64, device=device ), - prefix_lens_cpu=torch.tensor([prefix_len], dtype=torch.int64), + prefix_lens_cpu=torch.tensor([total_prefix_len], dtype=torch.int64), seq_lens=torch.tensor([fill_len], dtype=torch.int64, device=device), seq_lens_cpu=torch.tensor([fill_len], dtype=torch.int64), last_loc=last_loc, diff --git a/python/sglang/srt/disaggregation/decode_hicache_mixin.py b/python/sglang/srt/disaggregation/decode_hicache_mixin.py index 5f0dac9aaa55..6e05e837797b 100644 --- a/python/sglang/srt/disaggregation/decode_hicache_mixin.py +++ b/python/sglang/srt/disaggregation/decode_hicache_mixin.py @@ -60,7 +60,9 @@ class HiCacheRestoreResult(Enum): class DecodeHiCachePreallocMixin: """HiCache hooks for ``DecodePreallocQueue``: issue prefetch + reserve tokens.""" - def _build_decode_prefix_match(self, req: Req, result: Any) -> DecodePrefixMatch: + def _build_decode_prefix_match( + self, req: Req, result: Any, *, max_prefix_len: Optional[int] = None + ) -> DecodePrefixMatch: """Convert a ``match_prefix_for_req`` result into ``DecodePrefixMatch``. Performs the optional L3 storage hit length query when decode-side @@ -78,7 +80,7 @@ def _build_decode_prefix_match(self, req: Req, result: Any) -> DecodePrefixMatch last_host_node ): matched_len = l1_prefix_len + l2_host_hit_length - suffix_tokens = req.origin_input_ids[matched_len:] + suffix_tokens = req.origin_input_ids[matched_len:max_prefix_len] last_hash = self.tree_cache.get_last_hash_value(last_host_node) prefix_keys = ( self.tree_cache.get_prefix_hash_values(last_host_node) @@ -218,6 +220,7 @@ def _try_hicache_queue_load_back(self, dr: DecodeRequest) -> bool: dr.req.origin_input_ids, cow_mamba=False, include_req=True, + max_prefix_len=pm.decode_prefix_len, ) new_indices, restored_node = self.tree_cache.init_load_back( InitLoadBackParams( diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 9d5736074d54..daf066315614 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -43,7 +43,7 @@ from sglang.srt.mem_cache.l2_transfer import L2Transfer, L2TransferEngine from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool from sglang.srt.mem_cache.utils import get_storage_hash_str -from sglang.srt.runtime_context import get_parallel +from sglang.srt.runtime_context import get_memory, get_parallel from sglang.srt.utils import get_device_module logger = logging.getLogger(__name__) @@ -822,6 +822,9 @@ def start_writing(self) -> None: completion = self.l2_transfer_engine.submit_device_to_host( self._l2_transfers(host_indices, device_indices, pool_transfers) ) + self.mem_pool_device_allocator.set_hicache_transfer_done_event( + (id(self), "write"), completion.finish_event + ) self.ack_write_queue.append( HiCacheAck( @@ -917,6 +920,11 @@ def _move_op_indices( ) -> tuple[torch.Tensor, torch.Tensor, Optional[List[PoolTransfer]]]: return (*self.move_indices(op.host_indices, op.device_indices), None) + def _move_load_operation( + self, op: CacheOperation + ) -> tuple[torch.Tensor, torch.Tensor, Optional[List[PoolTransfer]]]: + return self._move_op_indices(op) + def _l2_transfers( self, host_indices: torch.Tensor, @@ -947,7 +955,7 @@ def start_loading(self) -> int: producer_id = self.layer_done_counter.update_producer() op = CacheOperation.merge_ops(self.load_queue) - host_indices, device_indices, pool_transfers = self._move_op_indices(op) + host_indices, device_indices, pool_transfers = self._move_load_operation(op) self.load_queue.clear() producer_event = self.layer_done_counter.events[producer_id] producer_event.start_event.record() @@ -966,6 +974,9 @@ def start_loading(self) -> int: on_layer_done=producer_event.complete, layer_num=self.layer_num, ) + self.mem_pool_device_allocator.set_hicache_transfer_done_event( + (id(self), "load"), completion.finish_event + ) self.ack_load_queue.append( HiCacheAck( @@ -1093,13 +1104,22 @@ def _page_transfer(self, operation: PrefetchOperation) -> int: # Get one batch token, and update the completed_tokens if succeed extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys) - hit_pages = self._page_transfer_kv_batch( - operation, - batch_hashes, - batch_host_indices, - extra_info, - kv_derived_transfers, - ) + try: + hit_pages = self._page_transfer_kv_batch( + operation, + batch_hashes, + batch_host_indices, + extra_info, + kv_derived_transfers, + ) + except Exception: + if not get_memory().enable_unified_memory: + raise + logger.exception( + "HiCache prefetch transfer failed for request %s", + operation.request_id, + ) + hit_pages = 0 # Check termination if hit_pages != len(batch_hashes): all_success = False @@ -1160,10 +1180,13 @@ def prefetch_io_aux_func(self): while not self.storage_stop_event.is_set(): try: operation = self.prefetch_buffer.get(block=True, timeout=1) - if operation is None: - continue + except Empty: + continue + if operation is None: + continue + try: self._page_transfer(operation) - + finally: self.prefetch_sync_queue.put( PrefetchAck( rid=operation.request_id, @@ -1171,8 +1194,6 @@ def prefetch_io_aux_func(self): operation=operation, ) ) - except Empty: - continue def prefetch_rate_limited(self) -> bool: """ @@ -1194,6 +1215,24 @@ def prefetch_rate_limited(self) -> bool: # todo: more sophisticated rate limiting based on storage backend performance return False + def alloc_prefetch_host_buffers( + self, operation: StorageOperation, need_size: int + ) -> Optional[torch.Tensor]: + """Allocate the host bounce for a storage hit.""" + return self.mem_pool_host.alloc(need_size) + + def can_fit_prefetch_host_buffers( + self, operation: StorageOperation, need_size: int + ) -> bool: + """Whether a prefetch bounce can fit when its host pools are empty.""" + return need_size <= self.mem_pool_host.size + + def free_prefetch_host_buffers( + self, operation: StorageOperation, host_indices: torch.Tensor + ) -> None: + """Roll back a hit-sized host bounce before transfer ownership moves.""" + self.mem_pool_host.free(host_indices) + def _storage_hit_query(self, operation) -> tuple[list[str], int]: last_hash = operation.last_hash tokens_to_fetch = operation.token_ids @@ -1228,10 +1267,22 @@ def prefetch_thread_func(self): operation = self.prefetch_queue.get(block=True, timeout=1) if operation is None: continue - if operation.is_terminated(): + try: + if operation.is_terminated(): + hash_value, storage_hit_count = [], 0 + else: + hash_value, storage_hit_count = self._storage_hit_query( + operation + ) + except Exception: + if not get_memory().enable_unified_memory: + raise + logger.exception( + "HiCache storage query failed for request %s", + operation.request_id, + ) hash_value, storage_hit_count = [], 0 - else: - hash_value, storage_hit_count = self._storage_hit_query(operation) + storage_hit_count_tensor = torch.tensor( storage_hit_count, dtype=torch.int ) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 6978b369a75d..89aa04874da5 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -145,6 +145,7 @@ def match_prefix_for_req( *, cow_mamba: bool = False, include_req: bool = False, + max_prefix_len: Optional[int] = None, ): if token_ids is None: token_ids = req.origin_input_ids + req.output_ids @@ -155,6 +156,10 @@ def match_prefix_for_req( # this request's SWA ring. No-op for other layouts. reprefill_tail = tree_cache.swa_reprefill_tail_tokens() key_limit = max(0, len(token_ids) - reprefill_tail) if reprefill_tail else None + if max_prefix_len is not None: + key_limit = ( + max_prefix_len if key_limit is None else min(key_limit, max_prefix_len) + ) match_result = tree_cache.match_prefix( MatchPrefixParams( diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index 0c4ec17ce821..7d11cda4f5df 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -16,7 +16,7 @@ from __future__ import annotations import abc -from typing import TYPE_CHECKING, Protocol +from typing import TYPE_CHECKING, Hashable, Protocol import torch @@ -178,9 +178,11 @@ def merge_and_sort_free(self): def translate_kv_indices_for_transfer( self, kv_indices: torch.Tensor ) -> torch.Tensor: - """Token ids as the PD transfer engine addresses them. Identity here - because a static pool's ids index its registered buffers directly; - virtual-id pools must override.""" + """Token ids as device transfer engines address them. + + Identity here: a static pool's token ids index its registered buffers + directly. Virtual-id pools must override. + """ return kv_indices def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): @@ -191,6 +193,10 @@ def load_cpu_copy( ): raise NotImplementedError() + def set_hicache_transfer_done_event(self, transfer_key: Hashable, event) -> None: + """Record an asynchronous HiCache transfer completion event if needed.""" + return + def alloc_extend(self, *args, **kwargs): raise NotImplementedError("alloc_extend is only for paged allocator") diff --git a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py index db4b2c3b0077..11776bd18c00 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py +++ b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py @@ -19,7 +19,7 @@ import logging import math from abc import abstractmethod -from typing import Callable, List, Optional, Sequence, Tuple, TypeGuard +from typing import Callable, Hashable, List, Optional, Sequence, Tuple, TypeGuard import torch from torch.profiler import record_function @@ -641,13 +641,25 @@ def free_full_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None 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.""" - return + if full_indices.numel() == 0: + return + assert full_indices.numel() == swa_indices.numel() + full_pages = full_indices.to(torch.int64) // self.page_size + swa_pages = swa_indices.to(torch.int64) // self.page_size + self.swa_attn_allocator.bind(full_pages, swa_pages) def clear_full_to_swa_mapping(self, full_indices: torch.Tensor) -> None: - # Paired with set_full_to_swa_mapping: shared mode has no mapping tensor. - return + if full_indices.numel() == 0: + return + full_pages = torch.unique(full_indices.to(torch.int64) // self.page_size) + swa = self.swa_attn_allocator + physical_pages = swa.virtual_to_physical[full_pages].clone() + swa.virtual_to_physical.index_fill_(0, full_pages, -1) + live = physical_pages > 0 + physical_pages = physical_pages[live] + full_pages = full_pages[live] + still_owned = swa.physical_to_virtual[physical_pages] == full_pages + swa.physical_to_virtual.index_fill_(0, physical_pages[still_owned], -1) # -- free-group -- @@ -722,6 +734,10 @@ def clear(self) -> None: # -- Lazy compaction hooks -- + def set_hicache_transfer_done_event(self, transfer_key: Hashable, event) -> None: + self.full_attn_allocator.set_hicache_transfer_done_event(transfer_key, event) + self.swa_attn_allocator.set_hicache_transfer_done_event(transfer_key, event) + def set_latest_forward_done_event(self, event: Optional[torch.cuda.Event]) -> None: """Forward the per-batch `forward_done` event to BOTH sub-allocators.""" with record_function("UnifiedSWAAlloc.set_latest_forward_done_event"): @@ -935,7 +951,8 @@ def can_reserve( def _compaction_allowed(self) -> bool: return all( - allocator.disagg_move_gate is None or allocator.disagg_move_gate() + allocator._pending_hicache_load_pages == 0 + and (allocator.disagg_move_gate is None or allocator.disagg_move_gate()) for allocator in (self.full_attn_allocator, self.swa_attn_allocator) ) diff --git a/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py b/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py index 8eb09ec21aca..81da67c95c50 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py +++ b/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py @@ -29,6 +29,7 @@ Callable, Dict, Generic, + Hashable, List, Optional, Sequence, @@ -425,6 +426,9 @@ def __init__( # write-race check. Single slot: at most ONE forward in flight per call # site; `_flush` materializes the write-set lazily, avoiding a sync here. self._inflight_forward: Optional[Tuple[torch.cuda.Event, torch.Tensor]] = None + self._hicache_transfer_done_events: Dict[Hashable, torch.cuda.Event] = {} + # HiCache reserves physical pages before a completion event exists. + self._pending_hicache_load_pages = 0 # Per-call move cap on NON-urgent `_flush`: bounds work per `on_idle()` so # a large backlog doesn't block ZMQ IPC. Urgent retries are uncapped. @@ -527,10 +531,30 @@ def clear(self) -> None: self.live_page_count = 0 self._inflight_forward = None self._latest_forward_done_event = None + self._hicache_transfer_done_events.clear() + self._pending_hicache_load_pages = 0 def clear_inverse_history(self) -> None: self._inverse_history.clear() + def set_hicache_transfer_done_event(self, transfer_key: Hashable, event) -> None: + self._hicache_transfer_done_events[transfer_key] = event + if transfer_key == "load" or ( + isinstance(transfer_key, tuple) and transfer_key[-1] == "load" + ): + if self._pending_hicache_load_pages > 0: + self._pending_hicache_load_pages = 0 + self._capacity_epoch += 1 + + def _wait_hicache_transfers(self) -> None: + if not self._hicache_transfer_done_events: + return + current_stream = torch.cuda.current_stream() + events = tuple(self._hicache_transfer_done_events.values()) + self._hicache_transfer_done_events.clear() + for event in events: + current_stream.wait_event(event) + # -- size reporting -- def _allocated_pages(self) -> int: @@ -725,6 +749,8 @@ def _peer_drainable_hole_bytes(self) -> int: neighbor = self._growth_side_neighbor() if neighbor is None or not neighbor.lazy_compaction: return 0 + if neighbor._pending_hicache_load_pages > 0: + 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. @@ -923,6 +949,8 @@ def _alloc_bind_fast_or_slow( with record_function("MultiEndedAlloc._alloc_bind_fast_or_slow"): if N == 0: return torch.empty(0, dtype=torch.int64, device=self.device) + if self.lazy_compaction and self._free_phys_pages.numel() > 0: + self._wait_hicache_transfers() # FAST PATH: eager, or lazy with no current holes. if not self.lazy_compaction or self._free_phys_pages.numel() == 0: @@ -1119,6 +1147,14 @@ def translate_write_loc_for_kernel( # -- alloc -- + def _expand_pages_to_tokens(self, pages: torch.Tensor) -> torch.Tensor: + if self.page_size == 1: + return pages + return ( + pages[:, None] * self.page_size + + torch.arange(self.page_size, device=self.device) + ).reshape(-1) + def alloc(self, need_size: int) -> Optional[torch.Tensor]: """Allocate `need_size` virtual TOKEN ids (id-owner only). Returns token-granular, page-structured ids, or None on shortfall. @@ -1149,13 +1185,7 @@ def alloc(self, need_size: int) -> Optional[torch.Tensor]: if phys_pages is None: self.free_virtual_ids = torch.cat([v_pages, self.free_virtual_ids]) return None - if self.page_size == 1: - return v_pages # v_pages already IS the token id list - # Expand page ids to token ids: (P, 1) * S + (S,) -> (P, S) -> (P*S,). - return ( - v_pages[:, None] * self.page_size - + torch.arange(self.page_size, device=self.device) - ).reshape(-1) + return self._expand_pages_to_tokens(v_pages) def alloc_with_virtual(self, virtual_pages: torch.Tensor) -> None: """Take physical PAGES for caller-supplied virtual PAGE ids (not token @@ -1173,6 +1203,63 @@ def alloc_with_virtual(self, virtual_pages: torch.Tensor) -> None: "physical room (the composite's byte-budget check should have caught this)" ) + def alloc_physical(self, need_size: int) -> Optional[torch.Tensor]: + """Reserve physical token slots without assigning virtual page ids.""" + if need_size <= 0: + return torch.empty(0, dtype=torch.int64, device=self.device) + assert need_size % self.page_size == 0, ( + f"MultiEndedAllocator({self.sub_pool_name!r}).alloc_physical: need_size=" + f"{need_size} must be a multiple of page_size={self.page_size}" + ) + if need_size > self.available_size() and not _relieve_for_alloc( + self, need_size + ): + return None + if self.lazy_compaction and self._free_phys_pages.numel() > 0: + self._wait_hicache_transfers() + physical_pages = self.take_physical_pages(need_size // self.page_size) + if physical_pages is None: + return None + if self.lazy_compaction: + self._pending_hicache_load_pages += int(physical_pages.shape[0]) + return self._expand_pages_to_tokens(physical_pages) + + def cancel_physical_reservation(self, free_index: torch.Tensor) -> None: + """Roll back a HiCache physical allocation before its H2D is submitted.""" + self.free_physical(free_index) + + def free_physical(self, free_index: torch.Tensor) -> None: + """Release physical token slots reserved by :meth:`alloc_physical`.""" + if free_index is None or free_index.numel() == 0: + return + assert free_index.numel() % self.page_size == 0, ( + f"MultiEndedAllocator({self.sub_pool_name!r}).free_physical: " + f"{free_index.numel()} tokens must contain whole pages of " + f"size {self.page_size}" + ) + physical_pages = ( + free_index.detach().to(torch.int64)[:: self.page_size] // self.page_size + ) + if self.lazy_compaction: + num_pages = int(physical_pages.shape[0]) + assert num_pages <= self._pending_hicache_load_pages, ( + f"MultiEndedAllocator({self.sub_pool_name!r}) released {num_pages} " + f"HiCache pages with only {self._pending_hicache_load_pages} pending" + ) + self._pending_hicache_load_pages -= num_pages + self._wait_hicache_transfers() + if self.forward_stream is not None and not self.lazy_compaction: + torch.cuda.current_stream().wait_stream(self.forward_stream) + virtual_pages = self.physical_to_virtual[physical_pages] + bound_mask = virtual_pages >= 0 + self.virtual_to_physical.index_fill_(0, virtual_pages[bound_mask], -1) + self.physical_to_virtual.index_fill_(0, physical_pages, -1) + if self.lazy_compaction: + self._free_phys_pages = torch.cat([self._free_phys_pages, physical_pages]) + self.live_page_count -= int(physical_pages.shape[0]) + return + self._compact_pending(physical_pages) + # -- paged alloc surface -- def alloc_extend( @@ -1325,6 +1412,7 @@ def free( self._free_lazy(free_index, pages=_pages) return # --- EAGER path --- + self._wait_hicache_transfers() # Near-no-op in normal mode (sampling's CPU sync already drained # forward_stream); in overlap mode it does the serializing. if self.forward_stream is not None: @@ -1783,6 +1871,9 @@ def _flush(self, *, urgent: bool) -> int: if self.disagg_move_gate is not None and not self.disagg_move_gate(): # Holes stay in the free list; the next flush picks them up. return 0 + if self._pending_hicache_load_pages > 0: + return 0 + self._wait_hicache_transfers() self._stats_n_flush_calls += 1 with record_function("MultiEndedAlloc._flush"): self._drain_pending_reuse(urgent=urgent) diff --git a/python/sglang/srt/mem_cache/buffer_mode/pipeline.py b/python/sglang/srt/mem_cache/buffer_mode/pipeline.py index c5f5e0561340..f397b1b0318e 100644 --- a/python/sglang/srt/mem_cache/buffer_mode/pipeline.py +++ b/python/sglang/srt/mem_cache/buffer_mode/pipeline.py @@ -48,6 +48,7 @@ PoolTransfer, SidecarPoolSpec, ) +from sglang.srt.mem_cache.pool_host.base import HostKVCache from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.unified_cache.cache_action import RebuildFullToSWAMapping from sglang.srt.mem_cache.unified_cache.components import ( @@ -68,6 +69,13 @@ logger = logging.getLogger(__name__) +def _minimum_full_transfer_tokens( + host_pool: HostKVCache, prefetch_threshold: int +) -> int: + page_size = host_pool.page_size + return max(page_size, -(-prefetch_threshold // page_size) * page_size) + + class _UnifiedBackupIntent(msgspec.Struct): """Buffer-mode backup intent, unpinned while queued. @@ -92,6 +100,7 @@ class _UnifiedBufferBackupEntry(msgspec.Struct): host_indices: torch.Tensor aux_xfers: list[PoolTransfer] lock_params: DecLockRefParams + occupied_units: int class _StagedPrefetch(msgspec.Struct): @@ -212,12 +221,40 @@ def validate_buffer_only_stack( # window), so every window-carrying intent would be dropped as # oversize and SWA storage coverage would silently be zero. window_tokens = swa.full_window_pages * swa._swa_kv_pool_host.page_size - if swa._swa_kv_pool_host.size < 2 * window_tokens: + shared_domain = ( + swa._swa_kv_pool_host.shared_allocation_domain + if isinstance(swa._swa_kv_pool_host, HostKVCache) + else None + ) + if shared_domain is not None: + full_host_pool = ( + swa.cache.cache_controller.mem_pool_host.anchor_entry.host_pool + ) + min_full_tokens = _minimum_full_transfer_tokens( + full_host_pool, swa.cache.prefetch_threshold + ) + one_transfer_bytes = ( + window_tokens * swa._swa_kv_pool_host.size_per_token + + min_full_tokens * full_host_pool.size_per_token + ) + one_transfer = [ + (full_host_pool.pool_label, min_full_tokens), + (swa._swa_kv_pool_host.pool_label, window_tokens), + ] + enough_capacity = shared_domain.can_fit_many_then( + one_transfer, one_transfer, empty=True + ) + capacity = f"{shared_domain.capacity_bytes} shared bytes" + requirement = f"{2 * one_transfer_bytes} shared bytes" + else: + enough_capacity = swa._swa_kv_pool_host.size >= 2 * window_tokens + capacity = f"{swa._swa_kv_pool_host.size} SWA tokens" + requirement = f"{2 * window_tokens} SWA tokens" + if not enough_capacity: raise ValueError( - "--hicache-host-memory-mode buffer_only requires an SWA " - f"host pool of at least two trailing windows " - f"({2 * window_tokens} tokens; got " - f"{swa._swa_kv_pool_host.size}): one staging a write " + "--hicache-host-memory-mode buffer_only requires a host arena " + "large enough for two minimum Full/SWA transfers " + f"({requirement}; got {capacity}): one staging a write " "while one stays reserved for prefetch window allocs." ) @@ -297,8 +334,94 @@ def reset(self) -> None: self.anchor_locked_tokens_ = 0 self._anchor_lock_cap_skips = 0 + def _shared_host_domain(self): + cc = self._cache.cache_controller + anchor = cc.mem_pool_host.anchor_entry.host_pool + if not isinstance(anchor, HostKVCache): + return None + return anchor.shared_allocation_domain + + def _transfer_tokens(self, transfer: PoolTransfer) -> int: + if transfer.host_indices is not None: + return len(transfer.host_indices) + if transfer.keys is None: + return 0 + entry = self._cache.cache_controller.mem_pool_host.entry_map.get(transfer.name) + return 0 if entry is None else len(transfer.keys) * entry.host_pool.page_size + + def host_allocation_units( + self, + host_indices: Optional[torch.Tensor], + aux_xfers: Optional[list[PoolTransfer]], + ) -> int: + """Host usage expressed in anchor-token units for scheduler accounting.""" + if host_indices is None: + return 0 + return self._host_request_units(len(host_indices), aux_xfers) + + def _shared_host_requests( + self, kv_tokens: int, aux_xfers: Optional[list[PoolTransfer]] + ) -> Optional[list[tuple[str, int]]]: + if self._shared_host_domain() is None: + return None + return [ + (pool.pool_label, tokens) + for pool, tokens in self._host_staging_sizes(kv_tokens, aux_xfers) + ] + + def _host_staging_sizes(self, kv_tokens, aux_xfers): + cc = self._cache.cache_controller + yield cc.mem_pool_host.anchor_entry.host_pool, kv_tokens + for transfer in aux_xfers or (): + if transfer.indices_from_pool is not None: + continue + entry = cc.mem_pool_host.entry_map.get(transfer.name) + if entry is not None: + yield entry.host_pool, self._transfer_tokens(transfer) + + def _host_request_units( + self, kv_tokens: int, aux_xfers: Optional[list[PoolTransfer]] + ) -> int: + """Host staging in anchor-token accounting units.""" + cc = self._cache.cache_controller + anchor = cc.mem_pool_host.anchor_entry.host_pool + if self._shared_host_domain() is None: + return kv_tokens + num_bytes = sum( + tokens * pool.size_per_token + for pool, tokens in self._host_staging_sizes(kv_tokens, aux_xfers) + ) + return (num_bytes + anchor.size_per_token - 1) // anchor.size_per_token + + def _shared_backup_fits(self, requests, *, empty: bool = False) -> bool: + can_fit = self._shared_host_domain().can_fit_many_then( + requests, self._shared_load_reserve_requests(), empty=empty + ) + if not empty: + can_fit = self._cache.cache_controller._sync_shared_host_value( + int(can_fit), torch.distributed.ReduceOp.MIN + ) + return bool(can_fit) + + def _shared_load_reserve_requests(self) -> list[tuple[str, int]]: + cc = self._cache.cache_controller + anchor = cc.mem_pool_host.anchor_entry.host_pool + swa_entry = cc.mem_pool_host.entry_map.get(PoolName.SWA) + min_full_tokens = _minimum_full_transfer_tokens( + anchor, self._cache.prefetch_threshold + ) + requests = [(anchor.pool_label, min_full_tokens)] + if swa_entry is not None and self._swa_window_pages: + requests.append( + ( + swa_entry.host_pool.pool_label, + self._swa_window_pages * swa_entry.host_pool.page_size, + ) + ) + return requests + def is_idle(self) -> bool: - """No pending or in-flight buffer-mode transfer work.""" + """No queued, staged, or in-flight buffer-mode work or anchor pins.""" return not ( self.pending_hit_allocs or self.staged_prefetches @@ -307,6 +430,7 @@ def is_idle(self) -> bool: or self.inflight_backup_node_ids or self.ongoing_write_through or self.ongoing_backup + or self.anchor_locks ) # ---- backup pipeline (device -> staging -> storage) ---- @@ -434,11 +558,21 @@ def _backup_oversize( capacity (total for KV, total minus the loads-priority margin for aux pools — matching ``_aux_budget_blocked``'s admission ceiling): such an intent could never stage and would wedge the FIFO head.""" + if aux_xfers is None: + aux_xfers = self._build_aux_staging_transfers(node_id, hash_values) cc = self._cache.cache_controller + shared_requests = self._shared_host_requests(intent_tokens, aux_xfers) + if shared_requests is not None: + pool_tokens = cc.mem_pool_host.size + max_write_units = max( + int(HICACHE_WRITE_STAGING_POOL_FRACTION * pool_tokens), + pool_tokens - pool_tokens // 10, + ) + if self._host_request_units(intent_tokens, aux_xfers) > max_write_units: + return True + return not self._shared_backup_fits(shared_requests, empty=True) if intent_tokens > cc.mem_pool_host.size: return True - if aux_xfers is None: - aux_xfers = self._build_aux_staging_transfers(node_id, hash_values) for t in aux_xfers or (): entry = cc.mem_pool_host.entry_map.get(t.name) if entry is not None and ( @@ -522,9 +656,6 @@ def flush_pending_writes(self) -> None: self.write_backlog_tokens_ -= intent_tokens self._log_backup_dropped(intent_tokens) continue - if self.write_staged_tokens_ >= live_cap: - # Yield to live fetch demand; retry next round. - break device_value, comp_xfers = self._cache.tree_core.build_backup_spec( snapshot.node_id ) @@ -543,6 +674,17 @@ def flush_pending_writes(self) -> None: self.write_backlog_tokens_ -= intent_tokens self._log_backup_dropped(intent_tokens) continue + shared_domain = self._shared_host_domain() + staging_at_limit = ( + self.write_staged_tokens_ + + self._host_request_units(intent_tokens, sizing_xfers) + > live_cap + if shared_domain is not None + else self.write_staged_tokens_ >= live_cap + ) + if staging_at_limit: + # Yield to live fetch demand; retry next round. + break if self._aux_budget_blocked(intent, sizing_xfers): # An aux pool lacks staging headroom: yield at the gate # instead of failing the alloc inside cc.write; acks free @@ -575,6 +717,11 @@ def _launch_backup_intent( # allocate no additional staging, but must ride the same D2H operation # so their bytes are present when the storage write starts. aux_xfers.extend(cache._build_backup_sidecar(device_value, comp_xfers)) + device_value = ( + cache.token_to_kv_pool_allocator.translate_kv_indices_for_transfer( + device_value + ) + ) host_indices = cc.write( device_value, node_id=snapshot.node_id, @@ -586,13 +733,15 @@ def _launch_backup_intent( # NOTE: no commit_backup — the node must never appear # host-resident in buffer mode; staging slots live in the entry. lock_params = cache.inc_lock_ref(snapshot.node_id).to_dec_params() + occupied_units = self.host_allocation_units(host_indices, aux_xfers) self.ongoing_write_through[snapshot.node_id] = _UnifiedBufferBackupEntry( intent=intent, host_indices=host_indices, aux_xfers=aux_xfers, lock_params=lock_params, + occupied_units=occupied_units, ) - self.write_staged_tokens_ += len(host_indices) + self.write_staged_tokens_ += occupied_units self.write_backlog_tokens_ -= len(snapshot.hash_values) * cache.page_size return True @@ -613,9 +762,14 @@ def _aux_budget_blocked( aux = self._build_aux_staging_transfers( snapshot.node_id, snapshot.hash_values ) + cc = self._cache.cache_controller + shared_requests = self._shared_host_requests( + len(snapshot.hash_values) * self._cache.page_size, aux + ) + if shared_requests is not None: + return not self._shared_backup_fits(shared_requests) if not aux: return False - cc = self._cache.cache_controller for t in aux: entry = cc.mem_pool_host.entry_map.get(t.name) if entry is None: @@ -695,13 +849,21 @@ def finish_backup_ack(self, ack_id: int) -> None: indices_from_pool=staged.indices_from_pool, ) ) - operation_id = self._cache.cache_controller.write_storage( - entry.host_indices, - snapshot.key.token_ids, - snapshot.hash_values, - snapshot.prefix_keys, - extra_pools=storage_xfers or None, - ) + try: + operation_id = self._cache.cache_controller.write_storage( + entry.host_indices, + snapshot.key.token_ids, + snapshot.hash_values, + snapshot.prefix_keys, + extra_pools=storage_xfers or None, + ) + except Exception: + self._free_staging_now(entry.host_indices, entry.aux_xfers) + self.write_staged_tokens_ -= entry.occupied_units + self.inflight_backup_node_ids.discard(snapshot.node_id) + _untrack_content_refs(self.inflight_backup_hashes, snapshot.hash_values) + self._log_backup_dropped(len(snapshot.hash_values) * self._cache.page_size) + raise self.ongoing_backup[operation_id] = entry def finish_storage_write_ack(self, operation_id: int) -> None: @@ -717,7 +879,7 @@ def finish_storage_write_ack(self, operation_id: int) -> None: snapshot = intent.snapshot self._cache.storage_existence_cache.add(PoolName.KV, snapshot.hash_values) self._free_staging_now(entry.host_indices, entry.aux_xfers) - self.write_staged_tokens_ -= len(entry.host_indices) + self.write_staged_tokens_ -= entry.occupied_units self.inflight_backup_node_ids.discard(snapshot.node_id) _untrack_content_refs(self.inflight_backup_hashes, snapshot.hash_values) @@ -858,12 +1020,6 @@ def pop_prefix_ctx(self, req_id: str) -> None: def has_staged(self, req_id: str) -> bool: return req_id in self.staged_prefetches - @staticmethod - def _occupied_span(host_indices) -> int: - """Occupancy units a buffer-mode prefetch holds: granted at - hit-alloc, sized to the allocation (0 while still querying).""" - return len(host_indices) if host_indices is not None else 0 - def stage_completed_prefetch( self, req_id: str, @@ -894,6 +1050,9 @@ def stage_completed_prefetch( for transfer in operation.pool_transfers or () if transfer.indices_from_pool is not None ) + occupied_tokens = operation.buffer_host_occupied_units + if occupied_tokens is None: + occupied_tokens = self.host_allocation_units(host_indices, aux_xfers) if num_tokens == 0 or prefix_tokens is None: # Nothing usable fetched: recompute. @@ -902,7 +1061,7 @@ def stage_completed_prefetch( cc.append_host_mem_release( host_indices[:num_tokens], extra_pools=aux_xfers or None ) - cc.prefetch_tokens_occupied -= self._occupied_span(host_indices) + cc.prefetch_tokens_occupied -= occupied_tokens cache.prefetch_loaded_tokens_by_reqid[req_id] = 0 cache.prefetch_loaded_storage_start_by_reqid.pop(req_id, None) return True @@ -914,8 +1073,6 @@ def stage_completed_prefetch( # itself is the evidence, so feeding is sound even if this staged # prefetch is later dropped unconsumed. cache.storage_existence_cache.add(PoolName.KV, list(staged_hashes)) - occupied_tokens = self._occupied_span(host_indices) - self.staged_prefetches[req_id] = _StagedPrefetch( req_id=req_id, key_tokens=prefix_tokens + list(prefetch_key[:num_tokens].token_ids), diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index 67e785716c70..ed92990c33d8 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -36,6 +36,7 @@ from sglang.srt.mem_cache.pool_host import HostPoolGroup, PoolEntry from sglang.srt.mem_cache.pool_host.base import uses_shared_host_layout from sglang.srt.mem_cache.pool_host.mha import MHATokenToKVPoolHost +from sglang.srt.runtime_context import get_memory if TYPE_CHECKING: from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator @@ -82,6 +83,7 @@ def __init__( pool_transfers=pool_transfers, ) self.pool_transfers_done = not bool(pool_transfers) + self.buffer_host_occupied_units: Optional[int] = None def mark_terminate(self): with self._lock: @@ -405,7 +407,15 @@ def _alloc_shared_host_requests_with_reclaim( for entry, (_, need_size) in zip(entries, requests, strict=True) ] - allocated = domain.alloc_many(domain_requests) + def try_allocate(): + allocated = domain.alloc_many(domain_requests) + if self._shared_host_alloc_succeeded_on_all_ranks(allocated, requests): + return allocated + if allocated is not None: + self._free_shared_host_allocations(entries, allocated) + return None + + allocated = try_allocate() if allocated is not None: return allocated @@ -438,6 +448,9 @@ def _alloc_shared_host_requests_with_reclaim( pool.page_size * pool.size_per_token, requested_bytes - domain.free_bytes(), ) + shortfall_bytes = self._sync_shared_host_value( + shortfall_bytes, torch.distributed.ReduceOp.MAX + ) tokens = ( ( (shortfall_bytes + pool.size_per_token - 1) @@ -448,15 +461,99 @@ def _alloc_shared_host_requests_with_reclaim( // pool.page_size * pool.page_size ) - if entry.host_evict_fn(tokens) <= 0: + evicted = entry.host_evict_fn(tokens) + self._assert_shared_host_value_equal(evicted, "evicted tokens") + if evicted <= 0: continue made_progress = True - allocated = domain.alloc_many(domain_requests) + allocated = try_allocate() if allocated is not None: return allocated if not made_progress: return None + def _sync_shared_host_value(self, value: int, op) -> int: + if not self._shared_host_consensus_groups(): + return value + synced = torch.tensor(value, dtype=torch.int64, device="cpu") + self._all_reduce_shared_host_groups(synced, op) + return int(synced.item()) + + def _assert_shared_host_value_equal(self, value: int, label: str) -> None: + if not self._shared_host_consensus_groups(): + return + synced = torch.tensor([value, -value], dtype=torch.int64, device="cpu") + self._all_reduce_shared_host_groups(synced, torch.distributed.ReduceOp.MIN) + if synced[0].item() != -synced[1].item(): + raise RuntimeError( + f"Shared host {label} diverged across attention ranks: " + f"min={synced[0].item()}, max={-synced[1].item()}" + ) + + def _shared_host_alloc_succeeded_on_all_ranks( + self, + allocated: Optional[list[torch.Tensor]], + requests: list[tuple[PoolName, int]], + ) -> bool: + if not self._shared_host_consensus_groups(): + return allocated is not None + + pool_names = list(PoolName) + ordered_pool_ids = [0] * len(pool_names) + ordered_sizes = [0] * len(pool_names) + for index, (name, size) in enumerate(requests): + ordered_pool_ids[index] = pool_names.index(name) + 1 + ordered_sizes[index] = size + values = [ + int(allocated is not None), + len(requests), + sum( + size * self.mem_pool_host.entry_map[name].host_pool.size_per_token + for name, size in requests + ), + *ordered_pool_ids, + *ordered_sizes, + ] + synced = torch.tensor( + [*values, *(-value for value in values)], + dtype=torch.int64, + device="cpu", + ) + self._all_reduce_shared_host_groups(synced, torch.distributed.ReduceOp.MIN) + mins = synced[: len(values)].tolist() + maxs = (-synced[len(values) :]).tolist() + if mins[1:] != maxs[1:]: + raise RuntimeError( + "Shared host allocation request diverged across attention ranks: " + f"min={mins[1:]}, max={maxs[1:]}" + ) + return bool(mins[0]) + + def _shared_host_consensus_groups(self) -> list: + groups = [ + group + for group in (self.attn_cp_group, self.attn_tp_group) + if group is not None and torch.distributed.get_world_size(group=group) > 1 + ] + if ( + not groups + and self.tp_group is not None + and torch.distributed.get_world_size(group=self.tp_group) > 1 + ): + groups.append(self.tp_group) + return groups + + def _all_reduce_shared_host_groups(self, tensor: torch.Tensor, op) -> None: + for group in self._shared_host_consensus_groups(): + torch.distributed.all_reduce(tensor, op=op, group=group) + + @staticmethod + def _free_shared_host_allocations( + entries: list[PoolEntry], allocated: list[torch.Tensor] + ) -> None: + for entry, indices in zip(entries, allocated, strict=True): + entry.host_pool.free(indices) + def allocate_shared_host_transfers( self, kv_device_indices: torch.Tensor, @@ -498,11 +595,146 @@ def allocate_shared_host_transfers( return None return host_indices, resolved + def _shared_prefetch_requests( + self, operation: StorageOperation, need_size: int + ) -> tuple[list[tuple[PoolName, int]], list[PoolTransfer]]: + anchor = self.mem_pool_host.anchor_entry + pool_transfers = operation.pool_transfers or [] + requests = [(anchor.name, need_size)] + independent_transfers = [] + for transfer in pool_transfers: + if transfer.indices_from_pool is not None: + continue + if transfer.host_indices is not None: + raise AssertionError( + "Shared prefetch pools must be allocated atomically." + ) + entry = self.mem_pool_host.entry_map.get(transfer.name) + if entry is None: + continue + num_pages = len(transfer.keys or []) + if transfer.hit_policy == PoolHitPolicy.TRAILING_PAGES: + num_pages = min(num_pages or 1, need_size // self.page_size) + requests.append((transfer.name, num_pages * entry.host_pool.page_size)) + independent_transfers.append(transfer) + return requests, independent_transfers + + def can_fit_prefetch_host_buffers( + self, operation: StorageOperation, need_size: int, *, empty: bool = True + ) -> bool: + anchor = self.mem_pool_host.anchor_entry + if not uses_shared_host_layout(anchor.host_pool): + return super().can_fit_prefetch_host_buffers(operation, need_size) + domain = anchor.host_pool.shared_allocation_domain + requests, _ = self._shared_prefetch_requests(operation, need_size) + domain_requests = [ + (self.mem_pool_host.entry_map[name].host_pool.pool_label, size) + for name, size in requests + ] + fits = domain.can_fit_many_then(domain_requests, (), empty=empty) + if not empty: + fits = self._sync_shared_host_value( + int(fits), torch.distributed.ReduceOp.MIN + ) + return bool(fits) + + def prefetch_rate_limited(self) -> bool: + if self.host_memory_mode == "buffer_only" and uses_shared_host_layout( + self.mem_pool_host.anchor_entry.host_pool + ): + # Shared-arena occupancy includes the Full and SWA byte footprint. + return self.prefetch_tokens_occupied >= self.prefetch_capacity_limit + return super().prefetch_rate_limited() + + def alloc_prefetch_host_buffers( + self, operation: StorageOperation, need_size: int + ) -> Optional[torch.Tensor]: + """Atomically allocate every prefetch slice from a shared arena.""" + anchor = self.mem_pool_host.anchor_entry + if not uses_shared_host_layout(anchor.host_pool): + return super().alloc_prefetch_host_buffers(operation, need_size) + + requests, independent_transfers = self._shared_prefetch_requests( + operation, need_size + ) + + allocated = self._alloc_shared_host_requests_with_reclaim(requests) + if allocated is None: + return None + self._sync_trailing_keys( + operation.pool_transfers or [], + operation.hash_value, + need_size // self.page_size, + ) + for transfer, indices in zip(independent_transfers, allocated[1:], strict=True): + transfer.host_indices = indices + return allocated[0] + + def allocate_storage_hit( + self, + operation: StorageOperation, + hit_tokens: int, + *, + allow_partial: bool, + min_tokens: int, + evict_host: Callable[[int], int], + ) -> tuple[Optional[torch.Tensor], int]: + host_indices = self.alloc_prefetch_host_buffers(operation, hit_tokens) + shared = uses_shared_host_layout(self.mem_pool_host.anchor_entry.host_pool) + if host_indices is None and not shared: + evict_host(hit_tokens) + host_indices = self.alloc_prefetch_host_buffers(operation, hit_tokens) + if host_indices is not None or not allow_partial: + return host_indices, hit_tokens + + if shared: + low, high = 0, hit_tokens // self.page_size + while low < high: + mid = (low + high + 1) // 2 + if self.can_fit_prefetch_host_buffers( + operation, mid * self.page_size, empty=False + ): + low = mid + else: + high = mid - 1 + alloc_len = low * self.page_size + else: + available = self.mem_pool_host.available_size() + alloc_len = min(hit_tokens, available - available % self.page_size) + if alloc_len >= min_tokens: + host_indices = self.alloc_prefetch_host_buffers(operation, alloc_len) + return host_indices, alloc_len + + def free_prefetch_host_buffers( + self, operation: StorageOperation, host_indices: torch.Tensor + ) -> None: + """Roll back an unsubmitted prefetch's anchor and independent pools.""" + pool_transfers = operation.pool_transfers or [] + self.mem_pool_host.free(host_indices) + for transfer in pool_transfers: + if transfer.indices_from_pool is not None or transfer.host_indices is None: + continue + entry = self.mem_pool_host.entry_map.get(transfer.name) + if entry is not None: + entry.host_pool.free(transfer.host_indices) + transfer.host_indices = None + def _move_op_indices( self, op: CacheOperation ) -> tuple[torch.Tensor, torch.Tensor, Optional[list[PoolTransfer]]]: return self.move_hybrid_indices(op) + def _move_load_operation( + self, op: CacheOperation + ) -> tuple[torch.Tensor, torch.Tensor, Optional[list[PoolTransfer]]]: + # Resolve unified virtual ids only when the H2D transfer is submitted. + op.device_indices = ( + self.mem_pool_device_allocator.translate_kv_indices_for_transfer( + op.device_indices + ) + ) + return self.move_hybrid_indices(op) + def _move_write_operation( self, op: CacheOperation ) -> tuple[torch.Tensor, torch.Tensor, Optional[list[PoolTransfer]]]: @@ -813,8 +1045,16 @@ def _page_transfer_sidecar( transfers_nonkv, operation.hash_value, kv_completed_pages ) self._resolve_sidecar_nonkv_derived_pool_transfers(operation) - results = self.storage_backend.batch_get_v2(transfers_nonkv) - pool_hits = count_pool_hits(results) + try: + results = self.storage_backend.batch_get_v2(transfers_nonkv) + pool_hits = count_pool_hits(results) + except Exception: + if not get_memory().enable_unified_memory: + raise + logger.exception( + "HiCache sidecar prefetch failed for request %s", + operation.request_id, + ) # Emit PrefetchAck to prefetch_sync_queue, even the operation has been canceled by the # scheduler thread. The prefetch sync thread expects the same number of PrefetchAck objects # to perform all_reduce. @@ -996,7 +1236,7 @@ def _resolve_device_transfers( derived_transfers: list[PoolTransfer] = [] def rollback_allocated() -> None: - for prev_pool, prev_free_fn, prev_indices in newly_allocated: + for prev_pool, prev_free_fn, prev_indices in reversed(newly_allocated): prev_free_fn(prev_indices) prev_pool.device_indices = None diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index bfa2ec19e4ad..c19b69031625 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -281,11 +281,6 @@ def build_hybrid_swa_group( transfer_layer_num = len(full_layer_mapping | swa_layer_mapping) if _uses_unified_page_envelope_host(full_kv_pool, swa_kv_pool, use_mla): memory = get_memory() - if memory.hicache_host_memory_mode == "buffer_only": - raise ValueError( - "--hicache-host-memory-mode buffer_only does not support a " - "shared FULL/SWA unified host arena." - ) kv_host_pool, swa_host_pool = ( UnifiedPageEnvelopeHostPool.build_hybrid_swa_pool_pair( device_pools=(full_kv_pool, swa_kv_pool), @@ -322,6 +317,18 @@ def build_hybrid_swa_group( target_device_layer_num=swa_kv_pool.layer_num, draft_layer_num=len(mtp_swa_device_pools), ) + from sglang.srt.mem_cache.allocator.unified_sub_pool import MultiEndedAllocator + + uses_virtual_device_indices = isinstance(swa_attn_allocator, MultiEndedAllocator) + if swa_attn_allocator is None: + swa_device_alloc_fn = None + swa_device_free_fn = None + elif uses_virtual_device_indices: + swa_device_alloc_fn = swa_attn_allocator.alloc_physical + swa_device_free_fn = swa_attn_allocator.cancel_physical_reservation + else: + swa_device_alloc_fn = swa_attn_allocator.alloc + swa_device_free_fn = swa_attn_allocator.free return HostPoolGroup( [ build_pool_entry( @@ -341,12 +348,8 @@ def build_hybrid_swa_group( transfer_layer_num=transfer_layer_num + 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 - ), + device_alloc_fn=swa_device_alloc_fn, + device_free_fn=swa_device_free_fn, packed_draft_device_pools=mtp_swa_device_pools, ), ] diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index f7685bbfab49..fd4e7c80cebf 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/python/sglang/srt/mem_cache/kv_cache_builder.py @@ -147,14 +147,6 @@ def resolve_decode_retraction_backup(*, tp_worker: BaseTpWorker) -> str: backend = disagg.disaggregation_decode_retraction_backup unified_hybrid_swa = memory.enable_unified_memory and tp_worker.is_hybrid_swa - unified_decode_radix = ( - memory.enable_unified_memory and disagg.disaggregation_decode_enable_radix_cache - ) - if backend == "host_pool" and unified_decode_radix: - raise ValueError( - "Unified-memory H2D/D2H does not support host-pool decode " - "retraction with decode radix cache yet." - ) if backend == "host_pool" and unified_hybrid_swa: raise ValueError( "Unified-memory hybrid-SWA H2D/D2H does not support host-pool " diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 2f4d7d0d5e4a..511f8bdb4776 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -60,6 +60,7 @@ class LogicalHostPool: """ shared_allocation_domain = None + stores_page_envelope = False def __init__(self, size: int, page_size: int, layout: str = "layer_first"): if size % page_size != 0: diff --git a/python/sglang/srt/mem_cache/pool_host/base.py b/python/sglang/srt/mem_cache/pool_host/base.py index af95b5c2bf7d..66aab12117fa 100644 --- a/python/sglang/srt/mem_cache/pool_host/base.py +++ b/python/sglang/srt/mem_cache/pool_host/base.py @@ -135,6 +135,7 @@ class HostKVCache(abc.ABC): dcp_size = 1 dcp_rank = 0 shared_allocation_domain = None + stores_page_envelope = False def __init__( self, diff --git a/python/sglang/srt/mem_cache/pool_host/group.py b/python/sglang/srt/mem_cache/pool_host/group.py index 78bd774f60a1..34f541463cc8 100644 --- a/python/sglang/srt/mem_cache/pool_host/group.py +++ b/python/sglang/srt/mem_cache/pool_host/group.py @@ -227,6 +227,10 @@ def dsa_kv_cache_store_fp8(self): def size_per_token(self): return self.anchor_entry.host_pool.size_per_token + @property + def stores_page_envelope(self) -> bool: + return self.anchor_entry.host_pool.stores_page_envelope + def clear(self) -> None: for entry in self.entries: entry.host_pool.clear() diff --git a/python/sglang/srt/mem_cache/pool_host/unified.py b/python/sglang/srt/mem_cache/pool_host/unified.py index c58211093388..7824fe346f2a 100644 --- a/python/sglang/srt/mem_cache/pool_host/unified.py +++ b/python/sglang/srt/mem_cache/pool_host/unified.py @@ -293,15 +293,30 @@ def _can_reserve_without_compaction(self, page_counts: dict[str, int]) -> bool: extension_bytes += extension_pages * side.page_bytes return extension_bytes <= self._current_gap_bytes() - def _can_fit_packed(self, page_counts: dict[str, int]) -> bool: + def _can_fit_state( + self, + live_page_counts: dict[str, int], + free_logical_page_counts: dict[str, int], + request_page_counts: dict[str, int], + ) -> bool: used_bytes = 0 - for name, page_count in page_counts.items(): - side = self.sides[name] - if page_count > self._logical_free_page_count(side): + for name, page_count in request_page_counts.items(): + if page_count > free_logical_page_counts[name]: return False - used_bytes += (side.live_page_count + page_count) * side.page_bytes + side = self.sides[name] + used_bytes += (live_page_counts[name] + page_count) * side.page_bytes return used_bytes <= self._allocatable_bytes + def _can_fit_packed(self, page_counts: dict[str, int]) -> bool: + return self._can_fit_state( + {name: side.live_page_count for name, side in self.sides.items()}, + { + name: self._logical_free_page_count(side) + for name, side in self.sides.items() + }, + page_counts, + ) + def _begin_compaction(self) -> list: if self._layout_lease_state.depth: raise RuntimeError( @@ -533,6 +548,48 @@ def alloc_many( self.sides[name].free_logical_extents = extents return results + def can_fit_many_then( + self, + requests: Sequence[tuple[str, int]], + following_requests: Sequence[tuple[str, int]], + *, + empty: bool = False, + ) -> bool: + """Whether both allocation groups fit in order without mutating state.""" + with self.lock: + _, first_page_counts = self._request_page_counts(requests) + _, following_page_counts = self._request_page_counts(following_requests) + if empty: + live_page_counts = {name: 0 for name in self.sides} + free_logical_page_counts = { + name: side.page_num for name, side in self.sides.items() + } + else: + live_page_counts = { + name: side.live_page_count for name, side in self.sides.items() + } + free_logical_page_counts = { + name: self._logical_free_page_count(side) + for name, side in self.sides.items() + } + if not self._can_fit_state( + live_page_counts, free_logical_page_counts, first_page_counts + ): + return False + live_after_first = { + name: live_page_counts[name] + first_page_counts[name] + for name in self.sides + } + free_after_first = { + name: free_logical_page_counts[name] - first_page_counts[name] + for name in self.sides + } + return self._can_fit_state( + live_after_first, + free_after_first, + following_page_counts, + ) + def free(self, name: str, indices: torch.Tensor) -> int: with self.lock: side = self.sides[name] @@ -600,6 +657,8 @@ def translate_index(self, name: str, index: int) -> int: class UnifiedPageEnvelopeHostPool(HostKVCache): """Host mirror that transfers complete unified-memory page envelopes.""" + stores_page_envelope = True + def __init__( self, device_pool, diff --git a/python/sglang/srt/mem_cache/rust_tree_core/adapter.py b/python/sglang/srt/mem_cache/rust_tree_core/adapter.py index ff11a09a8072..5c72968c71c3 100644 --- a/python/sglang/srt/mem_cache/rust_tree_core/adapter.py +++ b/python/sglang/srt/mem_cache/rust_tree_core/adapter.py @@ -4,7 +4,7 @@ from __future__ import annotations from array import array -from typing import TYPE_CHECKING, Optional, Sequence +from typing import TYPE_CHECKING, Callable, Optional, Sequence import torch @@ -337,6 +337,20 @@ def __init__(self, params: CacheInitParams): ) self._page_size = params.page_size + self.swa_write_back_eviction_barrier_enabled = False + self._swa_backup_index_mapper: Optional[ + Callable[[torch.Tensor], torch.Tensor] + ] = None + allocator = params.token_to_kv_pool_allocator + if allocator is not None and ComponentType.SWA in self.tree_components: + from sglang.srt.mem_cache.allocator.unified_sub_pool import ( + MultiEndedAllocator, + ) + + if isinstance(allocator.swa_attn_allocator, MultiEndedAllocator): + self._swa_backup_index_mapper = ( + allocator.translate_swa_indices_for_transfer + ) self.is_eagle = ( params.is_eagle and ComponentType.MAMBA not in self.tree_components ) @@ -463,6 +477,11 @@ def evict_device_next_node( result = EvictDeviceNextNodeResult( node_id=binding_result.node_id, made_progress=binding_result.made_progress, + backup_kv=( + _cache_action_from_tagged(binding_result.backup_kv) + if binding_result.backup_kv is not None + else None + ), ) return _fill_evict_result(binding_result, result) @@ -633,6 +652,10 @@ def evict_excess_path_states( def set_hicache_enabled(self) -> None: self._binding.set_hicache_enabled() + def enable_swa_write_back_eviction_barrier(self) -> None: + self.swa_write_back_eviction_barrier_enabled = True + self._binding.enable_swa_write_back_eviction_barrier() + def set_host_memory_buffer_only(self) -> None: self._binding.set_host_memory_buffer_only() @@ -724,7 +747,22 @@ def build_backup_spec( self, node_id: NodeId ) -> tuple[torch.Tensor, dict[ComponentType, list[PoolTransfer]]]: device_value, comp_xfers = self._binding.build_backup_spec(node_id) - return device_value, _comp_xfers_from_binding(comp_xfers) + comp_xfers = _comp_xfers_from_binding(comp_xfers) + if self._swa_backup_index_mapper is not None: + # Full tree values are stable virtual IDs; stored SWA physical IDs are not. + for transfer in comp_xfers.get(ComponentType.SWA, ()): + full_device_value = torch.cat( + [ + self.get_component_device_value(source_id, ComponentType.FULL) + for source_id in transfer.nodes_to_load + ] + ) + assert transfer.device_indices is not None + assert transfer.device_indices.numel() == full_device_value.numel() + transfer.device_indices = self._swa_backup_index_mapper( + full_device_value + ).to(torch.int64) + return device_value, comp_xfers def build_storage_backup_spec( self, node_id: NodeId, pass_prefix_keys: bool diff --git a/python/sglang/srt/mem_cache/storage/umbp/umbp_store.py b/python/sglang/srt/mem_cache/storage/umbp/umbp_store.py index 81e4ad008819..cafd4fa39063 100644 --- a/python/sglang/srt/mem_cache/storage/umbp/umbp_store.py +++ b/python/sglang/srt/mem_cache/storage/umbp/umbp_store.py @@ -919,6 +919,7 @@ def _resolve_node_tags() -> List[str]: # Initialize before the optional constructor-time pool registration. self.registered_pools: dict = {} self._kv_anchor_is_logical = False + self._kv_anchor_is_page_envelope = False self._registered_regions: set = set() self.client = UMBPClient(cfg) @@ -1003,6 +1004,7 @@ def register_mem_pool_host(self, mem_pool_host: HostKVCache): # A logical anchor owns indices; side pools carry the data. self._kv_anchor_is_logical = self.mem_pool_host.kv_buffer is None + self._kv_anchor_is_page_envelope = self.mem_pool_host.stores_page_envelope self._zero_copy_registered = False # Side-pool registration needs the mode even for a logical anchor. @@ -1165,68 +1167,61 @@ def register_mem_host_pool_v2(self, host_pool: HostKVCache, host_pool_name): # ------------------------------------------------------------------ # Key suffix generation — mirrors MooncakeStore # ------------------------------------------------------------------ - def _get_mha_buffer_meta(self, keys, indices): - ptr_list, element_size_list = self.mem_pool_host.get_page_buffer_meta(indices) - key_list = [] - for key_ in keys: - key_list.append(f"{key_}_{self.mha_suffix}_k") - key_list.append(f"{key_}_{self.mha_suffix}_v") - assert len(key_list) == len(ptr_list) - return key_list, ptr_list, element_size_list - - def _get_mha_split_heads_buffer_meta(self, keys, indices): - ptr_list, element_size_list = ( - self.mem_pool_host.get_split_heads_page_buffer_meta( - indices, self.split_factor - ) + def _anchor_key_suffixes(self) -> tuple[str, ...]: + if self._kv_anchor_is_page_envelope: + return (f"_{self.mha_suffix}_kv",) + if self.is_mla_backend: + return (f"_{self.mla_suffix}_k",) + ranks = ( + self.mha_suffix + if self.storage_config and self.storage_config.should_split_heads + else (self.mha_suffix,) + ) + return tuple( + f"_{rank}_{component}" for rank in ranks for component in ("k", "v") ) - key_list = [] - for key_ in keys: - for suffix in self.mha_suffix: - key_list.append(f"{key_}_{suffix}_k") - key_list.append(f"{key_}_{suffix}_v") - assert len(key_list) == len(ptr_list) - return key_list, ptr_list, element_size_list - def _get_mla_buffer_meta(self, keys, indices): - ptr_list, element_size_list = self.mem_pool_host.get_page_buffer_meta(indices) - key_list = [] - for key_ in keys: - key_list.append(f"{key_}_{self.mla_suffix}_k") - assert len(key_list) == len(ptr_list) - return key_list, ptr_list, element_size_list + def _anchor_keys(self, keys) -> tuple[list[str], int]: + suffixes = self._anchor_key_suffixes() + return [f"{key}{suffix}" for key in keys for suffix in suffixes], len(suffixes) def _batch_preprocess(self, keys, host_indices): assert len(keys) > 0 assert len(keys) == len(host_indices) // self.mem_pool_host.page_size - if self.is_mla_backend: - return self._get_mla_buffer_meta(keys, host_indices) + key_list, _ = self._anchor_keys(keys) + if ( + not self._kv_anchor_is_page_envelope + and not self.is_mla_backend + and self.storage_config + and self.storage_config.should_split_heads + ): + ptr_list, element_size_list = ( + self.mem_pool_host.get_split_heads_page_buffer_meta( + host_indices, self.split_factor + ) + ) else: - if self.storage_config and self.storage_config.should_split_heads: - return self._get_mha_split_heads_buffer_meta(keys, host_indices) - else: - return self._get_mha_buffer_meta(keys, host_indices) + ptr_list, element_size_list = self.mem_pool_host.get_page_buffer_meta( + host_indices + ) + assert len(key_list) == len(ptr_list) + return key_list, ptr_list, element_size_list def _batch_postprocess(self, results: List[bool], is_set_operate=False): """Convert per-key-component results to per-page results. - For MHA: each page has K+V → group pairs. - For MLA: each page has K only. + Unified page envelopes and MLA have one object per page. Ordinary MHA + has K+V pairs, or one pair per split rank. """ - if self.is_mla_backend: + group_size = len(self._anchor_key_suffixes()) + if group_size == 1: return list(results) - else: - if self.storage_config and self.storage_config.should_split_heads: - group_size = self.split_factor * 2 - groups = [ - results[i : i + group_size] - for i in range(0, len(results), group_size) - ] - return [all(g) for g in groups] - else: - # Group K/V pairs - kv_pairs = zip(results[::2], results[1::2]) - return [k and v for k, v in kv_pairs] + result_count = len(results) + if not (self.storage_config and self.storage_config.should_split_heads): + result_count -= result_count % group_size + return [ + all(results[i : i + group_size]) for i in range(0, result_count, group_size) + ] # ------------------------------------------------------------------ # Zero-copy v1 interface @@ -1284,19 +1279,8 @@ def _compute_expanded_depths( depths_per_page = [prefix_len + i for i in range(len(keys))] # Expand to match the key_strs layout produced by _batch_preprocess. - expanded = [] - for d in depths_per_page: - if self.is_mla_backend: - expanded.append(d) # MLA: 1 key per page - elif self.storage_config and self.storage_config.should_split_heads: - # split heads: 2 keys per split rank, split_factor ranks per page - for _ in range(self.split_factor): - expanded.append(d) - expanded.append(d) - else: - expanded.append(d) # K - expanded.append(d) # V - return expanded + objects_per_page = len(self._anchor_key_suffixes()) + return [depth for depth in depths_per_page for _ in range(objects_per_page)] def batch_set_v1( self, @@ -1354,23 +1338,7 @@ def batch_exists( self, keys: List[str], extra_info: Optional[HiCacheStorageExtraInfo] = None ) -> int: """Return count of consecutive existing keys from start.""" - if self.is_mla_backend: - query_keys = [f"{key}_{self.mla_suffix}_k" for key in keys] - key_multiplier = 1 - else: - query_keys = [] - if self.storage_config and self.storage_config.should_split_heads: - for key in keys: - for suffix in self.mha_suffix: - query_keys.append(f"{key}_{suffix}_k") - query_keys.append(f"{key}_{suffix}_v") - key_multiplier = 2 * self.split_factor - else: - for key in keys: - query_keys.append(f"{key}_{self.mha_suffix}_k") - query_keys.append(f"{key}_{self.mha_suffix}_v") - key_multiplier = 2 - + query_keys, key_multiplier = self._anchor_keys(keys) hit_count = self.client.batch_exists_consecutive(query_keys) return hit_count // key_multiplier diff --git a/python/sglang/srt/mem_cache/unified_cache/components/base.py b/python/sglang/srt/mem_cache/unified_cache/components/base.py index 642d02f124fe..bf002ee30d7b 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/base.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/base.py @@ -79,6 +79,9 @@ class PreparePrefetchResult: alloc_failed: bool = False # The component's pre-allocated host buffer (None = skip the build). host_indices: Optional[torch.Tensor] = None + # Shared arenas allocate all pool slices atomically after the storage hit + # size is known. The component still emits a keys-only transfer. + deferred_host_allocation: bool = False class CacheTransferPhase(str, Enum): @@ -125,6 +128,7 @@ def __init__(self, cache: UnifiedRadixCache, params: CacheInitParams): # Populated when the component passed to TreeCore constructor. self.tree_core: Optional[UnifiedTreeCore] = None self.is_evict_device_ongoing = False + self._evict_device_backup_node_id: Optional[NodeId] = None # Per-session frontier nodes (the deepest registered node per cached # path), not physical tree leaves: a frontier node may have children. self._session_leaves: dict[str, set[UnifiedTreeNode]] = defaultdict(set) @@ -510,6 +514,7 @@ def evict_device_start(self, request_cnt: int) -> None: assert not self.is_evict_device_ongoing, ( f"{self.component_type} device eviction already in progress" ) + assert self._evict_device_backup_node_id is None self._evict_device_start(request_cnt) self.is_evict_device_ongoing = True @@ -536,6 +541,17 @@ def evict_device_end(self) -> None: ) self._evict_device_end() self.is_evict_device_ongoing = False + self._evict_device_backup_node_id = None + + def take_backup_before_device_eviction(self) -> Optional[NodeId]: + node_id = self._evict_device_backup_node_id + self._evict_device_backup_node_id = None + return node_id + + def request_backup_before_device_eviction(self, node_id: NodeId) -> None: + """Pause the walk so the Controller can back up an internal node.""" + assert self._evict_device_backup_node_id is None + self._evict_device_backup_node_id = node_id @abstractmethod def _evict_device_start(self, request_cnt: int) -> None: diff --git a/python/sglang/srt/mem_cache/unified_cache/components/swa.py b/python/sglang/srt/mem_cache/unified_cache/components/swa.py index f325aa1ef0a9..c389fc93ed33 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/swa.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/swa.py @@ -22,6 +22,7 @@ PoolTransfer, PoolTransferResult, ) +from sglang.srt.mem_cache.pool_host.base import uses_shared_host_layout from sglang.srt.mem_cache.unified_cache.cache_action import ( FreeComponentDeviceSlot, FreeComponentHostSlot, @@ -86,6 +87,7 @@ def __init__(self, cache: UnifiedRadixCache, params: CacheInitParams): ) // params.page_size # HiCache state: set to host SWA pool when HiCache enabled self._swa_kv_pool_host = None + self._evict_device_last_backup_id: Optional[NodeId] = None component_type = ComponentType.SWA @@ -701,6 +703,7 @@ def eviction_priority(self, is_leaf: bool) -> int: def _evict_device_start(self, request_cnt: int) -> None: """Begin the device-eviction walk from this component's LRU cursor.""" self._evict_device_request_cnt = request_cnt + self._evict_device_last_backup_id = None if self.tree_core.enable_session_radix_cache: lru = self.tree_core.lru_lists[self.component_type] lru.cursor_begin() @@ -733,6 +736,15 @@ def _evict_device_next_node( self._evict_device_cursor = ( lru.cursor_next() if enabled else lru.get_lru_no_lock() ) + x = self._evict_device_cursor + if ( + x is not None + and x.id == self._evict_device_last_backup_id + and not x.backuped + ): + self._evict_device_cursor = ( + lru.cursor_next() if enabled else lru.get_prev_no_lock(x) + ) if ( tracker[ct] >= self._evict_device_request_cnt or self._evict_device_cursor is None @@ -749,6 +761,15 @@ def _evict_device_next_node( lru.cursor_next() if enabled else lru.get_prev_no_lock(x) ) return x.id + if ( + self.tree_core.is_write_back + and self.tree_core.swa_write_back_eviction_barrier_enabled + and not x.backuped + ): + # Back up a dirty internal SWA node before its device value is lost. + self._evict_device_last_backup_id = x.id + self.request_backup_before_device_eviction(x.id) + return None if not enabled: x_next = lru.get_prev_no_lock(x) self.tree_core._evict_component_and_detach_lru( @@ -770,6 +791,7 @@ def _evict_device_end(self) -> None: if self.tree_core.enable_session_radix_cache: self.tree_core.lru_lists[self.component_type].cursor_end() self._evict_device_cursor = None + self._evict_device_last_backup_id = None def acquire_component_lock( self, @@ -1021,6 +1043,8 @@ def prepare_prefetch( # device-guaranteed, require a full window. return PreparePrefetchResult() num_tokens = num_pages * self.cache.page_size + if uses_shared_host_layout(self._swa_kv_pool_host): + return PreparePrefetchResult(deferred_host_allocation=True) host_indices = self.cache.host_pool_group.alloc( num_tokens, pool=PoolName.SWA, @@ -1052,12 +1076,27 @@ def build_hicache_transfers( if not unbacked_swa_nodes: return None unbacked_swa_nodes.reverse() + allocator = self._unified_allocator() + if allocator is not None: + full_values = [] + for unbacked_node in unbacked_swa_nodes: + full_value = unbacked_node.component_data[BASE_COMPONENT_TYPE].value + assert full_value is not None + assert len(full_value) == len( + unbacked_node.component_data[ct].value + ) + full_values.append(full_value) + device_indices = allocator.translate_swa_indices_for_transfer( + torch.cat(full_values) + ) + else: + device_indices = torch.cat( + [n.component_data[ct].value for n in unbacked_swa_nodes] + ).to(torch.int64) return [ PoolTransfer( name=PoolName.SWA, - device_indices=torch.cat( - [n.component_data[ct].value for n in unbacked_swa_nodes] - ).to(torch.int64), + device_indices=device_indices, nodes_to_load=[n.id for n in unbacked_swa_nodes], ) ] @@ -1116,10 +1155,18 @@ def build_hicache_transfers( ] if phase == CacheTransferPhase.PREFETCH: - assert host_indices is not None # Keys are unknowable at build time; placeholders carry the # count, _sync_trailing_keys fills the real trailing hashes. - num_pages = host_indices.numel() // self.tree_core.page_size + if host_indices is None: + assert uses_shared_host_layout(self._swa_kv_pool_host), ( + "deferred SWA prefetch allocation requires a shared host arena" + ) + num_pages = min( + self.full_window_pages, + prefetch_tokens // self.tree_core.page_size, + ) + else: + num_pages = host_indices.numel() // self.tree_core.page_size return [ PoolTransfer( name=PoolName.SWA, diff --git a/python/sglang/srt/mem_cache/unified_cache/storage_attachment.py b/python/sglang/srt/mem_cache/unified_cache/storage_attachment.py index 2c02d29b29bf..9a8b02b10ce7 100644 --- a/python/sglang/srt/mem_cache/unified_cache/storage_attachment.py +++ b/python/sglang/srt/mem_cache/unified_cache/storage_attachment.py @@ -383,15 +383,19 @@ def _release_pending_storage_ops(self) -> None: cache.buffer_pipeline.release_anchor_lock(req_id) controller.append_host_mem_release( host_indices=info.host_indices[:completed_tokens], - extra_pools=[ - x for xfers in info.comp_xfers.values() for x in xfers - ], + extra_pools=( + [x for xfers in info.comp_xfers.values() for x in xfers] + if info.operation.pool_transfers_done + else None + ), ) controller.prefetch_tokens_occupied = max( 0, controller.prefetch_tokens_occupied - cache._prefetch_occupied_span( - info.prefetch_key, info.host_indices + info.prefetch_key, + info.host_indices, + operation=info.operation, ), ) except Exception: diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index 796a39912cc5..6c1cdc223a3b 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py @@ -421,6 +421,7 @@ def __init__( self.write_through_threshold = 256 self.is_write_back = False self.has_swa_host_pool = False + self.swa_write_back_eviction_barrier_enabled = False self.enable_session_radix_cache = params.enable_session_radix_cache self.eviction_strategy = get_eviction_strategy( params.eviction_policy.lower(), params.eviction_policy_config @@ -1480,15 +1481,20 @@ def evict_device_next_node( # The walk reads running totals for its doneness check; the result # carries only this step's delta. updated_tracker = defaultdict(int, tracker) + component = self.components_by_type[component_type] self._begin_tracking_unbacked_tokens() try: - result.node_id = self.components_by_type[ - component_type - ].evict_device_next_node( + result.node_id = component.evict_device_next_node( updated_tracker, result.device_frees, result.host_frees ) finally: result.unbacked_tokens = self._finish_tracking_unbacked_tokens() + backup_node_id = component.take_backup_before_device_eviction() + if backup_node_id is not None: + assert result.node_id is None + result.backup_kv = self._build_backup_kv_action( + self.node_by_id(backup_node_id), write_back=True + ) for ct, n in updated_tracker.items(): delta = n - tracker.get(ct, 0) if delta: @@ -1536,36 +1542,32 @@ def evict_device_leaf( result.unbacked_tokens = self._finish_tracking_unbacked_tokens() def drop_subtree_no_host(self, node_id: NodeId) -> DropSubtreeNoHostResult: - """Write-back fallback when a D-leaf's D->H backup fails under host - memory pressure: drop the subtree rooted at the unbacked leaf so - device eviction keeps making progress instead of leaving its KV - unevictable until host space frees up.""" + """Drop an unlocked subtree when its write-back cannot reserve host KV.""" result = DropSubtreeNoHostResult(is_dropped=False) node = self.node_by_id(node_id) - assert self._is_device_leaf(node), f"node {node.id} is not a D-leaf" # A failed backup never issues the D->H copy, so the subtree root has # no host state and no in-flight DMA reading its device slots. assert not node.backuped and node.write_through_pending_id is None - if any(cd.host_lock_ref > 0 for cd in node.component_data): - return result - descendants: list[UnifiedTreeNode] = [] - stack = list(node.children.values()) + subtree: list[UnifiedTreeNode] = [] + stack = [node] while stack: cur = stack.pop() - if any( - cd.lock_ref > 0 or cd.host_lock_ref > 0 for cd in cur.component_data + if ( + cur.write_through_pending_id is not None + or cur.load_back_pending_id is not None + or any( + cd.lock_ref > 0 or cd.host_lock_ref > 0 for cd in cur.component_data + ) ): return result - descendants.append(cur) + subtree.append(cur) stack.extend(cur.children.values()) - for desc in reversed(descendants): - # Host-only by construction: a device descendant would contradict - # this node being a D-leaf, and D-leaves evict before ancestors. - assert desc.evicted and desc.backuped, f"node {desc.id} not host-only" - assert desc.write_through_pending_id is None + for desc in reversed(subtree[1:]): + if not desc.evicted and desc.backuped: + self.kv_events.record_remove(desc, medium=StorageMedium.CPU) self._release_all_component_layers( desc, - StorageMedium.CPU, + StorageMedium.CPU if desc.evicted else StorageMedium.GPU, result.tracker, result.device_frees, result.host_frees, @@ -2046,6 +2048,9 @@ def _is_host_leaf(self, node: UnifiedTreeNode) -> bool: def set_hicache_enabled(self) -> None: self.enable_hicache = True + def enable_swa_write_back_eviction_barrier(self) -> None: + self.swa_write_back_eviction_barrier_enabled = True + def set_host_memory_buffer_only(self) -> None: self.is_host_memory_buffer_only = True diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py index 4fabcd2dc255..5ec7f94950ec 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py @@ -49,6 +49,7 @@ class EvictDeviceNextNodeResult(BaseEvictionResult): """ node_id: Optional[NodeId] = None + backup_kv: Optional[BackupKV] = None made_progress: bool = False unbacked_tokens: int = 0 @@ -152,6 +153,7 @@ class UnifiedTreeCoreInterface(ABC): write_through_threshold: int is_write_back: bool has_swa_host_pool: bool + swa_write_back_eviction_barrier_enabled: bool # Whether the host tier stages one node per FIFO backup intent. is_host_memory_buffer_only: bool kv_events: KVCacheEventRecorder @@ -458,6 +460,11 @@ def set_hicache_enabled(self) -> None: """Mark the host tier (HiCache) as wired.""" ... + @abstractmethod + def enable_swa_write_back_eviction_barrier(self) -> None: + """Preserve dirty SWA before cache-mode write-back eviction.""" + ... + @abstractmethod def set_host_memory_buffer_only(self) -> None: """Mark the host tier as buffer-only: one node staged per backup intent.""" diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index c5497d62cf57..06070db949bc 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -283,6 +283,7 @@ def __init__( "declined_rate_limited": 0, "declined_anchor_lost": 0, "declined_device_covered": 0, + "declined_host_oversize": 0, "revoked_insufficient": 0, "revoked_full_miss": 0, "l3_demand_requests": 0, @@ -484,11 +485,11 @@ def init_hicache(self, server_args: ServerArgs, params: CacheInitParams) -> None # content would drop-newest and punch storage holes. write_backlog_cap=2 * self.token_to_kv_pool_allocator.size_full, ) + # State initialization + if self.buffer_pipeline is not None: self.cache_controller.host_write_staged_tokens_fn = lambda: ( self.buffer_pipeline.write_staged_tokens_ ) - - # State initialization self.write_through_threshold = ( 1 if get_memory().hicache_write_policy == "write_through" else 2 ) @@ -496,6 +497,13 @@ def init_hicache(self, server_args: ServerArgs, params: CacheInitParams) -> None self.cache_controller is not None and self.cache_controller.write_policy == "write_back" ) + # Preserve the SWA host window before device eviction makes it unrecoverable. + if ( + get_memory().enable_unified_memory + and self.host_memory_mode == "cache" + and self.tree_core.has_swa_host_pool + ): + self.tree_core.enable_swa_write_back_eviction_barrier() # Pre-seed the logical dropped-tokens series. if self.metrics_collector is not None and self.cache_controller is not None: reasons = ["host_pressure"] @@ -749,16 +757,40 @@ def _accumulate_tracker( def _evict_device_next_node( self, component_type: ComponentType, tracker: dict[ComponentType, int] ) -> tuple[Optional[NodeId], bool]: - """Advance the eviction walk one node, consuming its step result.""" - result = self.tree_core.evict_device_next_node(component_type, tracker) - self._free_values(result.device_frees, result.host_frees) - if self._tracks_write_through_unbacked_evictions(): - self._record_dropped_tokens( - result.unbacked_tokens, - reason="write_through_unbacked_eviction", - ) - self._accumulate_tracker(tracker, result.tracker) - return result.node_id, result.made_progress + """Advance the walk, completing any pre-eviction backup barriers.""" + while True: + result = self.tree_core.evict_device_next_node(component_type, tracker) + self._free_values(result.device_frees, result.host_frees) + if self._tracks_write_through_unbacked_evictions(): + self._record_dropped_tokens( + result.unbacked_tokens, + reason="write_through_unbacked_eviction", + ) + self._accumulate_tracker(tracker, result.tracker) + if result.backup_kv is None: + return result.node_id, result.made_progress + + assert result.node_id is None + assert self.buffer_pipeline is None, ( + "SWA write-back eviction barriers are cache-mode only" + ) + written = self._execute_and_commit_kv_backup( + result.backup_kv, write_back=True + ) + if written <= 0: + node_id = result.backup_kv.node_ids[0] + dropped = self._drop_subtree_no_host(node_id, tracker) + logger.warning( + "write_back: auxiliary backup failed under host pressure " + "(component=%s, node=%d, subtree_dropped=%s)", + component_type.name, + node_id, + dropped, + ) + # A protected subtree stays intact. The component skips a + # failed backup on resume so later victims can still progress. + continue + self.writing_check(write_back=True) def _evict_device_leaf( self, node_id: NodeId, tracker: dict[ComponentType, int] @@ -1376,14 +1408,17 @@ def _retraction_device_transfers( component_transfers: dict[ComponentType, list[PoolTransfer]] = {} if self.supports_swa(): - kv_cache = self.token_to_kv_pool_allocator.get_kvcache() assert self.sliding_window_size is not None window_start = max(0, num_tokens - self.sliding_window_size) window_start = window_start // self.page_size * self.page_size window_indices = self.req_to_token_pool.req_to_token[ req.kv.req_pool_idx, window_start:num_tokens ].to(torch.int64) - swa_indices = kv_cache.translate_loc_from_full_to_swa(window_indices) + swa_indices = ( + self.token_to_kv_pool_allocator.translate_swa_indices_for_transfer( + window_indices + ) + ) assert bool((swa_indices > 0).all()), ( f"unmapped SWA window positions for request {req.rid}" ) @@ -1582,8 +1617,15 @@ def _execute_kv_backup(self, node_id, device_value, comp_xfers, sidecar_xfers): return None aux_xfers = [x for xfers in comp_xfers.values() for x in xfers] aux_xfers.extend(sidecar_xfers) + transfer_device_value = ( + self.token_to_kv_pool_allocator.translate_kv_indices_for_transfer( + device_value + ) + ) return self.cache_controller.write( - device_value, node_id=node_id, extra_pools=aux_xfers or None + transfer_device_value, + node_id=node_id, + extra_pools=aux_xfers or None, ) def _track_write_through_node( @@ -1982,7 +2024,7 @@ def prefetch_from_storage( if prep.alloc_failed: alloc_failed = True break - if prep.host_indices is None: + if prep.host_indices is None and not prep.deferred_host_allocation: continue transfers = self.tree_core.build_hicache_transfers( ct, @@ -2313,7 +2355,9 @@ def _check_hybrid_prefetch_result( self.buffer_pipeline.release_anchor_lock(req_id) del self.ongoing_prefetch[req_id] self.cache_controller.prefetch_tokens_occupied -= ( - self._prefetch_occupied_span(prefetch_key, host_indices) + self._prefetch_occupied_span( + prefetch_key, host_indices, operation=operation + ) ) self.prefetch_loaded_tokens_by_reqid[req_id] = 0 self.prefetch_loaded_storage_start_by_reqid.pop(req_id, None) @@ -2498,7 +2542,7 @@ def release_aborted_request(self, rid: str) -> None: # Buffer mode granted occupancy at hit-alloc, sized to the bounce; # cache mode reserved the requested span at enqueue. self.cache_controller.prefetch_tokens_occupied -= self._prefetch_occupied_span( - prefetch_key, host_indices + prefetch_key, host_indices, operation=operation ) def _invalidate_absent_from_hit_query(self, operation) -> None: @@ -2535,11 +2579,18 @@ def _account_prefetch_outcome(self, operation, revoked: bool) -> None: def prefetch_outcome_stats_snapshot(self) -> dict: return self._prefetch_outcome_stats.copy() - def _prefetch_occupied_span(self, prefetch_key, host_indices) -> int: + def _prefetch_occupied_span( + self, prefetch_key, host_indices, *, operation=None + ) -> int: """Occupancy units held by a prefetch: cache mode reserves the requested span at enqueue; buffer mode grants at hit-alloc, sized to the allocation (0 while still querying / parked).""" if self.host_memory_mode == "buffer_only": + if ( + operation is not None + and operation.buffer_host_occupied_units is not None + ): + return operation.buffer_host_occupied_units return len(host_indices) if host_indices is not None else 0 return len(prefetch_key) @@ -2573,7 +2624,9 @@ def revoke_pending_prefetch(self, req_id: str) -> None: cc.prefetch_tokens_occupied = max( 0, cc.prefetch_tokens_occupied - - self._prefetch_occupied_span(prefetch_key, _host_indices), + - self._prefetch_occupied_span( + prefetch_key, _host_indices, operation=operation + ), ) def _drain_storage_control_queues_impl( @@ -2647,20 +2700,24 @@ def _try_alloc_storage_hit(operation) -> bool: self.revoke_pending_prefetch(req_id) return True alloc_len = hit_tokens - host_indices = cc.mem_pool_host.alloc(alloc_len) - if host_indices is None: - self.evict_host(alloc_len) - host_indices = cc.mem_pool_host.alloc(alloc_len) - if host_indices is None and not buffer_mode: - # Memory-pressure fallback: a shorter page-aligned prefix. - # (Cache mode only — buffer mode parks for the full hit.) - available_size = cc.mem_pool_host.available_size() - alloc_len = min( - hit_tokens, - available_size - (available_size % self.page_size), + if buffer_mode and not cc.can_fit_prefetch_host_buffers( + operation, alloc_len + ): + self._prefetch_outcome_stats["declined_host_oversize"] += 1 + logger.warning( + "HiCache buffer prefetch declined req=%s: the Full/sidecar " + "bounce cannot fit in an empty host arena", + req_id, ) - if alloc_len >= self.prefetch_threshold: - host_indices = cc.mem_pool_host.alloc(alloc_len) + self.revoke_pending_prefetch(req_id) + return True + host_indices, alloc_len = cc.allocate_storage_hit( + operation, + hit_tokens, + allow_partial=not buffer_mode, + min_tokens=self.prefetch_threshold, + evict_host=self.evict_host, + ) if host_indices is None: if buffer_mode: return False @@ -2678,8 +2735,25 @@ def _try_alloc_storage_hit(operation) -> bool: operation.host_indices = host_indices self.ongoing_prefetch[req_id] = info._replace(host_indices=host_indices) if buffer_mode: - cc.prefetch_tokens_occupied += alloc_len - cc.prefetch_buffer.put(operation) + operation.buffer_host_occupied_units = ( + self.buffer_pipeline.host_allocation_units( + host_indices, operation.pool_transfers + ) + ) + cc.prefetch_tokens_occupied += operation.buffer_host_occupied_units + try: + cc.prefetch_buffer.put(operation) + except Exception: + cc.free_prefetch_host_buffers(operation, host_indices) + operation.host_indices = None + if buffer_mode: + cc.prefetch_tokens_occupied -= operation.buffer_host_occupied_units + operation.buffer_host_occupied_units = None + self.buffer_pipeline.pop_prefix_ctx(req_id) + self.buffer_pipeline.release_anchor_lock(req_id) + self.ongoing_prefetch.pop(req_id, None) + operation.mark_terminate() + raise return True def _drain_and_alloc_storage_hit(): diff --git a/rust/sglang-radix-tree/src/components/swa.rs b/rust/sglang-radix-tree/src/components/swa.rs index 4789adbe3c97..15a22069f960 100644 --- a/rust/sglang-radix-tree/src/components/swa.rs +++ b/rust/sglang-radix-tree/src/components/swa.rs @@ -789,6 +789,14 @@ impl TreeComponent for SwaComponent { if !tree_core.device_lru_list(SWA).in_list(Some(x)) { break None; } + if tree_core.component_state(SWA).evict_device_last_backup == Some(x) + && !tree_core.arena.node(x).backuped() + { + cursor = tree_core + .device_lru_list(SWA) + .get_prev_no_lock(x, &tree_core.arena); + continue; + } assert!( tree_core.arena.has_device_value(x, SWA), "Swa eviction cursor on a valueless node {x}" @@ -803,6 +811,18 @@ impl TreeComponent for SwaComponent { if tree_core.evictable_device_leaves.contains(x) { break Some(x); } + if tree_core.is_write_back + && tree_core.swa_write_back_eviction_barrier_enabled + && !tree_core.arena.node(x).backuped() + { + // A later Full backup cannot recover SWA data after this + // internal node is tombstoned. Pause on the same cursor so + // the Controller can preserve the dirty path first. + tree_core.component_state_mut(SWA).evict_device_backup_node = Some(x); + tree_core.component_state_mut(SWA).evict_device_last_backup = Some(x); + cursor = Some(x); + break None; + } // Internal nodes are tombstoned inline (no IO). tree_core.evict_component_and_detach_lru_( x, @@ -927,7 +947,7 @@ impl TreeComponent for SwaComponent { _mamba_pool_idx: Option, host_indices: Option, _token_ids: Option<&[i64]>, - _prefetch_tokens: usize, + prefetch_tokens: usize, _last_hash: Option<&str>, ) -> Result>, TreeCoreRuntimeError> { // unified_kv keeps SWA as a device-only ring. @@ -1024,11 +1044,17 @@ impl TreeComponent for SwaComponent { }]) } CacheTransferPhase::Prefetch => { - let host_indices = host_indices.expect("SWA PREFETCH build requires host indices"); - let sw_pages = host_indices.numel() / tree_core.page_size; + let sw_pages = host_indices.as_ref().map_or_else( + || { + self.sliding_window_size + .div_ceil(tree_core.page_size) + .min(prefetch_tokens / tree_core.page_size) + }, + |indices| indices.numel() / tree_core.page_size, + ); Some(vec![PoolTransfer { name: PoolName::Swa, - host_indices: Some(host_indices), + host_indices, keys: Some(vec!["__placeholder__".to_string(); sw_pages]), hit_policy: PoolHitPolicy::TrailingPages, ..Default::default() diff --git a/rust/sglang-radix-tree/src/python_bindings.rs b/rust/sglang-radix-tree/src/python_bindings.rs index d34a2862b82c..681afe3da65c 100644 --- a/rust/sglang-radix-tree/src/python_bindings.rs +++ b/rust/sglang-radix-tree/src/python_bindings.rs @@ -751,6 +751,7 @@ fn tracker_to_py(tracker: HashMap) -> HashMap { #[pyclass(get_all)] pub struct EvictDeviceNextNodeResultBinding { node_id: Option, + backup_kv: Option>, made_progress: bool, tracker: HashMap, new_device_frees: Py, @@ -1226,6 +1227,10 @@ impl TreeCoreBinding { let made_progress = node_id.is_some() || !result.tracker.is_empty(); Ok(EvictDeviceNextNodeResultBinding { node_id, + backup_kv: result + .backup_kv + .map(|backup| cache_action_to_py(py, CacheAction::BackupKV(backup))) + .transpose()?, made_progress, tracker: tracker_to_py(result.tracker), new_device_frees: frees_to_py(py, result.device_frees)?, @@ -1386,6 +1391,11 @@ impl TreeCoreBinding { py.allow_threads(|| self.core().enable_hicache) } + /// Preserve dirty SWA data before cache-mode write-back eviction. + fn enable_swa_write_back_eviction_barrier(&self, py: Python<'_>) { + py.allow_threads(|| self.core().enable_swa_write_back_eviction_barrier()); + } + /// Mark the SWA host pool as wired (HiCache). fn set_has_swa_host_pool(&self, py: Python<'_>) { py.allow_threads(|| self.core().set_has_swa_host_pool()); @@ -2656,6 +2666,11 @@ macro_rules! tree_core_binding { self.inner.enable_hicache(py) } + /// Preserve dirty SWA data before cache-mode write-back eviction. + fn enable_swa_write_back_eviction_barrier(&self, py: Python<'_>) { + self.inner.enable_swa_write_back_eviction_barrier(py) + } + /// Mark the SWA host pool as wired (HiCache). fn set_has_swa_host_pool(&self, py: Python<'_>) { self.inner.set_has_swa_host_pool(py) diff --git a/rust/sglang-radix-tree/src/unified_tree_core.rs b/rust/sglang-radix-tree/src/unified_tree_core.rs index 44ab0acff089..6b5c366f44ab 100644 --- a/rust/sglang-radix-tree/src/unified_tree_core.rs +++ b/rust/sglang-radix-tree/src/unified_tree_core.rs @@ -356,7 +356,7 @@ pub struct PoolTransferResult { } /// A device->host backup work item for the cache to execute. -#[derive(Default)] +#[derive(Default, Debug)] pub struct BackupKV { /// Backup these nodes device->host in order, stopping at the first failure; the /// caller orders them parent-before-child for write-through and child-first for @@ -451,6 +451,11 @@ pub struct ComponentState { /// leaf may be freed: the leaf's parent for Full, the LRU predecessor for /// SWA and Mamba. pub(crate) evict_device_cursor: Option, + /// Internal node whose component value must be backed up before the walk + /// can tombstone it. The Controller consumes this request between steps. + pub(crate) evict_device_backup_node: Option, + /// A resumed, still-unbacked victim is skipped after failed host allocation. + pub(crate) evict_device_last_backup: Option, /// Token budget for the current eviction walk. pub(crate) evict_device_request_cnt: usize, } @@ -507,6 +512,7 @@ pub struct EvictionStepResult { pub tracker: HashMap, pub device_frees: HashMap>, pub host_frees: HashMap>, + pub backup_kv: Option, } /// The radix tree mechanism: owns the tree structure, per-node values, the @@ -548,6 +554,8 @@ pub struct UnifiedTreeCore { pub(crate) enable_external_cache_linker: bool, /// Whether the cache wired a host SWA pool (HiCache). pub(crate) has_swa_host_pool: bool, + /// Whether dirty internal SWA nodes must be backed up before eviction. + pub(crate) swa_write_back_eviction_barrier_enabled: bool, /// Whether tree mutations emit BlockStored/BlockRemoved events. pub(crate) enable_kv_cache_events: bool, /// Queued placement events, drained by take_events. @@ -660,6 +668,8 @@ impl UnifiedTreeCore { state.is_evict_device_ongoing = true; state.evict_device_request_cnt = request_cnt; state.evict_device_cursor = None; + state.evict_device_backup_node = None; + state.evict_device_last_backup = None; } /// Finish the component's device-eviction bookkeeping; panics if no walk @@ -672,6 +682,8 @@ impl UnifiedTreeCore { ); state.is_evict_device_ongoing = false; state.evict_device_cursor = None; + state.evict_device_backup_node = None; + state.evict_device_last_backup = None; } /// Add newly evictable device tokens to the component's evictable size. @@ -731,6 +743,7 @@ impl UnifiedTreeCore { enable_storage: false, enable_external_cache_linker: false, has_swa_host_pool: params.has_swa_host_pool, + swa_write_back_eviction_barrier_enabled: false, enable_kv_cache_events: params.enable_kv_cache_events, kv_event_queue: Vec::new(), namespaced_event_hashes: HashMap::new(), @@ -2079,6 +2092,15 @@ impl UnifiedTreeCore { &mut result.device_frees, &mut result.host_frees, ); + let backup_node = self + .component_state_mut(component_type) + .evict_device_backup_node + .take(); + if let Some(backup_node) = backup_node { + assert!(node_id.is_none()); + result.backup_kv = + Some(self.build_backup_kv_action_(self.arena.node(backup_node), true)); + } for (ct, total) in tracker { let delta = total - baseline.get(&ct).copied().unwrap_or(0); if delta > 0 { @@ -2136,10 +2158,7 @@ impl UnifiedTreeCore { Ok((None, result)) } - /// Write-back fallback when a D-leaf's D->H backup fails under host - /// memory pressure: drop the subtree rooted at the unbacked leaf so - /// device eviction keeps making progress instead of leaving its KV - /// unevictable until host space frees up. + /// Drop an unlocked subtree when its write-back cannot reserve host KV. pub fn drop_subtree_no_host( &mut self, node_id: NodeId, @@ -2148,47 +2167,37 @@ impl UnifiedTreeCore { let mut result = EvictionStepResult::default(); { let node = self.arena.node(node_id); - assert!( - self.is_evictable_device_leaf_(node), - "node {node_id} is not a D-leaf" - ); // A failed backup never issues the D->H copy, so the subtree root has // no host state and no in-flight DMA reading its device slots. assert!(!node.backuped() && node.write_through_pending_id.is_none()); - if node.is_host_locked() { - return Ok((false, result)); - } } - let mut descendants: Vec = Vec::new(); - let mut stack: Vec = self - .arena - .node(node_id) - .children - .values() - .copied() - .collect(); + let mut subtree: Vec = Vec::new(); + let mut stack = vec![node_id]; while let Some(cur_id) = stack.pop() { let cur = self.arena.node(cur_id); - if cur.is_device_locked() || cur.is_host_locked() { + if cur.is_device_locked() + || cur.is_host_locked() + || cur.write_through_pending_id.is_some() + || cur.is_load_back_pending() + { return Ok((false, result)); } - descendants.push(cur_id); + subtree.push(cur_id); stack.extend(cur.children.values().copied()); } - for &desc_id in descendants.iter().rev() { - { - let desc = self.arena.node(desc_id); - // Host-only by construction: a device descendant would contradict - // this node being a D-leaf, and D-leaves evict before ancestors. - assert!( - desc.evicted() && desc.backuped(), - "node {desc_id} not host-only" - ); - assert!(desc.write_through_pending_id.is_none()); - } + for &desc_id in subtree[1..].iter().rev() { + let desc = self.arena.node(desc_id); + let medium = if desc.evicted() { + StorageMedium::Cpu + } else { + if desc.backuped() { + self.record_remove_event_(desc_id, StorageMedium::Cpu); + } + StorageMedium::Gpu + }; self.release_all_component_layers_( desc_id, - StorageMedium::Cpu, + medium, &mut result.tracker, &mut result.device_frees, &mut result.host_frees, @@ -2720,6 +2729,11 @@ impl UnifiedTreeCore { self.enable_hicache = true; } + /// Preserve dirty internal SWA nodes before cache-mode write-back eviction. + pub fn enable_swa_write_back_eviction_barrier(&mut self) { + self.swa_write_back_eviction_barrier_enabled = true; + } + /// Mark the host tier as buffer-only; wired after the host pools are built. pub fn set_host_memory_buffer_only(&mut self) { self.is_host_memory_buffer_only = true; diff --git a/test/registered/unit/disaggregation/test_unified_memory_move_gate.py b/test/registered/unit/disaggregation/test_unified_memory_move_gate.py index c33f6654a10f..112014db8bf5 100644 --- a/test/registered/unit/disaggregation/test_unified_memory_move_gate.py +++ b/test/registered/unit/disaggregation/test_unified_memory_move_gate.py @@ -150,6 +150,7 @@ class _Peer: def __init__(self, gate): self.lazy_compaction = True self._free_phys_pages = [0, 1, 2, 3] # only len() is read + self._pending_hicache_load_pages = 0 self.entry_bytes_per_page = 512 self.disagg_move_gate = gate diff --git a/test/registered/unit/mem_cache/test_buffer_mode_sidecar.py b/test/registered/unit/mem_cache/test_buffer_mode_sidecar.py index 3a9fcae76dfb..8c121b21d967 100644 --- a/test/registered/unit/mem_cache/test_buffer_mode_sidecar.py +++ b/test/registered/unit/mem_cache/test_buffer_mode_sidecar.py @@ -35,7 +35,11 @@ class TestBufferModeSidecar(unittest.TestCase): def _swa_component(): return SimpleNamespace( full_window_pages=2, - _swa_kv_pool_host=SimpleNamespace(page_size=2, size=8), + _swa_kv_pool_host=SimpleNamespace( + page_size=2, + size=8, + shared_allocation_domain=None, + ), ) @staticmethod @@ -135,6 +139,7 @@ def test_write_stages_and_persists_dsv4_full_and_swa_sidecars(self): ] controller = MagicMock() + controller.mem_pool_host.anchor_entry.host_pool.shared_allocation_domain = None controller.mem_pool_host.entry_map = { PoolName.SWA: SimpleNamespace( host_pool=SimpleNamespace(page_size=page_size) @@ -166,6 +171,9 @@ def _write(device_value, *, node_id, extra_pools): cache.cache_controller = controller cache.page_size = page_size cache._build_backup_sidecar.return_value = sidecars + cache.token_to_kv_pool_allocator.translate_kv_indices_for_transfer.side_effect = ( + lambda indices: indices + ) pipeline = BufferModePipeline.__new__(BufferModePipeline) pipeline._cache = cache @@ -231,12 +239,13 @@ def test_completed_prefetch_keeps_dsv4_full_and_swa_sidecars_for_h2d(self): ) for spec in self._dsv4_specs() ] + host_indices = torch.arange(4, dtype=torch.int64) operation = SimpleNamespace( id=23, pool_transfers=sidecars, storage_start=0, + buffer_host_occupied_units=len(host_indices), ) - host_indices = torch.arange(4, dtype=torch.int64) req_id = "sidecar-prefetch" cache = MagicMock() diff --git a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py index b401f8ab7a44..27977d1beec5 100644 --- a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py +++ b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py @@ -458,9 +458,8 @@ def test_pop_preallocated_rechecks_budget_after_lock(self): scheduler.output_streamer = MagicMock() queue.scheduler = scheduler - # The 4-token match is locked, then capped to zero because the whole - # 8-token request is inside the SWA window. Admission rejection must - # still release the original matched-node lock. + # Admission rejection must release the matched-node lock, including + # when the SWA lock was already released for fresh tail allocation. queue._allocatable_token_budgets = MagicMock(return_value=3) preallocated, failed = queue.pop_preallocated() @@ -478,7 +477,6 @@ def test_pop_preallocated_rechecks_budget_after_lock(self): skip_swa=True, ) self.assertFalse(req.swa_prefix_lock_released) - queue._swa_tail_len.assert_called_with(8) queue._allocatable_token_budgets.assert_called_once() def test_hicache_restore_commit_hands_over_lock_with_receipt(self): diff --git a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py b/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py index baeb7a1f2346..92e74b6c61da 100644 --- a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py +++ b/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py @@ -201,6 +201,7 @@ def setUp(self): def _start_writing(controller): with mock.patch.object(transfer_module, "device_module", _FakeDeviceModule): controller.l2_transfer_engine = L2TransferEngine("kernel") + controller.mem_pool_device_allocator = mock.Mock() controller.start_writing() def test_hybrid_load_forwards_merged_pool_transfers(self): @@ -217,11 +218,12 @@ def test_hybrid_load_forwards_merged_pool_transfers(self): controller.load_queue = [op, op] controller.layer_done_counter = mock.MagicMock() controller.layer_done_counter.update_producer.return_value = 0 - controller._move_op_indices.side_effect = lambda op: ( + controller._move_load_operation.side_effect = lambda op: ( op.host_indices, op.device_indices, op.pool_transfers, ) + controller.mem_pool_device_allocator = mock.Mock() controller.mem_pool_host = _host_group_stub([], can_use_write_back_jit=False) controller._l2_transfers.side_effect = lambda *args: ( HybridCacheController._l2_transfers(controller, *args) @@ -242,7 +244,7 @@ def test_hybrid_load_forwards_merged_pool_transfers(self): self.assertEqual(HybridCacheController.start_loading(controller), 0) - merged_op = controller._move_op_indices.call_args.args[0] + merged_op = controller._move_load_operation.call_args.args[0] merged_transfer = merged_op.pool_transfers[0] self.assertEqual(merged_transfer.host_indices.tolist(), [0, 1, 0, 1]) self.assertEqual(merged_transfer.keys, ["page-key", "page-key"]) diff --git a/test/registered/unit/mem_cache/test_hybrid_pool_assembler.py b/test/registered/unit/mem_cache/test_hybrid_pool_assembler.py index 7de4d16d654f..2c39fc26c694 100644 --- a/test/registered/unit/mem_cache/test_hybrid_pool_assembler.py +++ b/test/registered/unit/mem_cache/test_hybrid_pool_assembler.py @@ -1,12 +1,18 @@ """Unit tests for hybrid HiCache pool assembly.""" import unittest +from queue import Queue from types import SimpleNamespace from unittest.mock import MagicMock, patch import torch from sglang.srt.mem_cache.base_prefix_cache import EvictParams +from sglang.srt.mem_cache.hicache_storage import PoolHitPolicy, PoolName, PoolTransfer +from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( + HybridCacheController, + PrefetchOperation, +) from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import ( _evict_mamba_for_device_alloc, _evict_swa_for_device_alloc, @@ -14,6 +20,7 @@ build_full_draft_pools, ) from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool +from sglang.srt.mem_cache.pool_host import HostPoolGroup, PoolEntry from sglang.srt.mem_cache.pool_host.mha import get_mha_host_pool_cls from sglang.srt.mem_cache.pool_host.unified import UnifiedPageEnvelopeHostPool from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools @@ -172,6 +179,86 @@ def _build_unified_host_pair(bundle): class TestUnifiedPageEnvelopeHostPool(CustomTestCase): + def test_sidecar_read_error_is_only_a_cache_miss_with_unified_memory(self): + from sglang.srt.runtime_context import publish, reset_context + from sglang.srt.server_args import ServerArgs + + class FailingStorage: + def batch_get_v2(self, transfers): + raise ValueError("sidecar read failed") + + self.addCleanup(reset_context) + for unified in (False, True): + with self.subTest(unified=unified): + reset_context() + publish( + ServerArgs(model_path="dummy", enable_unified_memory=unified), + role="tokenizer", + ) + cc = HybridCacheController.__new__(HybridCacheController) + cc.storage_backend = FailingStorage() + cc.prefetch_sync_queue = Queue() + operation = PrefetchOperation( + "r", [1], pool_transfers=[PoolTransfer(name=PoolName.SWA)] + ) + operation.hash_value = ["h0"] + if unified: + with self.assertLogs(level="ERROR"): + cc._page_transfer_sidecar(operation, kv_completed_pages=1) + ack = cc.prefetch_sync_queue.get_nowait() + self.assertIs(ack.operation, operation) + self.assertEqual(ack.pool_hits, {}) + else: + with self.assertRaisesRegex(ValueError, "sidecar read failed"): + cc._page_transfer_sidecar(operation, kv_completed_pages=1) + self.assertTrue(cc.prefetch_sync_queue.empty()) + + def test_shorter_prefetch_reserves_full_and_swa_without_mutating_probe_keys(self): + page_size = 4 + full_pool, swa_pool = _build_unified_host_pair( + _build_unified_swa_pool(page_size) + ) + self.addCleanup(full_pool.destroy) + self.addCleanup(swa_pool.destroy) + cc = HybridCacheController.__new__(HybridCacheController) + cc.page_size = page_size + cc.host_memory_mode = "cache" + cc.attn_cp_group = cc.attn_tp_group = cc.tp_group = None + cc.mem_pool_host = HostPoolGroup( + [ + PoolEntry(PoolName.KV, full_pool, None, None), + PoolEntry(PoolName.SWA, swa_pool, None, None), + ] + ) + hit_tokens = full_pool.available_size() + hashes = [str(i) for i in range(hit_tokens // page_size)] + transfer = PoolTransfer( + name=PoolName.SWA, + keys=hashes[-2:], + hit_policy=PoolHitPolicy.TRAILING_PAGES, + ) + operation = PrefetchOperation( + "r", list(range(hit_tokens)), pool_transfers=[transfer] + ) + operation.hash_value = hashes + self.assertIsNone(cc.alloc_prefetch_host_buffers(operation, hit_tokens)) + fitting = [ + n + for n in range(page_size, hit_tokens + 1, page_size) + if cc.can_fit_prefetch_host_buffers(operation, n, empty=False) + ] + self.assertTrue(fitting) + self.assertEqual(transfer.keys, hashes[-2:]) + length = max(fitting) + self.assertLess(length, hit_tokens) + host_indices = cc.alloc_prefetch_host_buffers(operation, length) + self.assertEqual(host_indices.numel(), length) + self.assertEqual(transfer.host_indices.numel(), 2 * page_size) + self.assertEqual( + transfer.keys, hashes[length // page_size - 2 : length // page_size] + ) + cc.free_prefetch_host_buffers(operation, host_indices) + def test_shared_arena_can_reuse_bytes_across_sides(self): page_size = 4 full_pool, swa_pool = _build_unified_host_pair( diff --git a/test/registered/unit/mem_cache/test_rust_tree_core_integration.py b/test/registered/unit/mem_cache/test_rust_tree_core_integration.py index 688e8bbff9d8..912cc63d38bc 100644 --- a/test/registered/unit/mem_cache/test_rust_tree_core_integration.py +++ b/test/registered/unit/mem_cache/test_rust_tree_core_integration.py @@ -1827,6 +1827,77 @@ def test_write_through_threshold_assignment_reaches_the_core(): assert any(isinstance(action, BackupKV) for action in result.cache_actions) +@pytest.mark.parametrize("backup_nodes", ["leaf", "parent_and_leaf", "parent"]) +def test_unified_swa_backup_when_full_already_has_a_host_copy(backup_nodes): + from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools + + bundle = init_unified_swa_pools( + device="cpu", + kv_cache_dtype=torch.float16, + head_num=1, + head_dim=4, + v_head_dim=4, + swa_head_num=1, + swa_head_dim=4, + swa_v_head_dim=4, + page_size=1, + start_layer=0, + end_layer=2, + swa_attention_layer_ids=[1], + full_attention_layer_ids=[0], + total_bytes=4096, + enable_memory_saver=False, + need_sort=False, + lazy_compaction=False, + ) + allocator = bundle.token_to_kv_pool_allocator + padding = allocator.alloc(4) + values = allocator.alloc(4) + core = _swa_tree_core(window=4, token_to_kv_pool_allocator=allocator) + core.set_hicache_enabled() + core.has_swa_host_pool = True + sources = [] + if backup_nodes != "leaf": + parent = _insert(core, [1, 2], values[:2].tolist()).last_device_node + sources.append((parent, values[:2])) + node = _insert(core, [1, 2, 3, 4], values.tolist()).last_device_node + node_values = values if backup_nodes == "leaf" else values[2:] + sources.append((node, node_values)) + for source_id, source_values in sources: + core.set_component_device_value( + source_id, + ComponentType.SWA, + allocator.translate_swa_indices_for_transfer(source_values), + ) + backed_up = {} + if backup_nodes == "parent": + backed_up[ComponentType.SWA] = [ + PoolTransfer( + name=PoolName.SWA, + host_indices=torch.tensor([200, 201]), + device_indices=allocator.translate_swa_indices_for_transfer( + node_values + ), + nodes_to_load=[node], + ) + ] + sources.pop() + core.commit_backup(node, torch.arange(100, 100 + node_values.numel()), backed_up) + allocator.free(padding) + + full_indices, transfers = core.build_backup_spec(node) + + assert full_indices.numel() == 0 + (swa_transfer,) = transfers[ComponentType.SWA] + assert swa_transfer.nodes_to_load == [source_id for source_id, _ in sources] + assert torch.equal( + swa_transfer.device_indices, + allocator.translate_swa_indices_for_transfer( + torch.cat([source_values for _, source_values in sources]) + ), + ) + + def test_swa_prefetch_commit_end_to_end(): from sglang.srt.mem_cache.unified_cache.components import CacheTransferPhase diff --git a/test/registered/unit/mem_cache/test_umbp_store.py b/test/registered/unit/mem_cache/test_umbp_store.py index 271b78a59fbf..2faac54b2055 100755 --- a/test/registered/unit/mem_cache/test_umbp_store.py +++ b/test/registered/unit/mem_cache/test_umbp_store.py @@ -36,6 +36,8 @@ class MockStorageConfig: class MockHostKVCache: """Mock HostKVCache that simulates page_first layout with real buffers.""" + stores_page_envelope = False + def __init__(self, num_pages=4, page_size=1, element_size=1024): self.layout = "page_first" self.page_size = page_size @@ -91,6 +93,7 @@ class MockLogicalHostPool: layout = "page_first" page_size = 1 kv_buffer = None + stores_page_envelope = False class MockHybridSidePool: diff --git a/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py b/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py index 05063d710797..a82ea6593205 100644 --- a/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py +++ b/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py @@ -63,7 +63,9 @@ def _paged_allocator(lazy: bool): _TOMBSTONE_METHODS = [ (mea.MultiEndedAllocator, "_free_lazy"), (mea.MultiEndedAllocator, "free"), + (mea.MultiEndedAllocator, "free_physical"), (mea.MultiEndedAllocator, "_commit_move_batch"), + (unified_hybrid_swa.UnifiedSWAAllocatorBase, "clear_full_to_swa_mapping"), (mea.FloatMultiEndedAllocator, "free"), (mea.FloatMultiEndedAllocator, "make_room"), (mea.FloatMultiEndedAllocator, "_relocate_to_positions"), diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 791aa4065aa5..0ac3750d7818 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -743,6 +743,36 @@ def _build_internal_chain(self, component_type, enable_session_radix_cache): self.assertIsNotNone(_device_value(cache, node_id, ComponentType.FULL)) return cache, first, second, leaf + def test_swa_write_back_host_full_preserves_pins_and_continues(self): + ct = ComponentType.SWA + for session in _session_radix_cache_test_values(): + for pinned in (False, True): + with self.subTest(session=session, pinned=pinned): + cache, first, second, leaf = self._build_internal_chain(ct, session) + cache.tree_core.is_write_back = True + cache.tree_core.has_swa_host_pool = True + cache.tree_core.enable_swa_write_back_eviction_barrier() + if pinned: + receipt = cache.inc_host_lock_ref(first).to_dec_params() + expected = _device_value(cache, first, ct).clone() + tracker = {ComponentType.FULL: 0, ct: 0} + # Only host allocation fails; tree walking and freeing are real. + with mock.patch.object( + cache, "_execute_and_commit_kv_backup", return_value=0 + ): + cache._evict_components({ComponentType.FULL: 0, ct: 2}, tracker) + self.assertGreaterEqual(tracker[ct], 2) + self.assertEqual(tracker[ct], tracker[ComponentType.FULL]) + if pinned: + self.assertTrue( + torch.equal(_device_value(cache, first, ct), expected) + ) + cache.dec_host_lock_ref(first, receipt) + self.assertEqual(tracker[ct], 4) + else: + self.assertEqual(tracker[ct], 6) + cache.sanity_check() + def _evict_for_alloc_after_first_drain(self, cache, component_type): capacity = {"available": 0} auxiliary_drains = {"count": 0} @@ -7976,6 +8006,43 @@ def _component_with_cache(component_type, cache): class TestUnifiedRadixCacheActionRouting(CustomTestCase): """CacheAction routing: each type forwards to the right Controller API.""" + def test_retraction_uses_physical_swa_transfer_indices(self): + cache = object.__new__(UnifiedRadixCache) + cache.is_swa_enabled = True + cache._sliding_window_size = 3 + cache.tree_core = mock.Mock(page_size=2) + cache.page_size = 2 + cache.sidecar_pool_specs = () + cache.req_to_token_pool = mock.Mock() + cache.req_to_token_pool.req_to_token = torch.tensor( + [[1, 2, 3, 4, 5, 0]], dtype=torch.int64 + ) + cache.token_to_kv_pool_allocator = mock.Mock() + cache.token_to_kv_pool_allocator.translate_kv_indices_for_transfer.side_effect = ( + lambda indices: indices + 10 + ) + cache.token_to_kv_pool_allocator.translate_swa_indices_for_transfer.side_effect = ( + lambda indices: indices + 100 + ) + req = mock.Mock(rid="req", seqlen=6) + req.kv.req_pool_idx = 0 + + full_indices, transfers = cache._retraction_device_transfers(req) + + self.assertTrue( + torch.equal(full_indices, torch.tensor([11, 12, 13, 14, 15, 16])) + ) + self.assertEqual(len(transfers), 1) + self.assertEqual(transfers[0].name, PoolName.SWA) + self.assertTrue( + torch.equal( + transfers[0].device_indices, + torch.tensor([103, 104, 105, 106]), + ) + ) + cache.token_to_kv_pool_allocator.translate_swa_indices_for_transfer.assert_called_once() + cache.token_to_kv_pool_allocator.translate_loc_from_full_to_swa.assert_not_called() + def test_backup_publish_node_ids_collects_component_nodes_once(self): comp_xfers = { ComponentType.SWA: [PoolTransfer(name=PoolName.SWA, nodes_to_load=[3, 4])], @@ -9546,15 +9613,21 @@ def test_storage_cleanup_releases_buffer_prefetch_anchor(self): anchor_node_id=99, prefetch_key=RadixKey(array("q", range(16))), host_indices=host_indices, - operation=mock.Mock(), + operation=mock.Mock( + buffer_host_occupied_units=4, pool_transfers_done=True + ), anchor_lock_params=None, comp_xfers={}, ) } cache.ongoing_backup = {} cache.host_memory_mode = "buffer_only" - cache._prefetch_occupied_span.side_effect = lambda key, indices: ( - UnifiedRadixCache._prefetch_occupied_span(cache, key, indices) + cache._prefetch_occupied_span.side_effect = ( + lambda key, indices, *, operation=None: ( + UnifiedRadixCache._prefetch_occupied_span( + cache, key, indices, operation=operation + ) + ) ) controller = cache.cache_controller controller.terminate_prefetch.return_value = (4, None)