Avoid device syncs when releasing paged SWA segments - #35953
Closed
weireweire wants to merge 7 commits into
Closed
weireweire wants to merge 7 commits into
weireweire wants to merge 7 commits into
Conversation
weireweire
requested review from
Ying1123,
alphabetc1,
hanming-lu,
hnyls2002,
huangtingwei9988,
hzh0425,
ispobock,
merrymercy,
xiezhq-hermann and
yizhang2077
as code owners
August 22, 2026 06:26
Contributor
Author
|
/tag-and-rerun-ci |
weireweire
force-pushed
the
perf/paged-swa-segment-free
branch
6 times, most recently
from
August 25, 2026 07:55
f41b008 to
6b997a2
Compare
weireweire
added a commit
to weireweire/sglang
that referenced
this pull request
Aug 25, 2026
Root cause: Paged SWA cleanup rediscovered physical pages with torch.unique and dynamic mapping filters, whose data-dependent output sizes can synchronize the scheduler thread with the CUDA stream. Deferred segment release also flattened different requests, allowing equal relative positions to be mistaken for a shared boundary page. Fix: Use fixed-stride page representatives and request-scoped segment release with the CPU-maintained SWA eviction frontier. Preserve request boundaries through free groups, validate page-dense mappings rebuilt by load-back and recovery, and route hybrid SWA, DSV4 HiSparse, and shared-ID unified cleanup through the fixed-geometry paths while retaining the unknown-geometry legacy fallback. Validation: - pre-commit run --all-files --show-diff-on-failure - Python syntax compilation for all changed modules and tests - git diff --check and added-line getattr audit
weireweire
added a commit
to weireweire/sglang
that referenced
this pull request
Aug 25, 2026
Root cause: Paged SWA cleanup rediscovered physical pages with torch.unique and dynamic mapping filters, whose data-dependent output sizes can synchronize the scheduler thread with the CUDA stream. Deferred segment release also flattened different requests, allowing equal relative positions to be mistaken for a shared boundary page. Fix: Use fixed-stride page representatives and request-scoped segment release with the CPU-maintained SWA eviction frontier. Preserve request boundaries through free groups, validate page-dense mappings rebuilt by load-back and recovery, and route hybrid SWA and DSV4 HiSparse cleanup through fixed-geometry paths while retaining the unknown-geometry legacy fallback. Validation: - pre-commit run --all-files --show-diff-on-failure - Python syntax compilation for all changed modules and tests - git diff --check and added-line getattr audit
weireweire
force-pushed
the
perf/paged-swa-segment-free
branch
from
August 25, 2026 09:18
6b997a2 to
d92f2b5
Compare
Collaborator
|
Validation result on the Qwen3.5 AgentX workload:
We observed no measurable performance regression from this isolated cherry-pick. The positive movement is still within our 3% run-to-run threshold, so we are not claiming a material or statistically significant improvement. |
Root cause: Paged SWA cleanup rediscovered physical pages with torch.unique and dynamic mapping filters, whose data-dependent output sizes can synchronize the scheduler thread with the CUDA stream. Deferred segment release also flattened different requests, allowing equal relative positions to be mistaken for a shared boundary page. Fix: Use fixed-stride page representatives and request-scoped segment release with the CPU-maintained SWA eviction frontier. Preserve request boundaries through free groups, validate page-dense mappings rebuilt by load-back and recovery, and route hybrid SWA and DSV4 HiSparse cleanup through fixed-geometry paths while retaining the unknown-geometry legacy fallback. Validation: - pre-commit run --all-files --show-diff-on-failure - Python syntax compilation for all changed modules and tests - git diff --check and added-line getattr audit
weireweire
force-pushed
the
perf/paged-swa-segment-free
branch
from
August 27, 2026 09:16
d92f2b5 to
b69be4b
Compare
Collaborator
|
[by Codex] The required PR Test Base failure appears caused by this PR. Both root failures exercise unified-memory allocation/release paths changed here:
The other red Base jobs are fast-fail cascades from these roots. Please audit the new free_segments/free_swa_segment and grouped Full/SWA component-release paths for slots that are not returned to the full allocator. Root jobs: |
Root cause: The segment-aware cleanup path queues free_segments calls separately from legacy free calls while a free group is active. Unified Mamba and shared-SWA allocators overrode free_group_end but only drained the legacy queue, leaking every Full slot from a finished request; the failing CI request retained 11 slots. Fix: Reset and drain the segment queue in both unified composite allocators. Preserve shared-SWA request boundaries and free its still-live suffix directly from the CPU-known eviction frontier. Validation: - 13 tests passed in test_paged_free_segment.py inside the SGLang CUDA container - uvx pre-commit run --all-files --show-diff-on-failure - Python compileall for changed source and test - git diff --check and added-line getattr audit
Resolve the free-group state refactor by adopting the new None/list sentinel for deferred segment queues. Keep grouped Full/SWA and Mamba segment release draining after the group switches back to immediate-free mode. Validation: python compileall and focused pre-commit checks pass.
Preserve deferred segment draining alongside the new full-only SWA release path. Both queues now participate in group lifecycle without bypassing deferral. Validation: compileall and focused pre-commit checks pass.
Root cause: the unified SWA allocator bypasses the static SWA parent constructor and free-group hooks, so the newly added full_free_group state was neither initialized nor drained. The Base CPU test failed on the first grouped free_full call with AttributeError; initializing only the field would leave a silent deferred-release leak. Fix: initialize and reset full_free_group with the allocator lifecycle, create a fresh queue at free_group_begin, and drain it after Base finishes ordinary and segment releases. Validation: compileall, diff checks, focused pre-commit, and full pre-commit pass. The exact runtime test could not be run locally because this checkout has no Torch/Numpy environment; PR CI provides the runtime confirmation.
Collaborator
|
@weireweire could you fix the conflicts? |
Contributor
Author
|
closing as supersede by @hnyls2002 's #36637 and follow up PRs |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
torch.unique.Motivation
Paged SWA cleanup previously rediscovered physical pages from token-slot tensors with
torch.uniqueand dynamic filtering. Their data-dependent output sizes can synchronize the scheduler thread with the CUDA stream.For a contiguous request segment, the touched pages are already determined by its logical start position, length, and allocator page size. Fixed-stride representatives and the CPU-maintained SWA eviction frontier provide fixed-shape release geometry without GPU-to-CPU size discovery.
Deferred releases must also preserve request boundaries: flattening segments from different requests can incorrectly treat equal relative positions as a shared boundary page.
This focused change is extracted from the broader integration work in #34515, which increase E2E throughput by 40%. And the E2E test include this change is done there.
CI States
Latest PR Test (Base): ❌ Run #33233754247
Latest PR Test (Extra): ❌ Run #33233754226
Latest PR Test (AMD ROCm 7.2): ❌ Run #33233754342