[SWA] Avoid shape-dependent operations for paged eviction

[paulzhang-tm]

Why is this change needed?

Out-of-window SWA frees on the paged SWATokenToKVPoolAllocator still run a data-dependent boolean filter (_release_swa) and a torch.unique in the paged allocator's free(), so every page-aligned segment that free_swa_out_of_window_slots releases synchronizes the scheduler stream with the in-flight forward. With the scheduler's WAR fence at target verify these are the only long stalls left in a saturated MTP decode step (4 to 5 per step at bs 750).

Summary

• SWATokenToKVPoolAllocator.free_swa(free_index, *, start_pos=None): given the segment's position in its request's kv row and page alignment, release one representative per page with fixed-shape ops. Every page must map onto one SWA page; an async assert guards it. Calls without start_pos, unaligned segments and page_size == 1 (already sync-free) keep the legacy path. Grouped frees collect the representatives and release them at free_group_end.
• PagedTokenToKVPoolAllocator.free_page_ids: release live page ids after asynchronous validity and uniqueness checks, without torch.unique.
• free_swa_out_of_window_slots passes start_pos=req.kv.swa_evicted_seqlen to every SWA allocator; the unified allocator already took that keyword, and the pure-SWA and HiSparse allocators now accept and ignore it, so the call site needs no allocator dispatch.
• No process-global "mapping may be partial" fallback. The only allocation that leaves a request's head unmapped is the disaggregated-decode tail preallocation (alloc_extend_swa_tail), and that path already starts the eviction floor at the page-aligned tail (swa_evicted_seqlen = fill_len - swa_tail_len), so out-of-window segments are always fully mapped.

Same call sites and timing; the end state is bitwise identical to the legacy path.

Tests

test/registered/unit/mem_cache/test_sync_free_swa_inplace.py: fixed path vs legacy on paged, grouped, mixed, page-size-1 and staged-release allocators; page-id validity and uniqueness asserts on host and device; segment alignment and mapping contracts; grouped out-of-window eviction without torch.unique; tail-only preallocation freeing from the tail floor; every SWA allocator accepts start_pos.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #33944690097
Latest PR Test (Extra): ❌ Run #33944690141
Latest PR Test (AMD ROCm 7.2): ❌ Run #33944690067

[hnyls2002]

Thanks @paulzhang-tm. Closing as superseded by #38159 (merged), which lands the same mechanism on main and covers a bit more ground:

• free_swa_segment(idx, *, start_pos) is the segment form (one representative per page, mapping[idx[::page_size]] // page_size, fixed shape); free_swa(idx) stays set-shaped, so there is no start_pos=None fallback to fall into.
• free_page_ids(page_ids) on the paged / token allocators is the same page-native primitive; grouped frees queue the page ids and flush once.
• Every kv-row caller takes the segment path, not only free_swa_out_of_window_slots: SWARadixCache node frees, the unified FreeComponentDeviceSlot action, and the unified allocator's own free_swa_segment. A partial last page is released whole instead of falling back to the syncing path.
• The full-page-mapped guard is the _SWA_PEER_MAPPED invariant (SGLANG_INVARIANT_CHECK), and the page-set check runs under SGLANG_DEBUG_MEMORY_POOL, rather than an always-on gather plus sort per free.

The bs 750 MTP observation (4 to 5 of these stalls per decode step) is a useful data point for this line; if you still see them on main with #38159, please open an issue with the config.