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 28f679bf8..81994802e 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 @@ -5088,3 +5088,65 @@ Regression test: attached catch-up ack `[53]`. `writing_check()` must not raise `KeyError`, must preserve both ack entries, and must not release the attached node behind the unattached head entry. + +### C120 — 2026-06-02 Chunked prefill should cache completed middle-chunk pages + +Finding: + +- With chunked prefill enabled, the CP HiCache path previously skipped + pre-forward prepared backup for every chunked request: + `[CP_HICACHE_FALLBACK][prepare_write_backup_skipped] reason=chunked_req`. +- `cache_unfinished_req(chunked=True)` still inserted scheduler-visible valid + tails. For CP HiCache this is the wrong granularity: the cache is managed at + page granularity, so inserting a non-page-aligned chunk tail creates a stale + tail that a later chunk must prune while the same request can still own or + lock those KV slots. +- Remote fallback storms showed this exact shape: many + `prepare_write_backup_skipped` and `insert_deferred_pending_backup_split` + events, often with `ongoing=0 pending=0`, meaning the split was blocked by + node protection/refs rather than an actual active transfer. + +Requirement: + +- Intermediate chunks are not disposable. They should write completed pages + into CP HiCache to maximize reuse and avoid doing all host backup only at the + final chunk. +- The current incomplete page must remain request-owned until a later chunk or + final insert makes it page-complete. CP HiCache should not create a radix + node for a sub-page valid tail. + +Correction: + +- `prepare_write_backup_for_req()` no longer skips chunked requests. For + chunked CP HiCache requests it floors the backup end to the nearest completed + page and reserves/writes only `[start, floor_to_page_len(len(keys)))`. +- `cache_unfinished_req(chunked=True)` now inserts only completed pages when CP + HiCache is active. If no new completed page exists beyond + `cache_protected_len`, it keeps `req.prefix_indices` pointing at the full + request-owned KV span and returns without mutating the radix tree. +- The sub-page tail is deliberately retained in `req.prefix_indices` but not in + the CP HiCache radix tree. This trades at most one page of temporary + non-cacheable tail for stable page-granular radix/host-cache state. + +Regression tests: + +- `test_cp_hicache_metadata.py::TestHiRadixCacheCPBackup::test_prepare_write_backup_for_req_chunked_reserves_completed_pages` + verifies a 10-token chunk with `page_size=4` reserves only the first 8 + completed tokens instead of warning and skipping backup. +- `test_cp_hicache_metadata.py::TestHiRadixCacheCPBackup::test_cp_cache_unfinished_chunked_inserts_only_completed_pages` + verifies `cache_unfinished_req(chunked=True)` inserts only 8 page-completed + tokens while preserving full 10-token `req.prefix_indices` for the scheduler. + +Validation: + +- Remote targeted regression tests on `g0034`: + `2 passed, 3 warnings`. +- Remote full metadata suite: + `PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py` + passed with `107 passed, 5 warnings`. + +Remaining validation gap: + +- Requires a fresh ETE run with chunked prefill enabled to verify that the + fallback storm disappears in production traffic. An already-running remote + process will not pick up this change until restarted. diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 1e3e65a88..a683ccb40 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -2325,14 +2325,6 @@ class HiRadixCache(RadixCache): rid=getattr(req, "rid", ""), ) return - if getattr(req, "is_chunked", 0) > 0: - self._warn_cp_hicache_fallback( - "prepare_write_backup_skipped", - "chunked_req", - rid=getattr(req, "rid", ""), - is_chunked=getattr(req, "is_chunked", 0), - ) - return if getattr(req, "cp_hicache_prepared_backup", None) is not None: prepared = getattr(req, "cp_hicache_prepared_backup", None) self._warn_cp_hicache_fallback( @@ -2345,20 +2337,40 @@ class HiRadixCache(RadixCache): token_ids = req.fill_ids keys = convert_to_bigram_key(token_ids) if self.is_eagle else token_ids + backup_end = len(keys) + if getattr(req, "is_chunked", 0) > 0 and self.page_size > 1: + # Chunked prefill should still populate CP HiCache, but only for + # pages whose KV has become fully available. Leaving the current + # sub-page tail out of the prepared write avoids creating a + # scheduler-visible stale tail that the next chunk has to prune + # while the same request still owns/locks it. + backup_end = floor_to_page_len(backup_end, self.page_size) + if backup_end <= 0: + logger.debug( + "[HiCache-write] skip chunked CP backup with no completed page: rid=%s key_len=%d", + getattr(req, "rid", ""), + len(keys), + ) + return + keys_for_probe = keys[:backup_end] + else: + keys_for_probe = keys existing_prefix_len = self._probe_existing_radix_prefix_len_no_split( RadixKey( - keys, + keys_for_probe, getattr(req, "extra_key", None), is_bigram=self.is_eagle, ) ) - raw_start = min(max(req.cache_protected_len, existing_prefix_len), len(keys)) - if len(keys) <= raw_start: + raw_start = min(max(req.cache_protected_len, existing_prefix_len), backup_end) + if backup_end <= raw_start: return start = floor_to_page_len(raw_start, self.page_size) + if backup_end <= start: + return kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, start : len(keys) + req.req_pool_idx, start:backup_end ].to(dtype=torch.int64, copy=True) if len(kv_indices) == 0: self._warn_cp_hicache_fallback( @@ -2366,7 +2378,7 @@ class HiRadixCache(RadixCache): "empty_kv_indices", rid=getattr(req, "rid", ""), start=start, - key_len=len(keys), + key_len=backup_end, ) return diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index d9b0d2b14..35b6513ef 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -635,6 +635,18 @@ class RadixCache(BasePrefixCache): keys = convert_to_bigram_key(token_ids) if self.is_eagle else token_ids if not getattr(self, "_uses_cp_hicache", False): keys = page_align_keys(keys, self.page_size) + elif chunked and self.page_size > 1: + # CP HiCache owns KV at page granularity. For an unfinished + # chunked request, cache only completed pages and keep the current + # sub-page tail in the request-owned KV slots until a later chunk + # or the final insert can make it page-complete. This preserves + # intermediate chunk caching without repeatedly creating/pruning + # stale valid-tail radix nodes. + completed_len = floor_to_page_len(len(keys), self.page_size) + if completed_len <= req.cache_protected_len: + req.prefix_indices = kv_indices + return + keys = keys[:completed_len] values = kv_indices[: len(keys)].to(dtype=torch.int64, copy=True) radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle) prepared_cp_backup = getattr(req, "cp_hicache_prepared_backup", None) 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 8f8bc9b47..3b5e5062d 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -1221,30 +1221,102 @@ class TestHiRadixCacheCPBackup(CustomTestCase): ) ) - def test_prepare_write_backup_for_req_chunked_skip_warns_fallback(self): + def test_prepare_write_backup_for_req_chunked_reserves_completed_pages(self): cache = HiRadixCache.__new__(HiRadixCache) cache.disable = False cache._uses_cp_hicache = True cache.cache_controller = FakeReserveWriteController([]) + cache.is_eagle = False + cache.page_size = 4 + cache.root_node = TreeNode() + cache.root_node.key = RadixKey([]) + cache.root_node.children = {} + 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.req_to_token_pool = types.SimpleNamespace( + req_to_token=torch.arange(16, dtype=torch.int64).view(1, 16) + ) + cache._reserve_write_cp_indices_no_collective = ( + lambda indices, node_id: make_write_reservation( + indices, node_id=node_id, host_start=300 + ) + ) req = types.SimpleNamespace( rid="rid-chunked", is_chunked=1, + fill_ids=list(range(10)), + req_pool_idx=0, + extra_key=None, + cache_protected_len=0, cp_hicache_prepared_backup=None, ) - with self.assertLogs( - "sglang.srt.mem_cache.hiradix_cache", level="WARNING" - ) as captured: - cache.prepare_write_backup_for_req(req) + cache.prepare_write_backup_for_req(req) - self.assertTrue( - any( - "[CP_HICACHE_FALLBACK][prepare_write_backup_skipped]" in message - and "reason=chunked_req" in message - for message in captured.output - ) + self.assertEqual(len(cache.cache_controller.submitted), 1) + self.assertEqual( + cache.cache_controller.submitted[0].physical_device_indices.tolist(), + list(range(8)), ) + self.assertEqual(req.cp_hicache_prepared_backup.logical_len, 8) + + def test_cp_cache_unfinished_chunked_inserts_only_completed_pages(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache.disable = False + cache._uses_cp_hicache = True + cache.is_eagle = False + cache.page_size = 4 + class Pool: + def __init__(self): + self.req_to_token = torch.arange(16, dtype=torch.int64).view(1, 16) + self.writes = [] + + def write(self, index, values): + self.writes.append((index, values.detach().cpu().tolist())) + + pool = Pool() + cache.req_to_token_pool = pool + cache.token_to_kv_pool_allocator = RecordingTokenAllocator() + inserted = [] + + def insert(params): + inserted.append(params) + return types.SimpleNamespace(prefix_len=8, pending_backup_deferred_node=None) + + cache.insert = insert + cache.match_prefix = lambda params: types.SimpleNamespace( + device_indices=torch.arange(8, dtype=torch.int64), + last_device_node=TreeNode(), + ) + cache._free_kv_indices_range = lambda *args, **kwargs: None + cache.dec_lock_ref = lambda node: None + cache.inc_lock_ref = lambda node: None + + req = types.SimpleNamespace( + fill_ids=list(range(10)), + req_pool_idx=0, + extra_key=None, + cache_protected_len=0, + last_node=TreeNode(), + prefix_indices=torch.empty((0,), dtype=torch.int64), + cp_hicache_prepared_backup=None, + ) + + cache.cache_unfinished_req(req, chunked=True) + + self.assertEqual(len(inserted), 1) + self.assertEqual(inserted[0].key.token_ids, list(range(8))) + self.assertEqual(inserted[0].value.tolist(), list(range(8))) + self.assertEqual( + pool.writes, + [((0, slice(0, 8, None)), list(range(8)))], + ) + self.assertEqual(req.cache_protected_len, 8) + self.assertEqual(req.prefix_indices.tolist(), list(range(10))) def test_write_backup_deterministic_eviction_avoids_reserve_all_reduce(self): cache = HiRadixCache.__new__(HiRadixCache)