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:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user