[Mamba] Support speculative decoding with extra_buffer_lazy

[YAMY1234]

Motivation

extra_buffer_lazy (#27118) reduces each running request's mamba checkpoint footprint from 2 slots to 1, but was guarded off for speculative decoding ("Follow-up: Support MTP"). This PR implements that follow-up. The freed slots matter because the mamba pool is shared: whatever the running requests do not pin becomes extra RadixCache capacity for cached checkpoints (quantified below).

Modifications

Non-spec lazy allocates the second slot and swaps the write pointer right at the boundary — each decode step advances exactly 1 token, so the boundary is known at prepare time. Under MTP a verify step advances 1..num_draft_tokens accepted tokens and a crossing is only known after verify, so the swap is deferred: prepare a spare slot first, swap only once the crossing is confirmed.

• ScheduleBatch.mamba_lazy_spec_prepare (scheduler phase, before each verify):
  • when the next boundary is reachable this step (within 2×num_draft_tokens), make sure the spare slot exists, without touching mamba_next_track_idx — the slot may end up never written;
  • record this step's write target on the batch; by result-processing time the request's pointer may already have moved under overlap, so it has to be recorded now.
• Post-verify state commit: unchanged — the existing mask-gated scatter writes the crossing state into the spare slot on a real crossing, and writes nothing otherwise (the spare slot is kept and reused next step).
• _mamba_lazy_spec_update (result processing, accept counts now known):
  • running, crossed into the spare slot → promote it: free the old checkpoint, repoint mamba_next_track_idx;
  • running, crossed with no spare slot available → the state was written in place into the checkpoint slot; bookkeeping only;
  • finishing → donate the checkpoint slot only if no scatter has written it (this step's in-place fallback) or may still write it (an in-flight verify with no spare slot inside the crossing window); otherwise skip the RadixCache insert.
• Allocation failures no longer evict (applied to the non-spec boundary prealloc as well): a transient slot is not worth a cached checkpoint — restoring the old evict-on-failure behavior collapses the reuse hit rate back to the extra_buffer column in the benchmark below. On failure the boundary is simply not tracked.
• Slot ownership is unchanged: allocation, holding, free and donation stay request-owned (the existing ping-pong buffer); the batch carries only a read-only per-iteration write plan.
• extra_buffer_lazy still rejects DFLASH/DSPARK (their verify bypasses prepare_mamba_track_for_verify) and PD disaggregation (the decode-side pool is not wired for lazy).

Accuracy Tests

Qwen3-Next-80B-A3B-Instruct, NEXTN (steps=3, topk=1, draft=4), TP4. "Pool" is the mamba state pool (--max-mamba-cache-size), shared by running requests and cached checkpoints.

Config	GSM8K (200q)	decode-cache KL	prefill-cache KL	accept len
lazy	0.960	0.000214	0.000334	2.400
extra_buffer (baseline)	0.965	0.000229	0.000334	2.356
lazy, small pool (68 slots / 16 concurrent reqs)	0.965	n/a*	n/a*	2.426
lazy, every spare-slot allocation forced to fail	0.960	n/a*	n/a*	2.426

KL compares logprobs of a cache-hit generation against a no-cache recompute (existing KL kit). *With the small pool, every cached entry is evicted before the kit's second turn (0/32 hits, identical for the extra_buffer baseline), so there is nothing to compare — the first two rows carry the KL evidence.

Checkpoint-boundary finish test. The risky corner of this feature is a request finishing right at a checkpoint boundary while the overlapped next verify may still write its checkpoint slot — a stale write landing after the slot enters the RadixCache would silently corrupt it. Setup: prompt of 500 tokens, checkpoint interval 64 (checkpoints at 448, 512, ..., 768, 832, ...); requests are forced (exact max_new_tokens) to finish at chosen offsets around the boundary at 832, then we generate a 32-token continuation from the finished sequence twice — reusing the cache vs after flushing it — and compare. The "exhausted" rows force every spare-slot allocation to fail during the test, emulating a pool with no free slot at any boundary (worst case). cached_tokens of the continuation reveals which checkpoint the finished request left in the cache:

finish position	mamba pool	checkpoint served to the continuation	greedy tokens	max |Δlogprob|
−2..+2 around the boundary	has free slots	768 — the previous safe checkpoint; the just-written 832 slot is deliberately not cached (the in-flight verify could still write it)	32/32 identical	3.7e-4
−2..+2 around the boundary	exhausted	448 — no safe decode-phase checkpoint exists, so the insert is skipped and only the last in-prompt checkpoint remains	32/32 identical	3.6e-4
+16 / +32 past the boundary	has free slots	832 — the crossing was confirmed while still running; the checkpoint advanced normally	32/32 identical	2.0e-4
+16 / +32 past the boundary	exhausted	832 — written in place as the fallback; safe to cache because the crossing was already settled	32/32 identical	2.0e-4

The logprob diff is the per-token log-ratio between the two runs — numerical-noise level throughout. The same grid with MTP off also passes, all token-identical (28 cells in total).

Speed Tests and Profiling

What lazy buys is memory. RadixCache reuse benchmark: 48 independent conversations × 3 turns; on turns 2–3 we report the fraction of the prompt served from cache (cached_tokens / prompt_tokens); max_running_requests=16 for both strategies.

mamba pool	extra_buffer_lazy	lazy + evict-on-failure	extra_buffer
unconstrained (1024 slots)	0.9677	—	0.9677
constrained (96 slots)	0.106–0.147 (3 runs)	0.0161	0.0161

With an unconstrained pool the strategies are identical. Under a constrained pool, lazy's 16 running requests pin 32 slots instead of 48, leaving 64 vs 48 slots for cached checkpoints — about 7–9× higher reuse in this setup. The middle column keeps lazy but restores the old evict-on-allocation-failure behavior: the evictions consume exactly the capacity lazy freed, collapsing reuse back to the extra_buffer level — which is why this PR removes the evict (a temporary slot is never worth a cached checkpoint). Accept length is unchanged in every configuration (2.36–2.43, including forced allocation failure): checkpoint tracking never touches the sampling path.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ✅ Run #29730013349
Latest PR Test (Extra): ✅ Run #29730013194

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[hanming-lu]

LGTM as long as CI passes