Fix unified SWA: size a non-owner's v2p by the id space it must address

[ch-wan]

Stack — fourth of four, the top. Depends on #37511 → #37512 → #37550 (its base branch). Two commits of its own. It is an independent bug fix that happens to sit here so the whole stack shares one CI run; it can be rebased onto main on request.

Motivation

--enable-unified-memory on a hybrid sliding-window model dies mid-serving: a device-side index assert, the scheduler gone, and the request batch returning nothing. The same model at the same commit with the static pool is fine.

The unified SWA composite mints one virtual page id per allocation and binds it on both sides:

new_virtual_pages = fa.free_virtual_ids[:num_pages].clone()      # the FULL side's ids
v_tokens = fa.alloc(need_size)
self.swa_attn_allocator.alloc_with_virtual(new_virtual_pages)    # bound on swa too

So the swa side's virtual_to_physical is indexed by the full side's ids — while it was sized by the swa side's own page count. Those two numbers are unrelated: which side gets more pages out of a shared byte budget depends on its per-page byte cost. gemma-4-31b-it is 10 full-attention layers to 50 sliding, so at --mem-fraction-static 0.8 on 2×H100:

full  v2p len = 2,065,601
swa   v2p len =   206,561     <- indexed by ids up to 2,065,600

Two things make it hard to see. It needs cumulative allocation churn before an id that far up is handed out, so a fresh server passes its first eval and dies on the second. And on GPU the offending access is a write — alloc_bind_inplace's tl.store(v2p_ptr + v, p), unchecked — so the visible symptom is somewhere else entirely: the later read in _swa_write_loc_unified is what trips the assert. Under CUDA_LAUNCH_BLOCKING=1 the stack resolves to that read; the CPU reference path raises IndexError from the write.

Modifications

• MultiEndedAllocator takes virtual_num_pages. virtual_to_physical is sized by it — the id owner's count, passed down for a non-owner — while physical_to_virtual keeps this pool's own page count, because that one is indexed by physical id. is_slot_allocated's bound moves to the virtual count for the same reason.
• The composite passes virtual_num_pages=self.full_attn_allocator.num_virtual_ids when building the swa side (both the 2-pool end pair and the tri-pool float).
• test_unified_swa_shared_virtual_ids.py (CPU): the table spans the owner's id space across three layer splits, and alloc/free churn past the swa side's page count binds cleanly.
• test_gemma4_unified_swa_virtual_ids.py (1 GPU, base-b, est 65s): the end-to-end cell that was missing — no registered test ran unified memory on a hybrid-SWA model. The existing unified e2e tests are MLA or ShortConv; the existing SWA e2e tests run the static pool.

Cost is one int64 row per id in the owner's space on the non-owner side: 16.5 MB at the gemma-4-31b shape above, against the 1.65 MB it had.

Accuracy Tests

The failure needs two conditions, and separating them took several wrong turns:

1. full.num_pages > swa.num_pages — a property of the model, true when the full side's per-page byte cost is the lower one. Read entry_bytes_per_page, not the layer counts: gemma-4-E2B is 7 full to 28 sliding but 2:1 in pages.
2. cumulative allocation churn past swa.num_pages — a property of the KV budget, since only enough allocation brings an id that high into play.

Shrinking the pool makes (2) immediate, which turns a 2B model on one GPU into a one-round repro — this is what the new e2e test runs:

--model-path google/gemma-4-E2B-it --enable-unified-memory \
  --attention-backend triton --mem-fraction-static 0.8 \
  --disable-radix-cache --max-total-tokens 60000

GSM8K over 200, one GPU, 65 s wall clock:

main   + unified   0.05   test RED
main   + static    0.87   <- the control: the model and backend are fine
branch + unified   0.865 / 0.870 / 0.875 / 0.880   test GREEN

Found in the wild on the shape where the swa table is small by itself — gemma-4-31b-it, TP2, default budget, GSM8K over 200 alternating a side:

round 1   main 0.98 alive     branch 0.98 alive
round 2   main 0.00 DEAD      branch 0.98 alive
round 3   (main gone)         branch 0.98 alive
rounds 4-6                    branch 0.98, 0.98, 0.98

