chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Apply CP HiCache page-floor policy before backup and insert

Browse Source 
Backed CP HiCache tail splits were floored during match, but write preparation and radix insertion could still observe the raw sub-page prefix. That could skip one token during pre-forward reservation or call CpHiCacheNodeMetadata.split() with a non-page boundary. The shared helper now floors backed partial-tail prefixes before probe, match, and insert.

Constraint: CP HiCache uses page-granular physical ownership while radix keys may keep valid-tail lengths.

Rejected: Let probe and insert keep token-precise tail prefixes | it reintroduces half-page ownership through a different entry point.

Confidence: high

Scope-risk: moderate

Directive: Any CP radix path that can split or reserve against backed HiCache metadata must use the same page-floor policy.

Tested: Remote red tests first for probe and insert failures in g0034 container.

Tested: Remote py_compile for hiradix_cache.py and test_cp_hicache_metadata.py.

Tested: Remote targeted C7 tests: 4 passed, 3 warnings.

Tested: Remote pytest test_cp_hicache_metadata.py test_cp_hicache_load_back_owner_lanes.py: 94 passed, 5 warnings.

Not-tested: Live ETE traffic with concurrent divergent tail prompts.

Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-05-29 07:10:53 +08:00
parent 8001c4ae8e
commit 2b1524bd8c
3 changed files with 134 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md
View File

@@ -660,6 +660,11 @@ Implemented C7 split policy:
- If the floored boundary is zero, the match returns the parent/root node and
sacrifices the sub-page prefix. This avoids half-page owner accounting and
keeps host/device/draft metadata page-granular.
- The same floor policy is applied before forward-time write reservation and
final radix insertion. `_probe_existing_radix_prefix_len_no_split()` floors
backed partial-tail hits so backup reservation starts at the same page boundary
as `match_prefix()`, and `insert()` floors backed partial-tail splits before
calling `CpHiCacheNodeMetadata.split()`.
### C8. Owner-lane capacity must be padded-page based end to end

28
python/sglang/srt/mem_cache/hiradix_cache.py
View File

@@ -2009,6 +2009,12 @@ class HiRadixCache(RadixCache):
prefix_len = self.key_match_fn(child.key, key)
if prefix_len <= 0:
break
if prefix_len < len(child.key):
prefix_len = self._cp_floor_backed_partial_split_len(
child, prefix_len
)
if prefix_len <= 0:
break
total_prefix_len += prefix_len
if prefix_len < len(child.key):
@@ -2021,6 +2027,19 @@ class HiRadixCache(RadixCache):
return total_prefix_len
def _cp_floor_backed_partial_split_len(
self, child: TreeNode, prefix_len: int
) -> int:
if (
not self._uses_cp_hicache
or self.page_size <= 1
or prefix_len <= 0
or prefix_len % self.page_size == 0
or not self._node_backuped(child)
):
return prefix_len
return prefix_len // self.page_size * self.page_size
def prepare_write_backup_for_req(self, req) -> None:
if self.disable or not self._uses_cp_hicache:
return
@@ -3356,7 +3375,9 @@ class HiRadixCache(RadixCache):
and self._node_backuped(child)
and prefix_len % self.page_size != 0
):
prefix_len = prefix_len // self.page_size * self.page_size
prefix_len = self._cp_floor_backed_partial_split_len(
child, prefix_len
)
if prefix_len == 0:
break
new_node = self._split_node(child.key, child, prefix_len)
@@ -3464,6 +3485,11 @@ class HiRadixCache(RadixCache):
total_prefix_length += prefix_len
else:
# partial match, split the node
prefix_len = self._cp_floor_backed_partial_split_len(
node, prefix_len
)
if prefix_len <= 0:
break
new_node = self._split_node(node.key, node, prefix_len)
# shared-prefix node should also reflect max priority
new_node.priority = max(new_node.priority, priority)

102
test/registered/unit/mem_cache/test_cp_hicache_metadata.py
View File

@@ -2437,6 +2437,108 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
self.assertIs(root.children[(0, 1, 2, 3)], node)
self.assertEqual(node.key.token_ids, list(range(6)))
def test_cp_prepare_probe_floors_backed_tail_partial_hit_to_page(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.disable = False
cache.page_size = 4
cache.pending_host_backups = {}
cache.cache_controller = types.SimpleNamespace(has_draft_hicache=False)
cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4])
cache.key_match_fn = lambda child_key, key: _key_match_paged(
child_key, key, page_size=4
)
cache.maybe_bigram_convert = lambda key, value=None: (key, value)
root = TreeNode()
root.key = RadixKey([])
root.children = {}
cache.root_node = root
node = TreeNode()
node.id = 143
node.parent = root
node.key = RadixKey(list(range(6)))
node.value = None
node.host_len = 6
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=6,
padded_len=8,
owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64),
host_indices=torch.tensor([50, 51, 52, 53], dtype=torch.int64),
page_owners=torch.tensor([0, 1], dtype=torch.int8),
page_size=4,
)
root.children[(0, 1, 2, 3)] = node
prefix_len = cache._probe_existing_radix_prefix_len_no_split(
RadixKey([0, 1, 2, 3, 4])
)
self.assertEqual(prefix_len, 4)
def test_cp_insert_floors_backed_tail_split_to_page_boundary(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.disable = False
cache.is_eagle = False
cache.page_size = 4
cache.pending_host_backups = {}
cache.cache_controller = types.SimpleNamespace(
has_draft_hicache=False,
write_policy="write_back",
)
cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4])
cache.key_match_fn = lambda child_key, key: _key_match_paged(
child_key, key, page_size=4
)
cache.maybe_bigram_convert = lambda key, value=None: (key, value)
cache._update_leaf_status = lambda node: None
cache._update_host_leaf_status = lambda node: None
cache._inc_hit_count = lambda *args, **kwargs: None
cache._record_store_event = lambda node: None
cache.evictable_size_ = 0
cache.enable_storage = False
cache.enable_kv_cache_events = False
root = TreeNode()
root.key = RadixKey([])
root.value = torch.empty((0,), dtype=torch.int64)
root.children = {}
cache.root_node = root
node = TreeNode()
node.id = 144
node.parent = root
node.key = RadixKey(list(range(6)))
node.value = torch.arange(6, dtype=torch.int64)
node.host_len = 6
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=6,
padded_len=8,
owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64),
host_indices=torch.tensor([50, 51, 52, 53], dtype=torch.int64),
page_owners=torch.tensor([0, 1], dtype=torch.int8),
page_size=4,
)
root.children[(0, 1, 2, 3)] = node
result = cache.insert(
InsertParams(
key=RadixKey([0, 1, 2, 3, 4]),
value=torch.arange(5, dtype=torch.int64),
)
)
self.assertEqual(result.prefix_len, 4)
parent = root.children[(0, 1, 2, 3)]
self.assertEqual(parent.key.token_ids, [0, 1, 2, 3])
self.assertEqual(parent.cp_hicache.logical_len, 4)
self.assertEqual(parent.cp_hicache.padded_len, 4)
old_tail = parent.children[(4, 5)]
self.assertEqual(old_tail.key.token_ids, [4, 5])
self.assertEqual(old_tail.cp_hicache.logical_len, 2)
self.assertEqual(old_tail.cp_hicache.padded_len, 4)
new_tail = parent.children[(4,)]
self.assertEqual(new_tail.key.token_ids, [4])
self.assertEqual(new_tail.value.tolist(), [4])
def test_non_cp_match_prefix_uses_root_when_no_host_backup_exists(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = False