[caihuali95]

This PR is part of a stack (oldest at bottom):

• [Perf] attention: Capture translated FA4 verify metadata for DSPARK #44973
• [Perf] speculative: Fuse draft write-window gathering and KV translation #44972
• [Perf] mamba: Fuse strided state updates and replay index translation #44971
• [Perf] unified-memory: Reduce allocation and SWA retirement overhead #44970
• [Refactor] attention: Separate decode metadata from pool ownership #44969
• [Feat] unified-memory: Fuse drafts with sliding-window and recurrent-state layers #43326
• [Feat] unified-memory: Support asymmetric KV rows #43325
• [Perf] gdn: Select FlashInfer prefill for page-major pools #43171
• -> [Perf] unified-memory: Free SWA segments without a host sync #43170
• [Fix] swa: Retain a whole window behind insertion pages #43168

Current head 2509d740d3, based on main 47eeddcd18. Stack predecessor: #43168. This PR targets main after its stack predecessor merged.

Motivation

UnifiedSWAAllocatorBase (--enable-unified-memory with a hybrid sliding-window model) filtered every segment free by swa liveness with a boolean mask (reps[swa_v2p > 0]). A mask index runs nonzero, which copies its count to the host and blocks the scheduler until the in-flight forward drains, so under the overlap scheduler the next batch is prepared only after the GPU went idle. In gpt-oss-20b serving traces (SM90, page size 256) the free of tokens that left the window (free_swa_segment) and the request release (_release_page_reps at free_group_end) issued about 2.6 such syncs per prefill. Prefill cycles ran +1.3..+3.3% (speculative) and +2.5..+7.0% (non-speculative) longer than with the paged SWA allocator at matched batch size, all of it GPU idle.

Modifications

Two-pool allocation reclaim flushes the deferred FULL-only segment queue before sizing eviction. A regression test fills that queue alone and verifies allocation succeeds without evicting the tree.

The paged SWA allocator frees the same segments with a fixed shape: its callers split a kv row by host bookkeeping, and it trusts that split. The unified callers make the same split: free_kv_row_segments sends [swa_dead_lo, swa_evicted_seqlen) through free_full_segment, the tree picks FreeDeviceKV or FreeDeviceKVFullOnly per node, and the out-of-window free releases only tokens above swa_dead_lo, with a cursor that only moves up. So free_swa_segment and free_segment now free every page they are handed on the swa side, and free_full_segment releases the full side only, through its own deferred list for the free group (full_free_group), which flush_deferred_full_frees also drains when Mamba state needs memory. Each checks its premise with expect, as the paged allocator's swa.peer_mapped does: asynchronous, and loud from SGLANG_INVARIANT_CHECK=warn up. The two invariants are declared in allocator/swa.py beside the paged allocator's: the registry refuses a name it already holds, and a test fixture that patches sys.modules makes the next test import unified_hybrid_swa again. Page size 1 takes the same path. The set-shaped free() and free_swa() keep their filter; RecoverSWAWithLockedFull relies on it.

A second out-of-window free over the same range was filtered to a no-op. No caller issues one; under strict checks it is now reported instead.

Tests (CPU):

• test_segment_frees_dispatch_no_data_dependent_shape records, with a TorchDispatchMode, every op whose output shape depends on values (a boolean-mask index, nonzero, unique, item). Before: three mask indexes at page size 1 and at 4. After: none, and both sides' capacity comes back.
• test_a_double_ratchet_is_reported (replaces test_double_ratchet_is_filtered_not_crashed) and test_a_full_only_free_of_a_bound_swa_page_is_reported fail before (nothing raises) and pass after.
• The test that reclaims full-attention KV for Mamba state (on a model with both Mamba state and sliding-window layers) and the eviction fixture name the new list.
• test_importing_the_module_again_registers_nothing_twice runs the module again as such a fixture does: a duplicate-invariant error before, none after (test_lmcache_component_cursors.py errors in 3 of its 4 cases with the invariants in unified_hybrid_swa and passes with them in allocator/swa.py).
• The unified-memory CPU tests pass with SGLANG_INVARIANT_CHECK=2 exactly as without it, so none of them frees an already released sliding-window page through a segment free.

Accuracy Tests

Earlier combined stack b550cacde6 adds a test-fixture correction in #43326: the Gemma4 read-stream fixture declares its FULL address space, and the HiSparse registry test uses real CacheInitParams for the retained-window calculation. Both complete test files pass locally with GPUs hidden: 4 Gemma4 tests and 26 registry tests. Pre-commit passed. The five descendant patches are unchanged (git range-diff); this follow-up changes no production code. The earlier 455-test run did not cover these two files; GitHub CI has not yet been revalidated for this head.

Validation before the test-fixture follow-up, on combined stack 4cf27a551c: 455 tests and 921 subtests passed on B300, plus the two widened-history regressions against Rust TreeCore. The new FULL-fallback metadata cases cover eager execution and CUDA graph replay. New allocation-reclaim, explicit-prefill and JIT-reuse regressions fail on the previous code and pass with the fixes. These are combined-stack checks, not independent model reruns of each PR.

