chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Make the duplicate-reservation ack gate O(1) for fresh node ids

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-06-11 10:34:11 +00:00
parent 698a3e2431
commit 47cb9b9d11
2 changed files with 50 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
python/sglang/srt/managers/cache_controller.py
View File

@@ -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(

22
test/registered/unit/mem_cache/test_cp_hicache_metadata.py
View File

@@ -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):