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