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) <noreply@anthropic.com>
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user