Gdn ucache backend by ameynaik-hub · Pull Request #49766 · vllm-project/vllm

ameynaik-hub · 2026-07-24T22:30:29Z

[Perf] GDN ReplaySSM spec decode: opt-in fused CuTeDSL "ucache" backend (fp16 state/cache)

Purpose

Add an opt-in second backend for GDN ReplaySSM speculative decode:
gated_delta_rule_mtp_ucache_flush — a single fused CuTeDSL kernel (SM100 /
B200) that executes one GDN spec-decode step for the whole batch:

verify: for every request, replay the delta-rule recurrence from the frozen
SSM checkpoint through the cached ring history (k/u/g), then through the
T=4 speculative tokens, emitting each token's attention output;
flush (same launch, routed per-row by hist_len ≥ flush_min): additionally
fold the ring history into the checkpoint in place.

vs. the Triton baseline (gdn_replayssm_spec_decode), which launches its kernel
twice per layer per step (an IS_FLUSH=False verify pass + an IS_FLUSH=True
flush pass, each gridding the full batch with early-out for the other's rows).

Where the pieces live

piece	location
FlashInfer kernel (REQUIRED)	FlashInfer flashinfer-ai/flashinfer#4081 (`flashinfer-ai/flashinfer`) — must include the ring revision (`RING_SLOTS = 32`); pre-ring kernel builds are rejected at engine init
kernel file the adapter loads	`flashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache_flush.py` (import-self-contained; JIT-compiles via CuTeDSL at first use)
vLLM backend (this PR)	4 commits on top of ReplaySSM PR #47576

How the kernel is wired in

The adapter (vllm/model_executor/layers/fla/ops/gdn_ucache_spec.py) loads the
kernel module at first use, resolved in this order:

VLLM_GDN_UCACHE_MODULE=/abs/path/to/gdn_decode_bf16_wy_ucache_flush.py
else an installed FlashInfer providing flashinfer.gdn_kernels.gdn_decode_bf16_wy_ucache_flush

Availability is validated at engine init (part of resolve_gdn_spec_backend's
checks): a missing module — or a pre-ring kernel build (no RING_SLOTS = 32 in
the module source) — raises a ValueError at startup with fix instructions, not a
mid-serving error. The loader re-verifies the imported module (W_RING == 16,
RING_SLOTS == 32) with a RuntimeError on first use.

Minimal enable:

--use-replayssm-spec --replayssm-buffer-len 16 \
--additional-config '{"gdn_spec_backend": "flashinfer_ucache"}' \
--speculative-config '{"method": "qwen3_next_mtp", "num_speculative_tokens": 3}'
# + VLLM_GDN_UCACHE_MODULE=... if flashinfer.gdn_kernels isn't installed

No dtype flags needed — see Dtypes.

Drop-in with the Triton backend: shared ring, shared cursors

The backend uses the same circular-ring state and cursor model as the Triton
backend — it is a drop-in replacement, switchable per boot with no change to
the cache layout:

Same state layout: per-request 32-slot physical ring
(next_pow2(buffer_len + num_spec) = 32) for u/k/g, byte-identical page
tuple to the Triton backend. The live history window is at most 16 rows
(W_RING), located at [cache_base, cache_base + write_pos) mod 32.
Same cursors, same commit: the block-keyed write_pos / cache_base /
is_flush cursors and the shared commit_gdn_replayssm_spec /
reset_gdn_replayssm_spec_cursors kernels. The builder commits once per
step, outside the launches (rollback = the cursor simply advances by the
accepted count); all GDN layers share one cursor set; the kernel treats
cursors as read-only. The builder gathers both cursors into fixed-address
request-keyed buffers (CUDA-graph safe) and fills pad rows
(hist=0, base=0, idx=−1) every spec step.
Race-free appends by construction: the kernel always appends new tokens
past the live window at (cache_base + write_pos + s) & 31 — for flushing
and non-flushing rows alike. Window (≤16) + appends (≤8) always fit in 32
slots, so a flushing row can never overwrite ring slots that sibling CTAs of
the same request are still reading. No inter-CTA synchronization is needed.

Accuracy

Evals tie (re-validated on the ring revision): MMLU-Pro, HMMT, and GPQA
match the Triton fp32 baseline within noise for every arm (Triton bf16/fp16,
ucache bf16/fp16).
Acceptance length parity: 3.54–3.60 across all arms and batch sizes (any
state corruption would collapse AL).
A teacher-forced state-divergence study showed fp16 checkpoint ≈ fp32 within
mantissa-accumulation noise over 32k-token contexts (fp16's smaller ulp at the
state's dynamic range; the fp32 g running sum is kept).

Dtypes: SSM state, cache, input

Defaults (zero configuration):

tensor	default	alternatives	knob
input IO (q/k/v/a/b activations)	bf16	fp16 (whole-fp16 mode)	`GDN_UCACHE_IO_DTYPE=fp16`
SSM state (checkpoint pool)	fp16	bf16	`--mamba-ssm-cache-dtype` + `GDN_UCACHE_STATE_DTYPE`
u/k cache (ring buffers)	fp16	bf16, fp32	`GDN_UCACHE_RING_DTYPE` (vLLM-side override: `VLLM_REPLAYSSM_RING_DTYPE`)
g cache (cumulative log-decay)	fp32 always	—	(mantissa-sensitive running sum)

fp32 SSM state is NOT supported — the CuTeDSL kernel has no fp32-checkpoint
read path. On this backend --mamba-ssm-cache-dtype auto resolves to fp16
(the Qwen3.5 config updater prefers the kernel state dtype over the HF config's
mamba_ssm_dtype=float32 when this backend is selected).
Consistency by construction: the kernel compiles its dtypes from env at
import and asserts pool/ring dtypes on first call; vLLM allocates pools from the
same envs (the adapter setdefaults both to fp16). Both sides default fp16 →
they agree with no flags.
bf16 mode (explicit): --mamba-ssm-cache-dtype bfloat16 +
GDN_UCACHE_STATE_DTYPE=bf16 GDN_UCACHE_RING_DTYPE=bf16.

Limitations (validated at engine init — violations raise `ValueError`)

constraint	value	why
ring geometry	physical 32 slots (`RING_SLOTS`, matches Triton's pow2 ring), live window ≤ `W_RING = 16` → `--replayssm-buffer-len 16` exactly	window is one 16-row MMA tile; `flush_min = 17 − T` (13 for T=4)
verify window T	`1 + num_speculative_tokens ∈ {4, 8}`	native tile sizes; T=4 is the validated production point
head dims	`linear_key_head_dim == linear_value_head_dim == 128`	tile geometry (Qwen3.5: H=16, HV=64, K=V=128)
SSM state dtype	fp16 or bf16 — no fp32 checkpoint	kernel has no fp32 state read path
hardware	CuTeDSL SM100 (B200); init gate SM90+	tensor-core layouts
batch shape	uniform T-token verify rows (`total_spec == B·T` asserted)	persistent ring format cannot fall back per-step
spec mode	chain MTP only — no tree/Eagle speculation	the kernel replays one contiguous logical window; non-contiguous (tree) accepts would need an index-map read path
config	requires `--use-replayssm-spec`; incompatible with non-spec `--use-replayssm`	reuses the ReplaySSM ring page tuple `[2]=u [3]=k [4]=g`
kernel module	loadable at init AND a ring build (`RING_SLOTS == 32`)	pre-ring kernels use an incompatible flat layout

Test Result

dataset	protocol	records	tr_fp32	uc_fp16	delta	AL (both arms)
MMLU-Pro	1000 q, pass@1	1000	0.8530	0.8630	+1.00 pp	3.38
GPQA-Diamond	198 q, avg@4	792	0.8447	0.8384	−0.63 pp	3.38–3.40
HMMT Feb-2025	30 q, avg@16	480	0.8208	0.8083	−1.25 pp	3.48–3.49
AIME 2024	30 q, avg@16	480	0.8958	0.8958	±0.00 pp	3.49
AIME 2025	30 q, avg@16	480	0.8938	0.9083	+1.46 pp	3.50

AIME accuracy — nvidia/Qwen3.5-122B-A10B-NVFP4 (1×B200, TP1)

Protocol: 30 problems per set, avg@16 (480 records per cell), NVIDIA eval preset (max_tokens=64000, max_model_len=66560, integer-answer judging). Speculative arms use MTP num_speculative_tokens=3 with ReplaySSM; acceptance length matched at 3.49–3.50. Mean
answer length ≈ 20–23K tokens (p90 up to ~43K, ~3% of records reach the 64K cap — the longest-generation evals in the suite).

dataset	standard AR, fp32 state (no spec)	ReplaySSM-spec, Triton fp32 state	ReplaySSM-spec, ucache fp16 state + fp16 cache, bf16 in
AIME 2024	0.8854	0.8958	0.8958
AIME 2025	0.8896	0.8938	0.9083

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

…dard and speculative decode (Mamba2 + GDN) Signed-off-by: Johnny-Liou <a897111@gmail.com>

…decode Signed-off-by: Johnny-Liou <a897111@gmail.com>

…kernel Signed-off-by: Johnny-Liou <a897111@gmail.com>

… dedicated replayssm methods, drop flag params Signed-off-by: Johnny-Liou <a897111@gmail.com>

…-token flushes instead of raising Signed-off-by: Johnny-Liou <a897111@gmail.com>

…kernel Signed-off-by: Johnny-Liou <a897111@gmail.com>