[yuan-luo]

Motivation

Linear-attention models (KDA, GDN / Gated DeltaNet) keep one compact recurrent state per sequence instead of a growing KV cache. SGLang already caches these states in the radix tree (MambaRadixCache) so a shared prefix's state can be restored on a hit instead of recomputed.

The problem is capacity. The active mamba state pool stores one bf16 state per cached prefix, and that pool is sized to the running set, not to the cache. Measured on a linear-attn model: the token-KV pool holds millions of tokens, but the mamba pool saturates at a small number of checkpoints and is length-independent (one full state per cached sequence). As the number of distinct cached prefixes grows past the pool size, prefix reuse collapses — the KV tokens are still resident, but the states get evicted and recomputed. This is the dominant cost for workloads with many distinct short/medium prefixes (multi-doc RAG, many short sessions, agentic fan-out).

Modifications

This PR introduces an optional, dedicated int8 checkpoint pool for the radix-cached linear-attn states, behind --enable-int8-mamba-checkpoint (off by default).

• The radix-cached temporal (SSM) state is stored as int8 with a per-(head, k-channel) symmetric scale (Int8CheckpointStore, already in tree). The conv1d window stays bf16 (tiny).
• Each cached state takes ~1 byte/element instead of 2, so at roughly the same memory the pool holds ~2× more distinct prefixes, moving the reuse-collapse knee out by 2×.

Why int8, and why a separate pool

• Quality-safe by construction. A state is quantized once on store and dequantized once on a hit — it never re-enters the recurrence, so there is no error accumulation.
• int8 beats fp8 for this data: the linear-attn recurrent state is ~uniformly distributed (an accumulated sum of rank-1 updates), so fp8-e4m3 wastes bits on an exponent it doesn't need; int8-per-channel spends all 8 bits on mantissa precision (≈0.5% vs ≈0.7% decode error at 1 byte).
• Separate pool because the recurrence kernels need bf16 working memory — running requests cannot read int8 slots directly. So we keep an active bf16 pool (kernel-facing, sized to the running set) and an int8 checkpoint pool (radix-cached states). A donate quantizes active → int8; a cache hit copy-on-writes dequantize int8 → active.

Strategy-agnostic donate hook (spec-KDA-ready)

Both scheduler strategies converge on the same point — "an active slot becomes the cached mamba_value": no_buffer via a copy, extra_buffer via the ping-pong track buffer (the spec path). The quantize hook lives at that common convergence point, so a single implementation covers both, including the future spec path.

Design

Today (baseline) the radix-cached linear-attn states live in the same bf16 MambaPool as the running requests: a cached node's mamba_value is just a slot in that one pool, and on a hit it is copy_from'd bf16→bf16 into a fresh slot. So the cache and the working set compete for one pool, and every cached prefix costs a full bf16 state — the pool saturates at a small number of checkpoints and prefix reuse collapses. This PR splits them: a small bf16 working pool (running set, kernel-facing) + a separate, compact int8 checkpoint pool (the cache). The expensive, growing part — the cache — is the one that gets compressed.

	baseline (shared bf16 pool)	this PR (bf16 working + int8 cache)
pools	one bf16 pool (running + cache)	small bf16 working pool + separate int8 cache pool
where cached states live	in the bf16 pool (mamba_value = bf16 slot)	in the int8 pool (mamba_value = int8 slot)
bytes per cached prefix	a full bf16 state	int8 (≈ half)
cache capacity	pool_size − running_set, bf16-expensive → small	independent int8 pool → ~2× at the same bytes
cache hit	copy_from bf16→bf16 (no quant)	dequantize int8→bf16 (COW, forward stream)
working-pool size	grows with the cache (coupled)	stays small (sized to concurrency)

Advantages of the int8 design:

