[Fix] Snapshot SWA mappings before deferred frees

[LingZ315]

Motivation

Fixes #36081.

SGLang v0.5.18 can report a one-page SWA pool leak after grouped cache frees. SWATokenToKVPoolAllocator.free_swa() previously deferred Full-pool indices and resolved their Full-to-SWA mapping only in free_group_end(). Tombstone recovery can replace that mapping in the same free group, so the flush frees the replacement SWA page while the old physical page becomes unreachable.

I minimized this on an RTX PRO 5000 using the exact v0.5.18 Docker source revision. The unpatched allocator ended with 2 physical SWA pages in use but only 1 mapped page, reproducing the one-page discrepancy from the issue. The patched allocator preserves the invariant.

Modifications

• Keep Full indices batched while a free group is open.
• Before any Full-to-SWA mapping write (or at free_group_end()), snapshot and detach the queued ownership so a tombstone recovery cannot replace it.
• Defer the captured physical SWA indices until free_group_end(), preserving one batched lookup/clear on the steady-state path.
• Extend the unified radix-cache recovery test to cover grouped frees for both token and paged allocators.
• Assert that mapped SWA pages equal physically allocated SWA pages after recovery.

Accuracy Tests

This change does not affect model outputs.

Validated on an NVIDIA RTX PRO 5000:

• Exact v0.5.18 minimized reproduction, before fix: fails with 2 physical pages used vs. 1 mapped.
• Exact v0.5.18 minimized reproduction, after fix: passes.
• Targeted recovery test on current main: 2 passed, 17 skipped.
• test_swa_unittest.py: 13 passed.
• Full+SWA unified-cache configurations: 904 tests passed, 359 skipped.
• Black 26.1.0 and Ruff 0.15.1 checks passed.

The host does not provide /usr/local/cuda/bin/nvcc, so optional HiCache JIT kernels used their existing fallback in the broader unified-cache run.

Speed Tests and Profiling

Allocator-bookkeeping microbenchmark on an NVIDIA RTX PRO 5000 72GB Blackwell (PyTorch 2.11.0+cu130), with 50 warmups and 500 interleaved A/B samples. It mirrors grouped free_swa() bookkeeping and excludes the identical physical allocator release. These steady-state cases contain no mapping replacement, so the baseline is correct and directly comparable. Values are medians per free group; wall time includes Python launch and synchronization overhead.

page size	calls/group	pages/call	GPU before/after (us)	GPU delta	wall before/after (us)	wall delta
256	1	1	117.3 / 117.4	+0.1%	122.6 / 122.7	+0.1%
256	8	1	149.5 / 149.9	+0.3%	154.7 / 155.2	+0.3%
256	8	16	215.6 / 216.3	+0.3%	221.1 / 221.7	+0.2%
256	64	1	458.8 / 459.6	+0.2%	464.2 / 464.9	+0.2%

The page-size-1 cases (1, 8, and 64 calls/group) had wall-time deltas of +0.5%, -0.1%, and +0.4%. Across all seven cases, the final implementation stayed within -0.1% to +0.5%, i.e. measurement noise. It achieves this by inserting a snapshot barrier only before a mapping update or group end instead of performing a gather/clear on every free_swa() call.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations. (N/A: internal allocator correctness fix.)
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ⏳ Run #32713782397
Latest PR Test (Extra): ⏳ Run #32713782164
Latest PR Test (AMD ROCm 7.2): ⏳ Run #32713782533

[haydn-jones]

Anecdotal evidence: seems to have corrected the issue on my end. DSV4 Flash, 4xB200, DP=TP=EP=4 with dp attention.