Prevent duplicate CP HiCache write acks from crashing split drain

CP per-layer write-through can leave duplicated ready ack ids in the
ack queue while pending-split insertion drains write visibility. The
pre-scan observed both duplicates as registered before the first
completion removed radix state, so the second completion raised KeyError
and killed the scheduler.

The write-check path now records completed ack ids until matching queue
entries are drained. Duplicate completed acks are removed with a warning,
while genuinely unknown or unregistered acks still fail fast.

Constraint: CP HiCache pending-split drain may call writing_check outside the normal idle event path.
Rejected: Blind pop(..., None) for all unknown acks | would silently hide truly unattached or corrupt ack state.
Confidence: high
Scope-risk: moderate
Directive: Do not remove the completed-ack short-term memory unless duplicate and stale ack queues are proven impossible across TP MIN synchronization.
Tested: Local py_compile for hiradix_cache.py and test_cp_hicache_metadata.py.
Tested: Remote cjy-glm5-new py_compile and full test_cp_hicache_metadata.py, 121 passed.
Not-tested: Full ETE prefill workload after restarting service.
Co-authored-by: OmX <omx@oh-my-codex.dev>

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -10,7 +10,7 @@ import threading
 import time
 from dataclasses import dataclass, field
 from queue import Empty
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple
 
 import torch
 
@@ -2877,13 +2877,36 @@ class HiRadixCache(RadixCache):
                 self.write_backup(node)
 
     def writing_check(self, write_back=False):
-        def complete_write_ack_node(ack_id: int) -> TreeNode:
+        completed_write_ack_ids = getattr(self, "_completed_write_ack_ids", None)
+        if completed_write_ack_ids is None:
+            completed_write_ack_ids = set()
+            self._completed_write_ack_ids = completed_write_ack_ids
+
+        def complete_write_ack_node(ack_id: int) -> Optional[TreeNode]:
             if ack_id in self.pending_host_backups:
                 backuped_node = self._commit_pending_backup(ack_id)
                 self.ongoing_write_through.pop(ack_id, None)
+                completed_write_ack_ids.add(ack_id)
                 return backuped_node
-            backuped_node = self.ongoing_write_through.pop(ack_id)
+            backuped_node = self.ongoing_write_through.pop(ack_id, None)
+            if backuped_node is None:
+                if ack_id not in completed_write_ack_ids:
+                    raise KeyError(ack_id)
+                # A duplicated ready ack can pass the pre-scan while the first
+                # occurrence still owns radix state, then become already-consumed
+                # by the time this loop reaches the second occurrence.  Keep
+                # genuinely unattached acks deferred in the pre-scan, but make
+                # duplicate completion idempotent instead of crashing scheduler.
+                logger.warning(
+                    "[HiCache-write] writing_check skipped already-consumed write ack: "
+                    "ack_id=%d pending=%s ongoing=%s",
+                    ack_id,
+                    list(self.pending_host_backups.keys()),
+                    list(self.ongoing_write_through.keys()),
+                )
+                return None
             self.dec_node_lock_ref(backuped_node)
+            completed_write_ack_ids.add(ack_id)
             return backuped_node
 
         if write_back:
@@ -2893,7 +2916,7 @@ class HiRadixCache(RadixCache):
                     finish_event.synchronize()
                     for ack_id in ack_list:
                         backuped_node = complete_write_ack_node(ack_id)
-                        if self.enable_storage:
+                        if backuped_node is not None and self.enable_storage:
                             self.write_backup_storage(backuped_node)
                 self.cache_controller.ack_write_queue.clear()
                 assert len(self.ongoing_write_through) == 0
@@ -2918,6 +2941,7 @@ class HiRadixCache(RadixCache):
             return (
                 ack_id in self.pending_host_backups
                 or ack_id in self.ongoing_write_through
+                or ack_id in completed_write_ack_ids
             )
 
         finish_count = 0
@@ -2966,10 +2990,17 @@ class HiRadixCache(RadixCache):
             finish_event.synchronize()
             for ack_id in ack_list:
                 backuped_node = complete_write_ack_node(ack_id)
+                if backuped_node is None:
+                    continue
                 released_nodes.append(ack_id)
                 if self.enable_storage:
                     self.write_backup_storage(backuped_node)
             finish_count -= 1
+        if completed_write_ack_ids:
+            remaining_ack_ids = set()
+            for ack in self.cache_controller.ack_write_queue:
+                remaining_ack_ids.update(ack.node_ids)
+            completed_write_ack_ids.intersection_update(remaining_ack_ids)
         if released_nodes:
             logger.debug(
                 "[HiCache-write] writing_check released %d write locks: node_ids=%s remaining_ongoing=%d",

test/registered/unit/mem_cache/test_cp_hicache_metadata.py

@@ -2174,6 +2174,72 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
         self.assertEqual(cache.ongoing_write_through, {53: attached_node})
         self.assertEqual(dec_locked, [])
 
+    def test_writing_check_ignores_duplicate_ready_ack_for_same_node(self):
+        cache = HiRadixCache.__new__(HiRadixCache)
+        cache._uses_cp_hicache = True
+        cache.tp_world_size = 1
+        cache.enable_storage = False
+        cache.pending_host_backups = {}
+
+        class ReadyEvent:
+            def query(self):
+                return True
+
+            def synchronize(self):
+                pass
+
+        node = TreeNode()
+        node.id = 4614
+        dec_locked = []
+        cache.ongoing_write_through = {4614: node}
+        cache.dec_node_lock_ref = lambda released_node: dec_locked.append(released_node.id)
+        cache.cache_controller = types.SimpleNamespace(
+            ack_write_queue=[
+                HiCacheAck(ReadyEvent(), ReadyEvent(), [4614]),
+                HiCacheAck(ReadyEvent(), ReadyEvent(), [4614]),
+            ],
+        )
+
+        cache.writing_check()
+
+        self.assertEqual(cache.ongoing_write_through, {})
+        self.assertEqual(cache.cache_controller.ack_write_queue, [])
+        self.assertEqual(dec_locked, [4614])
+
+    def test_writing_check_drops_stale_duplicate_ack_before_valid_ack(self):
+        cache = HiRadixCache.__new__(HiRadixCache)
+        cache._uses_cp_hicache = True
+        cache.tp_world_size = 1
+        cache.enable_storage = False
+        cache.pending_host_backups = {}
+        cache._completed_write_ack_ids = {4614}
+
+        class ReadyEvent:
+            def query(self):
+                return True
+
+            def synchronize(self):
+                pass
+
+        node = TreeNode()
+        node.id = 4615
+        dec_locked = []
+        cache.ongoing_write_through = {4615: node}
+        cache.dec_node_lock_ref = lambda released_node: dec_locked.append(released_node.id)
+        cache.cache_controller = types.SimpleNamespace(
+            ack_write_queue=[
+                HiCacheAck(ReadyEvent(), ReadyEvent(), [4614]),
+                HiCacheAck(ReadyEvent(), ReadyEvent(), [4615]),
+            ],
+        )
+
+        cache.writing_check()
+
+        self.assertEqual(cache.ongoing_write_through, {})
+        self.assertEqual(cache.cache_controller.ack_write_queue, [])
+        self.assertEqual(dec_locked, [4615])
+        self.assertEqual(cache._completed_write_ack_ids, set())
+
     def test_write_backup_cp_failfast_on_unplanned_reservation_failure(self):
         cache = HiRadixCache.__new__(HiRadixCache)
         cache._uses_cp_hicache = True