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.
2026-06-02 07:58:02 +08:00
parent 8be4a3a8b5
commit ce3a20d11b
17 changed files with 2268 additions and 16 deletions
@@ -109,6 +109,7 @@ from sglang.srt.managers.cache_controller import HiCacheController
 from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
 from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata
 from sglang.srt.mem_cache.memory_pool_host import (
+    HostKVCache,
    MHATokenToKVPoolHost,
    MLATokenToKVPoolHost,
    NSATokenToKVPoolHost,
@@ -161,6 +162,50 @@ class FakeHostPool:
        return len(indices)


+class ContiguousPreferredHostPool(FakeHostPool):
+    def __init__(self, alloc_result):
+        super().__init__(alloc_result)
+        self.contiguous_alloc_calls = []
+
+    def alloc_contiguous_preferred(self, need_size):
+        self.contiguous_alloc_calls.append(need_size)
+        if self.alloc_result is None:
+            return None
+        return self.alloc_result[:need_size].clone()
+
+
+class DummyHostKVCacheForAlloc(HostKVCache):
+    def get_size_per_token(self):
+        return 1
+
+    def init_kv_buffer(self):
+        return None
+
+    def load_to_device_per_layer(
+        self, device_pool, host_indices, device_indices, layer_id, io_backend
+    ) -> None:
+        pass
+
+    def backup_from_device_per_layer(
+        self, device_pool, host_indices, device_indices, layer_id, io_backend
+    ) -> None:
+        pass
+
+    def backup_from_device_all_layer(
+        self, device_pool, host_indices, device_indices, io_backend
+    ) -> None:
+        pass
+
+    def get_data_page(self, index, flat: bool = True) -> torch.Tensor:
+        return torch.empty((0,), dtype=torch.uint8)
+
+    def get_dummy_flat_data_page(self) -> torch.Tensor:
+        return torch.empty((0,), dtype=torch.uint8)
+
+    def set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None:
+        pass
+
+
 class FakeDevicePool:
    device = "cpu"
    layer_num = 1
@@ -931,6 +976,46 @@ class TestHiCacheControllerCPWrite(CustomTestCase):
        self.assertEqual(controller.ack_write_queue[0].node_ids, [79])
        self.assertIn("all-layer backup fallback", "\n".join(logs.output))

+    def test_cp_reserve_write_uses_contiguous_preferred_host_alloc(self):
+        host_pool = ContiguousPreferredHostPool(
+            torch.tensor([100, 101, 102, 103], dtype=torch.int64)
+        )
+        draft_host_pool = ContiguousPreferredHostPool(
+            torch.tensor([200, 201, 202, 203], dtype=torch.int64)
+        )
+        controller = self.make_controller(
+            host_pool,
+            cp_rank=1,
+            draft_host_pool=draft_host_pool,
+            draft_mem_pool_device=FakeDevicePool("draft"),
+        )
+        logical_locs = torch.arange(4, 20, dtype=torch.int64)
+
+        reservation = controller.reserve_write_cp(logical_locs, node_id=179)
+
+        self.assertEqual(reservation.metadata.host_indices.tolist(), [100, 101, 102, 103])
+        self.assertEqual(
+            reservation.metadata.draft_host_indices.tolist(), [200, 201, 202, 203]
+        )
+        self.assertEqual(host_pool.contiguous_alloc_calls, [4])
+        self.assertEqual(draft_host_pool.contiguous_alloc_calls, [4])
+        self.assertEqual(host_pool.alloc_calls, [])
+        self.assertEqual(draft_host_pool.alloc_calls, [])
+
+    def test_host_alloc_contiguous_preferred_skips_fragmented_fifo_prefix(self):
+        host_pool = DummyHostKVCacheForAlloc.__new__(DummyHostKVCacheForAlloc)
+        host_pool.page_size = 4
+        host_pool.lock = __import__("threading").RLock()
+        host_pool.free_slots = torch.tensor(
+            [100, 101, 102, 103, 8, 9, 10, 11, 12, 13, 14, 15],
+            dtype=torch.int64,
+        )
+
+        selected = host_pool.alloc_contiguous_preferred(8)
+
+        self.assertEqual(selected.tolist(), [8, 9, 10, 11, 12, 13, 14, 15])
+        self.assertEqual(host_pool.free_slots.tolist(), [100, 101, 102, 103])
+
    def test_cp_reserve_zero_owned_queues_no_ack_until_submit(self):
        host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
        controller = self.make_controller(host_pool, cp_rank=3)