• ~2× more distinct prefixes cached at the same memory — directly lifts the reuse-collapse knee (256 → 512 in the A/B below), so high-distinct-prefix workloads (multi-doc RAG, many concurrent sessions, agentic fan-out) keep hitting the cache instead of recomputing the state.
• Cache size is decoupled from the working set. The baseline forces one pool to hold both, so a bigger cache means a bigger bf16 pool — which also drags the KV-cache budget down. Here the bf16 pool stays small (sized to concurrency) while the int8 cache grows independently, giving a big cache and a high KV budget (see the memory table) — something the single-pool design cannot do.
• Spend the 2× either way — as more cached prefixes, or as the same cache at ~half the mamba HBM (freeing HBM for KV → longer context / higher concurrency).
• Quality-safe and isolated — quantize-once-on-store / dequantize-once-on-hit (the state never re-enters the recurrence; within the serving stack's own noise on GSM8K), off by default, and the baseline bf16 path is byte-for-byte unchanged when disabled.

Flags

• --enable-int8-mamba-checkpoint (bool, default off): store radix-cached linear-attn states as int8.
• --int8-mamba-ckpt-size (int, optional): number of int8 checkpoint slots; defaults to 2 × max-mamba-cache-size.

Validation

Validated on Qwen3.5-35B-A3B (GDN-hybrid) on 4×H20-3e. The int8 mem-cache path is GDN/KDA-agnostic (both flow through the same MambaPool + MambaRadixCache + chunk_gated_delta_rule path), so this validates the integration model-independently.

Capacity + latency/throughput — A/B

Via the committed benchmark benchmark/bench_linear_attention/bench_int8_checkpoint_reuse.py (Qwen3.5-35B-A3B, extra_buffer, --max-mamba-cache-size 256 → active 256 / int8 512, prefix ~1000 tokens, parallel 8), run against the server with vs. without --enable-int8-mamba-checkpoint:

K (distinct prefixes)	reuse_frac (bf16 → int8)	probe p50 ms (bf16 → int8)	probe throughput rps (bf16 → int8)
128 (both fit)	0.940 → 0.940	149.9 → 152.4	46.9 → 50.6
384 (bf16 collapsed)	0.020 → 0.940	204.1 → 150.5	36.9 → 49.0
640 (both exceed pool)	0.000 → 0.048	204.1 → 207.5	37.7 → 36.7

The reuse-collapse knee moves 256 → 512 (≈2×). In the band 256 < K ≤ 512 — where the baseline has collapsed but int8 still serves prefixes from cache — int8 delivers ~47× more prefix reuse, ~26% lower probe-prefill latency, and ~33% higher throughput (K=384), because the probe is a cheap cache hit instead of a full prefix recompute. Below the bf16 pool size (K=128) int8 is on par (no overhead); above the int8 pool size (K=640) both collapse, as expected. 0 errors / crashes / leaks in both runs.

Memory efficiency — same cache, less mamba HBM, KV decoupled

The other way to spend the 2× density: cache the same number of distinct prefixes using far less mamba HBM, and keep the KV budget high. int8 decouples the cache from the active pool, so a small active pool drives the KV budget while a large int8 pool holds the cache (Qwen3.5-35B-A3B, per card):

config	cache cap	mamba HBM	max_total_num_tokens
bf16 single pool 1280	~1280	ssm 18.76 + conv 0.44 = 19.2 GB	8.32M
int8 (256 active + 1024 ckpt)	~1024	3.85 (active) + 4.16 (int8) = 8.0 GB	9.93M

bf16 forces a trade-off — a big cache needs a big pool, which drops the KV budget (8.32M). int8 gives a big cache and the high KV budget of a small pool (9.93M), caching ~1024 prefixes at ~4 MB each vs bf16's ~15 MB. Caveat (follow-up): the int8 pool's HBM is not yet subtracted from the KV-budget computation (int8's max_total equals the 256-pool value; the pool currently rides on GPU headroom). It should be accounted; even when accounted the mamba footprint (8 vs 19.2 GB) wins.

Quality / precision — does int8 degrade output?

Three levels of evidence, each isolating the int8 effect from the serving stack's noise. Headline: no systematic task-accuracy degradation; the perturbation is small, bounded, and does not compound with cache depth.

1. Mechanism (why it's bounded). The state is quantized once on store and dequantized once on a hit — it never re-enters the recurrence as a quant→dequant loop. On chained cache extensions the dequantized state does continue the recurrence, but the gated-delta rule is contractive (per-channel decay α<1 plus a negative-semidefinite delta-rule Jacobian), so an injected quantization error is attenuated geometrically. Measured on a faithful fp32 gated-delta recurrence driven through the real codec: the accumulated state error saturates (does not grow with depth) — ~6.6e-3 at 1 checkpoint → ~1.3e-2 by 32 chained checkpoints → flat through 128 (seq 2048). The per-(head,k-channel) scale matches the per-k-channel decay, so the dominant large state entries keep ~bf16 precision and the error concentrates on small entries that barely affect the readout o = Sᵀq.

2. Logit-level isolation (real model, deterministic). On prefix-extension cache hits where int8 dequant actually fires (Qwen3.5-35B-A3B, TP1 greedy → run-to-run deterministic, so the noise floor is exactly 0), the next-token distribution shifts by KL ≈ 6e-3 (max ~0.5 nats on a tail token), with greedy top-1 unchanged in every probe. For reference, loading a bf16 cached state instead of recomputing is itself not bit-exact (KL ≈ 5e-4); int8 is ~10× that already-accepted baseline and, crucially, never diverges from the bf16 path.

3. Task accuracy (GSM8K, int8 dequant confirmed firing). Qwen3.5-35B-A3B, 200 questions, 5-shot, greedy, extra_buffer, with the int8 dequant path instrumented and confirmed firing 599× (≈ every question × 3 runs). WARM = cache reused from the int8 pool; COLD = flush_cache per question = full bf16 recompute (deterministic baseline); a bf16-reuse server is the control.

config	COLD (bf16, deterministic)	WARM accuracy band (3 runs)	WARM run-to-run agreement
int8 (dequant ×599)	0.805	0.832 [0.830–0.835]	0.86
bf16 (control)	0.805	0.842 [0.815–0.855]	0.89

int8's WARM band sits inside bf16's; the ~1pp mean gap is well within the reuse path's run-to-run churn. Both WARM means exceed the COLD baseline (0.805), which confirms the warm-vs-cold spread is trajectory-shuffling noise, not a quality change — quantization cannot raise accuracy. The cache-reuse path is itself non-deterministic under overlap scheduling (even bf16 warm-vs-warm agrees on only ~89% of answers, and the reused-token count varies run-to-run), which is why accuracy is reported as a band over a deterministic cold baseline rather than a single number.

Methodology note (important for reproducing). int8 is only exercised when the cached prefix is genuinely reused from the int8 pool, which requires --mamba-scheduler-strategy extra_buffer: no_buffer only snapshots state at the full-sequence leaf, so a fixed-prefix / divergent-question workload (few-shot GSM8K, or the divergent-suffix benchmark) reuses ~0 mamba state and int8 dequant almost never fires (measured: 2 dequant calls in a 50-question no_buffer run vs 599 with extra_buffer). Any int8 quality claim must instrument the dequant count to confirm the path is actually live — a naive GSM8K run under the default no_buffer measures bf16, not int8.

Correctness — both scheduler strategies

• extra_buffer: store / match / int8-dequant-COW correct; reuse stable across repeated probes.
• no_buffer: int8 store/match/COW also correct (an appended-suffix probe hits cleanly, cached 663/666). NOTE: no_buffer shows low prefix reuse on a divergent-suffix workload by strategy design (it only snapshots state at the full-sequence leaf, so a divergent suffix lands on a _split_node with mamba_value=None) — this is unchanged from baseline no_buffer, not specific to int8. Mid-sequence prefix reuse requires extra_buffer's track-buffer.

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.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ✅ Run #27668615858
Latest PR Test (Extra): ✅ Run #28150073323

[gemini-code-assist]

Code Review

This pull request introduces an int8-compressed checkpoint pool (via Int8CheckpointStore and MambaCheckpointPool) for caching linear-attention (Mamba) recurrent states, which effectively doubles the cached-prefix capacity at fixed memory. It integrates this pool with the built-in MambaRadixCache, updates the scheduler's invariant checking and pool statistics to handle the two-pool setup, and adds corresponding CLI arguments and unit tests. The review feedback is highly constructive, suggesting improvements for device placement of index tensors, performing quantization math in float32 to avoid precision loss, using .view(-1) for robust tensor flattening, making memory calculations more explicit, and gracefully skipping CUDA tests if the fla package is missing.

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.

[merrymercy]

This folder test/srt is deprecated. please never add files under this folder.
if the test manual, move it under test/manual

[yuan-luo]

Addressed in #28813.

[merrymercy]

this should go to extra, we do not want to waste 400 s for this feature.

[yuan-luo]

Addressed in #28813.