[AMD][Bugfix] Mamba radix cache: avoid the degenerate single-token (M=1) recompute that corrupts the first token

[ntgiang71096]

Motivation

On a hybrid Gated-DeltaNet + MoE model (e.g. Qwen3.5/3.6) with the mamba radix prefix cache
(--mamba-scheduler-strategy extra_buffer), a cache hit on a prompt whose total length is
≡ 1 (mod 64) produces a wrong, silent first generated token.

Mechanism. The mamba radix cache stores the GDN recurrent state only at periodic checkpoints
(every chunk = 64 tokens), so checkpoints land at multiples of 64. On a hit, the engine restores the
nearest checkpoint and recomputes the remaining "tail" — i.e. the number of tokens past that
checkpoint (tail = extend_input_len, the recompute length). When the prompt length is ≡ 1 (mod 64),
the nearest checkpoint sits exactly one token back, so the tail is 1 (extend_input_len == 1).

Root cause (localized; the corrupting kernel is upstream aiter). That single-token (M=1) recompute
issues an M=1 MoE forward, which on gfx950 selects MoE tile block_m=16; the aiter ck_moe_stage1
GEMM returns NaN at that tile for a finite fp8 input (the 32/64/128 tiles are correct) → flat
logits → garbage first token. The GDN/mamba recurrence is bit-exact correct (verified). The
precise reason inside CK is still under investigation — it reproduces only under the live server's
warmed state, not in an isolated replay (see Related issues) — but the trigger and the fix are
verified. This PR is the correctness-safe sglang-side mitigation: it removes the tail == 1
trigger so the engine never issues the M=1 MoE, robust even on an un-fixed aiter.

The defect is the extra_buffer mamba-radix path. Every tail == 1 cache hit is corrupted —
24/24 (100%) of templated re-sends across six diverse prompts — while a no_buffer baseline
(which does not cache mamba state) never reaches the path and never corrupts.

The fact of corruption is reliable; the specific garbage text is not bit-stable run-to-run,
even at temperature 0. The aiter MoE reduction is itself non-deterministic, and the corrupted
state has near-tied (flat) top logits, so the wrong argmax flips between forwards; a clean state stays
stable because its logit gap is large. The clean, low-noise probe is therefore the first-token
argmax (quantified under Accuracy Tests).

Why it's silent / impact. Standard evals grade a final answer after a long generation that
re-grounds on the correct prompt, and the trigger is rare (~1.5% of prompt lengths), so the bug stays
latent. It bites hardest where the first token is load-bearing and there is no recovery step:

• Determinism contract broken: a cached/retried temp-0 re-send of a prompt of length ≡ 1 (mod 64)
  (65, 129, 193, …) returns a different first token than the original call (breaks response caches,
  "regenerate = same answer", eval reproducibility).
• Single-token outputs (classifier/router/yes-no gate, max_tokens=1): the garbage token is the
  answer — no recovery.
• First-token-logprob uses (logprob reranking, MCQ A/B/C/D): the first-token distribution is
  corrupted (measured — see Accuracy Tests).
