[Fix] Allocate SWA KV pools inside the memory-saver region

[guapisolo]

Summary

release_memory_occupation now frees the KV cache of MHA, MLA and DSA sliding-window models.

Symptom & Reproduction

• Symptom: with --enable-memory-saver, releasing kv_cache leaves the KV cache of SWA models resident, and every DSA model keeps its index-K buffers.
• Engine measurement: an H200 TP2 engine of a hybrid SWA model used 85.3 GiB after load, 32.9 GiB after releasing everything, and 85.3 GiB after resuming only the weights, so only the weights were pausable.
• Reproduction: build a pool through its KVCacheConfigurator builder with the flag on, pause the torch_memory_saver kv_cache tag, and read torch.cuda.mem_get_info; expected: the pool's bytes are freed; observed on the merge base: 0 GiB freed for all three SWA pool kinds.

Root Cause

1. pause(GPU_MEMORY_TYPE_KV_CACHE) unmaps only buffers allocated inside a memory-saver region.
2. _build_hybrid_swa_kv_pool never passes the flag, so SWAKVPool gives both sub-pools False.
3. _build_hybrid_mla_swa_kv_pool hard-codes False in the explicit sub-pool kwargs it hands SWAKVPool.
4. IndexKeyCache.__init__ allocates index-K outside memory_saver_adapter.region, in every DSA pool.

Fix

Each pool builder now passes get_exec().features.enable_memory_saver, as every other pool builder already does. The MHA builder uses a new SWAKVPool argument that is forwarded to both sub-pools and defaults to False. The hybrid MLA/DSA builder sets the flag in its explicit common kwargs, because explicit sub-pool kwargs bypass that argument. IndexKeyCache allocates index-K inside the GPU_MEMORY_TYPE_KV_CACHE region, like the MLA buffers of the same pool. With the flag off the adapter is a no-op, so allocation is unchanged. Unchanged: the NPU SWA path, LayerSplitIndexKeyCache.remote_buffer and the DSA compress-tail buffers.

Verification

• pre-commit run --files on the three changed files at this head: passed.
• Pause/resume on one H200 with torch_memory_saver 0.0.9.post1: pools built through the real builders (48 layers, 12 full and 36 SWA), freed GiB per tree. Run before the rebase at 29bc585196, whose three-file diff equals this head's except one deleted comment:

tree	MHA SWA (21.0)	hybrid MLA (5.9)	hybrid DSA (6.3, index-K 0.39)
base	0	0	0
first commit only	21.2	0	0
builder change only	21.2	6.0	6.0
full change	21.2	6.0	6.4

• After resume every paused buffer kept its data_ptr, read back zero, and accepted writes.
• CI at 34ba83fa52 (this head minus the comment-only commit): PR Test Base passed; /rerun-test of the SWA and DSA pool unit tests passed at 29bc585196.
• Not run: a full engine with a hybrid MLA/DSA model, or the NPU path.

Review Focus

• index_key_cache.py: index-K now joins the kv_cache tag for every DSA model; resume zero-fills it, and flush_cache() runs right after the pause, so stale index-K is never read.
• _build_hybrid_mla_swa_kv_pool: its explicit sub-pool kwargs bypass the SWAKVPool argument, so the flag must stay in common.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #36200923877
Latest PR Test (Extra): ❌ Run #36200923762
Latest PR Test (AMD ROCm 10): ❌ Run #36200923733

[guapisolo]

/tag-and-rerun-ci

[guapisolo]

/rerun-test test_dsa_pool_host_unit.py test_dsa_layer_split_broadcast.py test_swa_unittest.py registered/attention/unittests/dsa/test_dsa.py test_swa_pool_v_head_dim.py

[github-actions]

Results for /rerun-test test_dsa_pool_host_unit.py test_dsa_layer_split_broadcast.py test_swa_unittest.py registered/attention/unittests/dsa/test_dsa.py test_swa_pool_v_head_dim.py:

🚀 1-gpu-5090 (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/mem_cache/test_dsa_pool_host_unit.py

🚀 4-gpu-b200 (2 tests): ✅ View workflow run

cd test/ && python3 registered/unit/mem_cache/test_dsa_layer_split_broadcast.py
cd test/ && python3 registered/attention/unittests/dsa/test_dsa.py

🚀 1-gpu-h100 (2 tests): ✅ View workflow run

cd test/ && python3 registered/unit/mem_cache/test_swa_unittest.py
cd test/ && python3 registered/attention/unittests/dsa/test_dsa.py

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/mem_cache/test_swa_pool_v_head_dim.py

[guapisolo]

/tag-and-rerun-ci

[guapisolo]

cc cmt

CI record for this PR:

• PR Test Base passed at 34ba83f (the code under review): 55 jobs succeeded, 10 skipped, pr-test-finish green. https://github.com/sgl-project/sglang/actions/runs/36065717895
• Lint passed at 34ba83f. https://github.com/sgl-project/sglang/actions/runs/36065717158
• Targeted /rerun-test at 29bc585 (before the rebase; identical three-file diff), all passed:
  • test_dsa_pool_host_unit.py on 1-gpu-5090: 3 tests OK. https://github.com/sgl-project/sglang/actions/runs/35973273719
  • test_dsa_layer_split_broadcast.py and test_dsa.py on 4-gpu-b200: 1 and 16 tests OK (7 skipped). https://github.com/sgl-project/sglang/actions/runs/35973287126
  • test_swa_unittest.py and test_dsa.py on 1-gpu-h100: 24 and 16 tests OK (5 skipped). https://github.com/sgl-project/sglang/actions/runs/35973300011
  • test_swa_pool_v_head_dim.py on CPU: 1 test OK. https://github.com/sgl-project/sglang/actions/runs/35973312648

The current head c7db69c only removes a code comment on top of 34ba83f; its CI run stopped at the maintenance-mode rebase check and ran no tests.