Enforce one in-flight CP write ack per node at the producer

dfa168abe9 made duplicate ack completion idempotent at the consumer
(writing_check), which is correct and lock-balance-safe but masks the
producer invariant breach: two acks can only coexist for one radix
registration when an ack is orphaned in ack_write_queue after a
rollback cleared ongoing_write_through/pending_host_backups,
re-opening the _node_host_write_pending guard for a fresh
registration. _rollback_pending_backup (write_backup's exception
path) was the one rollback that did not scrub the node's acks — and
it also left a half-submitted layer-write state alive, which later
forwards would keep driving, writing D2H into host slots the rollback
had already evicted.

Close the invariant structurally:
- reserve_write_cp refuses (HiCacheWriteFailure, the existing
  skip-this-round path both callers already handle) when the node
  still has an undrained final ack, computed by scanning the small
  ack queue — no new state to keep in sync.
- _rollback_pending_backup now cancels the pending layer-write state
  (write-stream sync so in-flight per-layer copies finish before the
  host slots are evicted) and scrubs the node's queued acks via the
  scrub factored out of _rollback_prepared_cp_backup.
- The consumer-side duplicate guard from dfa168abe9 is kept as
  defense in depth.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

python/sglang/srt/managers/cache_controller.py

@@ -964,6 +964,37 @@ class HiCacheController:
                 device_indices, self.page_size, "physical_device_indices"
             )
 
+    def node_has_undrained_write_ack(self, node_id: int) -> bool:
+        """Whether a final write ack for this node is still in ack_write_queue.
+
+        The radix side drains acks via writing_check; until then the node's
+        previous backup owns radix/host state. Creating a second layer-write
+        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.
+        """
+        ack_write_queue = getattr(self, "ack_write_queue", None)
+        if not ack_write_queue:
+            return False
+        return any(node_id in ack.node_ids for ack in ack_write_queue)
+
+    def cancel_layer_write_state(self, node_id: int) -> bool:
+        """Drop a pending per-layer write state during backup rollback.
+
+        Without this, an exception between reservation and final ack leaves a
+        half-submitted state in pending_layer_writes: the next forward's layer
+        hooks keep driving it, eventually writing D2H into host slots the
+        rollback already evicted and appending an ack for a node the radix
+        side no longer has registered. Synchronize the write stream so any
+        already-issued per-layer copies finish before the caller frees the
+        host slots.
+        """
+        state = self.pending_layer_writes.pop(node_id, None)
+        if state is None:
+            return False
+        self.write_stream.synchronize()
+        return True
+
     def reserve_write_cp(
         self,
         device_indices: torch.Tensor,
@@ -972,6 +1003,17 @@ class HiCacheController:
     ) -> HiCacheWriteReservation | HiCacheWriteFailure:
         from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata
 
+        if node_id >= 0 and self.node_has_undrained_write_ack(node_id):
+            # Refusing here (instead of crashing later in writing_check) keeps
+            # the failure on the existing HiCacheWriteFailure path: the caller
+            # skips this backup round and may retry after the ack drains.
+            logger.warning(
+                "[CacheCtrl-write] reserve_write_cp refused: node_id=%d still "
+                "has an undrained write ack; skipping duplicate backup",
+                node_id,
+            )
+            return HiCacheWriteFailure(required_host_slots=0)
+
         page_size = self.page_size
         layout = self.cp_shared_kv_layout
         logical_len = len(device_indices)

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -2090,9 +2090,39 @@ class HiRadixCache(RadixCache):
             self.dec_node_lock_ref(node)
         return node
 
+    def _remove_undrained_write_acks(self, node_id: int) -> bool:
+        """Scrub this node's final write acks out of the controller queue.
+
+        Every rollback that clears the radix-side registration
+        (ongoing_write_through / pending_host_backups) MUST also scrub the
+        node's acks: an orphaned ack re-opens the _node_host_write_pending
+        guard for a fresh registration while the stale ack is still queued,
+        producing two acks for one registration (one inc_lock_ref, two
+        completions) and crashing writing_check on the second.
+        """
+        retained_acks = []
+        removed = False
+        for ack in self.cache_controller.ack_write_queue:
+            if node_id not in ack.node_ids:
+                retained_acks.append(ack)
+                continue
+            ack.finish_event.synchronize()
+            remaining_node_ids = [nid for nid in ack.node_ids if nid != node_id]
+            if remaining_node_ids:
+                retained_acks.append(ack._replace(node_ids=remaining_node_ids))
+            removed = True
+        if removed:
+            self.cache_controller.ack_write_queue = retained_acks
+        return removed
+
     def _rollback_pending_backup(self, node_id: int) -> TreeNode:
         pending = self.pending_host_backups.pop(node_id)
         node = pending.node
+        # Cancel the half-submitted per-layer state (so later layer hooks
+        # cannot write into the host slots evicted below) and scrub any
+        # already-appended final ack before releasing the registration.
+        self.cache_controller.cancel_layer_write_state(node_id)
+        self._remove_undrained_write_acks(node_id)
         self.cache_controller.evict_cp_host(pending.metadata)
         if node.cp_hicache is pending.metadata:
             node.cp_hicache = None
@@ -2271,21 +2301,7 @@ class HiRadixCache(RadixCache):
                 f"while per-layer writes are still pending; reason={reason}"
             )
 