• Structured output / tool calling (first token is {/[/a tool name): garbage → parse failure
  (plausible consequence of the measured first-token corruption; not separately harness-tested).

A concrete, replicable example (AMD MI350X, Qwen3.6-35B-A3B-FP8, TP1, temp 0, thinking off).
This is the entire request — a normal 193-token prompt (= 3×64 + 1, so ≡ 1 mod 64):

system: You are a precise assistant. Answer the user's question with a single word.

user: I am double-checking a short geography quiz before I submit it. So far I have reviewed
quiz item numbers 1 2 3 4 5 … 38 39. The one fact I still want to confirm is this question: what
is the capital of France? Reply with only the city name, nothing else.

→ templated prompt length = 193 tokens (≡ 1 mod 64 — the trigger)

Send it once (cache miss), then send the identical request again. The second call is a cache
hit reporting cached_tokens = 192 (the nearest checkpoint, 3×64) — so 192 of 193 tokens are
restored and exactly one token (the tail) is recomputed:

call	cache	cached_tokens	recompute tail	output
1st — fresh	miss	0	(full prefill)	Paris ✓
2nd…5th — identical re-send	HIT	192	1	niestety, nie mogę odpowiedzieć… · \/\/\/\/… · svp, please provide the question… · ú — 4/4 wrong, each different

Same request, same temp 0 — the only difference is that the second call was served from the cache. The
corruption is reliable (every re-send wrong) yet the garbage differs each time (temp-0 greedy over a
flat-logit state + the non-deterministic MoE — not a sampling effect). With longer generations the
recurrent-state corruption often carries forward through the whole response, not just the first token.

Modifications

One file, MambaRadixCache._match_post_processor
(python/sglang/srt/mem_cache/mamba_radix_cache.py). Before the copy-on-write (COW) state-restore
logic: when a hit would leave a recompute tail < 2, retreat the matched prefix to the nearest
earlier checkpoint ancestor (tail becomes chunk + 1 ≥ 2); if there is no earlier checkpoint, fall
back to a fresh compute (prefix 0). The KV prefix (value[:best_value_len]) and the COW restore source
(last_node.mamba_value) move to the same boundary in lock-step, so KV and mamba state stay
consistent.

MIN_SAFE_GDN_EXTEND = 2  # tail==1 is the degenerate single-token (M=1) extend
# Retreat to the previous GDN checkpoint so the recompute is always >=2 tokens (never M=1), moving the
# KV prefix and the COW restore source to the same boundary in lock-step. Root fix: upstream aiter.
if cow_mamba and req is not None and best_value_len > 0 and last_node.mamba_value is not None:
    fill = getattr(req, "full_untruncated_fill_ids", None)
    # empty array("q") (e.g. a PD-disagg decode match before fill_ids is populated) is falsy ->
    # input_len=None -> we cannot determine the tail -> do NOT retreat (else we'd discard real hits).
    input_len = len(fill) if fill else None
    if input_len is None and hasattr(req, "get_fill_ids"):
        gf = req.get_fill_ids()
        input_len = len(gf) if gf else None
    if input_len is not None and (
        input_len - sum(len(v) for v in value[:best_value_len]) < MIN_SAFE_GDN_EXTEND
    ):
        retreat_node, retreat_len = last_node, best_value_len
        while retreat_len > 0:
            retreat_node = retreat_node.parent
            retreat_len -= 1
            if retreat_node is not None and retreat_node.mamba_value is not None:
                break
        if retreat_len > 0 and retreat_node is not None and retreat_node.mamba_value is not None:
            last_node, best_value_len = retreat_node, retreat_len   # tail becomes chunk+1 (>=2)
        else:
            last_node, best_value_len = self.root_node, 0           # no earlier ckpt -> fresh

Design notes:

• Gated on cow_mamba, not on enable_mamba_extra_buffer — i.e. the same condition that guards
  the COW state-restore later in this method. Any strategy that performs the mamba COW restore is
  covered (e.g. int8); a strategy that doesn't set cow_mamba (e.g. no_buffer, which doesn't cache
  mamba state) never enters the guard, so it is inert there. Correctness-safe in all cases.
• Bails when input length is unknown (empty full_untruncated_fill_ids, e.g. a PD-disaggregated
  decode match before fill_ids is populated). Treating it as 0 would spuriously retreat every decode
  hit and discard real cache hits.
• Lock-step KV + mamba: both downstream uses derive from (last_node, best_value_len) — the KV
  slice and the COW source — so they can never desync.
• Backends / scope. The corruption is confirmed only on AMD gfx950 + aiter (cause: the aiter
  ck_moe_stage1 block_m=16 GEMM, a ROCm-only Composable Kernel). The trigger (tail == 1) is
  backend-independent — it comes from the generic extra_buffer cache — and whether other backends'
  MoE kernels also mishandle the degenerate M=1 shape is untested (we do not claim it is safe
  there). The retreat is therefore applied unconditionally on the mamba-COW path, not arch-gated:
  it is a routing-only change, always correct on any backend (worst case ≤ chunk extra recompute
  tokens on ~1.5% of full re-sends), and avoids a degenerate single-token shape regardless of vendor.
  If maintainers prefer to limit the change to the only confirmed-affected backend, gating the retreat
  to ROCm is a one-line change.

Accuracy Tests

Primary gate — deterministic, by construction. The fix changes only a routing decision (which
checkpoint to restore from), based on lengths. The test asserts the structural property: a
L ≡ 1 mod 64 re-send no longer produces a 1-token extend (extend_input_len ≥ 2, or a fresh
cached=0), with a tail = 2 control that must still hit the cache (guards against
over-triggering). Verified on current main at L=705 (705 = 11×64 + 1 → tail=1, was garbage) and
L=706 (= 11×64 + 2 → tail=2, correct).

Why the fix is correct (decisive isolation). Build L ∈ {193,194,195,196} sharing the same
checkpoint@192 (cached=192 for all): tail=1 (193) is corrupt; tail=2/3/4 (194/195/196) are
correct. So the stored checkpoint is correct — only the single-token recompute is wrong, and for
tail ≥ 2 the later recomputed tokens re-ground. The clean correctness probe (per Motivation) is the
first-token argmax: on tail=1 hits it matched fresh 0/40 against a 40/40 fresh-vs-fresh
floor before the fix; the retreat removes the path.

No regression. For tail ≥ 2, fresh, and decode the code path is unchanged by construction
(the retreat only fires on tail < 2), so those cases are unaffected. Corroborating: GSM8K (n=40) =
0.95 unchanged, radix cache hits 38/40; cache-HIT vs fresh next-token distributions are within the
MoE non-determinism floor (Jensen-Shannon over top-20 logprobs, temp 0).

Reproduce:

# Server — AMD MI350X / gfx950, Qwen3.6-35B-A3B-FP8, TP1:
python -m sglang.launch_server --model-path <Qwen3.6-35B-A3B-FP8> --tp 1 \
    --mamba-scheduler-strategy extra_buffer --page-size 64 --trust-remote-code
# Client — send the 193-token prompt from Motivation twice (temp 0, enable_thinking=false):
#   1st call  -> cache miss -> "Paris"
#   2nd call  -> cache HIT (cached_tokens=192, tail=1) -> garbage   [bug; fixed by this PR]
# Automated: tests/.../test_qwen35_radix_tail1_fidelity_mi35x.py

Scope of what was executed: the only backend exercised is AMD MI350X (gfx950 / ROCm / aiter) —
validated on Qwen3.6-35B-A3B-FP8, TP1, current main (d282bdd8f1); the fail/pass split is verified
at L=705/L=706. A broader tail sweep (tail=1 fails; tail 2..63 and ≥64 pass) was run on an earlier
container image; its results are consistent with the current-main L=705/L=706 verification but the
full sweep was not re-run on current main. The 397B / TP8 lane is not yet executed, and no
non-AMD backend was tested (the fix is applied generically — see the Backends / scope note above).

Speed Tests and Profiling

Negligible and bounded to the rare trigger: ≤ chunk (≤64) extra recompute tokens only on a full
re-send whose length is ≡ 1 mod chunk (~1/64 of lengths); a single-node prompt at that length falls
back to a full recompute for that one request. Zero effect on tail ≥ 2 hits or normal serving.
No kernel changes.

Checklist

• Format with pre-commit (all hooks pass; black reformatted the 3 changed files)
• Unit tests added (test_tail1_retreat_logic_cpu.py — CPU retreat-logic;
  test_qwen35_radix_tail1_fidelity_mi35x.py — E2E structural gate)
• Accuracy results provided (no-regression; deterministic correctness gate)
• Documentation (n/a — internal cache routing fix)

Related issues

The kernel root cause is filed upstream as ROCm/aiter#4032 — the ck_moe_stage1 block_m=16 GEMM
returns NaN at M=1 on gfx950 (clean fp8 input → NaN output). Confirmed in the live server: clamping
block_size_M ≥ 32 in aiter/fused_moe.py drives the tail=1 corruption to 0%, identifying the tile
as the cause. The why-inside-CK is still open (it reproduces only under the live server's warmed state,
not in an isolated replay), which is why this sglang-side trigger removal is the verifiable, shippable
change.

---

CI States

Latest PR Test (Base): ❌ Run #29419886515
Latest PR Test (Extra): ❌ Run #29419886372

[gemini-code-assist]

Code Review

This pull request introduces a retreat mechanism in the Mamba radix cache to prevent single-token (M=1) MoE forwards, which can produce NaNs on gfx950 GPUs. It also adds CPU-based unit tests and GPU-based integration tests to verify this logic. The reviewer identified a critical performance regression where perfect cache hits (tail == 0) would trigger an unnecessary retreat and recomputation. To resolve this, the reviewer suggested updating the retreat condition to only trigger when the tail is exactly 1, and provided corresponding updates for both the CPU and GPU test suites.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[ntgiang71096]

Hi @merrymercy @hnyls2002 this is ready for review whenever you have a moment. Could you kick off CI (/tag-and-rerun-ci)? Thank you!

[HaiShaw]

/tag-and-rerun-ci

[ntgiang71096]

Hi @hanming-lu @hnyls2002 — sorry to ping directly, but this one's green now and should be a quick review if you have a moment.

Small (+44, routing-only) correctness fix in mamba_radix_cache.py (_match_post_processor): on a radix cache hit where the prompt length is ≡ 1 (mod 64), the 1-token recompute (extend_input_len == 1) issues an M=1 MoE that returns NaN on gfx950 → silent wrong first token. The fix just retreats one checkpoint so the recompute is never M=1. The GDN recurrence itself is bit-exact (verified), and it's validated 24/24 on MI350X. Tests added (CPU decision-logic + MI35x fidelity), and build-test + base-a-test-cpu are green.

cc @fzyzcjy — you wrote _match_post_processor, so you'd catch anything off faster than anyone 🙂 It's the companion to the extra_buffer enablement in #30359 (just merged). Thanks!

[hnyls2002]

Thanks @ntgiang71096! The file this PR patches, python/sglang/srt/mem_cache/mamba_radix_cache.py, was removed from main in #40313 ("Remove swa and mamba radix cache"), so the patch no longer applies. The underlying M=1 gfx950 MoE tile issue may still be relevant against the new unified-cache implementation — if so, please reopen against the new tree and we can take another look.