[HiCache] Preserve restored SWA lock ownership - #36350
paulzhang-tm wants to merge 5 commits into
Conversation
|
/rerun-group radix_cache/unified_radix_tree hicache |
|
Results for 🚀 🚀 🚀 🚀 🚀 🚀 ⛔ |
b57b819 to
3cd2f33
Compare
Replay the acquiring request’s SWA skip set during early release so a restored tombstone’s lock remains owned by the request that acquired it. Treat device and host locks as unconditional ownership barriers during cascade eviction, including on lower-priority Mamba state, and cover both restoration races with regression tests.
3cd2f33 to
0cf22e2
Compare
There was a problem hiding this comment.
The release_window_lock skip-set replay fix looks good for me.
For the cascade change, I remember the assert is necessary to keep cc @hzh0425
The cascade takes Mamba out with SWA unless it's locked, and a request locks SWA and Mamba together at best_match_node, so how does a node end up Mamba-locked with its SWA device-evicted? The new test hand-builds that state in _build_core() instead of reaching it through match to evict. @paulzhang-tm Could you provide a concrete sequence to trigger the assert?
|
ref: #36848 is refactoring the |
Replace the hand-built TestUnifiedTreeCoreRestoredSWALocks fixture with three UnifiedRadixCacheSuite tests that reach the same state through the production path in the unified_kv + HiCache layout, where SWA has no host pool: a Full-evicted leaf is loaded back as Full+Mamba only, match_prefix selects it, the request lock takes Mamba and skips SWA, an overlapping insert restores SWA under the held Full lock, and SWA eviction cascades into the still-locked Mamba. On the merge-base the cascade test trips `assert cd.lock_ref == 0` in _should_cascade_evict_component and the early-release test observes the second request's SWA lock being released; both pass with the fix. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Fold the three test_ring_mode_* tests and their helper into a single linear test, test_swa_evict_cascade_spares_mamba_locked_while_swa_was_missing, that walks the unified_kv + HiCache sequence step by step: write-through, Full demote, load-back of Full+Mamba without SWA, match and request lock that skips the missing SWA, overlapping insert restoring SWA under the held Full lock, and the SWA eviction whose cascade meets the Mamba lock. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Reach the Mamba-locked, SWA-missing node with the real HiCache controller instead of tree-core calls: init_hicache, write-through on insert, evict, match_prefix, load_back + loading_check, inc_lock_ref. The only setup step that is not a cache call declares the layout under test (an SWA component with no host pool, as the unified_kv assembler leaves it), since no stack strategy builds HiCache with Mamba and without an SWA host pool today. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
|
@ispobock check test_swa_evict_cascade_spares_mamba_locked_while_swa_was_missing |
|
The test drives production APIs but on a pool config the assembler never builds: |
Review asked for a concrete trigger for the relaxed cascade check. The
Mamba-locked, SWA-missing node needs HiCache with Mamba present and no SWA
host pool, and no server-args configuration builds that layout:
{FULL, SWA, MAMBA} only matches _MambaSwaStrategy, which always allocates an
SWA host pool, so load-back restores SWA before the request lock is taken;
the only SWA-no-host-pool layout is DeepSeek V4 unified_kv, which has no
Mamba. The assertion is therefore a correct defensive invariant.
Restore _should_cascade_evict_component and its Rust counterpart to their
original form, put the original cascade tests back, and drop the cascade
regression test. The skip-set replay in release_window_lock and its tests
stay.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
|
#36848 is already merged, |
|
The issue seems to be gone on the main branch. I may close this PR for now. Feel free to reopen it if the issue persists. |
3 participants
Motivation
HiCache stores Full KV, SWA KV, and Mamba state as separate cache components that may be evicted and restored independently. A request whose lock skipped a tombstoned SWA node could later early-release an SWA lock that another request acquired after the node was restored:
release_window_lockwalked the window without consulting the acquiring request's skip set, so it decremented a lock it never owned.Scope after review
An earlier revision also relaxed the cascade-eviction assertion (
assert cd.lock_ref == 0on a strictly lower-priority component) so that SWA eviction would spare a locked Mamba value. Review asked for a concrete trigger. The state needs HiCache with Mamba present and no SWA host pool: only then does the SWA match validator accept tombstones, only then does a load-back restore Full and Mamba without SWA, and only then can the request lock take Mamba while skipping SWA. No server-args configuration yields that layout.{FULL, SWA, MAMBA}matches only_MambaSwaStrategy, which always builds an SWA host pool, so load-back restores SWA before the request lock is taken. The only layout without an SWA host pool is DeepSeek V4 unified_kv ({FULL, SWA}, plus a C128 sidecar on NPU with priority 2 and no locks), which has no Mamba. Node splits zero the new parent's Mamba lock, andswa_retain_floorkeeps the Mamba checkpoint node inside the retained SWA window, so no other path produces the state either. The assertion is therefore a correct defensive invariant, and this revision restores it unchanged in both the Python and Rust tree cores.Modifications
dec_swa_lock_only->release_window_lock), in the Python and Rust tree cores.TestDecSwaLockSkipfake intest_mamba_donated_alloc_ratio.pyand the Rustcomponents/swa.rstests to cover the skip-set plumbing.Accuracy Tests
Not applicable; this changes lock bookkeeping without changing model computation.
Focused unit tests (
test_unified_radix_cache_unittest.py -k "FULL_SWA_MAMBA or TestUnifiedTreeCore"plustest_mamba_donated_alloc_ratio.py):Speed Tests and Profiling
Not applicable; the change adds a constant-time skip check to an existing release path.
Checklist
-Robot
🤖 Generated with Claude Code
CI States
Latest PR Test (Base): ❌ Run #33767805627
Latest PR Test (Extra): ❌ Run #33767805402
Latest PR Test (AMD ROCm 7.2): ❌ Run #33767805697