Skip to content

[Fix] deterministic: replay hybrid Mamba prefix instead of restoring donated state - #48425

Open
ListenJ wants to merge 2 commits into
sgl-project:mainfrom
ListenJ:fix/deterministic-hybrid-mamba-cache
Open

ListenJ wants to merge 2 commits into
sgl-project:mainfrom
ListenJ:fix/deterministic-hybrid-mamba-cache

Conversation

@ListenJ

@ListenJ ListenJ commented Nov 4, 2026 •

Copy link
Copy Markdown

[Fix] deterministic: replay hybrid Mamba prefix instead of restoring donated state

Motivation

Hybrid linear-attention (Gated DeltaNet) models return different greedy logprobs
and tokens for the same prompt depending on whether it was served cold or from
the radix cache. Reproduced on SGLang 0.5.23 with Qwen/Qwen3.5-4B (hybrid GDN,
24 GDN layers + 8 full-attention layers) on an RTX 4090:

  • 5 prompts x (1 cold + 4 radix hits), temperature=0, seed=0, max_tokens=192, logprobs=true
  • default (radix on): 2/5 prompts token-divergent cold-vs-hit (positions 60/113), logprob maxdiff 0.0642–1.3649; hit repeats mutually identical; two independent runs identical
  • --disable-radix-cache: 0/5 divergent, logprob maxdiff 0.0
  • --enable-deterministic-inference (radix on): still 3/5 divergent (positions 22/45/142), maxdiff 0.0691–1.3964
  • vLLM with the same model/prompts is identical (negative control); upstream #47213 reports the same class on Qwen3.5-9B

Root cause (consistent with #47213 / #41351 / #33515): a radix hit can restore a
Mamba checkpoint that was produced by an earlier, differently shaped prefill
(decode-time tracking), and the persisted state went through a BF16 boundary while
the cold path kept FP32 internally. The restored state is not numerically
interchangeable with a cold prefill, so greedy logprobs/tokens drift even under
--enable-deterministic-inference. #46351 / FlashInfer #2830 fix the separate
FA2 tile/batch-composition issue, not this recurrent-state path (maintainer note
in #46351: "hybrid GDN (Qwen3.5) drift ... is a separate path, covered by neither fix").

Change

In arg_groups/overrides.py::_mamba_radix_cache_resolution, when the strategy is
left on auto, the model supports the Mamba extra_buffer radix cache, and
deterministic inference is enabled:

declared["mamba_radix_cache_strategy"] = "no_buffer"
declared["disable_overlap_schedule"] = True
return declared

no_buffer replays the recurrent computation instead of restoring donated Mamba
state, so a cache hit follows the same recurrent path as a cold prefill for the
hybrid state (no_buffer requires overlap scheduling off). Attention KV cache
handling is unchanged.

If the user explicitly opts into --mamba-radix-cache-strategy extra_buffer
together with --enable-deterministic-inference, the choice is kept but a loud
warning explains that cold/cache-hit logprobs can differ (#47213), instead of
silently breaking the deterministic contract.

Behavior / cost

  • Only active with --enable-deterministic-inference; default mode unchanged.
  • Deterministic + hybrid Mamba models trade Mamba prefix-state reuse for
    cold/hit equivalence (cost: prefix reuse for the recurrent state).
  • Explicit strategies are not overridden; incompatible combinations still fail
    fast through the existing validate_mamba_no_buffer (page_size=1, no overlap,
    no trtllm_mha).

Tests

  • New test/registered/unit/managers/test_mamba_radix_cache_deterministic.py:
    deterministic auto -> no_buffer + overlap off; non-deterministic auto ->
    extra_buffer; explicit strategy preserved; explicit extra_buffer warns under
    deterministic; --disable-radix-cache takes precedence. 5/5 pass locally.
  • Patch applies cleanly to current main (patch -p1 --dry-run, offset -9).

Validation (pre-PR matrix, RTX 4090 24GB / driver 610.57.04 / CUDA 13.4 / SGLang 0.5.23)

20 prompts (needle/math/code/JSON/translation/table/compare/error-check), greedy
(temperature=0, seed=0, max_tokens=192, logprobs=true), 1 cold + 4 cache-hit
repeats per prompt:

configuration token divergence cold-vs-hit logprob maxdiff
baseline, deterministic 3/20 1.647
patched, deterministic 0/20 0.0
patched, deterministic (independent server restart) 0/20 0.0
patched, deterministic, triton attention backend 0/20 0.0
patched, deterministic, --disable-radix-cache control 0/20 0.0
  • patched cold outputs are character-identical to baseline cold outputs (quality unchanged)
  • functional check (thinking disabled, longer outputs): needle/math/JSON-style answers all correct
  • Qwen3.5-0.8B: baseline 2/20 -> patched 0/20 (maxdiff 2.497 -> 0.0)
  • local 4GB box (0.8B): baseline 1/10 -> patched 0/10
  • latency cost of replaying the prefix under deterministic mode in this benchmark: ~3-4%

Default (non-deterministic) mode is intentionally out of scope here and still
shows drift (consistent with #41351); it needs the deeper recurrent-state fix
(#47213/#33515 direction), not this deterministic-contract workaround.
Concurrency and long-context variants are left for follow-up testing.


CI States

Latest PR Test (Base): ❌ Run #41874359113
Latest PR Test (Extra): ❌ Run #41874359261
Latest PR Test (AMD ROCm 10): ❌ Run #41874358649

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

@ListenJ
ListenJ marked this pull request as ready for review November 4, 2026 16:59

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@github-actions

This comment has been minimized.

@ListenJ

ListenJ commented Nov 6, 2026

Copy link
Copy Markdown
Author

CI status note for reviewers:

  • lint is green on the latest commit (f8713ba).
  • The remaining red checks are the fork-PR gate: pr-gate fails with "Missing required label 'run-ci'", and the test jobs were skipped rather than failed (e.g. run 41874359113). The bot comment referencing 13e0b0e predates the lint fix and is stale.
  • Since the last push we re-checked against current main (7a82e41): python/sglang/srt/arg_groups/overrides.py has not changed since 3ead93f7, so this patch still applies cleanly; the new unit test passes locally (5/5, including via test/run_suite.py --hw cpu --suite base-a-test-cpu).

Could a maintainer add the run-ci label when convenient so the test matrix can run? Happy to address anything reviewers find.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant