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

Cache completed chunked prefill pages in CP HiCache

Browse Source 
Chunked prefill was skipping prepared CP HiCache backup entirely while still inserting unfinished valid-tail radix state. That combination caused middle chunks to miss reusable host cache and repeatedly hit stale-tail split fallback paths.\n\nThis keeps CP HiCache page-granular: middle chunks now reserve/write only completed pages, and unfinished radix inserts expose only page-complete prefixes while preserving the request-owned sub-page tail in prefix_indices for later chunks.\n\nConstraint: CP HiCache radix and host residency are managed at page granularity, but chunked prefill can end on a sub-page tail.\nRejected: Continue skipping chunked backup | loses middle-chunk reuse and triggers fallback storms.\nRejected: Insert sub-page chunk tails into CP HiCache | creates stale-tail prune/split conflicts against request-owned KV.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not reintroduce chunked_req backup skip; keep CP HiCache chunked inserts page-complete unless the radix/host cache contract becomes sub-page aware.\nTested: python -m py_compile python/sglang/srt/mem_cache/radix_cache.py python/sglang/srt/mem_cache/hiradix_cache.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py\nTested: g0034 targeted chunked CP HiCache regression tests, 2 passed\nTested: g0034 PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py, 107 passed\nNot-tested: Fresh chunked-prefill ETE traffic run after process restart
This commit is contained in:
laoyao0822
2026-06-02 06:22:17 +08:00
parent 5bd68768d9
commit 6a25c312c7
4 changed files with 182 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

62
docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md
View File

@@ -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.

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

@@ -2325,14 +2325,6 @@ class HiRadixCache(RadixCache):
rid=getattr(req, "rid", "<unknown>"),
)
return
if getattr(req, "is_chunked", 0) > 0:
self._warn_cp_hicache_fallback(
"prepare_write_backup_skipped",
"chunked_req",
rid=getattr(req, "rid", "<unknown>"),
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", "<unknown>"),
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", "<unknown>"),
start=start,
key_len=len(keys),
key_len=backup_end,
)
return

12
python/sglang/srt/mem_cache/radix_cache.py
View File

@@ -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)

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

@@ -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)