From 47cb9b9d116d1c5c112e894110ec94de80f2db6b Mon Sep 17 00:00:00 2001 From: leavelet Date: Thu, 11 Jun 2026 10:34:11 +0000 Subject: [PATCH] Make the duplicate-reservation ack gate O(1) for fresh node ids MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit node_has_undrained_write_ack scanned ack_write_queue on every reservation, including the per-request-per-chunk prepare path whose node ids are freshly minted and can never be queued. Track the max node id ever appended to the queue: any queued id is <= max by construction (no monotonicity assumption needed for correctness), so fresh ids exit on one integer compare and the scan remains only for re-reservations of old ids (the rare write_backup fallback). All three ack_write_queue append sites now go through a single _append_write_ack funnel that maintains the max — the zero-owned-rank and non-CP write acks previously would have bypassed it, allowing false negatives on ranks owning no pages of a node. Co-Authored-By: Claude Fable 5 --- .../sglang/srt/managers/cache_controller.py | 33 +++++++++++++++++-- .../mem_cache/test_cp_hicache_metadata.py | 22 +++++++++++-- 2 files changed, 50 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index aa99d651f..b374e654f 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -439,6 +439,11 @@ class HiCacheController: self.ack_load_queue: List[HiCacheAck] = [] self.ack_write_queue: List[HiCacheAck] = [] self.pending_layer_writes: Dict[int, HiCacheLayerWriteState] = {} + # Highest node_id ever appended to ack_write_queue. TreeNode ids are + # strictly monotonic, so a reservation whose node_id exceeds this can + # never collide with a queued ack — the common (prepare-path) case + # short-circuits to one integer compare instead of a queue scan. + self._max_write_ack_node_id: int = -1 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) @@ -808,6 +813,10 @@ class HiCacheController: self.ack_write_queue.clear() self.ack_load_queue.clear() self.pending_layer_writes.clear() + # HiRadixCache.reset() restarts TreeNode.counter from 0; with the queue + # cleared, resetting the max keeps the duplicate-reservation fast path + # active for the re-minted ids instead of falling back to queue scans. + self._max_write_ack_node_id = -1 if self.enable_storage: self.prefetch_thread.join() self.backup_thread.join() @@ -928,12 +937,24 @@ class HiCacheController: 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, node_ids)) + self._append_write_ack(HiCacheAck(start_event, finish_event, node_ids)) + + def _append_write_ack(self, ack: HiCacheAck) -> None: + """Single funnel for ack_write_queue appends. + + Keeps _max_write_ack_node_id (the node_has_undrained_write_ack fast + path) covering every enqueued node id; appending around this helper + would re-open false negatives in the duplicate-reservation gate. + """ + self.ack_write_queue.append(ack) + for node_id in ack.node_ids: + if node_id > self._max_write_ack_node_id: + self._max_write_ack_node_id = node_id def _append_completed_write_ack(self, node_id: int) -> None: event = device_module.Event() event.record() - self.ack_write_queue.append(HiCacheAck(event, event, [node_id])) + self._append_write_ack(HiCacheAck(event, event, [node_id])) def _append_completed_load_ack(self, node_id: int) -> None: event = device_module.Event() @@ -972,7 +993,13 @@ class HiCacheController: state for the same node while its ack is undrained would put two acks for one radix registration into the queue (one inc_lock_ref, two completions) — the invariant is at most one in-flight ack per node. + + Cost: node ids are strictly monotonic, so fresh ids (the per-request + prepare path) exit on the integer compare below; the queue scan runs + only for re-reservations of old node ids (rare write_backup fallback). """ + if node_id > getattr(self, "_max_write_ack_node_id", -1): + return False ack_write_queue = getattr(self, "ack_write_queue", None) if not ack_write_queue: return False @@ -1254,7 +1281,7 @@ class HiCacheController: state.ack_appended = True reservation = state.reservation self.pending_layer_writes.pop(reservation.node_id, None) - self.ack_write_queue.append( + self._append_write_ack( HiCacheAck(state.start_event, state.finish_event, [reservation.node_id]) ) logger.debug( 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 f8d21f684..be3347d3e 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -2176,6 +2176,17 @@ class TestHiRadixCacheCPBackup(CustomTestCase): self.assertEqual(cache.ongoing_write_through, {53: attached_node}) self.assertEqual(dec_locked, []) + @staticmethod + def _make_exploding_queue(): + class _ExplodingQueue(list): + def __iter__(self): + raise AssertionError("fresh node ids must not scan the ack queue") + + def __bool__(self): + return True + + return _ExplodingQueue() + def test_reserve_write_cp_refuses_node_with_undrained_ack(self): """Producer-side invariant: at most one in-flight ack per node_id. A reservation for a node whose previous final ack is still queued is @@ -2190,6 +2201,7 @@ class TestHiRadixCacheCPBackup(CustomTestCase): ack_write_queue=[ HiCacheAck(MagicMock(), MagicMock(), [4614]), ], + _max_write_ack_node_id=4614, ) controller.node_has_undrained_write_ack = types.MethodType( HiCacheController.node_has_undrained_write_ack, controller @@ -2201,9 +2213,15 @@ class TestHiRadixCacheCPBackup(CustomTestCase): ) self.assertIsInstance(result, HiCacheWriteFailure) self.assertEqual(result.required_host_slots, 0) - # Different node id is not gated by the queue scan. + # An old (<= max) but absent node id is resolved by the queue scan. self.assertFalse( - HiCacheController.node_has_undrained_write_ack(controller, 4615) + HiCacheController.node_has_undrained_write_ack(controller, 4613) + ) + # Fresh ids (> max ever acked, i.e. every prepare-path reservation) + # take the O(1) fast path: no queue scan at all. + controller.ack_write_queue = self._make_exploding_queue() + self.assertFalse( + HiCacheController.node_has_undrained_write_ack(controller, 9999) ) def test_rollback_pending_backup_scrubs_acks_and_cancels_state(self):