From 315eaaff56921ac4de74f338b9996ea6928df813 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Sun, 24 May 2026 00:45:30 +0800 Subject: [PATCH] Preserve draft KV across CP HiCache hits Cache-hit prefill can skip draft forward for the prefix while PD transfer still reads draft KV for that same prefix. CP HiCache therefore needs to persist draft/MTP KV alongside target KV instead of relying on whatever remains in the draft GPU pool. Constraint: CP HiCache is host-only here; storage backends remain unsupported for CP shared KV. Constraint: CP shared KV must keep owner-page semantics and avoid falling back to full KV on every rank. Rejected: Recompute cached-prefix draft KV during prefill | loses the HiCache benefit and reintroduces the large hidden/KV footprint. Rejected: Change PD transfer to skip draft prefix KV | decode still needs draft cache continuity for MTP acceptance. Confidence: medium Scope-risk: moderate Directive: Keep target and draft CP HiCache metadata/load/write/evict paths in lockstep; changing one without the other can silently reduce MTP accept length. Tested: Remote g0034 container /sgl-workspace/sglang-tai: python3 -m pytest -q test/registered/unit/managers/test_hicache_controller_cp.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py => 58 passed, 3 warnings Not-tested: Full multi-node HiCache+MTP serving benchmark and accept-length recovery. --- python/sglang/srt/disaggregation/prefill.py | 34 +- .../sglang/srt/managers/cache_controller.py | 293 +++++++++++++++--- python/sglang/srt/managers/scheduler.py | 4 + python/sglang/srt/mem_cache/hiradix_cache.py | 126 +++++--- python/sglang/srt/models/deepseek_nextn.py | 11 + .../managers/test_hicache_controller_cp.py | 180 ++++++++++- .../mem_cache/test_cp_hicache_metadata.py | 127 +++++++- 7 files changed, 681 insertions(+), 94 deletions(-) diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index a04489e17..809f596f1 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -853,9 +853,18 @@ class SchedulerDisaggregationPrefillMixin: ) return prefill_queue = getattr(self, "disagg_prefill_bootstrap_queue", None) + has_draft_pool = ( + getattr(prefill_queue, "draft_token_to_kv_pool", None) is not None + ) + prefix_len = len(getattr(req, "prefix_indices", ())) + host_hit_length = int(getattr(req, "host_hit_length", 0) or 0) + draft_prefix_overlap = max(0, min(end_idx, prefix_len) - start_idx) _cp_draft_shared_kv_debug( "prefill_send_kv_chunk rid=%s room=%s start_idx=%s end_idx=%s " - "last_chunk=%s page_size=%s pages=%s state_pages=%s has_draft_pool=%s", + "last_chunk=%s page_size=%s pages=%s state_pages=%s " + "has_draft_pool=%s prefix_len=%s host_hit_length=%s " + "cache_protected_len=%s extend_input_len=%s fill_len=%s " + "origin_input_len=%s already_computed=%s draft_prefix_overlap=%s", req.rid, req.bootstrap_room, start_idx, @@ -864,6 +873,27 @@ class SchedulerDisaggregationPrefillMixin: page_size, _seq_summary(page_indices), _seq_summary(state_indices), - getattr(prefill_queue, "draft_token_to_kv_pool", None) is not None, + has_draft_pool, + prefix_len, + host_hit_length, + getattr(req, "cache_protected_len", None), + getattr(req, "extend_input_len", None), + len(req.fill_ids), + len(req.origin_input_ids), + getattr(req, "already_computed", None), + draft_prefix_overlap, ) + if has_draft_pool and draft_prefix_overlap > 0: + _cp_draft_shared_kv_debug( + "prefill_send_cachehit_draft_prefix rid=%s room=%s " + "draft_prefix_overlap=%s prefix_len=%s host_hit_length=%s " + "start_idx=%s end_idx=%s note=transfer_reads_draft_pool_for_cached_prefix", + req.rid, + req.bootstrap_room, + draft_prefix_overlap, + prefix_len, + host_hit_length, + start_idx, + end_idx, + ) req.disagg_kv_sender.send(page_indices, state_indices) diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 16d94655a..a17bb296b 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -70,10 +70,11 @@ class LayerDoneCounter: def __init__(self, num_layers: int): self.num_layers = num_layers # extra producer and consumer counters for overlap mode - self.num_counters = 3 + self.num_counters = 5 self.events = [LayerLoadingEvent(num_layers) for _ in range(self.num_counters)] self.producer_index = -1 self.consumer_index = -1 + self.consumer_indices: List[int] = [] def update_producer(self): self.producer_index = (self.producer_index + 1) % self.num_counters @@ -84,17 +85,26 @@ class LayerDoneCounter: ) return self.producer_index - def set_consumer(self, index: int): - self.consumer_index = index + def set_consumer(self, index): + if isinstance(index, (list, tuple, set)): + self.consumer_indices = [int(i) for i in index if int(i) >= 0] + self.consumer_index = ( + self.consumer_indices[0] if self.consumer_indices else -1 + ) + return + self.consumer_index = int(index) + self.consumer_indices = [self.consumer_index] if self.consumer_index >= 0 else [] def wait_until(self, threshold: int): - if self.consumer_index < 0: + if not self.consumer_indices: return - self.events[self.consumer_index].wait(threshold) + for consumer_index in self.consumer_indices: + self.events[consumer_index].wait(threshold) def reset(self): self.producer_index = -1 self.consumer_index = -1 + self.consumer_indices = [] class CacheOperation: @@ -275,6 +285,8 @@ class HiCacheController: pp_size: int = 1, enable_storage_metrics: bool = False, cp_shared_kv_layout: Optional[CpSharedKVLayout] = None, + draft_mem_pool_host: Optional["HostKVCache"] = None, + draft_mem_pool_device=None, ): self.tp_group = tp_group self.mem_pool_device_allocator = token_to_kv_pool_allocator @@ -285,6 +297,10 @@ class HiCacheController: mem_pool_device = mem_pool_device.full_kv_pool self.mem_pool_device = mem_pool_device self.cp_shared_kv_layout = cp_shared_kv_layout + + self.has_draft = False + self.mem_pool_device_draft = None + self.mem_pool_host_draft = None self.uses_cp_hicache = cp_shared_kv_layout is not None if self.uses_cp_hicache and not isinstance( self.mem_pool_device, NSATokenToKVPool @@ -317,20 +333,28 @@ class HiCacheController: self.layer_num = self.mem_pool_device.layer_num self.layer_done_counter = LayerDoneCounter(self.layer_num) self.mem_pool_device.register_layer_transfer_counter(self.layer_done_counter) + self.draft_mem_pool_host = None + self.draft_mem_pool_device = None if write_policy not in [ "write_through", "write_through_selective", "write_back", + "write_behind", ]: raise ValueError(f"Invalid write policy: {write_policy}") # self.write_queue = PriorityQueue[CacheOperation]() self.load_queue: List[CacheOperation] = [] self.write_queue: List[CacheOperation] = [] + self.draft_load_queue: List[CacheOperation] = [] + self.draft_write_queue: List[CacheOperation] = [] self.ack_load_queue: List[HiCacheAck] = [] self.ack_write_queue: List[HiCacheAck] = [] + if draft_mem_pool_host is not None or draft_mem_pool_device is not None: + self.attach_draft_pool(draft_mem_pool_device, draft_mem_pool_host) + self.stop_event = threading.Event() self.write_buffer = TransferBuffer(self.stop_event) self.load_buffer = TransferBuffer( @@ -354,6 +378,27 @@ class HiCacheController: # Preserve the historical error shape on init for unknown backends. raise ValueError(f"Failed to create storage backend: {e}") from e + @property + def has_draft_hicache(self) -> bool: + return ( + self.draft_mem_pool_host is not None + and self.draft_mem_pool_device is not None + ) + + def attach_draft_pool(self, draft_mem_pool_device, draft_mem_pool_host) -> None: + if draft_mem_pool_device is None and draft_mem_pool_host is None: + return + if draft_mem_pool_device is None or draft_mem_pool_host is None: + raise ValueError( + "draft_mem_pool_device and draft_mem_pool_host must be provided together" + ) + self.draft_mem_pool_device = draft_mem_pool_device + self.draft_mem_pool_host = draft_mem_pool_host + if hasattr(draft_mem_pool_device, "register_layer_transfer_counter"): + draft_mem_pool_device.register_layer_transfer_counter( + self.layer_done_counter + ) + def _start_storage_threads(self): """Start storage prefetch/backup threads and their queues. @@ -652,6 +697,8 @@ class HiCacheController: self.write_queue.clear() self.load_queue.clear() + self.draft_write_queue.clear() + self.draft_load_queue.clear() self.write_buffer.clear() self.load_buffer.clear() self.ack_write_queue.clear() @@ -706,6 +753,10 @@ class HiCacheController: self.write_queue.append( CacheOperation(host_indices, device_indices, node_id, priority) ) + if self.has_draft and not self.uses_cp_hicache: + self.draft_write_queue.append( + CacheOperation(host_indices, device_indices, node_id, priority) + ) self.start_writing() logger.info( "[CacheCtrl-write] write non-CP submitted: node_id=%d len=%d", @@ -715,12 +766,28 @@ class HiCacheController: return host_indices def start_writing(self) -> None: - if len(self.write_queue) == 0: + if len(self.write_queue) == 0 and len(self.draft_write_queue) == 0: return - op = CacheOperation.merge_ops(self.write_queue) - host_indices, device_indices = self.move_indices(op) + op = CacheOperation.merge_ops(self.write_queue) if self.write_queue else None + draft_op = ( + CacheOperation.merge_ops(self.draft_write_queue) + if self.draft_write_queue + else None + ) + node_ids = op.node_ids if op is not None else draft_op.node_ids + if op is not None: + host_indices, device_indices = self.move_indices(op, self.mem_pool_host) + else: + host_indices = device_indices = None + if draft_op is not None: + draft_host_indices, draft_device_indices = self.move_indices( + draft_op, self.draft_mem_pool_host + ) + else: + draft_host_indices = draft_device_indices = None self.write_queue.clear() + self.draft_write_queue.clear() start_event = device_module.Event() finish_event = device_module.Event() @@ -728,19 +795,31 @@ class HiCacheController: start_event.record() with device_module.stream(self.write_stream): start_event.wait(self.write_stream) - self.mem_pool_host.backup_from_device_all_layer( - self.mem_pool_device, host_indices, device_indices, self.io_backend - ) + if op is not None: + self.mem_pool_host.backup_from_device_all_layer( + self.mem_pool_device, host_indices, device_indices, self.io_backend + ) + if draft_op is not None: + self.draft_mem_pool_host.backup_from_device_all_layer( + self.draft_mem_pool_device, + draft_host_indices, + draft_device_indices, + self.io_backend, + ) finish_event.record() # NOTE: We must save the host indices and device indices here, # this is because we need to guarantee that these tensors are # still alive when the write stream is executing. - if host_indices.is_cuda: + if host_indices is not None and host_indices.is_cuda: host_indices.record_stream(self.write_stream) - if device_indices.is_cuda: + if device_indices is not None and device_indices.is_cuda: device_indices.record_stream(self.write_stream) + if draft_host_indices is not None and draft_host_indices.is_cuda: + draft_host_indices.record_stream(self.write_stream) + if draft_device_indices is not None and draft_device_indices.is_cuda: + draft_device_indices.record_stream(self.write_stream) - self.ack_write_queue.append(HiCacheAck(start_event, finish_event, op.node_ids)) + self.ack_write_queue.append(HiCacheAck(start_event, finish_event, node_ids)) def _append_completed_write_ack(self, node_id: int) -> None: event = device_module.Event() @@ -757,10 +836,10 @@ class HiCacheController: host_indices: torch.Tensor, device_indices: torch.Tensor, ) -> None: - validate_page_aligned_token_indices(host_indices, self.page_size, "host_indices") validate_page_aligned_token_indices( device_indices, self.page_size, "physical_device_indices" ) + validate_page_aligned_token_indices(host_indices, self.page_size, "host_indices") def _write_cp( self, @@ -786,6 +865,11 @@ class HiCacheController: logical_len=logical_len, owned_positions=owned_positions, host_indices=torch.empty((0,), dtype=torch.int64), + draft_host_indices=( + torch.empty((0,), dtype=torch.int64) + if self.has_draft_hicache + else None + ), ) ) @@ -800,31 +884,75 @@ class HiCacheController: owned_positions.numel(), ) return HiCacheWriteFailure(required_host_slots=len(physical_device_indices)) + draft_host_indices = None + if self.has_draft_hicache: + draft_host_indices = self.draft_mem_pool_host.alloc( + len(physical_device_indices) + ) + if draft_host_indices is None: + self.mem_pool_host.free(host_indices) + logger.info( + "[CacheCtrl-write] _write_cp FAILED (draft host full): node_id=%d logical_len=%d owned=%d", + node_id, + logical_len, + owned_positions.numel(), + ) + return HiCacheWriteFailure( + required_host_slots=len(physical_device_indices) + ) try: self._validate_cp_hicache_page_indices(host_indices, physical_device_indices) + if draft_host_indices is not None: + self._validate_cp_hicache_page_indices( + draft_host_indices, physical_device_indices + ) except Exception: self.mem_pool_host.free(host_indices) + if draft_host_indices is not None: + self.draft_mem_pool_host.free(draft_host_indices) raise self.write_queue.append( CacheOperation(host_indices, physical_device_indices, node_id, priority) ) + if draft_host_indices is not None: + self.draft_write_queue.append( + CacheOperation( + draft_host_indices, physical_device_indices, node_id, priority + ) + ) self.start_writing() logger.info( - "[CacheCtrl-write] _write_cp submitted: node_id=%d logical_len=%d owned=%d physical=%d", + "[CacheCtrl-write] _write_cp submitted: node_id=%d logical_len=%d owned=%d physical=%d draft=%s", node_id, logical_len, owned_positions.numel(), len(physical_device_indices), + draft_host_indices is not None, ) return HiCacheWriteResult( metadata=CpHiCacheNodeMetadata( logical_len=logical_len, owned_positions=owned_positions, host_indices=host_indices.cpu(), + draft_host_indices=( + draft_host_indices.cpu() if draft_host_indices is not None else None + ), ) ) + def set_draft_kv_pool(self, draft_device_pool, draft_host_pool) -> None: + """Register draft KV pools so L2 ops piggyback draft transfers.""" + self.has_draft = True + self.mem_pool_device_draft = draft_device_pool + self.mem_pool_host_draft = draft_host_pool + self.attach_draft_pool(draft_device_pool, draft_host_pool) + logger.info( + "HiCache draft KV registered: %s (host %d slots)", + type(draft_device_pool).__name__, + draft_host_pool.size, + ) + def load( self, host_indices: torch.Tensor, @@ -840,6 +968,10 @@ class HiCacheController: self.load_queue.append( CacheOperation(host_indices, device_indices, node_id, priority) ) + if self.has_draft and not self.uses_cp_hicache: + self.draft_load_queue.append( + CacheOperation(host_indices, device_indices, node_id, priority) + ) return device_indices def load_cp(self, nodes_to_load, node_id: int = -1) -> Optional[torch.Tensor]: @@ -849,6 +981,7 @@ class HiCacheController: return None host_chunks = [] + draft_host_chunks = [] physical_chunks = [] offset = 0 for node in nodes_to_load: @@ -862,9 +995,31 @@ class HiCacheController: self.cp_shared_kv_layout.logical_locs_to_physical(selected_logical_locs) ) host_chunks.append(node.cp_hicache.host_indices) + if self.has_draft_hicache: + draft_host_indices = getattr( + node.cp_hicache, "draft_host_indices", None + ) + if draft_host_indices is None: + self.mem_pool_device_allocator.free(device_indices) + raise RuntimeError( + "CP HiCache draft KV restore requested but node metadata " + "does not contain draft_host_indices" + ) + draft_host_chunks.append(draft_host_indices) if not host_chunks: - self._append_completed_load_ack(node_id) + # Keep CP load ACK rows identical across ranks. A zero-owned rank + # still queues a zero-length op with the logical node id so a later + # batched start_loading() merges the same node_ids as owning ranks. + self.load_queue.append( + CacheOperation( + torch.empty((0,), dtype=torch.int64), + torch.empty( + (0,), dtype=device_indices.dtype, device=device_indices.device + ), + node_id, + ) + ) return device_indices host_indices = torch.cat(host_chunks) @@ -874,6 +1029,17 @@ class HiCacheController: except Exception: self.mem_pool_device_allocator.free(device_indices) raise + draft_host_indices = None + if self.has_draft_hicache: + draft_host_indices = torch.cat(draft_host_chunks) + try: + self._validate_cp_hicache_page_indices( + draft_host_indices, physical_device_indices + ) + except Exception: + self.mem_pool_device_allocator.free(device_indices) + raise + self.load_queue.append( CacheOperation( host_indices, @@ -881,10 +1047,19 @@ class HiCacheController: node_id, ) ) + if draft_host_indices is not None: + self.draft_load_queue.append( + CacheOperation( + draft_host_indices, + physical_device_indices, + node_id, + ) + ) return device_indices - def move_indices(self, op: CacheOperation): + def move_indices(self, op: CacheOperation, mem_pool_host=None): + mem_pool_host = mem_pool_host or self.mem_pool_host host_indices, device_indices = op.host_indices, op.device_indices # move indices to GPU if using kernels, to host if using direct indexing if self.io_backend == "kernel": @@ -892,15 +1067,15 @@ class HiCacheController: host_indices = host_indices.to(self.device, non_blocking=True) return host_indices, device_indices elif self.io_backend == "direct": - if self.mem_pool_host.layout == "layer_first": + if mem_pool_host.layout == "layer_first": device_indices = device_indices.cpu() host_indices, idx = host_indices.sort() return host_indices, device_indices.index_select(0, idx) - elif self.mem_pool_host.layout == "page_first_direct": + elif mem_pool_host.layout == "page_first_direct": return host_indices, device_indices.cpu() else: raise ValueError( - f"Unsupported layout {self.mem_pool_host.layout!r} for io backend 'direct'" + f"Unsupported layout {mem_pool_host.layout!r} for io backend 'direct'" ) elif self.io_backend == "kernel_ascend": return host_indices, device_indices.cpu() @@ -908,40 +1083,76 @@ class HiCacheController: raise ValueError(f"Unsupported io backend") def start_loading(self) -> int: - if len(self.load_queue) == 0: + if len(self.load_queue) == 0 and len(self.draft_load_queue) == 0: return -1 producer_id = self.layer_done_counter.update_producer() - op = CacheOperation.merge_ops(self.load_queue) - host_indices, device_indices = self.move_indices(op) + op = CacheOperation.merge_ops(self.load_queue) if self.load_queue else None + draft_op = ( + CacheOperation.merge_ops(self.draft_load_queue) + if self.draft_load_queue + else None + ) + node_ids = op.node_ids if op is not None else draft_op.node_ids + if op is not None: + host_indices, device_indices = self.move_indices(op, self.mem_pool_host) + else: + host_indices = device_indices = None + if draft_op is not None: + draft_host_indices, draft_device_indices = self.move_indices( + draft_op, self.draft_mem_pool_host + ) + else: + draft_host_indices = draft_device_indices = None self.load_queue.clear() + self.draft_load_queue.clear() producer_event = self.layer_done_counter.events[producer_id] producer_event.start_event.record() with device_module.stream(self.load_stream): producer_event.start_event.wait(self.load_stream) - for i in range(self.layer_num): - self.mem_pool_host.load_to_device_per_layer( - self.mem_pool_device, - host_indices, - device_indices, - i, - self.io_backend, - ) - producer_event.complete(i) + draft_layer_num = ( + self.draft_mem_pool_device.layer_num if draft_op is not None else 0 + ) + for i in range(max(self.layer_num, draft_layer_num)): + if draft_op is not None and i < draft_layer_num: + if len(draft_host_indices) > 0: + self.draft_mem_pool_host.load_to_device_per_layer( + self.draft_mem_pool_device, + draft_host_indices, + draft_device_indices, + i, + self.io_backend, + ) + if op is not None and i < self.layer_num: + if len(host_indices) > 0: + self.mem_pool_host.load_to_device_per_layer( + self.mem_pool_device, + host_indices, + device_indices, + i, + self.io_backend, + ) + producer_event.complete(i) + elif op is None and i < self.layer_num: + producer_event.complete(i) # NOTE: We must save the host indices and device indices here, # this is because we need to guarantee that these tensors are # still alive when the load stream is executing. - if host_indices.is_cuda: + if host_indices is not None and host_indices.is_cuda: host_indices.record_stream(self.load_stream) - if device_indices.is_cuda: + if device_indices is not None and device_indices.is_cuda: device_indices.record_stream(self.load_stream) + if draft_host_indices is not None and draft_host_indices.is_cuda: + draft_host_indices.record_stream(self.load_stream) + if draft_device_indices is not None and draft_device_indices.is_cuda: + draft_device_indices.record_stream(self.load_stream) self.ack_load_queue.append( HiCacheAck( start_event=producer_event.start_event, finish_event=producer_event.finish_event, - node_ids=op.node_ids, + node_ids=node_ids, ) ) return producer_id @@ -957,6 +1168,16 @@ class HiCacheController: self.mem_pool_host.free(host_indices) return len(host_indices) + def evict_cp_host(self, metadata, backup_only: bool = True) -> int: + if not backup_only: + raise ValueError("Other eviction policies are not supported yet.") + + freed = self.evict_host(metadata.host_indices, backup_only=backup_only) + draft_host_indices = getattr(metadata, "draft_host_indices", None) + if self.has_draft_hicache and draft_host_indices is not None: + self.draft_mem_pool_host.free(draft_host_indices) + return freed + def prefetch( self, request_id: str, diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 1a7c0b426..57a6d9341 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -966,6 +966,10 @@ class Scheduler( _cp_draft_pool_summary(draft_token_to_kv_pool), _cp_draft_pool_summary(self.token_to_kv_pool_allocator.get_kvcache()), ) + if draft_token_to_kv_pool is not None and hasattr( + self.tree_cache, "attach_draft_kv_pool" + ): + self.tree_cache.attach_draft_kv_pool(draft_token_to_kv_pool) if ( self.disaggregation_mode == DisaggregationMode.DECODE diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 856204279..c830b914a 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -57,6 +57,7 @@ class CpHiCacheNodeMetadata: logical_len: int owned_positions: torch.Tensor host_indices: torch.Tensor + draft_host_indices: Optional[torch.Tensor] = None def __post_init__(self): if self.logical_len < 0: @@ -67,11 +68,23 @@ class CpHiCacheNodeMetadata: self.host_indices = self.host_indices.to( device="cpu", dtype=torch.int64 ).detach().clone() + if self.draft_host_indices is not None: + self.draft_host_indices = self.draft_host_indices.to( + device="cpu", dtype=torch.int64 + ).detach().clone() if self.owned_positions.numel() != self.host_indices.numel(): raise ValueError( "owned_positions and host_indices must have same length, got " f"{self.owned_positions.numel()} and {self.host_indices.numel()}" ) + if ( + self.draft_host_indices is not None + and self.owned_positions.numel() != self.draft_host_indices.numel() + ): + raise ValueError( + "draft_host_indices and owned_positions must have same length, got " + f"{self.draft_host_indices.numel()} and {self.owned_positions.numel()}" + ) if self.owned_positions.numel() > 0: if torch.any(self.owned_positions < 0) or torch.any( self.owned_positions >= self.logical_len @@ -96,19 +109,88 @@ class CpHiCacheNodeMetadata: logical_len=split_len, owned_positions=self.owned_positions[parent_mask], host_indices=self.host_indices[parent_mask], + draft_host_indices=( + self.draft_host_indices[parent_mask] + if self.draft_host_indices is not None + else None + ), ), CpHiCacheNodeMetadata( logical_len=self.logical_len - split_len, owned_positions=self.owned_positions[child_mask] - split_len, host_indices=self.host_indices[child_mask], + draft_host_indices=( + self.draft_host_indices[child_mask] + if self.draft_host_indices is not None + else None + ), ), ) class HiRadixCache(RadixCache): + def _create_token_to_kv_pool_host(self, kv_cache, server_args: ServerArgs): + from sglang.srt.mem_cache.memory_pool_host import ( + MHATokenToKVPoolHost, + MLATokenToKVPoolHost, + NSATokenToKVPoolHost, + ) + + if isinstance(kv_cache, MHATokenToKVPool): + return MHATokenToKVPoolHost( + kv_cache, + server_args.hicache_ratio, + server_args.hicache_size, + self.page_size, + server_args.hicache_mem_layout, + allocator_type=server_args.hicache_storage_backend, + ) + if isinstance(kv_cache, NSATokenToKVPool): + return NSATokenToKVPoolHost( + kv_cache, + server_args.hicache_ratio, + server_args.hicache_size, + self.page_size, + server_args.hicache_mem_layout, + allocator_type=server_args.hicache_storage_backend, + ) + if isinstance(kv_cache, MLATokenToKVPool): + return MLATokenToKVPoolHost( + kv_cache, + server_args.hicache_ratio, + server_args.hicache_size, + self.page_size, + server_args.hicache_mem_layout, + allocator_type=server_args.hicache_storage_backend, + ) + raise ValueError("HiRadixCache only supports MHA, MLA, and NSA KV pools") + + def attach_draft_kv_pool(self, draft_token_to_kv_pool) -> None: + if not self._uses_cp_hicache or draft_token_to_kv_pool is None: + return + if ( + getattr(self.cache_controller, "draft_mem_pool_device", None) + is draft_token_to_kv_pool + ): + return + draft_token_to_kv_pool_host = self._create_token_to_kv_pool_host( + draft_token_to_kv_pool, self._server_args + ) + self.cache_controller.attach_draft_pool( + draft_token_to_kv_pool, draft_token_to_kv_pool_host + ) + self.draft_token_to_kv_pool_host = draft_token_to_kv_pool_host + logger.info( + "[HiCache-draft] attached CP draft KV host pool: pool=%s host_pool=%s page_size=%d", + draft_token_to_kv_pool.__class__.__name__, + draft_token_to_kv_pool_host.__class__.__name__, + self.page_size, + ) + def __init__(self, params: CacheInitParams, server_args: ServerArgs): self._enable_metrics_flag = params.enable_metrics + self._server_args = server_args if not server_args.disable_hicache_numa_detect: bind_to_closest_numa_node_cuda() @@ -130,43 +212,10 @@ class HiRadixCache(RadixCache): "CP shared KV HiCache host integration requires NSATokenToKVPool." ) self.kv_cache = params.token_to_kv_pool_allocator.get_kvcache() - - from sglang.srt.mem_cache.memory_pool_host import ( - MHATokenToKVPoolHost, - MLATokenToKVPoolHost, - NSATokenToKVPoolHost, + self.token_to_kv_pool_host = self._create_token_to_kv_pool_host( + self.kv_cache, server_args ) - if isinstance(self.kv_cache, MHATokenToKVPool): - self.token_to_kv_pool_host = MHATokenToKVPoolHost( - self.kv_cache, - server_args.hicache_ratio, - server_args.hicache_size, - self.page_size, - server_args.hicache_mem_layout, - allocator_type=server_args.hicache_storage_backend, - ) - elif isinstance(self.kv_cache, NSATokenToKVPool): - self.token_to_kv_pool_host = NSATokenToKVPoolHost( - self.kv_cache, - server_args.hicache_ratio, - server_args.hicache_size, - self.page_size, - server_args.hicache_mem_layout, - allocator_type=server_args.hicache_storage_backend, - ) - elif isinstance(self.kv_cache, MLATokenToKVPool): - self.token_to_kv_pool_host = MLATokenToKVPoolHost( - self.kv_cache, - server_args.hicache_ratio, - server_args.hicache_size, - self.page_size, - server_args.hicache_mem_layout, - allocator_type=server_args.hicache_storage_backend, - ) - else: - raise ValueError(f"HiRadixCache only supports MHA and MLA yet") - self.tp_group = params.tp_cache_group self.tp_world_size = torch.distributed.get_world_size(group=self.tp_group) self.pp_rank = params.pp_rank @@ -1261,7 +1310,12 @@ class HiRadixCache(RadixCache): self._record_remove_event(x) if physical_count > 0: - num_evicted += self.cache_controller.evict_host(host_indices) + if self._uses_cp_hicache and hasattr( + self.cache_controller, "evict_cp_host" + ): + num_evicted += self.cache_controller.evict_cp_host(x.cp_hicache) + else: + num_evicted += self.cache_controller.evict_host(host_indices) x.host_len = 0 x.cp_hicache = None x.host_value = None diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index 6f96807fe..f349d9194 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -184,6 +184,12 @@ class DeepseekModelNextN(nn.Module): local_input_ids = pad_cp_local_input_ids_for_embedding( forward_batch, local_input_ids ) + self._debug_cp_draft_shared_kv( + "local_embedding_path " + f"full_tokens={full_num_tokens} " + f"local_tokens={local_num_tokens} " + f"padded_tokens={padded_token_count}" + ) hidden_states = self.embed_tokens(local_input_ids) if hidden_states.shape[0] != local_num_tokens: @@ -230,6 +236,11 @@ class DeepseekModelNextN(nn.Module): if hidden_states is None: # Conservative compatibility fallback: embed full input # so all TP ranks all-reduce the same shape, then CP-split. + self._debug_cp_draft_shared_kv( + "full_embedding_fallback " + f"full_tokens={input_ids.shape[0]} " + f"local_tokens={local_num_tokens}" + ) hidden_states = cp_split_and_rebuild_data( forward_batch, self.embed_tokens(input_ids) ) diff --git a/test/registered/unit/managers/test_hicache_controller_cp.py b/test/registered/unit/managers/test_hicache_controller_cp.py index 6616c8fdb..e898462f0 100644 --- a/test/registered/unit/managers/test_hicache_controller_cp.py +++ b/test/registered/unit/managers/test_hicache_controller_cp.py @@ -5,8 +5,72 @@ from unittest.mock import patch import torch -if "sgl_kernel" not in sys.modules: - sys.modules["sgl_kernel"] = types.ModuleType("sgl_kernel") +try: + import sgl_kernel # noqa: F401 + import sgl_kernel.kvcacheio # noqa: F401 +except (ImportError, RuntimeError): + if "sgl_kernel" not in sys.modules: + sys.modules["sgl_kernel"] = types.ModuleType("sgl_kernel") + sys.modules["sgl_kernel"].__file__ = "sgl_kernel_stub.py" + sys.modules["sgl_kernel"].__path__ = [] + if not hasattr(sys.modules["sgl_kernel"], "__getattr__"): + + def _sgl_kernel_getattr(name): + if name.startswith("__"): + raise AttributeError(name) + fn = lambda *args, **kwargs: None + setattr(sys.modules["sgl_kernel"], name, fn) + return fn + + sys.modules["sgl_kernel"].__getattr__ = _sgl_kernel_getattr + if "sgl_kernel.quantization" not in sys.modules: + quantization_stub = types.ModuleType("sgl_kernel.quantization") + quantization_stub.__file__ = "sgl_kernel_quantization_stub.py" + + def _quantization_getattr(name): + if name.startswith("__"): + raise AttributeError(name) + fn = lambda *args, **kwargs: None + setattr(quantization_stub, name, fn) + return fn + + quantization_stub.__getattr__ = _quantization_getattr + for name in ( + "ggml_dequantize", + "ggml_moe_a8", + "ggml_moe_a8_vec", + "ggml_moe_get_block_size", + "ggml_mul_mat_a8", + "ggml_mul_mat_vec_a8", + ): + setattr(quantization_stub, name, lambda *args, **kwargs: None) + sys.modules["sgl_kernel.quantization"] = quantization_stub + for name in ( + "sgl_per_token_group_quant_8bit", + "sgl_per_token_group_quant_fp8", + "sgl_per_token_quant_fp8", + "fp8_blockwise_scaled_mm", + "fp8_scaled_mm", + "silu_and_mul", + ): + if not hasattr(sys.modules["sgl_kernel"], name): + setattr(sys.modules["sgl_kernel"], name, lambda *args, **kwargs: None) + _sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT") + for _schema in ( + "sgl_per_token_group_quant_8bit(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()", + "sgl_per_token_group_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()", + "sgl_per_token_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s) -> ()", + "fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? bias=None) -> Tensor", + "fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> Tensor", + ): + try: + _sgl_kernel_lib.define(_schema) + except RuntimeError as exc: + if ( + "already" not in str(exc).lower() + and "duplicate" not in str(exc).lower() + ): + raise if "sgl_kernel.kvcacheio" not in sys.modules: kvcacheio_stub = types.ModuleType("sgl_kernel.kvcacheio") for name in ( @@ -56,12 +120,14 @@ class FakeHostPool: def backup_from_device_all_layer( self, device_pool, host_indices, device_indices, io_backend ): - self.backups.append((host_indices.clone(), device_indices.clone())) + self.backups.append((host_indices.clone(), device_indices.clone(), device_pool)) def load_to_device_per_layer( self, device_pool, host_indices, device_indices, layer_id, io_backend ): - self.loads.append((host_indices.clone(), device_indices.clone(), layer_id)) + self.loads.append( + (host_indices.clone(), device_indices.clone(), layer_id, device_pool) + ) def free(self, indices): self.frees.append(indices.clone()) @@ -72,6 +138,10 @@ class FakeDevicePool: device = "cpu" layer_num = 1 + def __init__(self, name="target", layer_num=1): + self.name = name + self.layer_num = layer_num + def register_layer_transfer_counter(self, counter): self.counter = counter @@ -80,11 +150,13 @@ class FakeAllocator: def __init__(self, alloc_result=None): self.alloc_result = alloc_result self.alloc_calls = [] + self.frees = [] self.cp_size = 4 self.cp_rank = 1 + self.device_pool = FakeDevicePool() def get_kvcache(self): - return FakeDevicePool() + return self.device_pool def alloc(self, need_size): self.alloc_calls.append(need_size) @@ -92,6 +164,10 @@ class FakeAllocator: return None return self.alloc_result[:need_size].clone() + def free(self, indices): + self.frees.append(indices.clone()) + return len(indices) + class HostIndicesTensor(torch.Tensor): @staticmethod @@ -166,7 +242,14 @@ class TestHiCacheControllerCPWrite(CustomTestCase): self.addCleanup(self.device_module_patcher.stop) self.addCleanup(self.nsa_pool_patcher.stop) - def make_controller(self, host_pool, allocator=None, cp_rank=1): + def make_controller( + self, + host_pool, + allocator=None, + cp_rank=1, + draft_host_pool=None, + draft_mem_pool_device=None, + ): allocator = allocator or FakeAllocator() controller = HiCacheController( token_to_kv_pool_allocator=allocator, @@ -178,6 +261,8 @@ class TestHiCacheControllerCPWrite(CustomTestCase): cp_shared_kv_layout=CpSharedKVLayout( page_size=4, cp_size=4, cp_rank=cp_rank ), + draft_mem_pool_host=draft_host_pool, + draft_mem_pool_device=draft_mem_pool_device, ) controller.layer_done_counter = DummyLayerDoneCounter() return controller @@ -200,7 +285,7 @@ class TestHiCacheControllerCPWrite(CustomTestCase): controller = self.make_controller(host_pool, cp_rank=1) logical_locs = torch.tensor([8, 9, 10], dtype=torch.int64) - with self.assertRaisesRegex(ValueError, "physical_device_indices.*whole pages"): + with self.assertRaisesRegex(ValueError, "host_indices.*whole pages"): controller.write(logical_locs, node_id=21) def test_cp_write_rejects_non_contiguous_owned_physical_page(self): @@ -242,6 +327,53 @@ class TestHiCacheControllerCPWrite(CustomTestCase): self.assertEqual(result.required_host_slots, 4) + def test_cp_write_with_draft_pool_backs_target_and_draft_locs(self): + host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64)) + draft_host_pool = FakeHostPool( + torch.tensor([200, 201, 202, 203], dtype=torch.int64) + ) + draft_device_pool = FakeDevicePool("draft") + controller = self.make_controller( + host_pool, + cp_rank=1, + draft_host_pool=draft_host_pool, + draft_mem_pool_device=draft_device_pool, + ) + logical_locs = torch.arange(4, 20, dtype=torch.int64) + + result = controller.write(logical_locs, node_id=77) + + self.assertEqual(result.metadata.host_indices.tolist(), [100, 101, 102, 103]) + self.assertEqual( + result.metadata.draft_host_indices.tolist(), [200, 201, 202, 203] + ) + self.assertEqual(host_pool.alloc_calls, [4]) + self.assertEqual(draft_host_pool.alloc_calls, [4]) + self.assertEqual(host_pool.backups[0][1].tolist(), [4, 5, 6, 7]) + self.assertEqual(draft_host_pool.backups[0][1].tolist(), [4, 5, 6, 7]) + self.assertIs(draft_host_pool.backups[0][2], draft_device_pool) + self.assertEqual(len(controller.ack_write_queue), 1) + self.assertEqual(controller.ack_write_queue[0].node_ids, [77]) + + def test_cp_write_draft_allocation_failure_rolls_back_target_host(self): + host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64)) + draft_host_pool = FakeHostPool(None) + controller = self.make_controller( + host_pool, + cp_rank=1, + draft_host_pool=draft_host_pool, + draft_mem_pool_device=FakeDevicePool("draft"), + ) + logical_locs = torch.arange(4, 20, dtype=torch.int64) + + result = controller.write(logical_locs, node_id=78) + + self.assertIsNone(result.metadata) + self.assertEqual(result.required_host_slots, 4) + self.assertEqual(host_pool.frees[0].tolist(), [100, 101, 102, 103]) + self.assertEqual(host_pool.backups, []) + self.assertEqual(draft_host_pool.backups, []) + def test_generate_storage_config_constructs_config_at_runtime(self): controller = HiCacheController.__new__(HiCacheController) controller.mem_pool_device = FakeDevicePool() @@ -300,6 +432,39 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite): self.assertEqual(allocator.alloc_calls, [16]) self.assertEqual(host_pool.loads[0][1].tolist(), [20, 21, 22, 23]) + def test_cp_load_with_draft_pool_restores_target_and_draft_locs(self): + host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64)) + draft_host_pool = FakeHostPool( + torch.tensor([200, 201, 202, 203], dtype=torch.int64) + ) + draft_device_pool = FakeDevicePool("draft") + allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64)) + controller = self.make_controller( + host_pool, + allocator=allocator, + cp_rank=1, + draft_host_pool=draft_host_pool, + draft_mem_pool_device=draft_device_pool, + ) + node = TreeNode() + node.host_len = 16 + node.cp_hicache = CpHiCacheNodeMetadata( + logical_len=16, + owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64), + host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64), + draft_host_indices=torch.tensor([200, 201, 202, 203], dtype=torch.int64), + ) + + device_indices = controller.load_cp([node], node_id=14) + controller.start_loading() + + self.assertEqual(device_indices.tolist(), list(range(64, 80))) + self.assertEqual(host_pool.loads[0][1].tolist(), [20, 21, 22, 23]) + self.assertEqual(draft_host_pool.loads[0][1].tolist(), [20, 21, 22, 23]) + self.assertIs(draft_host_pool.loads[0][3], draft_device_pool) + self.assertEqual(len(controller.ack_load_queue), 1) + self.assertEqual(controller.ack_load_queue[0].node_ids, [14]) + def test_cp_load_zero_owned_returns_full_logical_locs_and_noop_ack(self): host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64)) allocator = FakeAllocator(alloc_result=torch.arange(64, 68, dtype=torch.int64)) @@ -313,6 +478,7 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite): ) device_indices = controller.load_cp([node], node_id=12) + controller.start_loading() self.assertEqual(device_indices.tolist(), [64, 65, 66, 67]) self.assertEqual(host_pool.loads, []) diff --git a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py index 6ef679017..117c2acf6 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -1,14 +1,80 @@ import sys +import types import unittest -from unittest.mock import MagicMock, patch +from unittest.mock import patch import torch -# Stub out sgl_kernel before any sglang import so this CPU unit test does not -# require CUDA extension libraries to be installed. -for _mod in ("sgl_kernel", "sgl_kernel.kvcacheio"): - if _mod not in sys.modules: - sys.modules[_mod] = MagicMock() +# Prefer the real sgl_kernel package when the test image provides it so custom +# Torch operators are registered. Fall back to stubs on local CPU-only hosts. +try: + import sgl_kernel # noqa: F401 + import sgl_kernel.kvcacheio # noqa: F401 +except (ImportError, RuntimeError): + if "sgl_kernel" not in sys.modules: + sys.modules["sgl_kernel"] = types.ModuleType("sgl_kernel") + sys.modules["sgl_kernel"].__file__ = "sgl_kernel_stub.py" + sys.modules["sgl_kernel"].__path__ = [] + if not hasattr(sys.modules["sgl_kernel"], "__getattr__"): + + def _sgl_kernel_getattr(name): + if name.startswith("__"): + raise AttributeError(name) + fn = lambda *args, **kwargs: None + setattr(sys.modules["sgl_kernel"], name, fn) + return fn + + sys.modules["sgl_kernel"].__getattr__ = _sgl_kernel_getattr + if "sgl_kernel.quantization" not in sys.modules: + quantization_stub = types.ModuleType("sgl_kernel.quantization") + quantization_stub.__file__ = "sgl_kernel_quantization_stub.py" + + def _quantization_getattr(name): + if name.startswith("__"): + raise AttributeError(name) + fn = lambda *args, **kwargs: None + setattr(quantization_stub, name, fn) + return fn + + quantization_stub.__getattr__ = _quantization_getattr + for _name in ( + "ggml_dequantize", + "ggml_moe_a8", + "ggml_moe_a8_vec", + "ggml_moe_get_block_size", + "ggml_mul_mat_a8", + "ggml_mul_mat_vec_a8", + ): + setattr(quantization_stub, _name, lambda *args, **kwargs: None) + sys.modules["sgl_kernel.quantization"] = quantization_stub + for _name in ( + "sgl_per_token_group_quant_8bit", + "sgl_per_token_group_quant_fp8", + "sgl_per_token_quant_fp8", + "fp8_blockwise_scaled_mm", + "fp8_scaled_mm", + "silu_and_mul", + ): + if not hasattr(sys.modules["sgl_kernel"], _name): + setattr(sys.modules["sgl_kernel"], _name, lambda *args, **kwargs: None) + _sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT") + for _schema in ( + "sgl_per_token_group_quant_8bit(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()", + "sgl_per_token_group_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()", + "sgl_per_token_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s) -> ()", + "fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? bias=None) -> Tensor", + "fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> Tensor", + ): + try: + _sgl_kernel_lib.define(_schema) + except RuntimeError as exc: + if ( + "already" not in str(exc).lower() + and "duplicate" not in str(exc).lower() + ): + raise +if "sgl_kernel.kvcacheio" not in sys.modules: + sys.modules["sgl_kernel.kvcacheio"] = types.ModuleType("sgl_kernel.kvcacheio") from sglang.srt.mem_cache.base_prefix_cache import EvictParams, MatchPrefixParams from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata, HiRadixCache @@ -70,6 +136,22 @@ class TestCpHiCacheNodeMetadata(CustomTestCase): self.assertEqual(child.owned_positions.tolist(), [0, 4]) self.assertEqual(child.host_indices.tolist(), [22, 23]) + def test_split_keeps_draft_host_indices_aligned_with_owned_positions(self): + metadata = CpHiCacheNodeMetadata( + logical_len=10, + owned_positions=torch.tensor([0, 2, 5, 9], dtype=torch.int64), + host_indices=torch.tensor([20, 21, 22, 23], dtype=torch.int64), + draft_host_indices=torch.tensor([120, 121, 122, 123], dtype=torch.int64), + ) + + parent, child = metadata.split(5) + + self.assertEqual(parent.host_indices.tolist(), [20, 21]) + self.assertEqual(parent.draft_host_indices.tolist(), [120, 121]) + self.assertEqual(child.owned_positions.tolist(), [0, 4]) + self.assertEqual(child.host_indices.tolist(), [22, 23]) + self.assertEqual(child.draft_host_indices.tolist(), [122, 123]) + def test_zero_owned_metadata_is_valid(self): metadata = CpHiCacheNodeMetadata( logical_len=64, @@ -150,6 +232,15 @@ class TestCpHiCacheNodeMetadata(CustomTestCase): host_indices=torch.tensor([9], dtype=torch.int64), ) + def test_draft_host_length_mismatch_raises(self): + with self.assertRaisesRegex(ValueError, "draft_host_indices.*same length"): + CpHiCacheNodeMetadata( + logical_len=4, + owned_positions=torch.tensor([1, 2], dtype=torch.int64), + host_indices=torch.tensor([9, 10], dtype=torch.int64), + draft_host_indices=torch.tensor([109], dtype=torch.int64), + ) + def test_out_of_range_positions_raise(self): with self.assertRaisesRegex(ValueError, r"\[0, logical_len\)"): CpHiCacheNodeMetadata( @@ -222,6 +313,11 @@ class FakeEvictDeviceController: return len(device_indices) +class FakeTokenAllocator: + def available_size(self): + return 0 + + class TestHiRadixCacheCPBackup(CustomTestCase): def test_node_backuped_uses_cp_metadata(self): cache = HiRadixCache.__new__(HiRadixCache) @@ -341,6 +437,7 @@ class TestHiRadixCacheCPBackup(CustomTestCase): cache.evictable_leaves = set() cache.evictable_host_leaves = set() cache.eviction_strategy = FakeEvictionStrategy() + cache.token_to_kv_pool_allocator = FakeTokenAllocator() cache.evictable_size_ = 4 cache.protected_size_ = 0 cache._record_remove_event = lambda node: (_ for _ in ()).throw( @@ -602,14 +699,18 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase): cache.root_node.children[1] = node cache.evictable_host_leaves.add(node) - def mark_not_done(done_tensor, op=None, group=None): - self.assertIs(group, cache.tp_group) - done_tensor.fill_(0) + all_done_states = iter([False, True]) + cache._cp_all_ranks_true = lambda done: next(all_done_states) + cache._cp_broadcast_node_ids = lambda node_ids, max_ids: node_ids[:max_ids] + cache._cp_filter_all_ranks_safe_node_ids = ( + lambda node_ids, is_safe, **_kwargs: [ + node_id + for node_id in node_ids + if is_safe(cache._cp_node_by_id(node_id)) + ] + ) - with patch("torch.distributed.all_reduce", side_effect=mark_not_done): - physical_freed = cache._evict_host_for_physical_slots( - 0, synchronize_across_ranks=True - ) + physical_freed = cache._cp_evict_host_for_physical_slots(0) self.assertEqual(physical_freed, 0) self.assertEqual(freed, [])