Retain SWA down to the last state checkpoint

[ispobock]

Motivation

A hybrid SWA + mamba model throws away most of its decode-region prefix reuse at the default --mamba-track-interval. Measured on one GPU with 32 prompts, counting how many get a decode-region cache hit on their second turn:

page_size  track_interval    prompts reusing the decode region
   128          256              16/32  ->  32/32
   128          512               5/32  ->  32/32
    64          256               9/32  ->  32/32

Counting tokens rather than prompts, on the first row, over the same 32 second-turn requests:

reused prefix           110208  ->  124416 tokens   (+12.9%)
recomputed               18435  ->    4227 tokens   (-77.1%)
mean reuse fraction      0.845  ->   0.965

The gain is not deeper matches on requests that already hit, it is requests that were losing the whole generated region getting it back: 16 of 32 improved by 768 to 896 tokens each, none regressed, and the ones that already hit were at 98.5% reuse to begin with.

The cause is that SWA frees out-of-window slots relative to the tail, while a prefix match lands on a state checkpoint, which sits behind the tail. evict_threshold = pre_len - max(window, page) therefore strands checkpoints that are still reachable in the tree: the match needs a full window of live SWA below the checkpoint, and part of that window has already been freed. swa_evicted_seqlen only moves forward, so nothing gets it back.

Today the only way to get full reuse is to set the interval equal to the page size, which doubles checkpoint density and the mamba pool footprint. This change makes the two independent: keep the interval at 256 for the cheaper checkpoint storage and still reuse every decode-region prefix.

Modifications

free_swa_out_of_window_slots takes an optional retain_floor and promises not to free past it. It stays ignorant of mamba; the caller decides where the floor is.

The floor is computed in one place, on the cache that knows about both components, and both call sites (decode-side eviction and the chunked-prefill insert path) ask it rather than each deriving it. BasePrefixCache.swa_retain_floor returns None, so a cache without a second state stream is unaffected.

Accuracy

test_unified_radix_cache_kl_hybrid_bitexact asserts prefill and decode score every token identically at a kl_div floor of 1e-9, which makes it the right instrument here: the change increases how much state gets reused, so a wrong retention floor shows up as a nonzero KL on exactly the prompts that newly reuse.

Same tree, same diff, the only variable being whether the floor is applied. hits is the number of prompts whose second turn hit the decode region, nonzero counts per-prompt KL above zero:

                     floor off              floor on
SM90 (H200)      16/32 hits, 0 nonzero   32/32 hits, 0 nonzero
SM100 (B200)     17/32 hits, 0 nonzero   32/32 hits, 0 nonzero

Every prompt that newly reuses reads exactly 0, so the retained window is correct rather than merely present. The grid above repeats this at page_size 64, where the page size no longer coincides with the mamba chunk size, and reuse is still complete with the floor on.

On memory: peak device usage is unchanged (134.2 GB against 130.3 GB, and the direction is noise), but that metric does not answer the question, because the pools are preallocated from mem-fraction-static and holding more slots inside them does not move the device peak. What the change actually costs is SWA pool occupancy, which I did not measure. The analytic bound is the checkpoint spacing, and the probe that located this measured the extra retention at 127 to 255 tokens per request against a 511-token window, so on the config above it is under 1% of the SWA pool. Worth a reviewer's judgement rather than my assertion.

A cache with no mamba component is unaffected, checked at the seam rather than by sampling: BasePrefixCache.swa_retain_floor returns None, the unified override early-returns None when mamba is off, retain_floor is read in exactly one guarded branch, the original threshold line is untouched, and both call sites go through swa_retain_floor. With retain_floor=None the helper is byte-identical to before.

TODO

• The prefill region is not covered yet. req.mamba_last_track_seqlen is cleared at the end of cache_unfinished_req once the checkpoint has been handed to the tree, so at the first decode step the floor reads None and no extra retention happens. Covering it needs the floor to come from the tree rather than from the request.
• Move the SWA pool sizing floor from window to window + interval, so the extra retention cannot bite under pool pressure.

---

CI States

Latest PR Test (Base): 🚫 Run #31718123937
Latest PR Test (Extra): ✅ Run #31718123679

[ispobock]

/tag-and-rerun-ci

[ispobock]

/rerun-test test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_swa.py test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py

[github-actions]

Results for /rerun-test test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_swa.py test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py:

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py

🚀 2-gpu-h100 (2 tests): ❌ View workflow run

cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_swa.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py

🚀 4-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py

[ispobock]

/rerun-test test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_swa.py test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py test/registered/unit/mem_cache/test_swa_eviction_boundary.py

[github-actions]

Results for /rerun-test test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_swa.py test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py test/registered/unit/mem_cache/test_swa_eviction_boundary.py:

🚀 1-gpu-h100 (2 tests): ✅ View workflow run

cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py
cd test/ && python3 registered/unit/mem_cache/test_swa_eviction_boundary.py

🚀 2-gpu-h100 (2 tests): ✅ View workflow run

cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_swa.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py

🚀 4-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py

[ispobock]

/tag-and-rerun-ci