[AMD][DSV4] Skip the paged SWA page return under the per-request ring

[ankith117]

Motivation

--attention-backend dsv4 with --page-size 256 crashes both TP ranks in the
warmup decode, immediately after the server reports it is ready:

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

Deterministic, 3/3 runs, both ranks identically. The HTTP layer binds and then
the scheduler dies underneath it, so the process exits via SIGQUIT rather than
failing at startup.

Under the per-request SWA ring, the paged swa_attn_allocator is vestigial:
available_size() returns the full-attn allocator's size and
swa_available_size() computes req_slots * ring_cost, so nothing consults the
paged allocator's internal counter. Ring slots are owned by the req slot for its
lifetime and recycled with it -- they are never lent out per free. Returning
pages to the paged allocator therefore credits capacity it never lent, its
available_size() climbs past size, and the assert at the end of
free_group_end fires.

The two halves reached main a day apart and each is correct alone:

• The _free_swa_pages -> free_page_ids path for page_size > 1 ([mem_cache] Free hybrid SWA pages by one representative per page on page_size > 1 #38159).
• The per-request SWA ring, re-landed gated ([AMD][DSV4] Reland unified-KV pool sizing and SWA ring accounting, fully gated #38192). The gating covers the
  allocation and translation paths (_swa_req_ring guards in alloc_extend,
  alloc_decode, available_size, swa_available_size) but not the two free
  sites.

Triggering it requires page_size > 1 and the ring simultaneously, which is
exactly DSV4 on ROCm -- the backend forces page_size to 256 and the ring comes
from the model's sliding_window. No CI lane covers that intersection, and the
reland landed the day after the free path changed underneath it.

Modifications

Skip the page return when the ring is active, at both free sites in
SWATokenToKVPoolAllocator:

• _free_swa_pages: return early in ring mode, after
  clear_full_to_swa_mapping. The mapping clear must still run -- skipping it
  leaves stale peer indices for translate_loc_from_full_to_swa, which would
  turn a clean crash into wrong KV reads.
• free_group_end: guard the grouped drain with not self._swa_req_ring. The
  list is still drained; only the page return is skipped.

test/registered/unit/mem_cache/test_swa_ring_page_return.py (base-a-test-cpu,
no GPU) covers both free paths in ring and non-ring mode. The two ring cases
fail on the pre-fix code with the same swa.py assert seen on the GPU and pass
on the fix; the non-ring cases guard against disabling the page return outright,
and one case pins the mapping clear that must still run in ring mode.

+7 -1 in swa.py, no API or behavior change off the ring path.

Accuracy Tests

DeepSeek-V4-Flash, 2x MI45x (gfx1250), --tp 2 --attention-backend dsv4 --page-size 256 --kv-cache-dtype fp8_e4m3, ROCm 10.1.

Booting is not sufficient evidence here: the warmup request is 19 tokens and
never wraps the 256-token ring. Two checks that do exercise it, both against a
reference build that predates #38159/#38192 (v0.5.19.dev20260906+g5bebe7a033):

• 5 greedy prompts (temperature=0), including two of 614 and 620 tokens that
  force the ring to wrap: output_ids bit-identical to the reference on all
  5. Greedy decoding is deterministic, so this is an exact-match oracle on the
  KV path, not a smoke test.
• 32 concurrent 600+-token requests, exercising ring-slot recycling across the
  256 slots: 32/32 correct answers, zero asserts or exceptions in the log.

End-to-end on the fixed build, GSM8K full set (1319 questions, 5-shot,
temperature=0, --parallel 128):

Accuracy:          0.936
Invalid:           0.000
Latency:           447.437 s
Output throughput: 265.407 token/s

Unit coverage reproduces the failure without hardware: the two ring-mode cases
in test_swa_ring_page_return.py fail on the pre-fix code at the same
swa.py:559 assert seen on the GPU, and pass on the fix (5 passed, 7 s, CPU).

Without the fix the same config cannot serve a single request.

Speed Tests and Profiling

No performance work; the change removes calls on the ring path and touches
nothing else. Worth noting the fix restores the intended sizing rather than
degrading it -- max_total_num_tokens is 44,482,560 on the fixed build vs
28,960,000 on the pre-#38192 build used as the accuracy reference (+54%), since
falling back to that build also gives up the unified-KV pool sizing.

Serving throughput on the fixed build, from the GSM8K run above: 265.407 token/s
output at --parallel 128.

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.

cc @hnyls2002 @ispobock @alphabetc1 (allocator CODEOWNERS), @yuttian1 (#38192)

---

CI States

Latest PR Test (Base): ❌ Run #34306364007
Latest PR Test (Extra): ❌ Run #34306363799
Latest PR Test (AMD ROCm 7.2): ❌ Run #34306363963

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-09-08T22:34:21.522633Z	153e825	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[HaiShaw]

cc @amd-danli103 @yuttian1

[1am9trash]

Add a guard to avoid non-existed page-id free.
Only impact the amd side.

# python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py
if self._unified_kv:
    self.swa_req_ring_size = self.unified_swa_ring_size

# python/sglang/srt/mem_cache/allocator/swa.py
ring_size = kvcache.swa_req_ring_size
self._swa_req_ring = ring_size is not None

if self._swa_req_ring:
    # ......

LGTM.

[amd-danli103]

#38192 gated alloc on the ring and #38159 added the page_size>1 page return; DSV4 on ROCm is that intersection, so free_page_ids over-credits. Skipping the return after clear_full_to_swa_mapping is the right cut.
LGTM, but please fix lint first.

[yuttian1]

Looked through the allocation and free paths. This fix makes sense—the ring owns these SWA pages, so they shouldn’t be returned to the paged allocator. The mapping is still cleared, and the non-ring path is unchanged. LGTM.

[HaiShaw]

Fix bug.