The measurements in this section are historical results from before this rebase, at the revisions/configurations stated below. Current combined-stack validation is linked above.

SM90, TP1, gpt-oss-20b, no speculation, radix cache, overlap scheduler, CUDA graphs on, --mem-fraction-static 0.8. Baseline: the default pool. Unified: --enable-unified-memory. Every server ran with the per-step pool checks and the allocator invariants on (SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY=1 SGLANG_CHECK_KV_PAGE_INVARIANTS=1 SGLANG_INVARIANT_CHECK=1), so a free that this change gets wrong stops the server or is reported. These ran on this commit over main 7224fce. GSM8K, all 1314 test questions, 5-shot, greedy:

page size	backend	baseline	unified
256	triton	0.502	0.515
256	fa3	0.553	0.558
1	triton	0.533	0.536
1	fa3	0.566	0.535

gpt-oss-20b's GSM8K score moves by several points between identical runs (up to ~6 in our runs of one configuration), and three of the four rows favour the unified pool.

Under the same checks, two load tests at page size 256: one with a shared prefix (one 4096-token system prompt, 192 requests with 128-token questions and 2048 output tokens, concurrency 192) and one with long outputs (100 requests, 512 input / 4096 output tokens, concurrency 128). Every request completed on both arms and both backends, and no server reported a leak, double free, invariant violation or illegal memory access.

GPU unit tests for unified memory: 4104 passed, 0 failed.

Speed Tests and Profiling

The measurements in this section are historical results from before this rebase, at the revisions/configurations stated below. Current combined-stack validation is linked above.

bench_serving with random prompts at concurrency 64, in two shapes: 2048 input / 512 output tokens (200 prompts) and 8192 input / 64 output tokens (100 prompts). gpt-oss-20b, no speculation, page size 256, radix cache, overlap scheduler, CUDA graphs on, --mem-fraction-static 0.55. Each entry is the unified pool's gap to the default pool on the same tree, before (main) and after this PR; the default pool runs twice and the gap is against the mean of the two runs, with their spread (A/A) as the noise floor. This table ran on this commit over an earlier main.

backend	shape	req/s, main -> this PR	mean TTFT	mean E2E	A/A (req/s / TTFT / E2E)
fa3	2048 / 512	-1.4% -> -0.7%	+18.2% -> -0.0%	+1.5% -> +0.8%	0.6 / 0.3 / 0.7
fa3	8192 / 64	-1.1% -> -0.6%	+7.2% -> +0.5%	+1.2% -> +0.6%	0.8 / 2.1 / 0.8
triton	2048 / 512	-0.9% -> -0.1%	+14.9% -> -0.3%	+1.0% -> +0.1%	0.3 / 0.1 / 0.4
triton	8192 / 64	-5.0% -> -0.3%	+8.3% -> +0.5%	+5.4% -> +0.3%	0.4 / 0.2 / 0.4

The triton rows compare with the first default-pool run only, and their A/A is main's: on this PR's tree the second run compiled kernels during the benchmark. On prefill batches with identical new and cached token counts in both arms' logs, unified prefill throughput is 1.000 (fa3) and 0.999 (triton) of the default pool's with this PR, with an interquartile range of 0.996-1.003; on main the median is 0.997 with an interquartile range of 0.920-1.065, the spread the host syncs add.

The same configuration on this commit over main 7224fce, unified pool against the default pool:

backend	shape	req/s	mean TTFT	mean E2E	A/A (req/s / TTFT / E2E)
fa3	2048 / 512	-0.3%	+0.2%	+0.4%	0.1 / 0.6 / 0.1
fa3	8192 / 64	-0.2%	+0.4%	+0.2%	0.1 / 0.6 / 0.2
triton	2048 / 512	+0.7%	+0.2%	-0.6%	1.6 / 0.6 / 1.4
triton	8192 / 64	-2.2%	-0.2%	+2.2%	3.9 / 0.8 / 4.0

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.
• 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 #39127973340
Latest PR Test (Extra): ❌ Run #39127973371
Latest PR Test (AMD ROCm 10): ❌ Run #39127973401

[ch-wan]

[bug] Two-pool evict_to_free_tokens still flushes the old 3-tuple, so this new list is invisible there.

_flush_deferred_free_group returns when not any(pending_groups). After this change, free_full_segment queues into full_page_reps_group instead of free_page_reps_group (or full_free_group at page size 1). A group that only has those full-only segment frees is treated as empty on the two-pool path (UnifiedSWAAllocatorBase.evict_to_free_tokens) and left queued, so reclaim_plan can over-evict or fail while the pages are already waiting to be released. The tri-pool override goes through this method and is fine; gpt-oss is two-pool.

Suggestion: Pass self.full_page_reps_group in the two-pool tuple as well, or share one pending-groups helper with this flush.