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

Overlap CP HiCache backup without exposing partial host state

Browse Source 
CP shared KV with HiCache and EAGLE needs host backup to overlap forward while keeping radix visibility synchronous. The change reserves host slots before forward, drives target and draft backup from explicit layer-end hooks, and commits host visibility only after the final target/draft ack. It also probes the final insertion prefix before early reservation so repeated EAGLE prompts do not prepare duplicate suffix backups that later rollback as insert_miss.

Constraint: CP ranks use independent shared-KV pools, so target/draft host state must remain atomically visible at the radix boundary.

Constraint: Fused MLA and NSA store paths can bypass store-side notifier hooks, so layer end is the safer backup progress boundary.

Rejected: Store-side backup notifier as the primary trigger | fused store and zero-local paths made notifier coverage fragile.

Rejected: Reserve from cache_protected_len alone | EAGLE bigram/page alignment can make final insertion find a longer existing prefix and force duplicate rollback work.

Confidence: medium

Scope-risk: moderate

Directive: Do not add per-layer CP collectives here; keep radix state synchronous and data transfer asynchronous/local-event driven.

Tested: local git diff --check

Tested: local py_compile for touched CP HiCache/cache-controller/deepseek/test files

Tested: remote pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/managers/test_hicache_controller_cp.py -q (115 passed, 5 warnings)

Not-tested: full GLM5 ETE server rerun after this commit
This commit is contained in:
laoyao0822
2026-05-27 09:50:47 +08:00
parent 03529319a1
commit f355fdd39e
11 changed files with 747 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
test/registered/unit/managers/test_hicache_controller_cp.py
View File

@@ -175,6 +175,10 @@ class FakeDevicePool:
def register_layer_backup_notifier(self, notifier):
self.layer_backup_notifiers.append(notifier)
def notify_layer_end_for_backup(self, layer_id):
for notifier in self.layer_backup_notifiers:
notifier(layer_id)
class TestPageFirstPerLayerBackupTaiKernel(CustomTestCase):
def test_mla_page_first_per_layer_backup_uses_tai_lf_pf_kernel(self):
@@ -677,12 +681,12 @@ class TestHiCacheControllerCPWrite(CustomTestCase):
self.assertEqual(host_pool.layer_backups, [])
self.assertEqual(controller.ack_write_queue, [])
controller.on_layer_kv_stored(0)
allocator.device_pool.notify_layer_end_for_backup(0)
self.assertEqual(host_pool.layer_backups[0][2], 0)
self.assertEqual(controller.ack_write_queue, [])
controller.on_layer_kv_stored(1)
allocator.device_pool.notify_layer_end_for_backup(1)
self.assertEqual([x[2] for x in host_pool.layer_backups], [0, 1])
self.assertEqual(len(controller.ack_write_queue), 1)
@@ -708,13 +712,13 @@ class TestHiCacheControllerCPWrite(CustomTestCase):
reservation = controller.reserve_write_cp(logical_locs, node_id=85)
controller.submit_write_cp_per_layer(reservation, catch_up_all_layers=False)
controller.on_layer_kv_stored(0, source="target")
allocator.device_pool.notify_layer_end_for_backup(0)
self.assertEqual([x[2] for x in host_pool.layer_backups], [0])
self.assertEqual(draft_host_pool.layer_backups, [])
self.assertEqual(controller.ack_write_queue, [])
controller.on_layer_kv_stored(0, source="draft")
draft_device_pool.notify_layer_end_for_backup(0)
self.assertEqual([x[2] for x in draft_host_pool.layer_backups], [0])
self.assertEqual(len(controller.ack_write_queue), 1)

331
test/registered/unit/mem_cache/test_cp_hicache_metadata.py
View File

