[Fix] SWA allocator: skip unmapped reps in the segment free path (page-0 double free)

[TianDi101]

Bug

SWATokenToKVPoolAllocator.free_group_end() asserts on all ranks a second after
the server goes live, whenever an external cache linker is attached under
unified KV:

File ".../mem_cache/allocator/swa.py", line 559, in free_group_end
  assert self.swa_attn_allocator.available_size() <= self.swa_attn_allocator.size
AssertionError

available_size() > size is an over-free: more SWA pages were returned than exist.

Root cause

SWAComponent.evict_component() deliberately passes the node's FULL indices,
and says why:

Pass full indices to free_swa so slots with no SWA pair are skipped. Freeing
swa_value directly would double free those entries since they all map to the
same sentinel slot.

But apply_component_action() calls free_swa_segment(), not free_swa(). The
two paths do not have the same contract:

• free_swa() -> _release_swa() filters swa_indices > 0.
• free_swa_segment() -> _free_swa_pages() does not filter. Its only guard is
  expect(_SWA_PEER_MAPPED, swa_tokens > 0, ...), which is signal-only and off by
  default: resolve_level() returns OFF unless SGLANG_INVARIANT_CHECK /
  SGLANG_ENABLE_ASYNC_ASSERT is set, and Bucket.FATAL_UNCONTAINABLE forbids a
  recover. In a normal run it neither detects nor repairs anything.

So a rep with no SWA peer reads 0, becomes swa_pages = 0 // page_size = 0, and
free_page_ids() is documented "Free exactly these pages; no page twice, no
dedup." Reserved page 0 therefore goes back on the free list once per unmapped
rep until available_size() passes size.

Page 0 is the padding slot (free_pages = torch.arange(1, num_pages + 1)), so
unguarded this hands the padding slot out as a real KV page; the assert is
catching genuine corruption.

Why it is a regression

Before _free_swa_pages() existed, free_swa() routed page_size > 1 through
_expand_to_full_pages() -> _release_swa(), which filters > 0. page_size > 1
was safe by construction; only the page_size == 1 branch carried the bare
expect, with the comment "A filter here would make the output shape
data-dependent, which costs a device-to-host sync." The fixed-shape segment path
routes page_size > 1 around that filter.

The allocator's own debug check already encodes the intended semantics --
torch.unique(ref[ref > 0] // ps) excludes unmapped entries -- so with
SGLANG_DEBUG_MEMORY_POOL=1 the current code fails that assert too.

Fix

Drop unmapped reps in _free_swa_pages(), exactly as _release_swa() already does
on the set-shaped path.

Reproducer

Standalone, no model required -- builds a real SWATokenToKVPoolAllocator
(page_size=256) and exercises free_swa_segment():

scenario	before	after
normal decode, every page has an SWA peer	PASS	PASS
externally loaded prefix, no SWA peers	FAIL [0,0,0,0], available 5120 > size 4096	PASS, frees nothing
half window-evicted, half live	FAIL [1,2,0,0], available 4608 > size 4096	PASS, frees only [1,2]

End-to-end

DeepSeek-V4-Pro FP4, TP8 + DP attention, EAGLE MTP, page_size 256, concurrency
256, agentic trace replay, external linker on all 8 ranks.

• Unpatched: 8/8 ranks assert ~1s after The server is fired up and ready to roll!
• Patched: full 3628s replay, 11450 requests, errors=0, zero asserts.

Note on the sync

The filter is a boolean index, so the output shape is data-dependent and it costs
a D2H sync -- the thing the fixed-shape path was built to avoid. In a follow-up I
can move the filter to the release points (free_group_end() plus the ungrouped
free) so it runs once per free group instead of once per call; that keeps
_free_swa_pages() fixed-shape. Filing the straightforward version first since it
is the one with end-to-end validation behind it.

Context

Found while running the AMD DeepSeek-V4 unified-KV external-linker arm from
#38269. That PR is not the cause and does not need to change:
its branch predates _free_swa_pages(), and the defect is on main independently.
The linker just makes the unmapped-slot case common, because it installs FULL
slots whose SWA component was skipped under unified KV. Window-evicted tails can
reach the same path without any linker.

---

CI States

Latest PR Test (Base): ❌ Run #34232788705
Latest PR Test (Extra): ❌ Run #34232788466
Latest PR Test (AMD ROCm 7.2): ❌ Run #34232788658

[TianDi101]

Closing in favour of a single PR. The same fix now rides on #38269, which is the PR that makes the direct external linker install FULL slots with no SWA peer under unified KV and therefore triggers this path. Landing them together keeps the trigger and the fix in one merge.

Commit on that branch: 5264e61. Root cause, the #38159 interaction, and the unit + end-to-end verification are written up at #38269 (comment).