[HiCache] Keep auxiliary load-back out of Full KV pending ownership (#36317)

Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>
This commit is contained in:
weireweire
2026-08-27 01:30:38 +08:00
committed by GitHub
co-authored by weireweire
parent e27a7fac77
commit 5263568bcb
2 changed files with 38 additions and 13 deletions
@@ -331,6 +331,34 @@ class TestUnifiedTreeCoreLoadBackPending(CustomTestCase):
self.assertTrue(UnifiedTreeCore._can_reclaim_full_host_duplicate(core, shared))
core._update_duplicate_tracking.assert_called_once_with(shared)
def test_auxiliary_load_does_not_reuse_full_pending_pin(self):
core, shared, anchor_a, anchor_b = self._build_core(is_write_back=True)
core.components_by_type[ComponentType.SWA] = mock.Mock()
self._commit_load_back(core, anchor_a, shared)
self.assertEqual(shared.load_back_pending_id, anchor_a.id)
kv_transfer = PoolTransfer(
name=PoolName.KV,
host_indices=torch.tensor([1], dtype=torch.int64),
nodes_to_load=[anchor_b.id],
)
swa_transfer = PoolTransfer(
name=PoolName.SWA,
host_indices=torch.tensor([2], dtype=torch.int64),
nodes_to_load=[shared.id],
)
UnifiedTreeCore.commit_load_back(
core,
anchor_b.id,
torch.tensor([3], dtype=torch.int64),
kv_transfer,
{ComponentType.SWA: [swa_transfer]},
)
self.assertEqual(shared.load_back_pending_id, anchor_a.id)
self.assertEqual(anchor_b.load_back_pending_id, anchor_b.id)
def _write_backup(cache, node, write_back: bool = False) -> int:
"""Back up one node's KV D->H via the tree's build+execute primitives."""