Repository navigation
[Fix] mamba: Release consumed lazy Mamba prefix locks before donation retry - #43559
Merged
ispobock merged 6 commits intoOct 18, 2026
Merged
Conversation
alphabetc1
requested review from
Ying1123,
hanming-lu,
hnyls2002,
huangtingwei9988,
hzh0425,
ispobock,
merrymercy,
xiezhq-hermann and
yizhang2077
as code owners
October 10, 2026 13:13
alphabetc1
changed the title
alphabetc1
changed the title
alphabetc1
marked this pull request as draft
October 10, 2026 14:49
alphabetc1
force-pushed
the
fix/mamba-lazy-peak-slot-budget
branch
from
October 10, 2026 14:49
43d6c96 to
9245134
Compare
alphabetc1
changed the title
alphabetc1
changed the title
alphabetc1
marked this pull request as ready for review
October 16, 2026 09:41
alphabetc1
added
run-ci
ispobock
approved these changes
Oct 18, 2026
finnff
added a commit
to finnff/sglang
that referenced
this pull request
Oct 19, 2026
…fresh Conflict in mamba.py: upstream sgl-project#43559's lazy lock release (req.lock, TreeLock) moved into _try_alloc_mamba_slot(req), so the donation callers pass req and skip the stash on None. Claude-Session: https://claude.ai/code/session_019zFzB8crbELHdkhezNhjU4
detain
added a commit
to detain/sglang
that referenced
this pull request
Oct 19, 2026
…sumed lazy Mamba prefix locks before donation retry Hand-port: our fork renamed _alloc_mamba_slot -> _try_alloc_mamba_slot (entry 197/sgl-project#37943, Optional return + _record_skipped_checkpoint) and folded upstream's three donation call sites into one needs_new_slot branch, so the release-before-evict block lands inside _try_alloc_mamba_slot(req=None) and the single donation call site passes req. Upstream's new donated-ratio test targets the old name; retargeted to _try_alloc_mamba_slot. The entry-197 exhaustion test needed the exec bag published (mamba_radix_cache_strategy="extra_buffer", lazy off => the new block is inert) because the release block reads get_exec(). (cherry picked from commit ba3fc94, hand-ported) (cherry picked from commit fb0cdfdccbd8b9fc1f0cc4d82f934fbedf352024)
nanjiangwill
pushed a commit
to modal-projects/sglang
that referenced
this pull request
Oct 29, 2026
… retry (sgl-project#43559) (cherry picked from commit ba3fc94) Claude-Session: https://claude.ai/code/session_01VDwTQn4gGFgVhzGXFs4aJx
nanjiangwill
pushed a commit
to modal-projects/sglang
that referenced
this pull request
Oct 29, 2026
… retry (sgl-project#43559) (cherry picked from commit ba3fc94) Claude-Session: https://claude.ai/code/session_01VDwTQn4gGFgVhzGXFs4aJx
nanjiangwill
pushed a commit
to modal-projects/sglang
that referenced
this pull request
Nov 1, 2026
… retry (sgl-project#43559) (cherry picked from commit ba3fc94)
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
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.
Dependency
Depends on #43567, now merged and included in this branch. Releasing the prefix lock alone can expose stale in-place checkpoint donation under slot pressure.
Problem and change
With lazy overlap, the next decode can consume the temporary tracking slot before the previous prefill donates its checkpoint. Donation then needs a replacement slot while the consumed prefix state is still locked, exhausting the pool and asserting.
On a failed lazy donation allocation, release this request's consumed Mamba prefix reference before eviction/retry. Preserve KV locks and other requests' references; record the release to avoid a double decrement. Keep the allocation assertion and existing capacity formulas unchanged.
The CPU regression covers full 3-slot and 4-slot per-request pools and the subsequent normal unlock. A registered GPU regression additionally uses exactly 4 Mamba slots and 1 running request: warm a 128-token prefix, then extend it to 255 tokens and generate 32 tokens with overlap and decode CUDA graphs enabled. The first decode takes the fourth slot before prefill donation, which must reclaim the consumed prefix instead of requiring a fifth slot. It reuses the existing Mamba KL test server and CI registration. Checkpoint invalidation is handled separately in #43567.
Validation
Exact base/head red-green check with the same new testcase on GB300: base
f05da43d44067e4e4f95302b6d3bd40f7deb3b1dplus the testcase fails inprepare_for_caching_req->_alloc_mamba_slotwithAssertionError: Can not alloc mamba cache(exit 1); PR headc29503969e33f5c8e103cde6abb1caba209abd99passes (exit 0). The base run uses the unmodified base runtime, not a selectively reverted fix.Added GPU regression validated on GB300 with Qwen3.5-0.8B: both tests in
test_mamba_cache_kl.pypass (49.7 seconds), including the existing checkpoint-reuse test with KL=0.Negative control: remove only this PR's prefix-lock release block while retaining [Unified Radix] Invalidate lazy Mamba checkpoints before in-place tracking #43567; the new test fails in prefill checkpoint donation with
AssertionError: Can not alloc mamba cache.Pre-commit passed for the added test.
This branch independently: 14 CPU tests and 12 subtests passed; pre-commit passed.
The two PRs merge without conflicts. Their combined runtime files are identical to the previously validated combined fix (43d6c96).
Combined E2E on GB300 with overlap and decode CUDA graphs: Qwen3.5-0.8B passes 24 requests with 1 running request / 4 slots and 96 with 4 / 16; Qwen3.5-4B + NEXTN passes 24 with 1 / 4. Original main asserts in these cases. Both single-request 4-slot runs exactly match their 8-slot references in output tokens and logprobs. Ordinary and NEXTN 3-slot runs also complete 24 requests each.
The earlier combined GPU results cover both fixes together. The new negative control isolates the prefix-lock release change while retaining checkpoint invalidation.
CI States
Latest PR Test (Base): ✅ Run #39032574328
Latest PR Test (Extra): ❌ Run #39032574243
Latest PR Test (AMD ROCm 10): ❌ Run #39032574315