@@ -91,14 +91,20 @@ for _schema in (
raise
from sglang.srt.managers.cache_controller import (
HiCacheAck,
HiCacheWriteFailure,
HiCacheWriteReservation,
)
from sglang.srt.mem_cache.base_prefix_cache import EvictParams, MatchPrefixParams
from sglang.srt.mem_cache.base_prefix_cache import (
EvictParams,
InsertParams,
MatchPrefixParams,
)
from sglang.srt.mem_cache.hiradix_cache import (
CpHiCacheNodeMetadata,
HiRadixCache,
PendingHiCacheBackup,
PreparedCpHiCacheBackup,
_compute_shared_hicache_token_capacities,
)
from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode
@@ -544,18 +550,25 @@ class FakeReserveWriteController:
self.results = list(results)
self.reservations = []
self.submitted = []
self.submit_kwargs = []
self.evicted_host_indices = []
def reserve_write_cp(self, device_indices, priority=None, node_id=-1):
self.reservations.append((device_indices.clone(), node_id))
result = self.results.pop(0)
return result(device_indices) if callable(result) else result
if not callable(result):
return result
try:
return result(device_indices, node_id=node_id)
except TypeError:
return result(device_indices)
def submit_write_cp_all_layer(self, reservation):
self.submitted.append(reservation)
def submit_write_cp_per_layer(self, reservation):
def submit_write_cp_per_layer(self, reservation, **kwargs):
self.submitted.append(reservation)
self.submit_kwargs.append(kwargs)
def evict_cp_host(self, metadata):
self.evicted_host_indices.append(metadata.host_indices.clone())
@@ -883,6 +896,221 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
self.assertIn(node.id, cache.pending_host_backups)
self.assertEqual(len(cache.cache_controller.submitted), 1)
def test_write_backup_rolls_back_local_success_when_peer_needs_host_eviction(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.tp_world_size = 2
cache.tp_group = object()
cache.cache_controller = FakeReserveWriteController(
[
lambda device_indices: make_write_reservation(
device_indices, node_id=130, host_start=90
),
lambda device_indices: make_write_reservation(
device_indices, node_id=130, host_start=100
),
]
)
evictions = []
cache._evict_host_for_physical_slots = lambda required, synchronize_across_ranks=False: (
evictions.append((required, synchronize_across_ranks)) or required
)
cache.ongoing_write_through = {}
cache.pending_host_backups = {}
cache.inc_node_lock_ref = lambda node: None
reduce_values = iter([4, 0])
def fake_all_reduce(tensor, op=None, group=None):
tensor.fill_(next(reduce_values))
node = TreeNode()
node.id = 130
node.value = torch.arange(16, dtype=torch.int64)
with patch("torch.distributed.all_reduce", side_effect=fake_all_reduce):
backed_len = cache.write_backup(node)
self.assertEqual(backed_len, 16)
# The first local success must be released before all ranks enter the
# collective host eviction/retry branch. Otherwise CP ranks can diverge:
# successful ranks submit backup while failing ranks enter eviction
# all_reduce, which matches the observed Gloo 4-vs-1 mismatch.
self.assertEqual(
cache.cache_controller.evicted_host_indices[0].tolist(),
list(range(90, 106)),
)
self.assertEqual(evictions, [(4, True)])
self.assertEqual(len(cache.cache_controller.submitted), 1)
self.assertEqual(
cache.cache_controller.submitted[0].metadata.host_indices.tolist(),
list(range(100, 116)),
)
def test_insert_attaches_prepared_cp_backup_without_catchup_copy(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.root_node = TreeNode()
cache.root_node.key = RadixKey([])
cache.evictable_size_ = 0
cache.get_child_key_fn = lambda key: key.token_ids[0]
cache.key_match_fn = lambda lhs, rhs: 0
cache.maybe_bigram_convert = lambda key, value: (key, value)
cache.is_eagle = False
cache.enable_storage = False
cache.enable_kv_cache_events = False
cache._update_leaf_status = lambda node: None
cache._update_host_leaf_status = lambda node: None
cache._record_store_event = lambda node: None
cache.ongoing_write_through = {}
cache.pending_host_backups = {}
locked_nodes = []
cache.inc_node_lock_ref = lambda node: locked_nodes.append(node)
cache.cache_controller = types.SimpleNamespace(write_policy="write_through")
value = torch.arange(16, dtype=torch.int64)
reservation = make_write_reservation(value, node_id=131, host_start=110)
prepared = PreparedCpHiCacheBackup(
node_id=131,
reservation=reservation,
metadata=reservation.metadata,
logical_len=16,
)
result = cache.insert(
InsertParams(
key=RadixKey(list(range(16))),
value=value,
cp_hicache_prepared_backup=prepared,
)
)
node = cache.root_node.children[0]
self.assertEqual(result.prefix_len, 0)
self.assertEqual(node.id, 131)
self.assertTrue(prepared.attached)
self.assertIs(cache.ongoing_write_through[131], node)
self.assertIs(cache.pending_host_backups[131].node, node)
self.assertIs(cache.pending_host_backups[131].metadata, reservation.metadata)
self.assertEqual(cache.pending_host_backups[131].logical_len, 16)
self.assertEqual(locked_nodes, [node])
def test_prepare_write_backup_for_req_registers_before_forward_without_catchup(
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(1, 16)
)
cache.cache_controller = FakeReserveWriteController(
[
lambda device_indices, node_id: make_write_reservation(
device_indices, node_id=node_id, host_start=130
)
]
)
req = types.SimpleNamespace(
rid="rid-prepare",
fill_ids=list(range(16)),
cache_protected_len=0,
req_pool_idx=0,
is_chunked=0,
cp_hicache_prepared_backup=None,
)
cache.prepare_write_backup_for_req(req)
prepared = req.cp_hicache_prepared_backup
self.assertIsNotNone(prepared)
self.assertEqual(prepared.logical_len, 16)
self.assertEqual(len(cache.cache_controller.submitted), 1)
self.assertIs(cache.cache_controller.submitted[0], prepared.reservation)
self.assertEqual(
cache.cache_controller.submit_kwargs,
[{"catch_up_all_layers": False}],
)
def test_prepare_write_backup_for_req_skips_existing_insert_prefix(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache.disable = False
cache._uses_cp_hicache = True
cache.is_eagle = False
cache.page_size = 1
cache.root_node = TreeNode()
cache.root_node.key = RadixKey([])
cache.root_node.children = {}
cache.get_child_key_fn = lambda key: key.token_ids[0]
def key_match(lhs, rhs):
matched = 0
for left, right in zip(lhs.token_ids, rhs.token_ids):
if left != right:
break
matched += 1
return matched
cache.key_match_fn = key_match
cache.maybe_bigram_convert = lambda key, value=None: (key, value)
cache.req_to_token_pool = types.SimpleNamespace(
req_to_token=torch.arange(16, dtype=torch.int64).view(1, 16)
)
cache.cache_controller = FakeReserveWriteController([])
existing = TreeNode()
existing.id = 201
existing.parent = cache.root_node
existing.key = RadixKey(list(range(16)))
existing.value = torch.arange(16, dtype=torch.int64)
cache.root_node.children[0] = existing
req = types.SimpleNamespace(
rid="rid-existing-prefix",
fill_ids=list(range(16)),
extra_key=None,
cache_protected_len=8,
req_pool_idx=0,
is_chunked=0,
cp_hicache_prepared_backup=None,
)
cache.prepare_write_backup_for_req(req)
self.assertIsNone(req.cp_hicache_prepared_backup)
self.assertEqual(cache.cache_controller.reservations, [])
self.assertEqual(cache.cache_controller.submitted, [])
def test_rollback_unattached_prepared_cp_backup_removes_orphan_ack(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
value = torch.arange(4, dtype=torch.int64)
reservation = make_write_reservation(value, node_id=132, host_start=120)
prepared = PreparedCpHiCacheBackup(
node_id=132,
reservation=reservation,
metadata=reservation.metadata,
logical_len=4,
)
evicted = []
class ReadyEvent:
def synchronize(self):
pass
cache.cache_controller = types.SimpleNamespace(
pending_layer_writes={},
ack_write_queue=[HiCacheAck(ReadyEvent(), ReadyEvent(), [132])],
evict_cp_host=lambda metadata: evicted.append(metadata),
)
cache._rollback_prepared_cp_backup(prepared, "test")
self.assertEqual(cache.cache_controller.ack_write_queue, [])
self.assertEqual(evicted, [reservation.metadata])
def test_write_backup_cp_retry_failure_leaves_node_device_only(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
@@ -1641,5 +1869,102 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
self.assertEqual(cache.ongoing_load_back, {})
class TestCPHiCacheLayerBackupNotifications(CustomTestCase):
def _call_store_index_fast_path(self, out_loc):
from sglang.srt.layers.attention.nsa import nsa_indexer
notifications = []
fused_calls = []
class FakePool:
page_size = 64
start_layer = 0
def get_index_k_with_scale_buffer(self, layer_id):
return torch.empty((1,), dtype=torch.uint8)
def notify_layer_kv_stored_for_backup(self, layer_id, source="kv"):
notifications.append((layer_id, source))
forward_batch = types.SimpleNamespace(token_to_kv_pool=FakePool())
indexer = types.SimpleNamespace()
def fake_fused_store(key, buf, loc, page_size):
fused_calls.append((key, buf, loc.clone(), page_size))
with (
patch.object(nsa_indexer, "_is_cuda", True),
patch.object(nsa_indexer, "_is_fp8_fnuz", False),
patch.object(nsa_indexer, "can_use_nsa_fused_store", return_value=True),
patch.object(
nsa_indexer, "fused_store_index_k_cache", side_effect=fake_fused_store
),
):
nsa_indexer.Indexer._store_index_k_cache(
indexer,
forward_batch,
layer_id=5,
key=torch.empty((max(out_loc.numel(), 1), 128), dtype=torch.float32),
out_loc_override=out_loc,
)
return notifications, fused_calls
def test_nsa_indexer_fused_store_does_not_notify_cp_hicache_layer_backup(self):
notifications, fused_calls = self._call_store_index_fast_path(
torch.tensor([1, 2, 3], dtype=torch.int64)
)
self.assertEqual(notifications, [])
self.assertEqual(len(fused_calls), 1)
self.assertEqual(fused_calls[0][2].tolist(), [1, 2, 3])
def test_nsa_indexer_empty_store_does_not_notify_cp_hicache_layer_backup(self):
notifications, fused_calls = self._call_store_index_fast_path(
torch.empty((0,), dtype=torch.int64)
)
self.assertEqual(notifications, [])
self.assertEqual(fused_calls, [])
def test_cp_shared_zero_local_index_store_does_not_notify_layer_backup(self):
from sglang.srt.layers.attention.nsa import nsa_indexer
notifications = []
class FakePool:
page_size = 64
def notify_layer_kv_stored_for_backup(self, layer_id, source="kv"):
notifications.append((layer_id, source))
forward_batch = types.SimpleNamespace(token_to_kv_pool=FakePool())
indexer = types.SimpleNamespace(nsa_enable_prefill_cp=True)
with (
patch.object(nsa_indexer, "nsa_use_prefill_cp", return_value=True),
patch.object(
nsa_indexer,
"get_cp_shared_kv_local_out_cache_loc",
return_value=torch.empty((0,), dtype=torch.int64),
),
patch.object(
nsa_indexer,
"get_cp_shared_kv_local_physical_out_cache_loc",
side_effect=AssertionError("must not require physical locs"),
),
):
handled = nsa_indexer.Indexer._store_cp_shared_local_index_k_cache(
indexer,
forward_batch=forward_batch,
layer_id=6,
local_key=torch.empty((0, 128), dtype=torch.float32),
act_quant=None,
)
self.assertTrue(handled)
self.assertEqual(notifications, [])
if __name__ == "__main__":
unittest.main()