For contrast, two hybrid-SWA models that do not reproduce it at their default budgets:

model	full / sliding	full pages	swa pages	runs to trigger
gemma-4-31b-it	10 / 50	2,065,601	206,561	~3
gemma-4-E2B-it	7 / 28	7,576,246	3,788,123	~63
gpt-oss-20b	12 / 12	equal	equal	never

The unit test fails both ways round on main: the width check on the (1, 5) split, and the churn check with the IndexError the CPU reference path raises where the Triton store would not.

test/registered/unit/mem_cache/: 2129 passed (this branch, after rebasing onto #37550 — both changes land in the same multi_ended_allocator.py).

Speed Tests and Profiling

Not a performance change. The extra table is 16.5 MB on the gemma-4-31b shape, allocated once at startup; nothing moves per step.

Since this is the top of the stack, it is also where "does unified cost anything against
the static pool" gets answered for the whole feature. Four hybrid shapes, same build,
only --enable-unified-memory differing -- schedulable KV is identical, and accuracy
and throughput are within run-to-run spread:

model	type	KV tokens	GSM8K	tput
gemma-4-E2B-it (TP1)	SWA	2913941 -> 2913939	0.8717 -> 0.8600	+2.00%
Qwen3.5-0.8B (TP1)	18 linear + 6 full	5032604 -> 5032608	0.5400 -> 0.5350	+1.67%
Kimi-Linear-48B (TP2)	MLA + KDA	4536315 -> 4536315	0.9075 -> 0.9100	-0.63%
Inkling-Small (TP4)	full + SWA + ShortConv	450605 -> 450605	0.975 -> 0.975	+1.21%

Inkling matches down to the budget it is derived from (rest=4.962 GB,
max_mamba_cache_size=207 on both sides), and at --mem-fraction-static 0.9 both sides
boot with 74778 tokens, so unified needs no extra headroom. Graph memory moves in both
directions; #37512 has that column and the breakdown.

Method, because the first pass got this wrong: run a multi-GPU A/B sequentially on the
same cards. The TP pool is sized from the minimum free memory across ranks, so one
foreign 1.2 GB process on one card of one GPU set shrinks that whole side's KV -- it read
as a 22.5% Inkling regression until both sides ran on the same four cards.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests. — one CPU unit test and one per-PR e2e, both verified red on main.
• Update documentation according to Write documentations. — no user-facing surface changes.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Notes for reviewers

The e2e test's arguments are load-bearing. Three substitutions made it pass on main, so they are not tuning:

• Raising --max-total-tokens hides the defect — the small pool is condition (2).
• Synthetic short prompts at a nominal 3.7x the churn passed on main; the 5-shot GSM8K traffic is what reaches it.
• It cannot use GSM8KMixin: that path scores this model 0.155 whatever the pool, because its chat template does not fit the model's reasoning config. Calling run_eval gives 0.87 on a healthy server.

Registered at base-b, not nightly. The lesson from the last regression in this area (#37307) was that the guard existed but was nightly-gated, so it never blocked the PR that broke it.

One claim I could not confirm. I described the id space as rotating — ids off the front of the owner's free list, freed ids to the back — and the code does read that way (torch.cat([free_virtual_ids, free_v_pages])). But probes on alloc() and alloc_with_virtual never fired during a real server run: the free list is consumed inside the Triton kernel on the alloc_extend path. So "churn advances the cursor" is the observed behaviour (the trigger is empirically cumulative, and the measured thresholds above bear that out), not a mechanism I have traced end to end.

---

CI States

Latest PR Test (Base): ⏳ Run #33682091141
Latest PR Test (Extra): ❌ Run #33682090790
Latest PR Test (AMD ROCm 7.2): ⏳ Run #33682090991

---

CI States

Latest PR Test (Base): 🚫 Run #33697315211
Latest PR Test (Extra): 🚫 Run #33697315009
Latest PR Test (AMD ROCm 7.2): 🚫 Run #33697315235

[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-02T07:07:59.203252Z	e2319f2	New commits

ℹ️ 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.

[ch-wan]

Latest PR Test (Base): ✅ Run #33682091141