[Perf] Fuse SWA page lookup and mapping clear

[Jialin]

Motivation

SWA segment frees currently launch separate Torch operations to look up peer pages and clear FULL-to-SWA mappings. Fuse that metadata work into one Triton kernel to reduce per-free scheduler dispatch overhead.

Modifications

• Preserve page ownership, deferred frees, and the per-request SWA-ring early return. CPU/NPU keep the Torch reference path.
• Retain the peer-mapped invariant and debug validation, and avoid specializing on the page count or representative-pointer alignment.
• Add kernel/allocator regressions and include the current test/registered/kernel/ taxonomy in the kernel CI path filter.

Accuracy Tests

• Fresh focused CUDA suite with SGLANG_IS_IN_CI=true and SGLANG_ENABLE_ASYNC_ASSERT=true: 139 tests and 244 subtests passed at 750ba1d20.
• The deliberate-invalid-input debug test now scopes invariant signaling OFF only around its invalid free, avoiding a fatal device assertion before Python validation. Allocator debug checks and the separate STRICT contract tests remain enabled. This follow-up changes tests only.
• CUDA memory checking: zero errors.
• Additional CUDA ring/debug/grouping matrix: 8/8 cases passed, checking mapping clears and immediate/deferred page accounting.
• All pre-commit hooks passed. NPU dispatch was emulated; AMD/NPU hardware and model-level accuracy tests were not run locally.

Validation commands

These command payloads ran through uv run --project <existing CUDA environment> --no-sync, with this checkout's python/ first on PYTHONPATH.

env -u SGLANG_INVARIANT_CHECK SGLANG_IS_IN_CI=true SGLANG_ENABLE_ASYNC_ASSERT=true \
  pytest -q --tb=short -x \
  test/registered/kernel/mem_cache/test_swa_page_free.py \
  test/registered/unit/mem_cache/test_swa_unittest.py \
  test/registered/unit/mem_cache/test_paged_free_segment.py \
  test/registered/unit/mem_cache/test_unified_free_no_host_sync.py \
  test/registered/unit/mem_cache/test_unified_tri_pool.py \
  test/registered/unit/mem_cache/test_swa_ring_page_return.py \
  test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py

compute-sanitizer --tool memcheck --report-api-errors no --error-exitcode 1 \
  python test/registered/kernel/mem_cache/test_swa_page_free.py

pre-commit run --files \
  .github/workflows/_pr-test-check-changes.yml \
  python/sglang/kernels/ops/memory/allocator.py \
  python/sglang/srt/mem_cache/allocator/swa.py \
  test/registered/kernel/mem_cache/test_swa_page_free.py \
  test/registered/unit/mem_cache/test_swa_unittest.py

Speed Tests and Profiling

For a paged segment (page_size > 1, debug checks off), the metadata path changes from eight GPU launches to one:

Operation per SWA segment free	Baseline	Fused
Gather peer mappings (aten::index)	1	0
Check mapped peers (aten::gt)	1	0
SWA/FULL page arithmetic (aten::floor_divide)	2	0
FULL page starts (aten::mul)	1	0
Page offsets (aten::arange)	1	0
Expanded mapping indices (aten::add)	1	0
Clear mappings (aten::index_fill_)	1	0
get_and_clear_swa_pages_kernel	0	1
Total metadata launches	8	1

The kernel reads one representative per distinct FULL page, returns its SWA page ID and peer-mapped flag, then clears that FULL page's mapping entries. A partial segment tail still clears the whole allocated page. The block barrier orders representative loads before mapping clears. The expanded index tensor is eliminated; its former unsqueeze/reshape views did not launch GPU kernels.

Eight back-to-back one-page frees therefore use 64 → 8 metadata launches, saving 56 launches without deferring ownership resolution. At page size 1, the same path changes from 3 → 1 launches per free.

This breakdown covers lookup and mapping clear, not the entire eviction or scheduler. Invariant-check kernels, FULL-page release arithmetic, and free-list updates remain outside it. Debug checks may synchronize, and first-use JIT compilation remains possible. No model-level serving-throughput result is claimed.

Model-free measurements

Measured on NVIDIA GB300 with PyTorch 2.11.0+cu130 and Triton 3.6.0 at 7e2774215 (runtime code is unchanged by the test-only follow-up). All 16 shapes (page_size = 1/16/64/128, pages per free = 1/8/64/1024) matched the Torch reference. The table shows one-page frees, with int64 inputs and debug checks off:

Page size	CPU enqueue baseline → fused (µs)	GPU graph-replay interval baseline → fused (µs)
1	41.681 → 25.200	11.360 → 7.264
16	124.321 → 23.808	19.552 → 7.264
64	120.833 → 23.280	19.584 → 7.296
128	124.257 → 23.296	19.584 → 7.296

Medians of 100 samples after 10 warmups, measured without the profiler. CPU time covers eager enqueue work, excluding explicit completion waits. GPU timing uses CUDA events around graph replay and includes graph-launch overhead; a 256 MiB cache flush occurs outside the timed interval. Mapping resets, invariant checks, and allocator free-list updates are excluded.

Benchmark script · All samples and results

Profiler screenshots and traces

Actual Perfetto screenshots of one paged free (page_size = 128): eight baseline GPU kernels versus one fused kernel. Each view is zoomed independently. These profiled spans illustrate dispatch structure; the unprofiled timings are in the table above. No Python stacks were captured. Trace metadata is reduced to generic CPU/GPU labels and relative timestamps.

Baseline:

Fused:

Click an image to open the full-size image. Download the complete eight-free traces and open them in Perfetto: baseline, 64 kernels · fused, 8 kernels.

Checklist

• Format code according to the pre-commit guide.
• Add tests according to the test guide.
• Document the distinct-FULL-pages contract; no user-facing documentation change is needed.
• Model-level accuracy and serving-speed benchmarks (not run; allocator validation is described above).
• Follow the SGLang code-style guidance.

Review and Merge Process

Ready for maintainer feedback. CI was requested on 750ba1d20 with /tag-and-rerun-ci; results are pending. Merge remains gated on review and passing CI.

---

CI States

Latest PR Test (Base): 🚫 Run #35373755609
Latest PR Test (Extra): 🚫 Run #35373755078
Latest PR Test (AMD ROCm 10): ❌ Run #35373755553

Original commits

• 3fb8e3407ab0245048926a61d025c6e3fa2d7ef5

[Jialin]

/tag-and-rerun-ci

[Jialin]

/tag-and-rerun-ci

[Jialin]

/rerun-failed-ci

09/10 10:24PM PT: after fixed trunk in #39019
09/11 9:42AM PT: retry to pick up the fix above ^
09/11 6:28PM PT: test process got killed
09/12 10:08PM PT: test got cancelled
09/13 12:08AM PT: perf test slightly exceed threshold

[Jialin]

/rerun-failed-ci

09/14 11:39AM PT: Add run-ci and run-ci-extra labels
09/14 1:52PM PT: Timeout
09/14 11:44PM PT: Extra retry is not triggered
09/15 1:30AM PT: Retry due to cancel
09/15 9:38AM PT: CUDA device unavailable
09/15 2:10PM PT: CUDA device unavailable
09/15 6:40PM PT: mooncake failure
09/16 11:33AM PT: cache_hit failure (confirm flakiness)

[Jialin]

/rerun-failed-ci