Repository navigation
fix(mamba): skip prefill donation when an in-flight verify aliases the keep slot - #64
Open
rchalamala wants to merge 1 commit into
Open
rchalamala wants to merge 1 commit into
rchalamala wants to merge 1 commit into
Conversation
…e keep slot Under extra_buffer_lazy + speculative decoding with the overlap scheduler, the first verify after a lazy final prefill is launched before the prefill result is processed. If the pending ping-pong slot allocation fails at verify prepare, the verify plan falls back to the keep slot as its scatter destination. Prefill result processing then donated/inserted that physical slot into the radix tree under the prefill tracked depth C while the already-launched verify commits the crossing state T > C into the same slot, so the tree records newer state under the stale depth. Apply the decode guard keep_may_be_written_in_flight recompute at the prefill caching decision: skip the finishing insert and the unfinished donation while the pending position is empty and a crossing is reachable. The first verify result repairs the request-side label, so the checkpoint is published later under the correct depth.
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 32ffd43643
ℹ️ 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".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
4 tasks done
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
1 participant
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.
Motivation
With lazy extra-buffer mamba checkpoints and speculative decoding, a lazy
final prefill holds its checkpoint in one physical ping-pong slot (keep == next,
second position −1). When the first verify's pending-slot allocation fails —
naturally, when the checkpoint pool is momentarily full of unlocked cached
checkpoints, because that allocation deliberately does not evict — the verify
scatters in place into the keep slot. Prefill result processing then donates
that physical slot to the radix tree under the prefill's tracked depth C,
while the already-launched verify commits newer checkpoint state at a tracking
boundary T > C into the same slot. The tree deterministically labels newer
state as the old depth, and a later prefix match is served a checkpoint for
the wrong position.
Runtime-confirmed on Qwen3-Next-80B-A3B + its NEXTN MTP head (2xB300,
extra_buffer_lazy, track interval 256, overlap scheduler, 255-token prompts):
forced pending-slot failure publishes the keep slot under C=192 while the
verify commits T=256 into it, in both the finishing and unfinished cases, and
a repeat request matches the corrupted node (cached_tokens=192) and produces
greedy output divergent from the flushed-cache control. Natural reproduction
without any test flag (12-slot pool, 40 distinct prompts): 140 natural
pending-slot allocation failures across 35 requests, 70 aliased donations,
every one with the matching crossing commit into the same slot.
New patch, not an upstream duplicate: sgl-project#37837 resets the
ReplaySSM ring cursor on the donate path and never touches the donation
decision or the in-flight verify's captured destination — disjoint files and
failure mode.
Modifications
guard's keep_may_be_written_in_flight recompute (new helper
_mamba_lazy_keep_may_be_written_in_flight) before the finishing tree insert
(is_insert=False) and before the unfinished donation, when lazy extra-buffer
spec is active, the pending ping-pong position is empty, and a tracking-boundary
crossing is reachable. The first verify result repairs the request-side
label, so the checkpoint publishes later under the correct depth T.
Conservative superset semantics match the decode guard (size-1 buffer and
non-overlap cases return False by construction), and the lazy-spec-inactive
path is byte-equivalent (early False, pure reads).
12 tests — predicate matrix (lazy/spec/window/buffer arms), finishing
suppression, unfinished-donation suppression, no-alias arms unchanged, and
two flow tests pinning the defect end to end.
Accuracy Tests
two flow tests fail with exactly the defect behavior (is_insert=True,
donation made) and the 7 predicate tests fail structurally (attributable
control; the 3 no-alias arms pass on both trees).
suppressed (prefill_alias_suppressed; correct publication under T=256 at
finish); natural arm: 140 failures, 70 suppressions, 0 aliased donations.
The already-released request's verify result is skipped by design in
process_batch_result_decode (confirmed by decode_skip_finished receipts).
bookkeeping); one pre-existing failure (test_spec_invalid_result_lifecycle)
reproduces identically on the unfixed tree — unrelated.
codespell — all clean; git merge-tree --write-tree vs the integration base
exits 0 with 0 conflicts; pairwise merge-tree vs the heads of
fix/spec-invalid-prefix-lifecycle, metrics/decode-step-gap,
metrics/fp8-range-observations, metrics/generated-token-timing,
fix/mla-hicache-host-dedup-txn, and fix/renorm-deterministic — all clean.
Speed Tests and Profiling
No hot-path cost: the predicate is one short-circuiting evaluation per
prefill result-processing call under lazy extra-buffer spec (pure reads, no
allocation, no synchronization), and early-False otherwise. No serving
performance impact is expected or measured; full combined-tree serving
acceptance remains the campaign follow-up.
Checklist
suppresses every aliased donation.
new patch, not a duplicate.
CI States
Latest PR Test (Base): ❌ Run #36218669347
Latest PR Test (Extra): ❌ Run #36218669275
Latest PR Test (AMD ROCm 10): ❌ Run #36218669322