fix(unified-cache): allow multiple concurrent load-back pins per node

[ffahleraz]

Motivation

Under heavy device-eviction↔load-back cycling (small-VRAM GPUs with --enable-hierarchical-cache, hybrid GDN/mamba models), the scheduler crashes with:

AssertionError: node 2498 pinned by load-back 2563, new anchor 2498

at UnifiedTreeCore.commit_load_back (unified_tree_core.py). Hit repeatedly serving nvidia/Qwen3.6-35B-A3B-NVFP4 + EAGLE + hicache ratio-12 on RTX 5090 (32GB) under a production-trace replay; 96GB GPUs never trip it (no deep evicted chains → no overlap).

Root cause

UnifiedTreeNode.load_back_pending_id is single-valued, but two live load-backs can legitimately pin the same node:

1. Request A load-backs anchored at a descendant; its Full-KV chain covers the evicted ancestor and pins it (load_back_pending_id = A).
2. Before A's ack, request B anchors at that ancestor to restore its independently-evicted mamba state (mamba_component.build_hicache_transfers pins [node.id]).
3. The single pin slot can't represent both live pins → assertion (new anchor == the node itself, matching the production message shape exactly).

The overlap is semantically safe: both in-flight DMAs only read the shared host slots, and their writes target disjoint destinations (Full-KV pool vs mamba pool). Only the bookkeeping cannot represent it.

Modifications

• load_back_pending_id: Optional[int] → load_back_pending_ids: set[int].
• commit_load_back adds the anchor id (idempotent for a node sitting in both the Full and an aux transfer list of one anchor).
• finish_load_back discards its anchor id along the root path; duplicate tracking updates once the last pin drains.
• Reclaim guards (_is_settled_full_host_duplicate, _can_reclaim_full_host_duplicate) and the sanity check test set-emptiness; split inheritance copies the set.

Test / repro

test/registered/unit/mem_cache/test_unified_loadback_multipin.py — a deterministic CPU repro built on the existing UnifiedRadixCacheSuite fixtures (FULL+MAMBA): backup → full evict → load-back anchored at the leaf (ack withheld) → load-back anchored at the pinned ancestor.

• Unpatched: fails in 0.3s with the exact production assertion.
• Patched: passes; pins drain to empty after both acks; sanity_check clean.
• Full inherited FULL+MAMBA suite under the new class: 107 tests green.
• Serving-level: the patched server sustains 400-span production-trace replays at hicache ratio-12 on RTX 5090 with zero assertions (previously crashed mid-run).

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings as needed.

---

CI States

Latest PR Test (Base): ❌ Run #31902335393
Latest PR Test (Extra): ❌ Run #31902335225

[200lz]

This looks like the right fix for concurrent ownership. One lifecycle edge case that may be worth making explicit: should a pin identify only the anchor ID, or the specific load-back generation/component as well? If an anchor/node ID can ever be reused while a delayed ack is still in flight, a generation-tagged pin could avoid an ABA-style release of a newer load-back.

[ffahleraz]

Good question — I checked this specifically. Anchor pins are keyed by node id, and UnifiedTreeNode.counter is a process-lifetime monotonic class attribute (unified_tree_core.py: self.id = UnifiedTreeNode.counter; UnifiedTreeNode.counter += 1) that is never rewound — UnifiedTreeCore.reset() rebuilds the node arena/root/LRUs but does not touch the counter. So a node id (and therefore an anchor id) can never be reused within a process, including across cache resets with delayed acks in flight — the ABA release can't occur. If id recycling were ever introduced, a generation tag would indeed be needed, but today monotonicity makes the anchor id itself the generation.

[xiezhq-hermann]

can you confirm this is a fix for write back policy right?

[xiezhq-hermann]

@ziang663 can you help take a look?

[ziang663]

I think the fix is correct. Retrieving the set allows the same node to be pinned by multiple requests, while there's an assumption that the components must be different, which I believe can be guaranteed.

[hzh0425]

@ffahleraz pls fix the lint

[weireweire]

Does swa / mamba really need load_back_pending_id ? It should be designed to protect full kv, because the full host duplicate reclaim can evict full-kv under transfer, as the full kv lock only effect anchor node.

But for swa / mamba, 1. they won't be reclaim by dup, 2. their lock will protect full window/state, so the real bug maybe just that they should not participate in the full kv's load_back_pending_id .

Wonder if this simple change can fix the issue you see: e16de90

[ffahleraz]

@weireweire you are right, thanks for #36317. Looking at main again, commit_load_back publishes the Full device value before the ack, so the second anchor's Full chain stops above the pinned ancestor and the only thing that pinned it under a different anchor was the mamba transfer, which #36317 now keeps out of load_back_pending_id. There is no case left for the multi-pin set, so I am dropping it.

@hzh0425 @ziang663 sorry for going quiet on the lint. The same workload had a second crash, Can not alloc mamba cache in prepare_for_caching_req when every cached mamba state is pinned, which is still on main. I put that fix in #36770 together with an end-to-end regression for the scenario from this PR, so closing this one as superseded by #36317. Would appreciate a look at #36770 when you have a moment.