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

Avoid CP HiCache owner-lane planning stalls during load-back

Browse Source 
CP HiCache load-back eviction planning previously recomputed per-owner page counts from node token tensors while scanning evictable leaves. Under shared-KV pressure this can put scheduler-side planning onto an expensive tensor-padding path and stall before load_back can complete.

This stores per-CP-size owner page counts on CP HiCache metadata and uses that CPU metadata for backed/resident CP nodes. Streaming abort handling also accepts int-like status codes so abort responses do not crash on .name/.value access. Temporary debug runbooks remain ignored.

Unnecessary prefill hot-path timing logs were removed before commit; owner-lane eviction now keeps warning-level output for slow planning, insufficient eviction, or remaining deficits only.

Constraint: CP shared-KV cache residency is page-owner based and already records page owners in CpHiCacheNodeMetadata.
Rejected: Keep verbose prefill/owner-lane timing logs | they proved the issue but add hot-path noise after validation.
Confidence: medium
Scope-risk: moderate
Directive: Do not reintroduce tensor-derived owner counting on CP HiCache backed nodes without measuring scheduler CPU/GPU sync cost.
Tested: python -m py_compile python/sglang/srt/mem_cache/hiradix_cache.py python/sglang/srt/entrypoints/openai/serving_base.py python/sglang/srt/entrypoints/openai/serving_chat.py python/sglang/srt/entrypoints/openai/serving_completions.py test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py
Tested: git diff --check
Not-tested: Local pytest collection is blocked by missing starlette dependency.
Not-tested: Full ETE after log cleanup; previous pre-cleanup ETE replay reached 136098.82 prompt tok/s without killing prefill.
This commit is contained in:
laoyao0822
2026-06-09 06:55:15 +08:00
parent 50fde834ae
commit 81eb138a26
6 changed files with 144 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py
View File

@@ -319,6 +319,25 @@ class TestCpHiCacheLoadBackOwnerLanes(CustomTestCase):
self.assertEqual(counts, (1, 1, 0, 0))
def test_device_victim_owner_counts_use_cp_metadata_without_tensor_padding(self):
allocator = _make_allocator(page_size=4, cp_size=4)
cache = _make_cache(allocator)
node = _make_node(14, 140, [0, 0, 2], value=torch.arange(4, 16, dtype=torch.int64))
import sglang.srt.mem_cache.hiradix_cache as hiradix_cache
old_pad = hiradix_cache.pad_token_locs_to_page_boundary
try:
hiradix_cache.pad_token_locs_to_page_boundary = lambda *args, **kwargs: (_ for _ in ()).throw(
AssertionError("metadata-backed counts must not touch tensor padding")
)
counts = cache._cp_load_back_node_owner_page_counts(node, cp_size=4)
finally:
hiradix_cache.pad_token_locs_to_page_boundary = old_pad
self.assertEqual(counts, (2, 0, 1, 0))
self.assertIs(counts, node.cp_hicache.owner_page_counts(4))
def test_load_back_plan_fails_closed_without_cp_metadata(self):
allocator = _make_allocator()
cache = _make_cache(allocator)