-        retained_acks = []
-        removed_ack = False
-        for ack in self.cache_controller.ack_write_queue:
-            if prepared.node_id not in ack.node_ids:
-                retained_acks.append(ack)
-                continue
-            ack.finish_event.synchronize()
-            remaining_node_ids = [
-                node_id for node_id in ack.node_ids if node_id != prepared.node_id
-            ]
-            if remaining_node_ids:
-                retained_acks.append(ack._replace(node_ids=remaining_node_ids))
-            removed_ack = True
-        if removed_ack:
-            self.cache_controller.ack_write_queue = retained_acks
+        self._remove_undrained_write_acks(prepared.node_id)
         logger.debug(
             "[HiCache-write] rollback unattached prepared CP backup: node_id=%d logical_len=%d reason=%s",
             prepared.node_id,

test/registered/unit/mem_cache/test_cp_hicache_metadata.py

@@ -3,7 +3,7 @@ import re
 import sys
 import types
 import unittest
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import torch
 
@@ -1320,7 +1320,9 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
         cache._uses_cp_hicache = True
         evicted = []
         cache.cache_controller = types.SimpleNamespace(
-            evict_cp_host=lambda metadata: evicted.append(metadata) or 2
+            evict_cp_host=lambda metadata: evicted.append(metadata) or 2,
+            cancel_layer_write_state=lambda node_id: False,
+            ack_write_queue=[],
         )
         dec_locked = []
         cache.dec_node_lock_ref = lambda node: dec_locked.append(node)
@@ -2174,6 +2176,104 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
         self.assertEqual(cache.ongoing_write_through, {53: attached_node})
         self.assertEqual(dec_locked, [])
 
+    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
+        refused on the normal HiCacheWriteFailure path instead of creating a
+        second layer-write state (which would enqueue a duplicate ack)."""
+        from sglang.srt.managers.cache_controller import (
+            HiCacheController,
+            HiCacheWriteFailure,
+        )
+
+        controller = types.SimpleNamespace(
+            ack_write_queue=[
+                HiCacheAck(MagicMock(), MagicMock(), [4614]),
+            ],
+        )
+        controller.node_has_undrained_write_ack = types.MethodType(
+            HiCacheController.node_has_undrained_write_ack, controller
+        )
+        result = HiCacheController.reserve_write_cp(
+            controller,
+            device_indices=torch.arange(4, dtype=torch.int64),
+            node_id=4614,
+        )
+        self.assertIsInstance(result, HiCacheWriteFailure)
+        self.assertEqual(result.required_host_slots, 0)
+        # Different node id is not gated by the queue scan.
+        self.assertFalse(
+            HiCacheController.node_has_undrained_write_ack(controller, 4615)
+        )
+
+    def test_rollback_pending_backup_scrubs_acks_and_cancels_state(self):
+        """The write_backup exception rollback must not orphan an ack or a
+        half-submitted layer-write state: an orphaned ack re-opens the
+        registration guard while still queued, producing the duplicate-ack
+        crash; a live state would keep writing into evicted host slots."""
+        cache = HiRadixCache.__new__(HiRadixCache)
+        node = TreeNode()
+        node.id = 4614
+        metadata = object()
+        dec_locked = []
+        cache.pending_host_backups = {
+            4614: PendingHiCacheBackup(
+                node=node,
+                metadata=metadata,
+                logical_len=8,
+                submitted=True,
+                locked=True,
+            )
+        }
+        cache.dec_node_lock_ref = lambda n: dec_locked.append(n.id)
+
+        finish_event = MagicMock()
+        cancelled = []
+        evicted = []
+        cache.cache_controller = types.SimpleNamespace(
+            ack_write_queue=[
+                HiCacheAck(MagicMock(), finish_event, [4614]),
+                HiCacheAck(MagicMock(), MagicMock(), [4615]),
+            ],
+            cancel_layer_write_state=lambda node_id: cancelled.append(node_id),
+            evict_cp_host=lambda meta: evicted.append(meta),
+        )
+        node.cp_hicache = metadata
+
+        returned = cache._rollback_pending_backup(4614)
+
+        self.assertIs(returned, node)
+        self.assertEqual(cache.pending_host_backups, {})
+        self.assertEqual(cancelled, [4614])
+        self.assertEqual(evicted, [metadata])
+        # The 4614 ack is scrubbed (after syncing its finish event); the
+        # unrelated 4615 ack is retained.
+        finish_event.synchronize.assert_called_once()
+        self.assertEqual(
+            [ack.node_ids for ack in cache.cache_controller.ack_write_queue],
+            [[4615]],
+        )
+        self.assertEqual(dec_locked, [4614])
+        self.assertIsNone(node.cp_hicache)
+        self.assertEqual(node.host_len, 0)
+
+    def test_remove_undrained_write_acks_splits_shared_ack(self):
+        """A grouped ack covering several nodes only loses the rolled-back
+        node id; other nodes' completion is preserved."""
+        cache = HiRadixCache.__new__(HiRadixCache)
+        cache.cache_controller = types.SimpleNamespace(
+            ack_write_queue=[
+                HiCacheAck(MagicMock(), MagicMock(), [4614, 4615]),
+            ],
+        )
+        removed = cache._remove_undrained_write_acks(4614)
+        self.assertTrue(removed)
+        self.assertEqual(
+            [ack.node_ids for ack in cache.cache_controller.ack_write_queue],
+            [[4615]],
+        )
+        self.assertFalse(cache._remove_undrained_write_acks(9999))
+
     def test_writing_check_ignores_duplicate_ready_ack_for_same_node(self):
         cache = HiRadixCache.__new__(HiRadixCache)
         cache._uses_cp_hicache = True