From 4e9b4d05c09ae3aa4f52ab56e3339b74892e9bf8 Mon Sep 17 00:00:00 2001 From: leavelet Date: Wed, 24 Jun 2026 01:13:59 +0000 Subject: [PATCH] CP HiCache: kill the O(victims*candidates) GPU-sync storm in the owner-lane eviction planner A CP4 prefill server hung 79s in _plan_cp_load_back_owner_lane_evictions (the [HiCache-load] slow-scan) and was killed by the detokenizer health-check. Root cause (verified from the b300 hang log + code): the L1 free-room watermark (--hicache-l1-free-room-ratio 0.25 over a 31507-page lane) hands the planner a ~7877-page single-owner deficit; the planner then rescans all ~2000 evictable candidates once per victim (~195 iters), and for EACH candidate every iteration it recomputes _cp_load_back_node_owner_page_counts -- which under the pooled shared-L2 path (no page_owners on CpSharedL2NodeMetadata) takes the device-tensor fallback and does cp_size per-owner .item() device->host syncs. That is ~195 * 2000 * 4 ~= 1.5M CUDA syncs on the synchronous load-back admission path, blocking the scheduler for ~79s. Fix (byte-identical victim selection, just fast): - Memoize the owner-count histogram per planning call ({node.id: counts}); the counts are invariant while node.value is fixed (the plan does not mutate values), so node.id is a safe key for the plan's duration. Threaded explicitly to both call sites (planner loop + ancestor-unlock helper). Turns O(victims*candidates) recomputes into O(distinct nodes). - Replace the cp_size per-owner sum().item() loop with one bincount().tolist() device->host sync. Net: ~79s -> ~1s; the eviction plan (victims, planned_freed) is unchanged. bincount == the per-owner loop proven over 8000 random vectors incl. the (-1)%cp==cp-1 zero-loc edge. New memo regression test; existing count-fn + planner tests pass (the one pre-existing unrelated EAGLE-tail failure is unchanged). (The 7877-page watermark magnitude is a separate, config-side issue: --hicache-l1-free-room-ratio 0.25 reserves ~25% of a ~2M-token lane -- ~30x more than a 64K chunk needs; lower it.) Co-Authored-By: Claude Opus 4.8 (1M context) --- python/sglang/srt/mem_cache/hiradix_cache.py | 82 ++++++++++++++----- .../test_cp_hicache_load_back_owner_lanes.py | 27 ++++++ 2 files changed, 87 insertions(+), 22 deletions(-) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 36e1eb721..2f6efa9f1 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -2496,28 +2496,52 @@ class HiRadixCache(RadixCache): return self._cp_device_node_is_load_back_victim_after_plan(node, set()) def _cp_load_back_node_owner_page_counts( - self, node: TreeNode, cp_size: int + self, node: TreeNode, cp_size: int, *, memo: Optional[dict] = None ) -> Tuple[int, ...]: + """Per-owner device page-count histogram for a node. + + The owner-lane eviction planner calls this O(victims * candidates) times over an + UNCHANGING node set (the plan does not mutate node.value), so the per-node counts + are invariant for the whole planning call. Pass a planner-scoped ``memo`` + ({node.id: counts}) to compute each node's counts exactly once -- node.id is a safe + key for the duration of one plan. Without ``memo`` the behaviour is unchanged. + + The histogram itself uses a single ``bincount().tolist()`` device->host sync instead + of ``cp_size`` per-owner ``.item()`` syncs -- under the pooled shared-L2 path (no + ``page_owners`` on the metadata) this fallback is taken for every candidate, so the + per-owner sync loop (re-run O(victims * candidates) times) was the dominant cost. + """ + if memo is not None: + cached = memo.get(node.id) + if cached is not None: + return cached metadata = getattr(node, "cp_hicache", None) if metadata is not None and not self._is_cp_shared_l2_metadata(metadata): - return metadata.owner_page_counts(cp_size) - - value = getattr(node, "value", None) - if value is None or int(value.numel()) == 0: - return tuple(0 for _ in range(cp_size)) - padded_value = pad_token_locs_to_page_boundary( - value, - self.page_size, - name=( - "CP HiCache load-back eviction device values " - f"node_id={getattr(node, 'id', '?')}" - ), - ) - - first_locs = padded_value[:: self.page_size] - logical_pages = torch.div(first_locs, self.page_size, rounding_mode="floor") - owners = torch.remainder(logical_pages - 1, cp_size) - return tuple(int((owners == owner).sum().item()) for owner in range(cp_size)) + counts = metadata.owner_page_counts(cp_size) + else: + value = getattr(node, "value", None) + if value is None or int(value.numel()) == 0: + counts = tuple(0 for _ in range(cp_size)) + else: + padded_value = pad_token_locs_to_page_boundary( + value, + self.page_size, + name=( + "CP HiCache load-back eviction device values " + f"node_id={getattr(node, 'id', '?')}" + ), + ) + first_locs = padded_value[:: self.page_size] + logical_pages = torch.div( + first_locs, self.page_size, rounding_mode="floor" + ) + owners = torch.remainder(logical_pages - 1, cp_size) + # one device->host sync (bincount over [0, cp_size)) == the per-owner + # sum().item() loop, byte-identical counts, cp_size-fewer syncs. + counts = tuple(torch.bincount(owners, minlength=cp_size).tolist()) + if memo is not None: + memo[node.id] = counts + return counts def _cp_load_back_ancestor_unlock_contribution( self, @@ -2525,6 +2549,8 @@ class HiRadixCache(RadixCache): deficits: List[int], planned_evicted_nodes: set, cp_size: int, + *, + memo: Optional[dict] = None, ) -> int: ancestor = getattr(node, "parent", None) while ancestor is not None and ancestor != getattr(self, "root_node", None): @@ -2535,7 +2561,9 @@ class HiRadixCache(RadixCache): continue if not self._cp_device_node_is_load_back_victim_base(ancestor): return 0 - counts = self._cp_load_back_node_owner_page_counts(ancestor, cp_size) + counts = self._cp_load_back_node_owner_page_counts( + ancestor, cp_size, memo=memo + ) contribution = sum( min(int(count), int(deficit)) for count, deficit in zip(counts, deficits) @@ -2558,6 +2586,10 @@ class HiRadixCache(RadixCache): candidate_nodes = set(getattr(self, "evictable_leaves", set())) initial_candidate_count = len(candidate_nodes) iteration = 0 + # Per-plan memo {node.id: owner_counts}: the planner rescans the same node set + # once per victim, but each node's owner counts are invariant for the whole plan, + # so compute the (device-sync) histogram once instead of O(victims * candidates). + owner_counts_memo: dict = {} while any(v > 0 for v in deficits): iteration += 1 @@ -2589,7 +2621,9 @@ class HiRadixCache(RadixCache): node, planned_evicted_nodes ): continue - counts = self._cp_load_back_node_owner_page_counts(node, cp_size) + counts = self._cp_load_back_node_owner_page_counts( + node, cp_size, memo=owner_counts_memo + ) contribution = sum( min(int(count), int(deficit)) for count, deficit in zip(counts, deficits) @@ -2598,7 +2632,11 @@ class HiRadixCache(RadixCache): if contribution <= 0: unlock_contribution = ( self._cp_load_back_ancestor_unlock_contribution( - node, deficits, planned_evicted_nodes, cp_size + node, + deficits, + planned_evicted_nodes, + cp_size, + memo=owner_counts_memo, ) ) if contribution <= 0 and unlock_contribution <= 0: diff --git a/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py b/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py index bb4b471b0..fdc648c7a 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py @@ -338,6 +338,33 @@ class TestCpHiCacheLoadBackOwnerLanes(CustomTestCase): self.assertEqual(counts, (2, 0, 1, 0)) self.assertIs(counts, node.cp_hicache.owner_page_counts(4)) + def test_owner_counts_memo_is_byte_identical_and_computed_once(self): + # The planner passes a per-plan {node.id: counts} memo so the (device-sync) owner + # histogram is computed once per node instead of O(victims*candidates). Verify + # (a) the memo result is byte-identical to the un-memoized result, and (b) the + # cached value is authoritative for the plan's duration -- mutating node.value does + # NOT change the memoized result (the planner never mutates value mid-plan, which + # is exactly the invariant the memo relies on), proving the cached path is taken. + allocator = _make_allocator(page_size=4, cp_size=4) + cache = _make_cache(allocator) + node = TreeNode(id=13) + node.value = torch.arange(4, 10, dtype=torch.int64) + + no_memo = cache._cp_load_back_node_owner_page_counts(node, cp_size=4) + memo = {} + first = cache._cp_load_back_node_owner_page_counts(node, cp_size=4, memo=memo) + self.assertEqual(first, no_memo) + self.assertEqual(first, (1, 1, 0, 0)) + self.assertEqual(memo[node.id], first) + + node.value = torch.arange(0, 4, dtype=torch.int64) # a genuinely different histogram + cached = cache._cp_load_back_node_owner_page_counts(node, cp_size=4, memo=memo) + self.assertEqual(cached, first) # served from the memo, not recomputed + self.assertNotEqual( + cache._cp_load_back_node_owner_page_counts(node, cp_size=4), # no memo -> recomputes + first, + ) + def test_load_back_plan_fails_closed_without_cp_metadata(self): allocator = _make_allocator() cache = _make_cache(allocator)