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

Stabilize CP HiCache residency under L1/L2 pressure

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_cp_hicache_metadata.py
+161
View File
@@ -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
test/registered/unit/mem_cache/test_cp_shared_kv_layout.py
+83 -1
View File
@@ -341,7 +341,7 @@ class TestCPSharedPagedAllocator(CustomTestCase):
self.assertIsNotNone(locs)
logical_pages = locs.view(-1, page_size)[:, 0] // page_size
self.assertEqual(logical_pages.tolist(), [9, 3, 1, 7, 4])
self.assertEqual(logical_pages.tolist(), [1, 3, 5, 7, 4])
self.assertEqual(
((logical_pages - 1) % cp_size).tolist(),
page_compute_owners,
@@ -392,6 +392,40 @@ class TestCPSharedPagedAllocator(CustomTestCase):
self.assertEqual(allocator.free_pages.tolist(), [2, 3, 4])
self.assertEqual(allocator.release_pages.tolist(), [])
def test_contiguous_owner_lane_selection_prefers_later_physical_run(self):
from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator
page_size = 64
cp_size = 4
allocator = CPSharedPagedTokenToKVPoolAllocator(
logical_size=page_size * 32,
physical_size=page_size * 8,
page_size=page_size,
dtype=torch.bfloat16,
device="cpu",
kvcache=None,
need_sort=False,
cp_size=cp_size,
cp_rank=0,
)
allocator.free_pages = torch.tensor(
[1, 9, 13, 17]
+ [page for page in range(2, 33) if page not in {9, 13, 17}],
dtype=torch.int64,
)
locs = allocator.alloc_pages_with_owners([0, 0, 0])
self.assertIsNotNone(locs)
logical_pages = locs.view(-1, page_size)[:, 0] // page_size
self.assertEqual(logical_pages.tolist(), [9, 13, 17])
self.assertEqual(
((logical_pages - 1) % cp_size).tolist(),
[0, 0, 0],
)
for selected_page in logical_pages.tolist():
self.assertNotIn(selected_page, allocator.free_pages.tolist())
def test_compute_owner_alloc_does_not_evict_lanes_when_first_try_succeeds(self):
from types import SimpleNamespace
@@ -663,6 +697,54 @@ class TestCPSharedPagedAllocator(CustomTestCase):
)
self.assertIsNotNone(locs)
def test_compute_owner_lane_eviction_uses_l1_free_room_target(self):
from sglang.srt.mem_cache.base_prefix_cache import EvictResult
from sglang.srt.mem_cache.common import _evict_for_compute_owner_lanes
page_size = 64
class FakeAllocator:
cp_size = 2
def __init__(self):
self.page_size = page_size
def available_size(self):
return 0
def compute_owner_lane_stats(self, _page_compute_owners):
return [1, 0], [0, 8], [1, 0]
def compute_owner_lane_capacity_pages(self):
return [8, 8]
class FakeTreeCache:
hicache_l1_free_room_ratio = 0.5
hicache_l1_free_room_trigger_ratio = 0.25
def __init__(self):
self.owner_deficits = []
def is_chunk_cache(self):
return False
def evictable_size(self):
return page_size * 8
def evict(self, params):
self.owner_deficits.append(list(params.owner_lane_deficits))
return EvictResult(num_tokens_evicted=0)
tree_cache = FakeTreeCache()
_evict_for_compute_owner_lanes(
tree_cache=tree_cache,
allocator=FakeAllocator(),
page_compute_owners=[0],
)
# required=1 page, available=0, target_room=ceil(8*0.5)=4 pages.
self.assertEqual(tree_cache.owner_deficits[0], [5, 0])
def test_compute_owner_capacity_wait_reports_owner_lane_deficits(self):
from types import SimpleNamespace