From 0043037f78c1510db1034e217d4580be2faf402d Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Fri, 29 May 2026 05:39:25 +0800 Subject: [PATCH] Keep CP HiCache owner lanes padded for valid-tail victims Load-back owner-lane eviction can need to evict a device-resident node whose radix value is a valid tail shorter than the physical tail page. The eviction planner now pads that value only to the page boundary before deriving owner counts, matching the existing write/load capacity contract without exposing padding to radix or scheduler lengths. Constraint: CP HiCache capacity remains page-owner based while radix node values remain valid-token based. Constraint: Avoid collectives; owner-lane capacity must be deterministic from local metadata and logical page ids. Rejected: Require device-resident victim values to be page-aligned | valid-tail cache nodes are now an intentional supported state. Rejected: Pad to cp_size pages | this would waste KV and violate the page-boundary-only contract. Confidence: medium Scope-risk: narrow Directive: If split-inside-tail support is added later, preserve page ownership/refcount semantics before sharing one padded physical page across radix nodes. Tested: local py_compile for hiradix_cache.py and touched CP HiCache tests. Tested: remote g0034 new C8 exact tests: 3 passed, 3 warnings. Tested: remote g0034 CP HiCache impacted suites: 146 passed, 5 warnings. Tested: remote g0034 CP shared KV C1-C5 suite: 122 passed, 5 warnings. Not-tested: full local pytest, blocked by missing runtime dependencies such as starlette. Not-tested: CUDA E2E runtime for this commit. Co-authored-by: OmX --- ..._prefill_cp_page_aligned_cache_contract.md | 19 ++++++ python/sglang/srt/mem_cache/hiradix_cache.py | 16 +++-- .../test_cp_hicache_load_back_owner_lanes.py | 10 +++ .../mem_cache/test_cp_hicache_metadata.py | 66 +++++++++++++++++++ 4 files changed, 104 insertions(+), 7 deletions(-) diff --git a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md index b58e543fc..56135aa07 100644 --- a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md +++ b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md @@ -656,6 +656,10 @@ Current state: - HiCache capacity planning counts `page_size` tokens per owner page. - Some entry points still pass scalar valid lengths or page-aligned-only `device_indices`. +- `HiRadixCache._cp_load_back_node_owner_page_counts()` still rejects a + device-resident CP node whose `node.value` is a valid tail shorter than the + padded physical page span. This can break load-back owner-lane eviction when + the victim that should free a lane is a valid-tail node. Correction: @@ -670,6 +674,21 @@ Tests: one page on owner 1. - Target and draft required vectors must match when draft HiCache is enabled. +Implemented C8 slice: + +- CP write admission derives required host capacity from the padded page owner + vector; a 100-token suffix on 64-token pages requires `(64, 64, 0, ...)`, + not a 100-token scalar and not CP-size padding. +- Target and draft admission share the same padded required vector; only their + available-capacity vectors differ. +- Host capacity snapshots count committed and pending valid-tail nodes by + `metadata.page_owners` / `metadata.padded_len`, so target/draft usage reflects + physical pages while `host_len` remains valid-token based. +- Load-back owner-lane eviction now pads a device-resident valid-tail victim to + its physical tail page before counting freed owner pages. A 6-token value on + 4-token pages frees owner pages `(1, 1, 0, 0)` when it spans logical pages 1 + and 2. + ### C9. Draft/EAGLE currently has safe fallbacks, not the final contract Current state: diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index b2812816b..706726f96 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -1033,14 +1033,16 @@ class HiRadixCache(RadixCache): value = getattr(node, "value", None) if value is None or int(value.numel()) == 0: return tuple(0 for _ in range(cp_size)) - if int(value.numel()) % self.page_size != 0: - raise RuntimeError( - "CP HiCache load-back eviction requires page-aligned device values: " - f"node_id={getattr(node, 'id', '?')} value_len={value.numel()} " - f"page_size={self.page_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 = value[:: self.page_size] + 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)) 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 ad7f0bdec..90286e72f 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 @@ -252,6 +252,16 @@ class TestCpHiCacheLoadBackOwnerLanes(CustomTestCase): self.assertEqual(plan.page_owners, [0, 1]) self.assertEqual(plan.required_by_owner, [1, 1, 0, 0]) + def test_device_victim_owner_counts_pad_valid_tail_to_physical_pages(self): + 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) + + counts = cache._cp_load_back_node_owner_page_counts(node, cp_size=4) + + self.assertEqual(counts, (1, 1, 0, 0)) + def test_load_back_plan_fails_closed_without_cp_metadata(self): allocator = _make_allocator() cache = _make_cache(allocator) 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 eeaec8700..b262e45de 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -1556,6 +1556,72 @@ class TestHiRadixCacheCPBackup(CustomTestCase): self.assertEqual(snapshot.pending_target, (64, 0)) self.assertEqual(snapshot.pending_draft, (64, 0)) + def test_cp_write_admission_uses_padded_owner_vector_for_valid_tail(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.page_size = 64 + cache.cache_controller = types.SimpleNamespace( + cp_shared_kv_layout=FakeCpLayout(cp_size=8, cp_rank=0, page_size=64), + has_draft_hicache=True, + draft_mem_pool_host=types.SimpleNamespace(size=1024), + ) + cache.token_to_kv_pool_host = types.SimpleNamespace(size=1024) + cache.pending_host_backups = {} + cache.root_node = TreeNode() + cache.root_node.children = {} + cache.evictable_host_leaves = set() + + admission = cache._cp_build_write_admission( + torch.arange(64, 164, dtype=torch.int64), + node_id=503, + phase="unit", + ) + + self.assertEqual(admission.required_by_owner, (64, 64, 0, 0, 0, 0, 0, 0)) + self.assertEqual( + admission.target_available_by_owner, + (1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024), + ) + self.assertEqual( + admission.target_available_by_owner, admission.draft_available_by_owner + ) + self.assertEqual(admission.deficit_by_owner, (0, 0, 0, 0, 0, 0, 0, 0)) + + def test_cp_capacity_snapshot_counts_padded_valid_tail_for_target_and_draft(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.cache_controller = types.SimpleNamespace( + cp_shared_kv_layout=FakeCpLayout(cp_size=2, cp_rank=0, page_size=64), + has_draft_hicache=True, + draft_mem_pool_host=types.SimpleNamespace(size=1024), + ) + cache.token_to_kv_pool_host = types.SimpleNamespace(size=1024) + cache.pending_host_backups = {} + root = TreeNode() + root.children = {} + cache.root_node = root + + committed = TreeNode() + committed.id = 504 + committed.parent = root + committed.key = RadixKey(list(range(100))) + committed.host_len = 100 + committed.cp_hicache = CpHiCacheNodeMetadata( + logical_len=100, + padded_len=128, + owned_positions=torch.arange(64, dtype=torch.int64), + host_indices=torch.arange(100, 164, dtype=torch.int64), + draft_host_indices=torch.arange(200, 264, dtype=torch.int64), + page_owners=torch.tensor([0, 1], dtype=torch.int8), + page_size=64, + ) + root.children[0] = committed + + snapshot = cache._cp_host_capacity_snapshot() + + self.assertEqual(snapshot.committed_target, (64, 64)) + self.assertEqual(snapshot.committed_draft, (64, 64)) + class TestHiRadixCacheCPSplitEvict(CustomTestCase): def test_split_node_splits_cp_metadata_by_owned_positions(self):