Share CP HiCache host budget across target and draft KV

CP HiCache previously let the target host pool and the draft/MTP host pool each consume the full --hicache-size budget. With EAGLE/MTP enabled this doubled per-rank host allocation and could kill scheduler ranks during startup before Python emitted a traceback. The cache now treats target KV and draft KV as one logical host-cache object: target and draft capacities are computed from one per-rank byte budget, draft may receive more token capacity when its per-token footprint is smaller, and draft attachment remains tied to target residency.

Constraint: --hicache-size is a per-rank host budget and must not be multiplied by attaching draft KV.

Rejected: Give draft another independent --hicache-size allocation | repeats the observed host OOM failure mode.

Rejected: Disable draft HiCache attachment under CP | avoids OOM but breaks target/draft cache-hit consistency for MTP.

Confidence: medium

Scope-risk: moderate

Directive: Keep target and draft KV as one logical HiCache object; do not let draft host allocation consume an independent full hicache-size budget.

Tested: python -m py_compile on modified scheduler/cache/test files

Tested: remote g0034 container PYTHONPATH=python python -m pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py -q (45 passed)

Not-tested: full multi-rank GLM5 server restart after clearing existing remote router/defunct process state

Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-05-27 04:28:50 +08:00
co-authored by OmX
parent d14c02b0dc
commit 8571fe0cd9
6 changed files with 1365 additions and 12 deletions
@@ -91,7 +91,11 @@ for _schema in (
raise
from sglang.srt.mem_cache.base_prefix_cache import EvictParams, MatchPrefixParams
from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata, HiRadixCache
from sglang.srt.mem_cache.hiradix_cache import (
CpHiCacheNodeMetadata,
HiRadixCache,
_compute_shared_hicache_token_capacities,
)
from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -117,6 +121,33 @@ class TestCpHiCacheImports(CustomTestCase):
)
class TestHiRadixCacheCPDraftHostPool(CustomTestCase):
def test_shared_budget_keeps_draft_at_least_target_capacity(self):
target_tokens, draft_tokens = _compute_shared_hicache_token_capacities(
total_host_bytes=1000,
target_size_per_token=6,
draft_size_per_token=2,
page_size=10,
)
self.assertEqual(target_tokens, 120)
self.assertEqual(draft_tokens, 140)
self.assertGreaterEqual(draft_tokens, target_tokens)
self.assertLessEqual(target_tokens * 6 + draft_tokens * 2, 1000)
def test_shared_budget_handles_equal_target_and_draft_size(self):
target_tokens, draft_tokens = _compute_shared_hicache_token_capacities(
total_host_bytes=1000,
target_size_per_token=6,
draft_size_per_token=6,
page_size=10,
)
self.assertEqual(target_tokens, 80)
self.assertEqual(draft_tokens, 80)
self.assertLessEqual(target_tokens * 6 + draft_tokens * 6, 1000)
class TestCpHiCacheNodeMetadata(CustomTestCase):
def test_split_zero_len_moves_all_positions_to_child(self):
metadata = CpHiCacheNodeMetadata(