Batch CP HiCache host admission before write reservation
CP shared-KV prefill batching can present multiple requests to HiCache write-through in one forward. Keeping host free-room admission per request causes repeated small L2 evictions and defeats the free-room policy, while request metadata and radix attachment still need per-request ownership.\n\nThis change adds a batch prepare path that aggregates required host tokens by CP owner lane, performs one host admission/eviction step for the batch, and then reserves/submits each request independently. SessionAwareCache and the scheduler prefer the batch API when available.\n\nConstraint: Radix nodes, host slots, draft slots, rollback, and ack semantics remain per request.\nRejected: Merge reservations or radix nodes across requests | would complicate rollback and writing_check ack ownership.\nRejected: Add collective synchronization for host admission | local owner-lane logic already provides deterministic capacity accounting.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not reintroduce per-request host free-room eviction in the scheduler path without profiling host-full workloads.\nTested: local py_compile for touched Python files\nTested: remote g0034 docker PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py => 117 passed, 5 warnings\nNot-tested: full ETE bs>1 HiCache run under production traffic
This commit is contained in:
@@ -110,6 +110,7 @@ from sglang.srt.mem_cache.hiradix_cache import (
|
||||
CpHiCacheCapacitySnapshot,
|
||||
CpHiCacheEvictionPlan,
|
||||
CpHiCacheNodeMetadata,
|
||||
CpWriteAdmission,
|
||||
HiRadixCache,
|
||||
PendingHiCacheBackup,
|
||||
PreparedCpHiCacheBackup,
|
||||
@@ -732,6 +733,110 @@ class TestCpHiCacheFreeRoom(CustomTestCase):
|
||||
self.assertEqual(admission.draft_available_by_owner, (64, 256))
|
||||
self.assertEqual(admission.deficit_by_owner, (64, 0))
|
||||
|
||||
def test_cp_host_write_batch_admission_uses_aggregate_required(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_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_from_required(
|
||||
(128, 0),
|
||||
node_id=603,
|
||||
phase="batch_unit",
|
||||
)
|
||||
|
||||
self.assertEqual(admission.required_by_owner, (128, 0))
|
||||
self.assertEqual(admission.target_available_by_owner, (96, 256))
|
||||
# required=128, target_room=128, available=96.
|
||||
self.assertEqual(admission.deficit_by_owner, (160, 0))
|
||||
|
||||
def test_prepare_write_backups_for_reqs_runs_one_batch_admission(self):
|
||||
cache = HiRadixCache.__new__(HiRadixCache)
|
||||
cache.disable = False
|
||||
cache._uses_cp_hicache = True
|
||||
cache.is_eagle = False
|
||||
cache.page_size = 1
|
||||
cache.req_to_token_pool = types.SimpleNamespace(
|
||||
req_to_token=torch.arange(16, dtype=torch.int64).view(2, 8)
|
||||
)
|
||||
cache.cache_controller = FakeReserveWriteController(
|
||||
[
|
||||
lambda device_indices, node_id: make_write_reservation(
|
||||
device_indices, node_id=node_id, host_start=700
|
||||
),
|
||||
lambda device_indices, node_id: make_write_reservation(
|
||||
device_indices, node_id=node_id, host_start=800
|
||||
),
|
||||
]
|
||||
)
|
||||
cache._probe_existing_radix_prefix_len_no_split = lambda key: 0
|
||||
cache._cp_zero_counts = lambda cp_size=None: (0,)
|
||||
cache._cp_add_counts = lambda lhs, rhs: (lhs[0] + rhs[0],)
|
||||
cache._cp_required_host_tokens_by_rank = lambda indices: (int(indices.numel()),)
|
||||
batch_admissions = []
|
||||
evicted_admissions = []
|
||||
|
||||
def build_batch_admission(required, *, node_id, phase):
|
||||
batch_admissions.append((required, node_id, phase))
|
||||
return CpWriteAdmission(
|
||||
node_id=node_id,
|
||||
phase=phase,
|
||||
required_by_owner=required,
|
||||
target_available_by_owner=(0,),
|
||||
draft_available_by_owner=(0,),
|
||||
deficit_by_owner=(1,),
|
||||
eviction_plan=CpHiCacheEvictionPlan(
|
||||
victims=(),
|
||||
planned_freed=(1,),
|
||||
remaining_deficit=(0,),
|
||||
),
|
||||
)
|
||||
|
||||
cache._cp_build_write_admission_from_required = build_batch_admission
|
||||
cache._evict_cp_host_for_write_admission = (
|
||||
lambda admission, **_: evicted_admissions.append(admission) or True
|
||||
)
|
||||
|
||||
req0 = types.SimpleNamespace(
|
||||
rid="batch-rid-0",
|
||||
fill_ids=list(range(4)),
|
||||
cache_protected_len=0,
|
||||
req_pool_idx=0,
|
||||
is_chunked=0,
|
||||
cp_hicache_prepared_backup=None,
|
||||
)
|
||||
req1 = types.SimpleNamespace(
|
||||
rid="batch-rid-1",
|
||||
fill_ids=list(range(6)),
|
||||
cache_protected_len=0,
|
||||
req_pool_idx=1,
|
||||
is_chunked=0,
|
||||
cp_hicache_prepared_backup=None,
|
||||
)
|
||||
|
||||
cache.prepare_write_backups_for_reqs([req0, req1])
|
||||
|
||||
self.assertEqual(len(batch_admissions), 1)
|
||||
self.assertEqual(batch_admissions[0][0], (10,))
|
||||
self.assertEqual(batch_admissions[0][2], "batch_prepare")
|
||||
self.assertEqual(len(evicted_admissions), 1)
|
||||
self.assertIsNotNone(req0.cp_hicache_prepared_backup)
|
||||
self.assertIsNotNone(req1.cp_hicache_prepared_backup)
|
||||
self.assertEqual(len(cache.cache_controller.submitted), 2)
|
||||
|
||||
|
||||
class FakeWriteFailure:
|
||||
metadata = None
|
||||
@@ -1030,6 +1135,20 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
|
||||
|
||||
self.assertEqual(calls, [req])
|
||||
|
||||
def test_session_aware_cache_forwards_cp_hicache_batch_prepare(self):
|
||||
calls = []
|
||||
|
||||
class Inner:
|
||||
def prepare_write_backups_for_reqs(self, reqs):
|
||||
calls.append(list(reqs))
|
||||
|
||||
wrapper = SessionAwareCache(Inner())
|
||||
reqs = [object(), object()]
|
||||
|
||||
wrapper.prepare_write_backups_for_reqs(reqs)
|
||||
|
||||
self.assertEqual(calls, [reqs])
|
||||
|
||||
def test_node_backuped_uses_cp_metadata(self):
|
||||
cache = HiRadixCache.__new__(HiRadixCache)
|
||||
cache._uses_cp_hicache = True
|
||||
|
||||
Reference in New Issue
Block a user