Stabilize CP HiCache residency under L1/L2 pressure

CP shared KV now keeps explicit L1 and host free-room targets so pressure is handled by planned eviction instead of repeated capacity-edge retries. The host allocator gains contiguous-preferred page reservation, L1 owner-lane allocation prefers contiguous physical pages, and CP HiCache metadata preserves pending backup safety for page-granular radix updates. Mooncake transfer stats and allocator microbenchmarks are included to make the remaining transfer bottlenecks measurable rather than inferred.

Constraint: CP shared KV uses decode CP size 1 with all prefill CP ranks participating in transfer, so L1/L2 cache residency must remain page-granular and avoid extra collectives.\nConstraint: Production HiCache can be hundreds of GB, so allocator metadata overhead must be visible before enabling aggressive contiguous allocation broadly.\nRejected: Evict only the exact deficit | this keeps the cache at the cliff and causes repeated evict/allocate pressure.\nRejected: Rely on allocator scans alone for contiguity | remote microbenchmarks show fragmented 220GB-equivalent host metadata can make contiguous-preferred scans multi-ms.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not increase L1/L2 free-room defaults or add new CP collectives without ETE evidence and transfer/allocator measurements.\nTested: python -m py_compile on touched runtime/test/benchmark files.\nTested: PYTHONPATH=. python -m pytest -q test/registered/unit/benchmark/test_cp_hicache_allocator_bench.py => 4 passed, 1 warning.\nTested: Remote g0034 log /mnt/beegfs/cjy/log/sglang_cp_hicache_20260601_233723.log shows active prefill process with L1/L2 free-room args, 702 HTTP 200 chat completions, 6272 prefill batches, and no fatal scheduler traceback in latest scan.\nTested: User-reported L1/L2 cache ETE validation passed on remote run.\nNot-tested: Full local pytest suite; local environment is missing several runtime dependencies.\nNot-tested: CUDA allocator microbenchmark during active production prefill process.\nNot-tested: Mooncake straggler fix; stats show transfer tail latency remains a separate bottleneck.
This commit is contained in:
laoyao0822
2026-06-02 07:58:02 +08:00
parent 8be4a3a8b5
commit ce3a20d11b
17 changed files with 2268 additions and 16 deletions
@@ -107,6 +107,8 @@ from sglang.srt.mem_cache.base_prefix_cache import (
MatchPrefixParams,
)
from sglang.srt.mem_cache.hiradix_cache import (
CpHiCacheCapacitySnapshot,
CpHiCacheEvictionPlan,
CpHiCacheNodeMetadata,
HiRadixCache,
PendingHiCacheBackup,
@@ -605,6 +607,132 @@ class TestCpHiCacheNodeMetadata(CustomTestCase):
self.assertEqual(child.page_owners.numel(), 0)
class TestCpHiCacheFreeRoom(CustomTestCase):
def test_free_room_deficit_does_not_evict_when_required_fits_trigger(self):
deficit = hiradix_cache._free_room_deficit(
required=64,
available=96,
capacity=256,
page_size=64,
target_ratio=0.5,
trigger_ratio=0.0,
)
self.assertEqual(deficit, 0)
def test_free_room_deficit_evicts_to_target_after_trigger(self):
deficit = hiradix_cache._free_room_deficit(
required=64,
available=96,
capacity=256,
page_size=64,
target_ratio=0.5,
trigger_ratio=0.25,
)
self.assertEqual(deficit, 96)
def test_free_room_deficit_rounds_room_to_page(self):
deficit = hiradix_cache._free_room_deficit(
required=64,
available=0,
capacity=100,
page_size=64,
target_ratio=0.01,
trigger_ratio=0.0,
)
self.assertEqual(deficit, 128)
def test_cp_host_write_admission_uses_trigger_target_room(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache.page_size = 64
cache.hicache_host_free_room_ratio = 0.5
cache.hicache_host_free_room_trigger_ratio = 0.25
cache._cp_required_host_tokens_by_rank = lambda _indices: (64, 0)
cache._cp_host_capacity_snapshot = lambda: CpHiCacheCapacitySnapshot(
target_capacity=(256, 256),
draft_capacity=None,
committed_target=(160, 0),
committed_draft=(0, 0),
pending_target=(0, 0),
pending_draft=(0, 0),
)
cache._plan_cp_host_evictions = lambda deficit: CpHiCacheEvictionPlan(
victims=(),
planned_freed=tuple(0 for _ in deficit),
remaining_deficit=tuple(0 for _ in deficit),
)
admission = cache._cp_build_write_admission(
torch.arange(64, dtype=torch.int64),
node_id=600,
phase="unit",
)
self.assertEqual(admission.target_available_by_owner, (96, 256))
self.assertEqual(admission.deficit_by_owner, (96, 0))
def test_cp_host_write_admission_does_not_evict_when_trigger_room_fits(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache.page_size = 64
cache.hicache_host_free_room_ratio = 0.5
cache.hicache_host_free_room_trigger_ratio = 0.0
cache._cp_required_host_tokens_by_rank = lambda _indices: (64, 0)
cache._cp_host_capacity_snapshot = lambda: CpHiCacheCapacitySnapshot(
target_capacity=(256, 256),
draft_capacity=None,
committed_target=(160, 0),
committed_draft=(0, 0),
pending_target=(0, 0),
pending_draft=(0, 0),
)
cache._plan_cp_host_evictions = lambda deficit: CpHiCacheEvictionPlan(
victims=(),
planned_freed=tuple(0 for _ in deficit),
remaining_deficit=tuple(0 for _ in deficit),
)
admission = cache._cp_build_write_admission(
torch.arange(64, dtype=torch.int64),
node_id=601,
phase="unit",
)
self.assertEqual(admission.target_available_by_owner, (96, 256))
self.assertEqual(admission.deficit_by_owner, (0, 0))
def test_cp_host_write_admission_uses_draft_room_deficit(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache.page_size = 64
cache.hicache_host_free_room_ratio = 0.25
cache.hicache_host_free_room_trigger_ratio = 0.25
cache._cp_required_host_tokens_by_rank = lambda _indices: (64, 0)
cache._cp_host_capacity_snapshot = lambda: CpHiCacheCapacitySnapshot(
target_capacity=(256, 256),
draft_capacity=(256, 256),
committed_target=(0, 0),
committed_draft=(192, 0),
pending_target=(0, 0),
pending_draft=(0, 0),
)
cache._plan_cp_host_evictions = lambda deficit: CpHiCacheEvictionPlan(
victims=(),
planned_freed=tuple(0 for _ in deficit),
remaining_deficit=tuple(0 for _ in deficit),
)
admission = cache._cp_build_write_admission(
torch.arange(64, dtype=torch.int64),
node_id=602,
phase="unit",
)
self.assertEqual(admission.target_available_by_owner, (256, 256))
self.assertEqual(admission.draft_available_by_owner, (64, 256))
self.assertEqual(admission.deficit_by_owner, (64, 0))
class FakeWriteFailure:
metadata = None
@@ -2592,6 +2720,39 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase):
class TestHiRadixCacheCPLoadBack(CustomTestCase):
def test_cp_load_back_plan_uses_l1_free_room_target(self):
class FreeRoomAllocator:
def compute_owner_lane_stats(self, _page_owners):
return [1, 0], [0, 8], [1, 0]
def compute_owner_lane_capacity_pages(self):
return [8, 8]
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.page_size = 64
cache.hicache_l1_free_room_ratio = 0.5
cache.hicache_l1_free_room_trigger_ratio = 0.25
cache.token_to_kv_pool_allocator = FreeRoomAllocator()
node = TreeNode(id=701)
node.host_len = 64
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=64,
padded_len=64,
owned_positions=torch.arange(64, dtype=torch.int64),
host_indices=torch.arange(64, dtype=torch.int64),
page_owners=torch.tensor([0], dtype=torch.int8),
page_size=64,
)
plan = cache._build_cp_load_back_plan([node], node_id=701)
self.assertEqual(plan.required_by_owner, [1, 0])
self.assertEqual(plan.available_by_owner, [0, 8])
# required=1 page, available=0, target_room=ceil(8*0.5)=4 pages.
self.assertEqual(plan.deficit_by_owner, [5, 0])
def test_cp_load_back_uses_host_len_not_host_value(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True