Stabilize EAGLE draft cache hits under CP HiCache

The failing runs showed EAGLE accept length collapsing when draft cache-hit suffixes used the new partial-current splice path.  This keeps target partial-current reuse enabled, but returns EAGLE/NextN draft cache-hit suffixes to the previous full-materialize path with an explicit fallback warning until the draft splice path has value-level ETE proof.\n\nThe same change set also tightens the page-granular CP HiCache contract for scheduler-visible hits and makes the prefill-to-decode EAGLE handoff observable without cloning hot-path metadata.  Exact non-page CP hits are floored to a page boundary for new scheduling decisions, while internal unfinished-request refresh keeps its exact accounting.\n\nConstraint: CP shared KV and HiCache operate at page granularity; exposing token-precise CP tails to scheduler-visible cache hits can force non-page partial materialization.\nConstraint: EAGLE/NextN draft has only one executable layer, so draft prefetch and draft partial-current splice need a separate correctness contract from target layers.\nRejected: Keep draft partial-current splice enabled | remote logs correlate it with avg accept length around 0.068 and median 0.\nRejected: Clone decode metadata tensors on transfer | slot ownership until process_prebuilt consumes them avoids extra hot-path copies.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not re-enable draft partial-current reuse without metadata/draft-KV value checks and ETE accept-length evidence.\nTested: g0034 container py_compile for touched modules.\nTested: g0034 container PYTHONPATH=python python -m pytest -q test/registered/unit/disaggregation/test_decode_queue_compaction.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -> 183 passed, 5 warnings, 2 subtests passed.\nNot-tested: Fresh ETE accept-length run after this exact commit; requires user-driven traffic restart.
This commit is contained in:
laoyao0822
2026-05-30 22:31:43 +08:00
parent 10296a5fef
commit b328baec7c
14 changed files with 920 additions and 21 deletions
@@ -2596,9 +2596,12 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(6)))))
self.assertEqual(result.device_indices.tolist(), [])
self.assertEqual(result.host_hit_length, 6)
# CP HiCache exposes page-granular host hits. A non-page exact tail is
# floored and the tail is recomputed by the incoming request.
self.assertEqual(result.host_hit_length, 4)
self.assertIs(result.last_device_node, root)
self.assertIs(result.last_host_node, node)
self.assertEqual(result.last_host_node.key.token_ids, [0, 1, 2, 3])
self.assertIsNot(result.last_host_node, node)
def test_cp_backed_tail_split_floors_to_page_boundary(self):
cache = HiRadixCache.__new__(HiRadixCache)
@@ -2811,6 +2814,47 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
self.assertEqual(result.last_device_node.key.token_ids, [0, 1, 2, 3])
self.assertNotIn((4, 5), result.last_device_node.children)
def test_cp_match_prefix_floors_exact_valid_tail_for_exact_key(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.device = "cpu"
cache.disable = False
cache.page_size = 4
cache.ongoing_write_through = {}
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: (key, None)
cache._update_leaf_status = lambda node: None
cache._update_host_leaf_status = lambda node: None
cache.enable_storage = False
cache.enable_kv_cache_events = False
root = TreeNode()
root.key = RadixKey([])
root.value = torch.empty((0,), dtype=torch.int64)
root.host_len = 0
root.children = {}
cache.root_node = root
node = TreeNode()
node.id = 149
node.parent = root
node.key = RadixKey(list(range(6)))
node.value = torch.arange(6, dtype=torch.int64)
node.host_value = None
node.host_len = 0
node.cp_hicache = None
root.children[(0, 1, 2, 3)] = node
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(6)))))
self.assertEqual(result.device_indices.tolist(), [0, 1, 2, 3])
self.assertEqual(result.host_hit_length, 0)
self.assertEqual(result.last_device_node.key.token_ids, [0, 1, 2, 3])
self.assertNotIn((4, 5), result.last_device_node.children)
def test_cp_insert_floors_backed_tail_split_to_page_boundary(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
@@ -2938,6 +2982,72 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
self.assertTrue(prepared.attached)
self.assertIs(cache.pending_host_backups[148].node, new_tail)
def test_cp_insert_replaces_exact_valid_tail_from_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.ongoing_write_through = {}
cache.cache_controller = types.SimpleNamespace(
has_draft_hicache=False,
write_policy="write_through",
)
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.protected_size_ = 0
cache.enable_storage = False
cache.enable_kv_cache_events = False
cache.inc_node_lock_ref = lambda node: None
root = TreeNode()
root.key = RadixKey([])
root.value = torch.empty((0,), dtype=torch.int64)
root.children = {}
cache.root_node = root
node = TreeNode()
node.id = 150
node.parent = root
node.key = RadixKey(list(range(6)))
node.value = torch.arange(6, dtype=torch.int64)
node.host_len = 0
node.cp_hicache = None
root.children[(0, 1, 2, 3)] = node
reservation = make_write_reservation(
torch.arange(4, 6, dtype=torch.int64), node_id=151, host_start=210
)
prepared = PreparedCpHiCacheBackup(
node_id=151,
reservation=reservation,
metadata=reservation.metadata,
logical_len=2,
)
result = cache.insert(
InsertParams(
key=RadixKey(list(range(6))),
value=torch.arange(6, dtype=torch.int64),
cp_hicache_prepared_backup=prepared,
)
)
self.assertEqual(result.prefix_len, 4)
parent = root.children[(0, 1, 2, 3)]
self.assertEqual(parent.key.token_ids, [0, 1, 2, 3])
new_tail = parent.children[(4, 5)]
self.assertIsNot(new_tail, node)
self.assertEqual(new_tail.key.token_ids, [4, 5])
self.assertTrue(prepared.attached)
self.assertIs(cache.pending_host_backups[151].node, new_tail)
def test_non_cp_match_prefix_uses_root_when_no_host_backup_exists(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = False
@@ -569,7 +569,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
forward_batch.out_cache_loc = torch.arange(64, dtype=torch.int64)
self.assertFalse(can_reuse_current_extend_kv(forward_batch))
def test_should_reuse_current_extend_kv_allows_draft_cache_hit_suffix(self):
def test_should_reuse_current_extend_kv_disables_draft_partial_cache_hit_suffix(
self,
):
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
@@ -593,8 +595,13 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
)
with envs.SGLANG_CP_SHARED_KV_CURRENT_REUSE.override(True):
self.assertTrue(runtime.should_reuse_current_extend_kv(forward_batch))
self.assertEqual(runtime._CURRENT_REUSE_FALLBACK_LOG_COUNTS, {})
with self.assertLogs(runtime.logger.name, level="WARNING") as logs:
self.assertFalse(runtime.should_reuse_current_extend_kv(forward_batch))
self.assertIn(
"[CP_SHARED_KV_FALLBACK][current_reuse]",
logs.output[0],
)
self.assertIn("draft_partial_current_reuse_disabled", logs.output[0])
forward_batch.spec_info = TargetSpecInfo()
forward_batch.cp_shared_kv_mla_prefetcher = object()