Repository navigation
Gdn ucache backend - #50380
Draft
ameynaik-hub wants to merge 23 commits into
Draft
Gdn ucache backend#50380ameynaik-hub wants to merge 23 commits into
ameynaik-hub wants to merge 23 commits into
Conversation
…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>
Signed-off-by: Johnny-Liou <a897111@gmail.com>
Signed-off-by: Johnny-Liou <a897111@gmail.com>
… state reconstruction Signed-off-by: Johnny-Liou <a897111@gmail.com>
Signed-off-by: Johnny-Liou <a897111@gmail.com>
…prefill instead of raising Signed-off-by: Johnny-Liou <a897111@gmail.com>
… x|B page width Signed-off-by: Johnny-Liou <a897111@gmail.com>
…ax_query_len Signed-off-by: Johnny-Liou <a897111@gmail.com>
Signed-off-by: Johnny-Liou <a897111@gmail.com>
…zed_layers` configs. (vllm-project#47318) Signed-off-by: Daniel Afrimi <dafrimi@nvidia.com> Signed-off-by: <dafrimi@nvidia.com> (cherry picked from commit 0a2965b)
Signed-off-by: Johnny-Liou <a897111@gmail.com>
…ugh the spec kernel Signed-off-by: Johnny-Liou <a897111@gmail.com>
…ltiples Signed-off-by: Johnny-Liou <a897111@gmail.com>
Signed-off-by: Johnny-Liou <a897111@gmail.com>
Opt-in `flashinfer_ucache` GDN cached-spec backend (single fused CuTeDSL verify+flush kernel) alongside PR vllm-project#47576's Triton path, selected via additional_config gdn_spec_backend. Default precision on this backend: fp16 SSM-state checkpoint + fp16 u/k ring caches with bf16 input IO. --mamba-ssm-cache-dtype auto resolves to fp16 (the Qwen3.5 config updater prefers the kernel state dtype over the HF-config mamba_ssm_dtype=float32, which the ucache kernel does not read), and the adapter setdefaults the kernel module's GDN_UCACHE_STATE/RING_DTYPE envs to fp16, so pool allocation and the compiled kernel dtypes agree with no flags (set the envs + an explicit dtype for bf16 mode). Backend selection validates its requirements at engine init, including that the ucache kernel module is actually loadable (VLLM_GDN_UCACHE_MODULE path or flashinfer.gdn_kernels) — misconfiguration fails loudly at init instead of a raw ImportError on the first spec-decode step. - gdn_ucache_spec.py: adapter (kernel module load, static max_num_seqs padding with pad-skip -1 sentinel rows, eager block-keyed hist_len commit kernel) - gdn_attn.py: ucache cursor metadata (hist_len gather + col0 state indices, fixed-address buffers, graph-replay pad-row fills) - mamba_utils.py: resolve_gdn_spec_backend + init checks, linear ring_slots=16 override (kernel W_RING), ring/ckpt dtype resolution - models/config.py: Qwen3.5 updater resolves 'auto' ckpt dtype to the ucache kernel state dtype when the backend is selected - qwen_gdn_linear_attn.py: backend dispatch - base.py / qwen3_5.py: thread vllm_config into the spec dtype calc
Micro-optimizations on the fused launch path, each removing a measured per-layer-per-step overhead: - env-gated zero-copy padding (VLLM_GDN_UCACHE_ZEROCOPY_PAD=1): as_strided bucket views instead of qkv/a/b staging copies (pad rows never dereferenced thanks to the -1 pad-skip sentinel) - bucket-aware padding: pad_to = the step's CUDA-graph bucket (spec_padded_rows) instead of static max_num_seqs; builder passes bucket-length PRE-PADDED hist/idx slices so the adapter skips its per-layer staging copies+fills - in-place spec output: the layer hands its core_attn_out slice as the kernel's direct STG target, skipping the per-layer slice-assign DtoD - strided a/b: pass chunk views directly (wrapper strided-a/b mode) instead of two per-layer .contiguous() copies - stable detached A_log/dt_bias: detach once per parameter so the wrapper's identity-keyed fp32->bf16 cast cache hits (was 2 cast kernels per layer per step) tests/kernels/test_replayssm_ucache_spec_gdn.py: kernel test (block- strided pool parity, 64-bit addressing past 2GiB, pad-skip rows).
Pairs with the FlashInfer ring kernel (ameyn/gdn-ucache-ring): the ucache backend now shares the Triton backend's ring layout and cursors end to end. - get_state_shape: drop the ring_slots=16 override — both backends allocate the identical pow2 ring (next_pow2(16+3) = 32 slots) - gdn_attn builder: ucache branch reuses commit_gdn_replayssm_spec / reset_gdn_replayssm_spec_cursors (block-keyed write_pos/cache_base, same flush cadence: the commit's arming rule wp+2T>L coincides with the kernel's P>=flush_min), then gathers write_pos AND cache_base into request-keyed fixed-address buffers for the captured kernel; the hist_len master buffer and its commit kernel are no longer used - adapter: cache_base flows through padding/staging like hist_len and into the kernel call; init assert requires a ring kernel build (mod.RING_SLOTS == 32) so pre-ring kernels fail loudly - commit_gdn_ucache_hist is retained only for the legacy kernel test and is slated for removal with that test's ring update Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…e ring check, ring-era kernel test Three review findings on the ring-cursor commit: - Pad-row staleness (eager mixed steps): the pad fill of the gathered spec buffers (hist=0, base=0, idx=pad sentinel) ran only inside the full-CG pure-spec-decode branch, but the layer passes bucket-length PRE-PADDED slices on eager mixed prefill+spec steps too — a stale tail row could reach the kernel carrying a reallocated block id plus old cursors (ghost CTAs appending to a live request's block, folding garbage into its checkpoint when the stale hist >= flush_min). The fill now runs at gather time, every spec step. (Pattern predates the ring commit — it extended it to cache_base — but fixed here for all three buffers.) - Init-time ring check: _ucache_kernel_available previously checked only that the kernel module file exists, so a pre-ring kernel passed engine init and failed at the FIRST spec step mid-serving (or, under python -O with the loader's assert stripped, silently corrupted state against ring cursors). The init gate now scans the module source for RING_SLOTS == 32 without importing it, and the loader's post-import check is a RuntimeError instead of an assert. - Ring-era kernel test: tests/kernels/test_replayssm_ucache_spec_gdn.py still allocated 16-deep flat rings and drove the legacy commit_gdn_ucache_hist protocol — impossible to pass against any kernel the loader accepts. Rewritten for the ring: 32-slot pools, block-keyed (write_pos, cache_base, is_flush) cursors driven by the shared commit_gdn_replayssm_spec / reset_gdn_replayssm_spec_cursors (incl. base wrap + both-cursor reset assertions), protocol equivalence vs the wrapper's restart_hist_on_flush=True commit, and the strided/null-page/2GB/pad-skip suites ported to (hist, base) windows. The consumerless commit_gdn_ucache_hist shim is removed.
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
mergify
Bot
added
performance
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
6 tasks done
6 tasks done
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
flashinfer-ai/flashinfer) — must include the ring revision (RING_SLOTS = 32); pre-ring kernel builds are rejected at engine initflashinfer/gdn_kernels/gdn_decode_bf16_wy_ucache_flush.py(import-self-contained; JIT-compiles via CuTeDSL at first use)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.pyflashinfer.gdn_kernels.gdn_decode_bf16_wy_ucache_flushDefaults (zero configuration):
GDN_UCACHE_IO_DTYPE=fp16--mamba-ssm-cache-dtype+GDN_UCACHE_STATE_DTYPEGDN_UCACHE_RING_DTYPE(vLLM-side override:VLLM_REPLAYSSM_RING_DTYPE)GDN_UCACHE_STATE_SR=1fp32-checkpoint read path. On this backend
--mamba-ssm-cache-dtype autoresolves to fp16 (the Qwen3.5 config updater prefers the kernel state
dtype over the HF config's
mamba_ssm_dtype=float32when this backend isselected).
import and asserts pool/ring dtypes on first call; vLLM allocates pools
from the same envs (the adapter
setdefaults both to fp16). Both sidesdefault fp16 → they agree with no flags.
--mamba-ssm-cache-dtype bfloat16+GDN_UCACHE_STATE_DTYPE=bf16 GDN_UCACHE_RING_DTYPE=bf16.Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.