diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index b0ea8192f..48aa66219 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -2440,19 +2440,21 @@ class HiRadixCache(RadixCache): return False def _cp_drain_write_acks_before_pending_split(self) -> None: - """Drain already-finished CP write acks before declaring split blocked. + """Do not drain CP write acks from the radix pending-split path. - This is intentionally only called on the rare stale-tail/pending-split - path. It may run the final write-visibility collective when a completed - ack is available, but it avoids adding a collective to the normal insert - hot path. + This helper is reached from radix insert, where CP ranks can legitimately + diverge: one rank may be checking a node-pending-backup split while + another is checking stale-tail pruning or has already returned to the + scheduler event loop. `writing_check()` may issue the final + write-visibility TP collective when a local ack is ready. Calling it + here can therefore interleave a scalar write-ack collective with the + scheduler's inflight-poll collective on other ranks and abort Gloo with + a message-size mismatch. + + Keep this path local and conservative: report the split as still blocked + and let the globally ordered scheduler/cache-event path drain write acks. """ - - if not self._uses_cp_hicache or not getattr( - self, "ongoing_write_through", None - ): - return - self.writing_check() + return def _cp_node_split_still_pending(self, node: TreeNode) -> bool: if not self._uses_cp_hicache or not self._node_host_write_pending(node): 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 9c3c0f461..12a80e154 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -2393,6 +2393,45 @@ class TestHiRadixCacheCPBackup(CustomTestCase): self.assertEqual(dec_locked, [4615]) self.assertEqual(cache._completed_write_ack_ids, set()) + def test_pending_split_probe_does_not_run_collective_writing_check(self): + """Pending-split probing runs inside radix insert, where only a subset of + CP ranks may take this branch. It must not opportunistically drain write + acks via writing_check(), because writing_check can issue a TP collective + and race against the scheduler's inflight-poll collective on other ranks. + """ + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.tp_world_size = 2 + cache.tp_group = object() + cache.enable_storage = False + cache.pending_host_backups = {} + node = TreeNode() + node.id = 10616 + cache.ongoing_write_through = {10616: node} + cache.dec_node_lock_ref = lambda released_node: None + + class ReadyEvent: + def query(self): + return True + + def synchronize(self): + pass + + cache.cache_controller = types.SimpleNamespace( + ack_write_queue=[HiCacheAck(ReadyEvent(), ReadyEvent(), [10616])], + ) + cache._cp_hicache_all_reduce = lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("pending split probe must not issue a collective") + ) + + cache._cp_drain_write_acks_before_pending_split() + + self.assertEqual(cache.ongoing_write_through, {10616: node}) + self.assertEqual( + [ack.node_ids for ack in cache.cache_controller.ack_write_queue], + [[10616]], + ) + def test_write_backup_cp_failfast_on_unplanned_reservation_failure(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True