[yuchenwang3]

Motivation

With Mamba extra_buffer caching and overlap scheduling, the next decode/verify batch can capture req_to_token before the previous prefill result publishes its cache insertion. That insertion can replace freshly computed KV slots with the tree's canonical slots and free the duplicates. The first dependent step then uses a different prefix mapping from later steps.

Processing the result after get_next_batch_to_run() is too late: batch preparation has already captured the mapping.

Related to #38140, which addresses the same ordering issue for plain radix caches and explicitly leaves Mamba cache timing unchanged. This patch covers tracked Mamba prefills.

Modifications

• Process a pending tracked Mamba prefill result before selecting/preparing the next batch.
• Consume each result once and retain the opportunistic unified-memory flush at the synchronization boundary.
• Keep overlap for untracked prefills, disabled caches, and decode requests in mixed batches.
• Leave DP/MLP-sync scheduling unchanged; that path needs a collective decision and is outside this patch's scope.

Latest validation (November 5)

Merged main c03b6b3df into head 1ccea14a1d9d77b959450e55a9da357652d2c001, preserving branch history. The conflict resolution uses upstream's early-result processing path for pending tracked Mamba cache publications and retains its opportunistic unified-memory flush. The regression fixture follows the current scheduler iteration and early-processing interfaces.

All 5 Mamba cache-boundary tests pass in a CPU source-method harness using the actual scheduler loop and predicates. The harness avoids optional server/GPU imports; normal test collection in the local environment stops at missing starlette. All applicable changed-file pre-commit hooks and git diff --check pass. No GPU or server run was performed on this head.

Previous validation (October 3)

Rebased onto main 10edb05e72afb79ecfdf7dbe7603f90ba05ad3ee at 0b8607ec03cb10d6fb0afee433d1a879667ecb90. This includes the required CI minimum base 8854857a991e06b3ef9c4caeae34c94fc3588336.

The Mamba prefill-publication, batch-result boundary, and chunked-abort CPU suites pass all 10 tests, including the existing subtests. All applicable changed-file pre-commit hooks pass. The earlier SM90 server smoke is historical; it has not been rerun on this rebased head.

Previous accuracy tests (pre-rebase revision)

Base: 9c8d4641ff6415a640352a950b98e312479ddf6a.

The new regression exercises the real overlap loop with a controlled cache publisher. On unmodified main it captures the old slots before publication, including when the existing post-selection synchronization is enabled. It also launches a dependent batch before handling a prefill that finishes or fails. These checks pass with the patch.

Source tested before this rebase on an SM90 CUDA GPU with Python 3.12, PyTorch 2.13.0+cu130, Triton 3.7.1, and sglang-kernel 0.4.7 (also passed on PyTorch 2.11 / Triton 3.6):

python -m pytest \
  test/registered/unit/managers/test_mamba_prefill_cache_boundary.py \
  test/registered/unit/managers/test_scheduler_decision_batch_params.py \
  test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py \
  test/registered/unit/managers/test_scheduler_chunked_abort_race.py \
  test/registered/unit/mem_cache/test_mamba_unittest.py -q

27 passed, 16 subtests passed. All applicable pre-commit hooks and git diff --check pass.

Full-server smoke with Accio-Lab/occamy-1.0, BF16, FA3, extra_buffer, TP=1, and overlap enabled: both main and the patch complete cold/warm requests and three fixed-length 128-token requests. The natural-stop case returns the same 9 tokens in all four cold/warm runs, with 384 cached prompt tokens on a hit. This single-request smoke checks for regressions; the controlled scheduler test above reproduces the mapping-publication failure.

Cold/warm logprobs are not bit-exact in this configuration (maximum difference 0.016763 on both main and the patch). This PR does not claim general cold/warm numerical determinism.

Speed Tests and Profiling

Same shared SM90 GPU and server configuration, 422 prompt tokens, 128 generated tokens, one cold-cache request followed by two warm-cache requests:

Request latency	Main	Patch
Cold	0.852 s	0.807 s
Warm 1	0.828 s	0.765 s
Warm 2	0.797 s	0.764 s

These are small-sample smoke timings on a shared GPU, not evidence of a speedup or a concurrent-serving throughput benchmark. The fix synchronizes at tracked Mamba prefill publication; the steady decode overlap path remains unchanged.

Checklist

• Format code with pre-commit.
• Add registered regression tests.
• Follow the SGLang code style.
• Provide focused correctness tests and end-to-end smoke timings, with the limits stated above.

---

CI States

Latest PR Test (Base): ❌ Run #42095571707
Latest PR Test (Extra): ❌ Run #42095571740
Latest PR Test (AMD ROCm 10): ⏳ Run #42095571365

Tip

Base CI only covers the core cases, and its workload may not exercise your change. Use /rerun-test to run the tests your change touches, or copy this prompt for an agent:

Find the registered tests this PR affects and run them with /rerun-test, following https://github.com/sgl-project/sglang/blob/main/.agents/skills/rerun-test/SKILL.md

AI assistance: OpenAI Codex performed this rebase and local validation.

[yuchenwang3]

@merrymercy Could you review the Mamba prefill cache-publication ordering fix, or help route it to the appropriate scheduler owner? The PR includes the controlled regression, 27 passing tests plus 16 subtests, and a full-server SM90 smoke, with the limitations recorded in the description. The CI gate for f273a50b44db1be65bb337dd81e2cddee1db3f10 reports a missing run-ci label; could a maintainer authorize CI if the scope looks suitable?

[yuchenwang3]

@hnyls2002 Could you review the Mamba prefill cache-publication ordering? The fix processes the pending prefill before the next batch captures its KV slots. Rebased onto 10edb05e72 at 0b8607ec03; the three focused CPU suites pass 10 tests, and changed-file pre-commit checks pass. The earlier SM90 smoke was not rerun on this revision. The minimum-base requirement is satisfied; could you authorize CI?

[yuchenwang3]

@xiezhq-hermann Could you take a look at the Mamba prefill cache ordering here? The pending prefill is processed before the next batch captures its KV slots. The current head still needs CI authorization.

[yuchenwang3]

Merged main at c03b6b3df into 1ccea14a1 and resolved the conflict. The Mamba publication check now uses upstream's early-result path, keeping its unified-memory flush. All five cache-boundary tests pass in a CPU source-method harness. Changed-file pre-commit checks pass. Normal local collection is blocked by missing starlette; I did not run a local GPU/server test for this merge. The branch is mergeable with run-ci present. Current-head CI has passed lint, title, the small-GPU job and seven CPU shards; remaining jobs are queued/running. Review is still needed.