From 4cfaaf3c5c14eaa8786c983b7ae83b66466b0a1f Mon Sep 17 00:00:00 2001 From: Johnny-Liou Date: Fri, 3 Jul 2026 17:34:49 -0400 Subject: [PATCH 01/23] [Kernel] ReplaySSM: cache SSM inputs instead of state for faster standard and speculative decode (Mamba2 + GDN) Signed-off-by: Johnny-Liou --- benchmarks/replayssm/e2e_decode_speedup.py | 238 ++++++ .../replayssm/e2e_spec_decode_throughput.py | 246 ++++++ ...ayssm_prefill_decode_equivalence_mamba2.py | 212 +++++ .../test_replayssm_spec_decode_mamba2.py | 709 +++++++++++++++++ .../test_replayssm_standard_decode_mamba2.py | 450 +++++++++++ tests/kernels/mamba/utils.py | 286 +++++++ .../test_fused_recurrent_packed_decode.py | 6 +- .../kernels/test_replayssm_spec_decode_gdn.py | 663 ++++++++++++++++ .../test_replayssm_standard_decode_gdn.py | 273 +++++++ ...eplayssm_teacher_decode_equivalence_gdn.py | 207 +++++ tests/v1/e2e/test_replayssm_decode.py | 46 ++ vllm/config/cache.py | 22 + vllm/config/vllm.py | 72 +- vllm/engine/arg_utils.py | 22 + .../model_executor/layers/fla/ops/__init__.py | 4 + .../fla/ops/fused_recurrent_replayssm.py | 290 +++++++ .../fla/ops/gdn_replayssm_spec_decode.py | 642 +++++++++++++++ vllm/model_executor/layers/mamba/abstract.py | 12 +- vllm/model_executor/layers/mamba/gdn/base.py | 10 +- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 186 ++++- .../layers/mamba/mamba_mixer2.py | 250 +++++- .../layers/mamba/mamba_utils.py | 294 +++++++ .../layers/mamba/ops/replayssm_config.py | 119 +++ ...tive_state_update_replayssm_output_only.py | 550 +++++++++++++ .../selective_state_update_replayssm_spec.py | 742 ++++++++++++++++++ ...state_update_replayssm_state_and_output.py | 373 +++++++++ vllm/model_executor/models/nemotron_h.py | 39 +- vllm/model_executor/models/qwen3_5.py | 30 +- vllm/v1/attention/backend.py | 9 + vllm/v1/attention/backends/gdn_attn.py | 167 ++++ vllm/v1/attention/backends/mamba_attn.py | 292 +++++++ vllm/v1/worker/gpu_model_runner.py | 1 + vllm/v1/worker/ubatch_utils.py | 9 + 33 files changed, 7420 insertions(+), 51 deletions(-) create mode 100644 benchmarks/replayssm/e2e_decode_speedup.py create mode 100644 benchmarks/replayssm/e2e_spec_decode_throughput.py create mode 100644 tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py create mode 100644 tests/kernels/mamba/test_replayssm_spec_decode_mamba2.py create mode 100644 tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py create mode 100644 tests/kernels/test_replayssm_spec_decode_gdn.py create mode 100644 tests/kernels/test_replayssm_standard_decode_gdn.py create mode 100644 tests/kernels/test_replayssm_teacher_decode_equivalence_gdn.py create mode 100644 tests/v1/e2e/test_replayssm_decode.py create mode 100644 vllm/model_executor/layers/fla/ops/fused_recurrent_replayssm.py create mode 100644 vllm/model_executor/layers/fla/ops/gdn_replayssm_spec_decode.py create mode 100644 vllm/model_executor/layers/mamba/ops/replayssm_config.py create mode 100644 vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py create mode 100644 vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_spec.py create mode 100644 vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_state_and_output.py diff --git a/benchmarks/replayssm/e2e_decode_speedup.py b/benchmarks/replayssm/e2e_decode_speedup.py new file mode 100644 index 000000000000..b6faccb45ead --- /dev/null +++ b/benchmarks/replayssm/e2e_decode_speedup.py @@ -0,0 +1,238 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end autoregressive decode benchmark: ReplaySSM vs the standard SSM kernel. + +Loads a hybrid SSM model, replicates one prompt across the batch, and times a +long greedy decode (CUDA graphs on) once with the standard kernel and once with +ReplaySSM, then reports the per-step / throughput speedup. Works for any hybrid +SSM model supported by vLLM (Mamba2 or GDN). The two modes run in separate +subprocesses so each gets a clean CUDA context. + +GDN models (Qwen3.5) default to the Triton prefill backend, which starts +instantly; FlashInfer is also fine but JIT-compiles via nvcc on first run (slow +startup). The prefill backend never affects the decode speedup measured here. + +The FlashInfer FP4-MoE autotuner is disabled by default (it is unstable under +CUDA-graph capture on the pre-release Blackwell FP4 path); pass +--no-disable-flashinfer-autotune for non-FP4 models. + +Examples: + python e2e_decode_speedup.py --model-id nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 + python e2e_decode_speedup.py --model-id Qwen/Qwen3.5-4B --buffer-len 16 + python e2e_decode_speedup.py --dtype auto --buffer-len 16 \ + --model-id nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 # B300 NVFP4 + python e2e_decode_speedup.py --dtype auto --buffer-len 16 \ + --model-id nvidia/Qwen3.5-122B-A10B-NVFP4 # B300 NVFP4 MoE +""" + +import argparse +import json +import os +import subprocess +import sys +import time + +DEFAULT_PROMPT = "My cat wrote all this CUDA code for a new language model and" + +MODE_LABEL = {"standard": "standard", "replayssm": "ReplaySSM"} + + +def parse_args(): + p = argparse.ArgumentParser( + description="E2E decode speedup: ReplaySSM vs the standard SSM kernel." + ) + p.add_argument("--model-id", default="nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16") + p.add_argument("--prompt", default=DEFAULT_PROMPT) + p.add_argument("--batch-size", type=int, default=256) + p.add_argument("--num-steps", type=int, default=1000) + p.add_argument("--warmup-steps", type=int, default=128) + p.add_argument("--repeats", type=int, default=1) + p.add_argument("--buffer-len", type=int, default=16, + help="ReplaySSM input-buffer length (16 for Mamba2 and GDN).") + p.add_argument("--dtype", default="bfloat16", + choices=["bfloat16", "float16", "float32", "auto"]) + p.add_argument("--gpu-memory-utilization", type=float, default=0.9) + p.add_argument("--max-model-len", type=int, default=None) + p.add_argument("--gdn-prefill-backend", default="triton", + choices=["triton", "flashinfer", "auto"], + help="GDN prefill kernel (GDN models only; Mamba2 ignores it). " + "'triton' (default) starts instantly; 'flashinfer'/'auto' " + "are also fine but JIT-compile via nvcc on first run " + "(slow startup). Decode speed is identical either way.") + p.add_argument("--disable-flashinfer-autotune", + action=argparse.BooleanOptionalAction, default=True, + help="Disable the FlashInfer FP4-MoE autotuner (default: on). " + "It is unstable under CUDA-graph capture on the " + "pre-release Blackwell FP4 path; pass " + "--no-disable-flashinfer-autotune for non-FP4 models.") + p.add_argument("--replayssm-route", default="output_only", + choices=["output_only", "state_and_output"], + help="Mamba2 cached route: output_only (cached_bc) or " + "state_and_output (cached_dot). GDN models ignore it.") + p.add_argument("--mamba-ssm-cache-dtype", default="auto", + choices=["auto", "float32", "float16", "bfloat16"], + help="SSM state dtype (both modes). 'auto' = config-driven; " + "'float32' = fp32 state, 'bfloat16' = s16 state.") + p.add_argument("--baseline-ssm-config", default="", + help="Pin the STANDARD baseline's SSM launch config as " + "'bsm,nw' via override_ssm_config (forces the in-process " + "engine so the override reaches the kernel). Empty = off.") + p.add_argument("--worker", choices=["standard", "replayssm"], default=None, + help=argparse.SUPPRESS) + return p.parse_args() + + +def resolve_max_model_len(args) -> int: + if args.max_model_len is not None: + return args.max_model_len + return args.num_steps + 256 + + +def run_worker(args): + # override_ssm_config is a module global; it only reaches the model if the + # engine runs in-process (default V1 spawns a separate EngineCore). Force it. + if args.worker == "standard" and args.baseline_ssm_config: + os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" + + import torch + + from vllm import LLM, SamplingParams + + mode = args.worker + max_model_len = resolve_max_model_len(args) + + llm_kwargs = dict( + model=args.model_id, + tensor_parallel_size=1, + dtype=args.dtype, + max_model_len=max_model_len, + trust_remote_code=True, + enable_prefix_caching=False, + enable_chunked_prefill=False, + max_num_seqs=args.batch_size, + max_num_batched_tokens=max(max_model_len, args.batch_size * 64), + enforce_eager=False, + disable_log_stats=True, + gpu_memory_utilization=args.gpu_memory_utilization, + # SSM state dtype (applies to both standard and ReplaySSM). + mamba_ssm_cache_dtype=args.mamba_ssm_cache_dtype, + # Skip the vision tower of multimodal hybrids (Qwen3.5 is a + # *ForConditionalGeneration model); ignored by text-only Mamba2 models. + language_model_only=True, + # GDN prefill kernel only; decode (and thus the speedup) is unaffected. + additional_config={"gdn_prefill_backend": args.gdn_prefill_backend}, + ) + if args.disable_flashinfer_autotune: + # FP4-MoE autotuner is unstable under CUDA-graph capture on Blackwell; + # re-enable (--no-disable-flashinfer-autotune) only for non-FP4 models. + llm_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False} + if mode == "replayssm": + llm_kwargs.update(use_replayssm=True, replayssm_buffer_len=args.buffer_len, + replayssm_route=args.replayssm_route) # route ignored by GDN models + + _ssm_cm = None + if mode == "standard" and args.baseline_ssm_config: + from vllm.model_executor.layers.mamba.ops.mamba_ssm import override_ssm_config + _bsm, _nw = (int(x) for x in args.baseline_ssm_config.split(",")) + _ssm_cm = override_ssm_config((_bsm, _nw)) + _ssm_cm.__enter__() # active through LLM() graph capture + decode + print(f"[{mode}] override_ssm_config -> (BLOCK_SIZE_M={_bsm}, num_warps={_nw})", + flush=True) + + llm = LLM(**llm_kwargs) + prompts = [args.prompt] * args.batch_size + + def timed_generate(n_tokens): + sp = SamplingParams( + n=1, temperature=0.0, ignore_eos=True, + min_tokens=n_tokens, max_tokens=n_tokens, + ) + if torch.cuda.is_available(): + torch.cuda.synchronize() + t0 = time.perf_counter() + outs = llm.generate(prompts, sp, use_tqdm=False) + if torch.cuda.is_available(): + torch.cuda.synchronize() + elapsed = time.perf_counter() - t0 + produced = min(len(o.outputs[0].token_ids) for o in outs) + assert produced == n_tokens, f"expected {n_tokens} tokens, got {produced}" + return elapsed + + timed_generate(args.warmup_steps) + + best = None + for _ in range(args.repeats): + elapsed = timed_generate(args.num_steps) + tok_s = args.batch_size * args.num_steps / elapsed + per_step_ms = elapsed / args.num_steps * 1e3 + print(f"[{mode}] {elapsed:.3f}s {tok_s:,.0f} tok/s {per_step_ms:.3f} ms/step", + flush=True) + if best is None or elapsed < best["elapsed_s"]: + best = {"mode": mode, "elapsed_s": elapsed, + "tok_s": tok_s, "per_step_ms": per_step_ms} + + print("RESULT_JSON " + json.dumps(best), flush=True) + if _ssm_cm is not None: + _ssm_cm.__exit__(None, None, None) + + +def run_one_mode(args, mode) -> dict: + cmd = [ + sys.executable, __file__, "--worker", mode, + "--model-id", args.model_id, "--prompt", args.prompt, + "--batch-size", str(args.batch_size), "--num-steps", str(args.num_steps), + "--warmup-steps", str(args.warmup_steps), "--repeats", str(args.repeats), + "--buffer-len", str(args.buffer_len), "--dtype", args.dtype, + "--gpu-memory-utilization", str(args.gpu_memory_utilization), + "--gdn-prefill-backend", args.gdn_prefill_backend, + "--replayssm-route", args.replayssm_route, + "--mamba-ssm-cache-dtype", args.mamba_ssm_cache_dtype, + "--baseline-ssm-config", args.baseline_ssm_config, + ] + cmd.append("--disable-flashinfer-autotune" if args.disable_flashinfer_autotune + else "--no-disable-flashinfer-autotune") + if args.max_model_len is not None: + cmd += ["--max-model-len", str(args.max_model_len)] + + result = None + proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + text=True, bufsize=1) + for line in proc.stdout: + sys.stdout.write(line) + sys.stdout.flush() + if line.startswith("RESULT_JSON "): + result = json.loads(line[len("RESULT_JSON "):]) + proc.wait() + if proc.returncode != 0: + raise RuntimeError(f"mode '{mode}' worker exited with {proc.returncode}") + if result is None: + raise RuntimeError(f"mode '{mode}' produced no RESULT_JSON line") + return result + + +def main(): + args = parse_args() + if args.worker is not None: + run_worker(args) + return + + print(f"model={args.model_id} batch_size={args.batch_size} " + f"steps={args.num_steps} buffer_len={args.buffer_len} dtype={args.dtype}") + + std = run_one_mode(args, "standard") + fla = run_one_mode(args, "replayssm") + speedup = std["per_step_ms"] / fla["per_step_ms"] + + print() + header = f"{'mode':<10}{'ms/step':>12}{'tok/s':>16}{'wall (s)':>12}" + print(header) + print("-" * len(header)) + for r in (std, fla): + print(f"{MODE_LABEL[r['mode']]:<10}{r['per_step_ms']:>12.3f}" + f"{r['tok_s']:>16,.0f}{r['elapsed_s']:>12.3f}") + print("-" * len(header)) + print(f"speedup (standard / ReplaySSM, per step): {speedup:.3f}x") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/replayssm/e2e_spec_decode_throughput.py b/benchmarks/replayssm/e2e_spec_decode_throughput.py new file mode 100644 index 000000000000..c5a8f309af6c --- /dev/null +++ b/benchmarks/replayssm/e2e_spec_decode_throughput.py @@ -0,0 +1,246 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end speculative-decode throughput: AR vs standard spec vs ReplaySSM spec. + +Real GSM8K prompts (chat-formatted) at a fixed batch size, CUDA graphs on, with +ignore_eos so every mode emits exactly --max-tokens tokens per sequence, making +tokens/s directly comparable. Reports throughput, mean acceptance length, and +the ReplaySSM-spec speedup over both baselines. + + ar : no speculative decoding + standard : vLLM native spec decoding (one recurrent state per draft token) + cache : ReplaySSM cached spec decoding (use_replayssm_spec) + +Each mode runs in its own subprocess for a clean CUDA context. + +The FlashInfer FP4-MoE autotuner is disabled by default (it is unstable under +CUDA-graph capture on the pre-release Blackwell FP4 path); pass +--no-disable-flashinfer-autotune for non-FP4 models. + +Examples (B300): + python e2e_spec_decode_throughput.py --batch-size 512 \ + --model-id nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 + python e2e_spec_decode_throughput.py --batch-size 512 \ + --model-id nvidia/Qwen3.5-122B-A10B-NVFP4 --spec-method qwen3_next_mtp \ + --moe-backend triton +""" + +import argparse +import json +import subprocess +import sys +import time + +MODE_LABEL = {"ar": "AR", "standard": "standard-spec", "cache": "ReplaySSM-spec"} + + +def parse_args(): + p = argparse.ArgumentParser( + description="E2E spec-decode throughput: AR vs standard vs ReplaySSM." + ) + p.add_argument("--model-id", + default="nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4") + p.add_argument("--tensor-parallel-size", type=int, default=1) + p.add_argument("--batch-size", type=int, default=512) + p.add_argument("--num-spec", type=int, default=3, + help="Draft tokens per step (spec window = num_spec + 1).") + p.add_argument("--spec-method", default="mtp") + p.add_argument("--buffer-len", type=int, default=16, + help="ReplaySSM buffer length (power of two, >= 1 + num_spec).") + p.add_argument("--max-tokens", type=int, default=256) + p.add_argument("--max-model-len", type=int, default=2048) + p.add_argument("--dtype", default="auto") + p.add_argument("--kv-cache-dtype", default="auto") + p.add_argument("--gpu-memory-utilization", type=float, default=0.9) + p.add_argument("--warmup-s", type=float, default=30.0, + help="Sustained full-batch decode before timing, to ramp the " + "GPU SM clock to its steady-state boost.") + p.add_argument("--enable-thinking", action="store_true", + help="Keep reasoning mode on (default off for short outputs).") + p.add_argument("--disable-flashinfer-autotune", + action=argparse.BooleanOptionalAction, default=True, + help="Disable the FlashInfer FP4-MoE autotuner (default: on). " + "It is unstable under CUDA-graph capture on the " + "pre-release Blackwell FP4 path; pass " + "--no-disable-flashinfer-autotune for non-FP4 models.") + p.add_argument("--modes", default="ar,standard,cache", + help="Comma-separated subset of {ar,standard,cache}.") + p.add_argument("--moe-backend", default=None, + help="Override the draft-model MoE backend (e.g. triton). On " + "Blackwell the draft's bf16 MoE hangs on flashinfer_trtllm; " + "pass triton. Main-model MoE is left at auto.") + p.add_argument("--worker", choices=["ar", "standard", "cache"], default=None, + help=argparse.SUPPRESS) + return p.parse_args() + + +def gsm8k_messages(batch_size): + from datasets import load_dataset + + questions = [r["question"] + for r in load_dataset("openai/gsm8k", "main", split="test")] + return [[{"role": "user", "content": questions[i % len(questions)]}] + for i in range(batch_size)] + + +def _sum_counter(metrics, name): + return sum( + m.value for m in metrics + if getattr(m, "name", None) == name and hasattr(m, "value") + ) + + +def run_worker(args): + import torch + + from vllm import LLM, SamplingParams + + mode = args.worker + spec_window = 1 if mode == "ar" else 1 + args.num_spec + + llm_kwargs = dict( + model=args.model_id, + tensor_parallel_size=args.tensor_parallel_size, + dtype=args.dtype, + kv_cache_dtype=args.kv_cache_dtype, + max_model_len=args.max_model_len, + max_num_seqs=args.batch_size, + trust_remote_code=True, + enable_prefix_caching=False, + enforce_eager=False, + disable_log_stats=False, + gpu_memory_utilization=args.gpu_memory_utilization, + seed=0, + # Avoid the FlashInfer GDN-prefill cutlass-DSL JIT stall on Blackwell + # (same default as the decode benchmark). + additional_config={"gdn_prefill_backend": "triton"}, + compilation_config={ + "max_cudagraph_capture_size": max(8, args.batch_size * spec_window) + }, + ) + if args.disable_flashinfer_autotune: + # FP4-MoE autotuner is unstable under CUDA-graph capture on Blackwell; + # re-enable (--no-disable-flashinfer-autotune) only for non-FP4 models. + llm_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False} + if mode != "ar": + spec_cfg = { + "method": args.spec_method, + "num_speculative_tokens": args.num_spec, + } + # Override only the draft MoE backend (the trtllm hang is in the draft). + if args.moe_backend: + spec_cfg["moe_backend"] = args.moe_backend + llm_kwargs["speculative_config"] = spec_cfg + if mode == "cache": + llm_kwargs["use_replayssm_spec"] = True + llm_kwargs["replayssm_buffer_len"] = args.buffer_len + + llm = LLM(**llm_kwargs) + messages = gsm8k_messages(args.batch_size) + chat_kwargs = {"enable_thinking": args.enable_thinking} + sp = SamplingParams(n=1, temperature=0.0, max_tokens=args.max_tokens, + ignore_eos=True, seed=0) + + def timed_chat(): + torch.cuda.synchronize() + t0 = time.perf_counter() + outs = llm.chat(messages, sp, chat_template_kwargs=chat_kwargs, + use_tqdm=False) + torch.cuda.synchronize() + return time.perf_counter() - t0, outs + + deadline = time.perf_counter() + args.warmup_s + while time.perf_counter() < deadline: + timed_chat() + + pre_acc = _sum_counter(llm.get_metrics(), "vllm:spec_decode_num_accepted_tokens") + pre_dft = _sum_counter(llm.get_metrics(), "vllm:spec_decode_num_drafts") + + elapsed, outs = timed_chat() + produced = min(len(o.outputs[0].token_ids) for o in outs) + assert produced == args.max_tokens, f"expected {args.max_tokens}, got {produced}" + tok_s = args.batch_size * args.max_tokens / elapsed + + accept_len = None + if mode != "ar": + drafts = _sum_counter(llm.get_metrics(), + "vllm:spec_decode_num_drafts") - pre_dft + accepted = _sum_counter(llm.get_metrics(), + "vllm:spec_decode_num_accepted_tokens") - pre_acc + accept_len = 1.0 + accepted / drafts if drafts else None + + result = {"mode": mode, "elapsed_s": elapsed, "tok_s": tok_s, + "accept_len": accept_len} + print(f"[{mode}] {elapsed:.2f}s {tok_s:,.0f} tok/s" + + (f" accept_len={accept_len:.2f}" if accept_len else ""), flush=True) + print("RESULT_JSON " + json.dumps(result), flush=True) + + +def run_one_mode(args, mode): + cmd = [ + sys.executable, __file__, "--worker", mode, + "--model-id", args.model_id, + "--tensor-parallel-size", str(args.tensor_parallel_size), + "--batch-size", str(args.batch_size), "--num-spec", str(args.num_spec), + "--spec-method", args.spec_method, "--buffer-len", str(args.buffer_len), + "--max-tokens", str(args.max_tokens), + "--max-model-len", str(args.max_model_len), "--dtype", args.dtype, + "--kv-cache-dtype", args.kv_cache_dtype, + "--gpu-memory-utilization", str(args.gpu_memory_utilization), + "--warmup-s", str(args.warmup_s), + ] + if args.enable_thinking: + cmd.append("--enable-thinking") + if args.moe_backend: + cmd += ["--moe-backend", args.moe_backend] + cmd.append("--disable-flashinfer-autotune" if args.disable_flashinfer_autotune + else "--no-disable-flashinfer-autotune") + + result = None + proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + text=True, bufsize=1) + for line in proc.stdout: + sys.stdout.write(line) + sys.stdout.flush() + if line.startswith("RESULT_JSON "): + result = json.loads(line[len("RESULT_JSON "):]) + proc.wait() + if proc.returncode != 0: + raise RuntimeError(f"mode '{mode}' worker exited with {proc.returncode}") + if result is None: + raise RuntimeError(f"mode '{mode}' produced no RESULT_JSON line") + return result + + +def main(): + args = parse_args() + if args.worker is not None: + run_worker(args) + return + + print(f"model={args.model_id} tp={args.tensor_parallel_size} " + f"batch_size={args.batch_size} num_spec={args.num_spec} " + f"buffer_len={args.buffer_len} max_tokens={args.max_tokens}") + + modes = [m for m in args.modes.split(",") if m] + results = {m: run_one_mode(args, m) for m in modes} + + print() + header = f"{'mode':<16}{'tok/s':>14}{'accept_len':>12}{'wall (s)':>12}" + print(header) + print("-" * len(header)) + for m in modes: + r = results[m] + al = f"{r['accept_len']:.2f}" if r["accept_len"] else "-" + print(f"{MODE_LABEL[m]:<16}{r['tok_s']:>14,.0f}{al:>12}{r['elapsed_s']:>12.2f}") + print("-" * len(header)) + if "cache" in results and "standard" in results: + print(f"ReplaySSM / standard : " + f"{results['cache']['tok_s'] / results['standard']['tok_s']:.2f}x") + if "cache" in results and "ar" in results: + print(f"ReplaySSM / AR : " + f"{results['cache']['tok_s'] / results['ar']['tok_s']:.2f}x") + + +if __name__ == "__main__": + main() diff --git a/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py b/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py new file mode 100644 index 000000000000..8fe571bdf145 --- /dev/null +++ b/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py @@ -0,0 +1,212 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Prefill == decode equivalence for the Mamba2 ReplaySSM kernels. + +The SSM recurrence is path-independent: vLLM's chunked prefill (the SSD kernel +``mamba_chunk_scan_combined_varlen``) and step-by-step decode must produce the +same per-position outputs and final state over a sequence. This file feeds one +set of inputs through the production dt flow (raw dt + softplus + a per-head +dt_bias, applied inside each kernel) to: + + * the exact fp32 step recurrence (``selective_state_update_ref``) -- the + ground truth, + * the chunked prefill kernel, + * the baseline decode kernel, + * both ReplaySSM decode routes, + +and checks all of them agree. Prefill (a chunked scan) and decode (a step +recurrence) are different code paths, so they differ numerically: the chunked +scan carries ~2e-2 (fp32) / ~4e-2 (bf16) vs the exact recurrence, far above the +near-exact decode. We therefore anchor every path on the exact recurrence at +SSD-level tolerances (the same regime as ``test_mamba_ssm_ssd.py``), which the +chunked scan sets, keyed off the activation dtype. + +State and activation/buffer precision are swept independently, including the +fp32-state + bf16-activation production config. +""" + +import pytest +import torch + +from tests.kernels.mamba.utils import selective_state_update_ref +from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_state_update +from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_output_only import ( # noqa: E501 + selective_state_update_replayssm_output_only, +) +from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_state_and_output import ( # noqa: E501 + selective_state_update_replayssm_state_and_output, +) +from vllm.model_executor.layers.mamba.ops.ssd_combined import ( + mamba_chunk_scan_combined_varlen, +) +from vllm.utils.torch_utils import set_random_seed +from vllm.v1.attention.backends.mamba2_attn import compute_varlen_chunk_metadata + + +def _prefill_tolerances(act_dtype: torch.dtype) -> tuple[float, float]: + # The chunked prefill scan, not the decode, sets these: it carries ~2e-2 + # (fp32) / ~6e-2 (bf16) vs the exact recurrence, while ReplaySSM decode is + # near-exact (~2e-6 fp32). Keyed off the activation dtype (the outputs are + # in act_dtype). Same regime as test_mamba_ssm_ssd.py. + if act_dtype == torch.float32: + return 1e-2, 3e-2 + return 6e-2, 1e-1 + + +def _run_prefill_decode_equivalence( + *, + route: str, + state_dtype: torch.dtype, + act_dtype: torch.dtype, + nheads: int, + headdim: int, + ngroups: int, + dstate: int, + seqlen: int, + chunk_size: int, + max_cache_len: int, + seed: int = 0, +) -> None: + """Prefill the whole sequence and decode it step by step; check both match + the exact fp32 recurrence (and each other). All paths use the production dt + flow (raw dt + softplus + a per-head dt_bias), so this also checks prefill + and decode apply the softplus/bias preprocessing consistently. ``state_dtype`` + is the recurrent-state precision; ``act_dtype`` the activation/buffer one.""" + device = "cuda" + rtol, atol = _prefill_tolerances(act_dtype) + set_random_seed(seed) + + # Production dt flow: raw dt + a per-head dt_bias (~-4 keeps softplus(dt + + # bias) small and well-conditioned). dt_bias is (nheads,) for prefill and + # (nheads, headdim) for the decode kernels/reference. + A = -torch.exp(torch.rand(nheads, device=device, dtype=act_dtype)) + dt = torch.randn(seqlen, nheads, device=device, dtype=act_dtype) + dt_bias = torch.rand(nheads, device=device, dtype=act_dtype) - 4 + dt_bias_hd = dt_bias.view(nheads, 1).expand(nheads, headdim) + X = torch.randn(seqlen, nheads, headdim, device=device, dtype=act_dtype) + B = torch.randn(seqlen, ngroups, dstate, device=device, dtype=act_dtype) + C = torch.randn(seqlen, ngroups, dstate, device=device, dtype=act_dtype) + A_bcast = A.view(nheads, 1, 1).expand(nheads, headdim, dstate) + + # Chunked prefill over the whole sequence (implicit batch=1, varlen). The + # kernel always returns the final state in fp32, so no state_dtype plumbing. + cu_seqlens = torch.tensor((0, seqlen), device=device).cumsum(0).to(torch.int32) + cu_chunk_seqlens, last_chunk_indices, seq_idx = compute_varlen_chunk_metadata( + cu_seqlens, chunk_size) + y_prefill = torch.empty(seqlen, nheads, headdim, device=device, dtype=act_dtype) + final_state_prefill = mamba_chunk_scan_combined_varlen( + X, dt, A, B, C, chunk_size, cu_seqlens=cu_seqlens, + cu_chunk_seqlens=cu_chunk_seqlens, last_chunk_indices=last_chunk_indices, + seq_idx=seq_idx, out=y_prefill, D=None, dt_bias=dt_bias, + dt_softplus=True) + + # Step paths: exact fp32 recurrence (ground truth), baseline, ReplaySSM. + # State follows state_dtype; caches follow act_dtype (dt_cache is fp32). + state_ref = torch.zeros( + 1, nheads, headdim, dstate, device=device, dtype=torch.float32) + state_base = torch.zeros( + 1, nheads, headdim, dstate, device=device, dtype=state_dtype) + state_dec = torch.zeros( + 1, nheads, headdim, dstate, device=device, dtype=state_dtype) + x_cache = torch.zeros( + 1, nheads, max_cache_len, headdim, device=device, dtype=act_dtype) + dt_cache = torch.zeros( + 1, nheads, max_cache_len, device=device, dtype=torch.float32) + B_cache = torch.zeros( + 1, ngroups, max_cache_len, dstate, device=device, dtype=act_dtype) + bc_pre = torch.empty( + 1, ngroups, max_cache_len, device=device, dtype=torch.float32) + write_pos = torch.zeros(1, dtype=torch.int32, device=device) + # No skip connection (D=0) on any path here; the D!=0 path is covered by the + # standard-decode suite. The baseline kernel needs a D tensor, not None. + D_zero = torch.zeros(nheads, headdim, device=device) + + y_ref = torch.empty(seqlen, nheads, headdim, device=device, dtype=torch.float32) + y_base = torch.empty(seqlen, nheads, headdim, device=device, dtype=act_dtype) + y_dec = torch.empty(seqlen, nheads, headdim, device=device, dtype=act_dtype) + for t in range(seqlen): + dt_t = dt[t].view(1, nheads, 1).expand(1, nheads, headdim) + is_flush = write_pos == max_cache_len - 1 + + y_ref[t] = selective_state_update_ref( + state_ref, X[t:t + 1].float(), dt_t.float(), A_bcast.float(), + B[t:t + 1].float(), C[t:t + 1].float(), dt_bias=dt_bias_hd.float(), + dt_softplus=True)[0] + + out_b = torch.empty(1, nheads, headdim, device=device, dtype=act_dtype) + selective_state_update( + state_base, X[t:t + 1], dt_t, A_bcast, B[t:t + 1], C[t:t + 1], + D=D_zero, dt_bias=dt_bias_hd, dt_softplus=True, out=out_b) + y_base[t] = out_b[0] + + out_d = torch.empty(1, nheads, headdim, device=device, dtype=act_dtype) + common = dict( + dt_bias=dt_bias_hd, dt_softplus=True, x_cache=x_cache, + dt_cache=dt_cache, B_cache=B_cache, write_pos=write_pos, + is_flush=is_flush, max_cache_len=max_cache_len, out=out_d) + if route == "output_only": + selective_state_update_replayssm_output_only( + state_dec, X[t:t + 1], dt_t, A_bcast, B[t:t + 1], C[t:t + 1], + bc_pre=bc_pre, **common) + else: + selective_state_update_replayssm_state_and_output( + state_dec, X[t:t + 1], dt_t, A_bcast, B[t:t + 1], C[t:t + 1], + **common) + y_dec[t] = out_d[0] + + write_pos = torch.where( + is_flush, torch.zeros_like(write_pos), write_pos + 1) + + # Every path computes the same recurrence; anchor each on the fp32 truth. + torch.testing.assert_close(y_prefill.float(), y_ref, rtol=rtol, atol=atol) + torch.testing.assert_close(y_base.float(), y_ref, rtol=rtol, atol=atol) + torch.testing.assert_close(y_dec.float(), y_ref, rtol=rtol, atol=atol) + # Headline: ReplaySSM decode matches the chunked prefill directly. + torch.testing.assert_close(y_dec.float(), y_prefill.float(), rtol=rtol, atol=atol) + # Final state too (the recurrence ends in state_ref after the loop). + torch.testing.assert_close( + final_state_prefill[0].float(), state_ref[0], rtol=rtol, atol=atol) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("route", ["output_only", "state_and_output"]) +@pytest.mark.parametrize( + "precision", + # fp32 state is the default; bf16/fp16 are reduced-footprint configs. fp16 + # appears as an activation dtype (s32_afp16, sfp16_afp16) and as a finer- + # mantissa state under bf16 activations (sfp16_a16). + [ + pytest.param((torch.float32, torch.float32), id="s32_a32"), + pytest.param((torch.float32, torch.bfloat16), id="s32_a16"), + pytest.param((torch.bfloat16, torch.bfloat16), id="s16_a16"), + pytest.param((torch.float32, torch.float16), id="s32_afp16"), + pytest.param((torch.float16, torch.float16), id="sfp16_afp16"), + pytest.param((torch.float16, torch.bfloat16), id="sfp16_a16"), + ], +) +@pytest.mark.parametrize( + "geometry", # (nheads, headdim, dstate, ngroups) + [ + pytest.param((8, 64, 64, 2), id="small"), + pytest.param((96, 80, 128, 8), id="nano4b"), + ], +) +def test_replayssm_prefill_decode_equivalence( + route: str, + precision: tuple[torch.dtype, torch.dtype], + geometry: tuple[int, int, int, int], +): + state_dtype, act_dtype = precision + nheads, headdim, dstate, ngroups = geometry + _run_prefill_decode_equivalence( + route=route, + state_dtype=state_dtype, + act_dtype=act_dtype, + nheads=nheads, + headdim=headdim, + ngroups=ngroups, + dstate=dstate, + seqlen=16, + chunk_size=8, + max_cache_len=4, + ) diff --git a/tests/kernels/mamba/test_replayssm_spec_decode_mamba2.py b/tests/kernels/mamba/test_replayssm_spec_decode_mamba2.py new file mode 100644 index 000000000000..85e3e954db92 --- /dev/null +++ b/tests/kernels/mamba/test_replayssm_spec_decode_mamba2.py @@ -0,0 +1,709 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Speculative-decode correctness for the Mamba2 ReplaySSM kernels. + +A spec/verify step takes the checkpoint state ``S_0`` + the committed ring buffer ++ a window of ``max_spec_len = 1 + num_speculative_tokens`` draft tokens, and +computes the recurrence OUTPUT at each window position (causal: draft ``s`` reads +the buffer up to its own position). It does NOT write the state per draft -- only +a flush step folds the *committed* history into the checkpoint, so unaccepted +drafts can be rolled back. + +The history window is ``L = B + max_spec_len`` (block ``B`` = ``buffer_len``): the +physical circular buffer is ``next_pow2(L)`` and ``max_cache_len`` passed to the +kernel is the logical ``L``. Verify launches the non-flush kernel; flush launches +the reconstruct kernel; both run every step with device-side row routing. + +Oracle (no separate spec reference needed): the already-verified ReplaySSM +STANDARD decode kernel (``selective_state_update_replayssm_output_only``) stepped +one token at a time over the SAME window from the SAME checkpoint+buffer, with +``is_flush=False`` so it reads the SAME fixed checkpoint at every window position. +Its per-position output must equal the spec kernel's. For the multi-step rollback +the ground truth is the baseline ``selective_state_update`` decode of the accepted +token stream (it writes the full state every step). + +Three checks per (precision, geometry, base_block, max_spec_len): + * single spec step output == standard-decode-of-window, swept over the buffer + fill level up to the max non-flush fill C = B - max_spec_len, + * a 40-step rollback with a random accepted count k per step: accepted-position + outputs track the baseline decode, and the committed checkpoint at each flush + matches the baseline state at the matching folded-token count, + * continuous-batching realism: sparse state slots via state_batch_indices with + NULL padding rows and per-row buffer fill levels. + +Precision: the checkpoint-readout dot is keyed on the activation dtype -- tf32x3 +(~fp32 parity, 3-pass) for fp32 activations, single-pass tf32 for bf16 -- and the +cache GEMM uses tf32x3, so fp32 is near-exact (rel ~1e-6, tight tolerance). At bf16 +activations the bf16 ``C``/``x``/``B`` operands dominate (rel ~5e-3), a small +EXPECTED gap that is bounded per step and non-accumulating across the decode. +""" + +import pytest +import torch + +from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_state_update +from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_output_only import ( # noqa: E501 + selective_state_update_replayssm_output_only, +) +from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_spec import ( # noqa: E501 + commit_replayssm_spec, + reset_replayssm_spec_cursors, + selective_state_update_replayssm_spec, +) +from vllm.utils.torch_utils import set_random_seed +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + +DEV = "cuda" + + +def _lbt(base_block: int, max_spec_len: int) -> tuple[int, int]: + """Logical flush threshold L = B + max_spec_len and physical pow2 buffer.""" + L = base_block + max_spec_len + buf = 1 << (L - 1).bit_length() + return L, buf + + +def _tolerances(both_fp32: bool) -> tuple[float, float]: + # fp32 activations use the tf32x3 checkpoint dot, matching the standard-decode + # oracle to ~1e-6; atol=1e-3 still flags a regression to single-pass tf32. bf16 + # stays loose: the bf16 C/x/B operands dominate (rel ~5e-3), unaffected by the + # checkpoint dot's precision. + if both_fp32: + return 1e-4, 1e-3 + return 6e-2, 2e-1 + + +def _tied_A(nheads: int, headdim: int, dstate: int) -> torch.Tensor: + # TIE_HDIM: A is scalar per head (stride(-1)==stride(-2)==0). + A = -torch.rand(nheads, device=DEV) - 1.0 + return A.view(nheads, 1, 1).expand(nheads, headdim, dstate) + + +def _scatter_packed_history( + post_conv_cache: torch.Tensor, + dt_cache: torch.Tensor, + slot: int, + x_hist: torch.Tensor, # (wp, H, P) + B_hist: torch.Tensor, # (wp, G, N) + dt_hist_raw: torch.Tensor, # (wp, H) raw dt + d_inner: int, + ngroups: int, + dstate: int, +) -> None: + """Fill the committed history [0, wp) into the packed circular caches at + post_origin=0 (so logical pos == physical pos). Channel map matches the + kernel's x_view/B_view: x[h,d]->h*P+d, B[g,n]->d_inner+g*N+n (C is not + cached; the kernel reads it fresh from conv_out). dt_cache stores RAW dt + (the kernel applies bias+softplus on read).""" + wp, H, P = x_hist.shape + G, N = ngroups, dstate + for p in range(wp): + post_conv_cache[slot, p, :d_inner] = x_hist[p].reshape(d_inner) + post_conv_cache[slot, p, d_inner : d_inner + G * N] = B_hist[p].reshape(G * N) + # C not cached; only x|B seeded. + dt_cache[slot, :, p] = dt_hist_raw[p].float() + + +def _pack_window_conv_out( + x_w: torch.Tensor, # (S, H, P) + B_w: torch.Tensor, # (S, G, N) + C_w: torch.Tensor, # (S, G, N) + d_inner: int, + ngroups: int, + dstate: int, + act_dtype: torch.dtype, +) -> torch.Tensor: + S = x_w.shape[0] + G, N = ngroups, dstate + conv_dim = d_inner + 2 * G * N + conv_out = torch.zeros(S, conv_dim, device=DEV, dtype=act_dtype) + conv_out[:, :d_inner] = x_w.reshape(S, d_inner) + conv_out[:, d_inner : d_inner + G * N] = B_w.reshape(S, G * N) + conv_out[:, d_inner + G * N :] = C_w.reshape(S, G * N) + return conv_out + + +def _standard_window_oracle( + *, + S0_slot: torch.Tensor, # (H, P, N) checkpoint for the row + x_all, + dt_all, + B_all, + C_all, + z_all, # (T_tot, ...) history+window raw inputs + A, + D, + dt_bias, + dt_softplus, + wp: int, + spec_len: int, + buffer_len: int, + act_dtype: torch.dtype, +) -> torch.Tensor: + """Step the standard output_only decode over history+window (is_flush=False + throughout, so the checkpoint stays fixed). Return the window outputs + (spec_len, H, P).""" + H, P, N = S0_slot.shape + G = B_all.shape[1] + state = S0_slot[None].clone() # (1, H, P, N) + x_cache = torch.zeros(1, H, buffer_len, P, device=DEV, dtype=act_dtype) + dt_cache = torch.zeros(1, H, buffer_len, device=DEV, dtype=torch.float32) + B_cache = torch.zeros(1, G, buffer_len, N, device=DEV, dtype=act_dtype) + bc_pre = torch.empty(1, G, buffer_len, device=DEV, dtype=torch.float32) + no_flush = torch.zeros(1, device=DEV, dtype=torch.int8) + win = [] + for t in range(wp + spec_len): + out_t = torch.empty(1, H, P, device=DEV, dtype=act_dtype) + selective_state_update_replayssm_output_only( + state, + x_all[t : t + 1], + dt_all[t : t + 1, :, None].expand(1, H, P), + A, + B_all[t : t + 1], + C_all[t : t + 1], + D=D, + z=z_all[t : t + 1] if z_all is not None else None, + dt_bias=dt_bias[:, None].expand(H, P), + dt_softplus=dt_softplus, + x_cache=x_cache, + dt_cache=dt_cache, + B_cache=B_cache, + write_pos=torch.tensor([t], device=DEV, dtype=torch.int32), + is_flush=no_flush, + max_cache_len=buffer_len, + bc_pre=bc_pre, + out=out_t, + ) + if t >= wp: + win.append(out_t.clone()) + return torch.cat(win, dim=0) + + +def _run_single_step( + *, + state_dtype, + act_dtype, + nheads, + headdim, + dstate, + ngroups, + buffer_len, # history block B + max_spec_len, + wp, + has_z, + dt_softplus=True, + seed=0, + perturb=False, +): + set_random_seed(seed) + H, P, N, G = nheads, headdim, dstate, ngroups + d_inner = H * P + spec_len = max_spec_len + L, buf = _lbt(buffer_len, max_spec_len) + both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32 + rtol, atol = _tolerances(both_fp32) + + A = _tied_A(H, P, N) + dt_bias = torch.rand(H, device=DEV) - 4.0 + D = torch.randn(H, P, device=DEV) + + T_tot = wp + spec_len + x = torch.randn(T_tot, H, P, device=DEV, dtype=act_dtype) + dt = torch.randn(T_tot, H, device=DEV, dtype=act_dtype) + B = torch.randn(T_tot, G, N, device=DEV, dtype=act_dtype) + C = torch.randn(T_tot, G, N, device=DEV, dtype=act_dtype) + z = torch.randn(T_tot, H, P, device=DEV, dtype=act_dtype) if has_z else None + + num_blocks = 2 + S0 = torch.randn(num_blocks, H, P, N, device=DEV, dtype=state_dtype) * 0.1 + + oracle = _standard_window_oracle( + S0_slot=S0[1], + x_all=x, + dt_all=dt, + B_all=B, + C_all=C, + z_all=z, + A=A, + D=D, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + wp=wp, + spec_len=spec_len, + buffer_len=buf, + act_dtype=act_dtype, + ) + + # spec path: prefill packed history, one verify call (non-flush). + state_spec = S0.clone() + cache_conv_dim = d_inner + G * N # x|B only (no C) + post_conv_cache = torch.zeros( + num_blocks, buf, cache_conv_dim, device=DEV, dtype=act_dtype + ) + dt_cache = torch.zeros(num_blocks, H, buf, device=DEV, dtype=torch.float32) + _scatter_packed_history( + post_conv_cache, dt_cache, 1, x[:wp], B[:wp], dt[:wp], d_inner, G, N + ) + if perturb: + # teeth: corrupt one history slot -> outputs must diverge from the oracle. + post_conv_cache[1, max(0, wp - 1), 0] += 5.0 + + conv_out = _pack_window_conv_out(x[wp:], B[wp:], C[wp:], d_inner, G, N, act_dtype) + dt_spec = dt[wp:].float() + z_spec = z[wp:] if has_z else None + write_pos = torch.zeros(num_blocks, dtype=torch.int32, device=DEV) + write_pos[1] = wp + post_origin = torch.zeros(num_blocks, dtype=torch.int32, device=DEV) + is_flush = torch.zeros(num_blocks, dtype=torch.int8, device=DEV) + qsl = torch.tensor([0, spec_len], device=DEV, dtype=torch.int32) + sbi = torch.tensor([1], device=DEV, dtype=torch.int32) + out_spec = torch.empty(spec_len, H, P, device=DEV, dtype=act_dtype) + selective_state_update_replayssm_spec( + state_spec, + post_conv_cache, + dt_cache, + conv_out, + dt_spec, + A, + write_pos=write_pos, + post_conv_state_pos=post_origin, + is_flush=is_flush, + query_start_loc=qsl, + state_batch_indices=sbi, + max_cache_len=L, + max_spec_len=max_spec_len, + d_inner=d_inner, + ngroups=G, + dstate=N, + D=D, + z=z_spec, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + out=out_spec, + ) + + torch.testing.assert_close(out_spec, oracle, rtol=rtol, atol=atol) + # non-flush verify must not touch the checkpoint. + torch.testing.assert_close(state_spec, S0, rtol=0, atol=0) + + +def _wp_set(base_block: int, max_spec_len: int) -> list[int]: + # Verify-path fills: 0 (empty buffer = pure checkpoint readout), the max + # non-flush fill C = B - max_spec_len (the tightest edge), and a midpoint. + C = base_block - max_spec_len + return sorted({0, max(0, C // 2), max(0, C)}) + + +_PRECISIONS = [ + pytest.param((torch.float32, torch.float32), id="s32_a32"), + pytest.param((torch.float32, torch.bfloat16), id="s32_a16"), + pytest.param((torch.bfloat16, torch.bfloat16), id="s16_a16"), +] +# (nheads, headdim, dstate, ngroups). The full precision x block x max_spec_len +# sweep runs on the small shapes (cheap compiles); the production Nemotron-3 +# Mamba2 shapes are exercised by test_spec_step_real_geometry. +_SMALL = pytest.param((8, 64, 64, 4), id="small") +_TINY = pytest.param((4, 64, 16, 1), id="tiny") +_REAL_GEOMETRIES = [ + pytest.param((96, 80, 128, 8), id="nano4b"), + pytest.param((128, 64, 128, 8), id="super120b"), + pytest.param((256, 64, 128, 8), id="ultra550b"), +] +_BASE_BLOCKS = [16, 32] # history block B (replayssm_buffer_len) +_MAX_SPEC_LENS = [2, 4, 6, 8] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("precision", _PRECISIONS) +@pytest.mark.parametrize("geometry", [_SMALL, _TINY]) +@pytest.mark.parametrize("base_block", _BASE_BLOCKS) +@pytest.mark.parametrize("max_spec_len", _MAX_SPEC_LENS) +@pytest.mark.parametrize("has_z", [False, True]) +def test_spec_step_matches_standard_decode( + precision, geometry, base_block, max_spec_len, has_z +): + # base_block >= max_spec_len always holds here. Sweep the non-flush fill level. + state_dtype, act_dtype = precision + nheads, headdim, dstate, ngroups = geometry + for wp in _wp_set(base_block, max_spec_len): + _run_single_step( + state_dtype=state_dtype, + act_dtype=act_dtype, + nheads=nheads, + headdim=headdim, + dstate=dstate, + ngroups=ngroups, + buffer_len=base_block, + max_spec_len=max_spec_len, + wp=wp, + has_z=has_z, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("precision", _PRECISIONS) +@pytest.mark.parametrize("geometry", _REAL_GEOMETRIES) +@pytest.mark.parametrize("max_spec_len", [2, 8]) +def test_spec_step_real_geometry(precision, geometry, max_spec_len): + # Production Nemotron-3 Mamba2 shapes (Nano-4B / Super-120B / Ultra-550B) at + # the deployable block B=16, both ends of the max_spec_len range. + state_dtype, act_dtype = precision + nheads, headdim, dstate, ngroups = geometry + for wp in _wp_set(16, max_spec_len): + _run_single_step( + state_dtype=state_dtype, + act_dtype=act_dtype, + nheads=nheads, + headdim=headdim, + dstate=dstate, + ngroups=ngroups, + buffer_len=16, + max_spec_len=max_spec_len, + wp=wp, + has_z=True, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +def test_spec_step_teeth(): + # A correct run passes; corrupting one cached history value must break the + # output match (guards against a vacuous oracle). + _run_single_step( + state_dtype=torch.float32, + act_dtype=torch.float32, + nheads=8, + headdim=64, + dstate=64, + ngroups=4, + buffer_len=16, + max_spec_len=4, + wp=5, + has_z=True, + ) + with pytest.raises(AssertionError): + _run_single_step( + state_dtype=torch.float32, + act_dtype=torch.float32, + nheads=8, + headdim=64, + dstate=64, + ngroups=4, + buffer_len=16, + max_spec_len=4, + wp=5, + has_z=True, + perturb=True, + ) + + +def _run_rollback( + *, + state_dtype, + act_dtype, + nheads, + headdim, + dstate, + ngroups, + buffer_len, # history block B + max_spec_len, + num_steps=40, + has_z=True, + seed=0, +): + """Drive the verify+commit loop from an empty buffer; accept a random k each + step. The baseline ``selective_state_update`` decode of the accepted stream is + the ground truth for outputs; the committed checkpoint at each flush must match + the baseline state at the matching folded-token count (n_folded bookkeeping).""" + set_random_seed(seed) + H, P, N, G = nheads, headdim, dstate, ngroups + d_inner = H * P + L, buf = _lbt(buffer_len, max_spec_len) + both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32 + rtol, atol = _tolerances(both_fp32) + + A = _tied_A(H, P, N) + dt_bias = torch.rand(H, device=DEV) - 4.0 + D = torch.randn(H, P, device=DEV) + + num_blocks = 2 + S0 = torch.randn(num_blocks, H, P, N, device=DEV, dtype=state_dtype) * 0.1 + state_spec = S0.clone() + state_base = S0.clone() # full pool; row in slot 1 + + cache_conv_dim = d_inner + G * N # x|B only (no C) + post_conv_cache = torch.zeros( + num_blocks, buf, cache_conv_dim, device=DEV, dtype=act_dtype + ) + dt_cache = torch.zeros(num_blocks, H, buf, device=DEV, dtype=torch.float32) + write_pos = torch.zeros(num_blocks, dtype=torch.int32, device=DEV) + post_origin = torch.zeros(num_blocks, dtype=torch.int32, device=DEV) + is_flush = torch.zeros(num_blocks, dtype=torch.int8, device=DEV) + sbi = torch.tensor([1], device=DEV, dtype=torch.int32) + reset_replayssm_spec_cursors( + write_pos, + post_origin, + is_flush, + torch.ones(1, dtype=torch.int8, device=DEV), + sbi, + L, + max_spec_len, + ) + + n_folded = 0 + snapshots = {0: state_base[1].clone()} + total_accepted = 0 + g = torch.Generator(device="cpu").manual_seed(seed + 1) + + for _ in range(num_steps): + spec_len = max_spec_len + x = torch.randn(spec_len, H, P, device=DEV, dtype=act_dtype) + dt = torch.randn(spec_len, H, device=DEV, dtype=act_dtype) + Bw = torch.randn(spec_len, G, N, device=DEV, dtype=act_dtype) + Cw = torch.randn(spec_len, G, N, device=DEV, dtype=act_dtype) + zw = torch.randn(spec_len, H, P, device=DEV, dtype=act_dtype) if has_z else None + + conv_out = _pack_window_conv_out(x, Bw, Cw, d_inner, G, N, act_dtype) + dt_spec = dt.float() + qsl = torch.tensor([0, spec_len], device=DEV, dtype=torch.int32) + out_spec = torch.empty(spec_len, H, P, device=DEV, dtype=act_dtype) + + wp_before = int(write_pos[1].item()) + flush_before = int(is_flush[1].item()) + selective_state_update_replayssm_spec( + state_spec, + post_conv_cache, + dt_cache, + conv_out, + dt_spec, + A, + write_pos=write_pos, + post_conv_state_pos=post_origin, + is_flush=is_flush, + query_start_loc=qsl, + state_batch_indices=sbi, + max_cache_len=L, + max_spec_len=max_spec_len, + d_inner=d_inner, + ngroups=G, + dstate=N, + D=D, + z=zw, + dt_bias=dt_bias, + dt_softplus=True, + out=out_spec, + ) + # a flush verify folds the committed history [0, wp_before) into S_0. + if flush_before and wp_before > 0: + n_folded += wp_before + + k = int(torch.randint(1, spec_len + 1, (1,), generator=g).item()) + + base_out = [] + for s in range(k): + ot = torch.empty(1, H, P, device=DEV, dtype=act_dtype) + selective_state_update( + state_base, + x[s : s + 1], + dt[s : s + 1, :, None].expand(1, H, P), + A, + Bw[s : s + 1], + Cw[s : s + 1], + D=D, + z=zw[s : s + 1] if has_z else None, + dt_bias=dt_bias[:, None].expand(H, P), + dt_softplus=True, + state_batch_indices=sbi, + out=ot, + ) + base_out.append(ot.clone()) + total_accepted += 1 + snapshots[total_accepted] = state_base[1].clone() + base_out = torch.cat(base_out, dim=0) + + # (a) accepted-position outputs track the baseline decode. + torch.testing.assert_close(out_spec[:k], base_out, rtol=rtol, atol=atol) + + commit_replayssm_spec( + write_pos, + post_origin, + is_flush, + torch.tensor([k], device=DEV, dtype=torch.int32), + sbi, + L, + max_spec_len, + ) + + # (b) committed checkpoint == baseline state at the folded-token count. + if n_folded in snapshots: + torch.testing.assert_close( + state_spec[1], snapshots[n_folded], rtol=rtol, atol=atol + ) + + # The buffer must have flushed at least once over 40 steps (else the state + # check above never ran) -- otherwise the test is vacuous. + assert n_folded > 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("precision", _PRECISIONS) +@pytest.mark.parametrize("base_block", _BASE_BLOCKS) +@pytest.mark.parametrize("max_spec_len", _MAX_SPEC_LENS) +def test_spec_rollback_tracks_baseline(precision, base_block, max_spec_len): + state_dtype, act_dtype = precision + _run_rollback( + state_dtype=state_dtype, + act_dtype=act_dtype, + nheads=8, + headdim=64, + dstate=64, + ngroups=4, + buffer_len=base_block, + max_spec_len=max_spec_len, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("precision", _PRECISIONS) +@pytest.mark.parametrize("with_padding", [False, True]) +def test_spec_continuous_batching(precision, with_padding): + """Sparse state slots via state_batch_indices with NULL padding rows and a + DIFFERENT buffer fill level per row, in a single verify call. Each active + row is checked against its own standard-decode window oracle; padding-row + outputs and unused state slots are left untouched.""" + state_dtype, act_dtype = precision + set_random_seed(0) + H, P, N, G = 4, 64, 16, 2 + d_inner = H * P + conv_dim = d_inner + 2 * G * N + base_block = 16 + max_spec_len = 4 + spec_len = max_spec_len + L, buf = _lbt(base_block, max_spec_len) + both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32 + rtol, atol = _tolerances(both_fp32) + + batch = 3 + padding = 2 if with_padding else 0 + padded_batch = batch + padding + total_slots = 12 + C = base_block - max_spec_len + wps = [0, C // 2, C] # staggered non-flush fills incl. the tight edge + + A = _tied_A(H, P, N) + dt_bias = torch.rand(H, device=DEV) - 4.0 + D = torch.randn(H, P, device=DEV) + + state_indices = (torch.randperm(total_slots - 1, device=DEV)[:batch] + 1).to( + torch.int32 + ) + sbi = torch.cat( + [ + state_indices, + torch.full((padding,), NULL_BLOCK_ID, dtype=torch.int32, device=DEV), + ] + ) + unused = torch.ones(total_slots, dtype=torch.bool, device=DEV) + unused[state_indices] = False + + S0 = torch.randn(total_slots, H, P, N, device=DEV, dtype=state_dtype) * 0.1 + state_spec = S0.clone() + post_conv_cache = torch.zeros( + total_slots, buf, d_inner + G * N, device=DEV, dtype=act_dtype + ) + dt_cache = torch.zeros(total_slots, H, buf, device=DEV, dtype=torch.float32) + write_pos = torch.zeros(total_slots, dtype=torch.int32, device=DEV) + post_origin = torch.zeros(total_slots, dtype=torch.int32, device=DEV) + is_flush = torch.zeros(total_slots, dtype=torch.int8, device=DEV) + + # per-row raw inputs (history + window); build the per-row oracle. + oracles = [] + conv_out = torch.zeros( + padded_batch * spec_len, conv_dim, device=DEV, dtype=act_dtype + ) + dt_spec = torch.zeros(padded_batch * spec_len, H, device=DEV, dtype=torch.float32) + z_pack = torch.zeros(padded_batch * spec_len, H, P, device=DEV, dtype=act_dtype) + for r in range(batch): + wp = wps[r] + slot = int(state_indices[r].item()) + T_tot = wp + spec_len + x = torch.randn(T_tot, H, P, device=DEV, dtype=act_dtype) + dt = torch.randn(T_tot, H, device=DEV, dtype=act_dtype) + Bv = torch.randn(T_tot, G, N, device=DEV, dtype=act_dtype) + Cv = torch.randn(T_tot, G, N, device=DEV, dtype=act_dtype) + zv = torch.randn(T_tot, H, P, device=DEV, dtype=act_dtype) + oracles.append( + _standard_window_oracle( + S0_slot=S0[slot], + x_all=x, + dt_all=dt, + B_all=Bv, + C_all=Cv, + z_all=zv, + A=A, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + wp=wp, + spec_len=spec_len, + buffer_len=buf, + act_dtype=act_dtype, + ) + ) + _scatter_packed_history( + post_conv_cache, + dt_cache, + slot, + x[:wp], + Bv[:wp], + dt[:wp], + d_inner, + G, + N, + ) + write_pos[slot] = wp + seg = slice(r * spec_len, (r + 1) * spec_len) + conv_out[seg] = _pack_window_conv_out( + x[wp:], Bv[wp:], Cv[wp:], d_inner, G, N, act_dtype + ) + dt_spec[seg] = dt[wp:].float() + z_pack[seg] = zv[wp:] + + qsl = torch.arange( + 0, (padded_batch + 1) * spec_len, spec_len, device=DEV, dtype=torch.int32 + ) + out_spec = torch.full( + (padded_batch * spec_len, H, P), 42.0, device=DEV, dtype=act_dtype + ) + selective_state_update_replayssm_spec( + state_spec, + post_conv_cache, + dt_cache, + conv_out, + dt_spec, + A, + write_pos=write_pos, + post_conv_state_pos=post_origin, + is_flush=is_flush, + query_start_loc=qsl, + state_batch_indices=sbi, + max_cache_len=L, + max_spec_len=max_spec_len, + d_inner=d_inner, + ngroups=G, + dstate=N, + D=D, + z=z_pack, + dt_bias=dt_bias, + dt_softplus=True, + out=out_spec, + ) + + for r in range(batch): + seg = slice(r * spec_len, (r + 1) * spec_len) + torch.testing.assert_close(out_spec[seg], oracles[r], rtol=rtol, atol=atol) + if with_padding: + pad = slice(batch * spec_len, padded_batch * spec_len) + assert torch.equal(out_spec[pad], torch.full_like(out_spec[pad], 42.0)) + # non-flush verify leaves every state slot untouched. + assert torch.equal(state_spec[unused], S0[unused]) + torch.testing.assert_close(state_spec, S0, rtol=0, atol=0) diff --git a/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py b/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py new file mode 100644 index 000000000000..7d730fefb18c --- /dev/null +++ b/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py @@ -0,0 +1,450 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Standard (autoregressive) decode correctness for the Mamba2 ReplaySSM kernels. + +ReplaySSM caches the recent SSM inputs ``(x, dt, B)`` in a small ring buffer and +reconstructs / reads out the recurrent state on the fly, writing the full state +back to HBM only when the buffer flushes. This file checks that, over a +multi-step decode, both ReplaySSM routes reproduce the exact SSM recurrence. + +For every step we assert, against trusted oracles driven one token at a time: + + * the ReplaySSM output matches its pure-PyTorch reference, which models the + kernel's exact arithmetic (including its bf16 reconstruction), at every + precision, + * when the state is fp32, the output also matches the baseline decode kernel + (``selective_state_update``); at bf16 state the baseline legitimately + differs -- it downcasts the fp32 state to bf16 every step, while ReplaySSM + accumulates a whole buffer in fp32 and is the more accurate path, + * when state and activations are both fp32, the output also matches the exact + elementwise ``selective_state_update_ref``, + * the cached inputs match the reference cache management, + * the checkpoint state matches the reference (and the baseline at fp32 state). + +State and activation/buffer precision are swept independently. Nemotron-3 +defaults to fp32 SSM state (``mamba_ssm_cache_dtype=float32``), but bf16 state is +also supported; the buffer dtype follows the activation dtype and ``dt_cache`` +is always fp32. +""" + +import pytest +import torch + +from tests.kernels.mamba.utils import ( + allocate_update_caches, + selective_state_update_ref, + selective_state_update_replayssm_output_only_ref, + selective_state_update_replayssm_state_and_output_ref, +) +from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_state_update +from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_output_only import ( # noqa: E501 + selective_state_update_replayssm_output_only, +) +from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_state_and_output import ( # noqa: E501 + selective_state_update_replayssm_state_and_output, +) +from vllm.utils.torch_utils import set_random_seed +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + + +def _tolerances(dtype: torch.dtype) -> tuple[float, float]: + # fp32 demands true fp32 parity. A correct fp32 reconstruction (tl.dot with + # input_precision="tf32x3"/"ieee") matches the elementwise baseline to + # ~1e-5, while a TF32 reconstruction drifts to ~1e-2. atol=1e-3 sits between, + # so it flags TF32 degradation yet passes a correct fp32 kernel. bf16 stays + # loose (bf16 rounding dominates and is unaffected by the matmul precision). + if dtype == torch.float32: + return 1e-4, 1e-3 + return 6e-2, 2e-1 + + +def _tied_A(nheads: int, headdim: int, dstate: int, device: str) -> torch.Tensor: + A = -torch.rand(nheads, device=device) - 1.0 + return A.view(nheads, 1, 1).expand(nheads, headdim, dstate) + + +def _tied_dt( + batch: int, + nheads: int, + headdim: int, + device: str, + dtype: torch.dtype, +) -> torch.Tensor: + dt = torch.randn(batch, nheads, device=device, dtype=dtype) + return dt.unsqueeze(-1).expand(batch, nheads, headdim) + + +def _tied_dt_bias(nheads: int, headdim: int, device: str) -> torch.Tensor: + dt_bias = torch.rand(nheads, device=device) - 4.0 + return dt_bias.view(nheads, 1).expand(nheads, headdim) + + +def _run_standard_decode( + *, + route: str, + state_dtype: torch.dtype, + act_dtype: torch.dtype, + batch: int, + nheads: int, + headdim: int, + ngroups: int, + dstate: int, + max_cache_len: int, + num_steps: int, + has_z: bool, + dt_softplus: bool, + use_dt_bias: bool, + desync_write_pos: bool = False, + seed: int = 0, +) -> None: + """Drive ``num_steps`` decode steps and check every path against the + trusted recurrence. ``state_dtype`` is the recurrent-state precision; + ``act_dtype`` is the activation/buffer precision.""" + device = "cuda" + both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32 + rtol, atol = _tolerances(torch.float32 if both_fp32 else torch.bfloat16) + set_random_seed(seed) + + # One state copy per path; all start identical at the same precision. + state0 = torch.randn( + batch, nheads, headdim, dstate, dtype=state_dtype, device=device) + state_anchor = state0.clone() + state_baseline = state0.clone() + state_cached = state0.clone() + state_ref = state0.clone() + + A = _tied_A(nheads, headdim, dstate, device) + dt_bias = _tied_dt_bias(nheads, headdim, device) if use_dt_bias else None + D = torch.randn(nheads, headdim, device=device) + + # Caches follow the activation dtype (dt_cache is forced to fp32 inside). + x_cache, dt_cache, B_cache, _ = allocate_update_caches( + batch, nheads, ngroups, headdim, dstate, max_cache_len, device, + act_dtype, act_dtype) + x_cache_ref, dt_cache_ref, B_cache_ref, _ = allocate_update_caches( + batch, nheads, ngroups, headdim, dstate, max_cache_len, device, + act_dtype, act_dtype) + bc_pre = torch.empty( + batch, ngroups, max_cache_len, device=device, dtype=torch.float32) + + if desync_write_pos: + # Rows start at different ring positions so they flush on different + # steps, exercising per-row write-position handling. + write_pos = (torch.arange(batch, device=device, dtype=torch.int32) + % max_cache_len) + else: + write_pos = torch.zeros(batch, dtype=torch.int32, device=device) + + for _ in range(num_steps): + x = torch.randn(batch, nheads, headdim, device=device, dtype=act_dtype) + dt = _tied_dt(batch, nheads, headdim, device, act_dtype) + B = torch.randn(batch, ngroups, dstate, device=device, dtype=act_dtype) + C = torch.randn(batch, ngroups, dstate, device=device, dtype=act_dtype) + z = torch.randn_like(x) if has_z else None + is_flush = write_pos == max_cache_len - 1 + + # Trusted recurrence (mutates state_anchor in place, returns output). + out_anchor = selective_state_update_ref( + state_anchor, x, dt, A, B, C, D=D, z=z, dt_bias=dt_bias, + dt_softplus=dt_softplus) + + # Upstream baseline decode kernel. + out_baseline = torch.empty_like(x) + selective_state_update( + state_baseline, x, dt, A, B, C, D=D, z=z, dt_bias=dt_bias, + dt_softplus=dt_softplus, out=out_baseline) + + # ReplaySSM kernel under test + its pure-PyTorch reference. + out_cached = torch.empty_like(x) + common = dict( + D=D, z=z, dt_bias=dt_bias, dt_softplus=dt_softplus, + x_cache=x_cache, dt_cache=dt_cache, B_cache=B_cache, + write_pos=write_pos, is_flush=is_flush, + max_cache_len=max_cache_len, out=out_cached) + if route == "output_only": + selective_state_update_replayssm_output_only( + state_cached, x, dt, A, B, C, bc_pre=bc_pre, **common) + out_ref = selective_state_update_replayssm_output_only_ref( + state_ref, x, dt, A, B, C, D=D, z=z, dt_bias=dt_bias, + dt_softplus=dt_softplus, x_cache=x_cache_ref, + dt_cache=dt_cache_ref, B_cache=B_cache_ref, + write_pos=write_pos, max_cache_len=max_cache_len) + else: + selective_state_update_replayssm_state_and_output( + state_cached, x, dt, A, B, C, **common) + out_ref = selective_state_update_replayssm_state_and_output_ref( + state_ref, x, dt, A, B, C, D=D, z=z, dt_bias=dt_bias, + dt_softplus=dt_softplus, x_cache=x_cache_ref, + dt_cache=dt_cache_ref, B_cache=B_cache_ref, + write_pos=write_pos, max_cache_len=max_cache_len) + + # The reference models the kernel's exact arithmetic (including its bf16 + # reconstruction), so the kernel must match it tightly at every + # precision. At fp32 this also flags any TF32 reconstruction drift. + torch.testing.assert_close(out_cached, out_ref, rtol=rtol, atol=atol) + # When the STATE is fp32 the baseline decode kernel is a valid oracle (it + # does not downcast the state per step), so ReplaySSM must match it. At + # bf16 state the baseline legitimately differs: it downcasts the fp32 + # state to bf16 every step, while ReplaySSM accumulates a whole buffer in + # fp32 and is the MORE accurate path -- so it is not a tight oracle there. + if state_dtype == torch.float32: + torch.testing.assert_close(out_cached, out_baseline, rtol=rtol, atol=atol) + # The exact elementwise reference is valid only when state AND + # activations are fp32; otherwise it downcasts the state (at readout for + # bf16 activations, or per step for bf16 state). + if both_fp32: + torch.testing.assert_close(out_cached, out_anchor, rtol=rtol, atol=atol) + + # Cached inputs match the reference cache management. + torch.testing.assert_close(x_cache, x_cache_ref, rtol=rtol, atol=atol) + torch.testing.assert_close(dt_cache, dt_cache_ref, rtol=rtol, atol=atol) + torch.testing.assert_close(B_cache, B_cache_ref, rtol=rtol, atol=atol) + + # Checkpoint state at flush matches the reference (and, when the state is + # fp32, the baseline kernel). + if bool(is_flush.any()): + torch.testing.assert_close( + state_cached[is_flush], state_ref[is_flush], + rtol=rtol, atol=atol) + if state_dtype == torch.float32: + torch.testing.assert_close( + state_cached[is_flush], state_baseline[is_flush], + rtol=rtol, atol=atol) + + write_pos = torch.where( + is_flush, torch.zeros_like(write_pos), write_pos + 1) + + +# State/activation precisions. fp32 state is the default; bf16/fp16 are the +# reduced-footprint configs. fp16 appears both as an activation dtype (fully-fp16 +# model sfp16_afp16, or fp16 act over fp32 state s32_afp16) and as a state dtype +# under bf16 activations (sfp16_a16): fp16 has a finer mantissa than bf16 at the +# same 2 bytes, so it is a more accurate state at no extra footprint. We still +# skip fp16 state under fp32 activations (the unused low-state/high-act mix). +_PRECISIONS = [ + pytest.param((torch.float32, torch.float32), id="s32_a32"), + pytest.param((torch.float32, torch.bfloat16), id="s32_a16"), + pytest.param((torch.bfloat16, torch.bfloat16), id="s16_a16"), + pytest.param((torch.float32, torch.float16), id="s32_afp16"), + pytest.param((torch.float16, torch.float16), id="sfp16_afp16"), + pytest.param((torch.float16, torch.bfloat16), id="sfp16_a16"), +] +# Small synthetic shapes for the full axis sweep (compile fast). +_SMALL_GEOMETRIES = [ + pytest.param((8, 64, 64, 4), id="small"), + pytest.param((4, 64, 16, 1), id="tiny"), +] +# Production Mamba2 shapes (nheads, headdim, dstate, ngroups), TP=1. +_REAL_GEOMETRIES = [ + pytest.param((96, 80, 128, 8), id="nano4b"), + pytest.param((128, 64, 128, 8), id="super120b"), + pytest.param((256, 64, 128, 8), id="ultra550b"), +] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("route", ["output_only", "state_and_output"]) +@pytest.mark.parametrize("precision", _PRECISIONS) +@pytest.mark.parametrize("max_cache_len", [1, 4, 16]) +@pytest.mark.parametrize("geometry", _SMALL_GEOMETRIES) +@pytest.mark.parametrize("has_z", [False, True]) +def test_replayssm_standard_decode_matches_reference( + route: str, + precision: tuple[torch.dtype, torch.dtype], + max_cache_len: int, + geometry: tuple[int, int, int, int], + has_z: bool, +): + state_dtype, act_dtype = precision + nheads, headdim, dstate, ngroups = geometry + _run_standard_decode( + route=route, + state_dtype=state_dtype, + act_dtype=act_dtype, + batch=4, + nheads=nheads, + headdim=headdim, + ngroups=ngroups, + dstate=dstate, + max_cache_len=max_cache_len, + num_steps=2 * max_cache_len + 1, + has_z=has_z, + dt_softplus=True, + use_dt_bias=True, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("route", ["output_only", "state_and_output"]) +@pytest.mark.parametrize("precision", _PRECISIONS) +@pytest.mark.parametrize("geometry", _REAL_GEOMETRIES) +def test_replayssm_standard_decode_real_geometry( + route: str, + precision: tuple[torch.dtype, torch.dtype], + geometry: tuple[int, int, int, int], +): + # Production Mamba2 shapes for the Nemotron-3 family (Nano-4B / Super-120B / + # Ultra-550B), at the production buffer length (8). All three precisions, + # including the bf16 state case. + state_dtype, act_dtype = precision + nheads, headdim, dstate, ngroups = geometry + _run_standard_decode( + route=route, + state_dtype=state_dtype, + act_dtype=act_dtype, + batch=4, + nheads=nheads, + headdim=headdim, + ngroups=ngroups, + dstate=dstate, + max_cache_len=8, + num_steps=17, + has_z=True, + dt_softplus=True, + use_dt_bias=True, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("route", ["output_only", "state_and_output"]) +@pytest.mark.parametrize( + "precision", + [ + pytest.param((torch.float32, torch.float32), id="s32_a32"), + pytest.param((torch.float32, torch.bfloat16), id="s32_a16"), + pytest.param((torch.float32, torch.float16), id="s32_afp16"), + ], +) +def test_replayssm_standard_decode_desync_write_pos( + route: str, + precision: tuple[torch.dtype, torch.dtype], +): + # Rows start at staggered ring positions, so they flush on different steps + # and hold genuinely different cached histories. + state_dtype, act_dtype = precision + _run_standard_decode( + route=route, + state_dtype=state_dtype, + act_dtype=act_dtype, + batch=4, + nheads=8, + headdim=64, + ngroups=4, + dstate=64, + max_cache_len=4, + num_steps=12, + has_z=True, + dt_softplus=True, + use_dt_bias=True, + desync_write_pos=True, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("route", ["output_only", "state_and_output"]) +@pytest.mark.parametrize("precision", _PRECISIONS) +@pytest.mark.parametrize("with_padding", [False, True]) +def test_replayssm_standard_decode_with_batch_indices( + route: str, + precision: tuple[torch.dtype, torch.dtype], + with_padding: bool, +): + # Sparse state allocation via state_batch_indices, with NULL_BLOCK_ID + # padding rows. The pure-PyTorch references do not model the sparse + # gather, so the anchor here is the upstream baseline decode kernel. + state_dtype, act_dtype = precision + device = "cuda" + both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32 + rtol, atol = _tolerances(torch.float32 if both_fp32 else torch.bfloat16) + set_random_seed(0) + + batch = 3 + padding = 2 if with_padding else 0 + padded_batch = batch + padding + total_state_slots = 16 + nheads = 4 + ngroups = 2 + headdim = 64 + dstate = 16 + max_cache_len = 4 + num_steps = 2 * max_cache_len + + state = torch.randn( + total_state_slots, nheads, headdim, dstate, dtype=state_dtype, + device=device) + state_baseline = state.clone() + state_cached = state.clone() + state_before = state.clone() + + state_indices = ( + torch.randperm(total_state_slots - 1, device=device)[:batch] + 1 + ).to(torch.int32) + state_batch_indices = torch.cat([ + state_indices, + torch.full((padding,), NULL_BLOCK_ID, dtype=torch.int32, device=device), + ]) + unused_states = torch.ones(total_state_slots, dtype=torch.bool, device=device) + unused_states[state_indices] = False + + A = _tied_A(nheads, headdim, dstate, device) + dt_bias = _tied_dt_bias(nheads, headdim, device) + D = torch.randn(nheads, headdim, device=device) + x_cache = torch.zeros( + total_state_slots, nheads, max_cache_len, headdim, device=device, + dtype=act_dtype) + dt_cache = torch.zeros( + total_state_slots, nheads, max_cache_len, device=device, + dtype=torch.float32) + B_cache = torch.zeros( + total_state_slots, ngroups, max_cache_len, dstate, device=device, + dtype=act_dtype) + bc_pre = torch.empty( + padded_batch, ngroups, max_cache_len, device=device, dtype=torch.float32) + write_pos = torch.zeros(padded_batch, dtype=torch.int32, device=device) + + for _ in range(num_steps): + x = torch.randn(padded_batch, nheads, headdim, device=device, + dtype=act_dtype) + dt = _tied_dt(padded_batch, nheads, headdim, device, act_dtype) + B = torch.randn(padded_batch, ngroups, dstate, device=device, + dtype=act_dtype) + C = torch.randn(padded_batch, ngroups, dstate, device=device, + dtype=act_dtype) + z = torch.randn_like(x) + is_flush = write_pos == max_cache_len - 1 + + out_baseline = torch.empty_like(x) + selective_state_update( + state_baseline, x, dt, A, B, C, D=D, z=z, dt_bias=dt_bias, + dt_softplus=True, state_batch_indices=state_batch_indices, + out=out_baseline) + + out_cached = torch.full_like(x, 42) + common = dict( + D=D, z=z, dt_bias=dt_bias, dt_softplus=True, x_cache=x_cache, + dt_cache=dt_cache, B_cache=B_cache, write_pos=write_pos, + is_flush=is_flush, max_cache_len=max_cache_len, + state_batch_indices=state_batch_indices, out=out_cached) + if route == "output_only": + selective_state_update_replayssm_output_only( + state_cached, x, dt, A, B, C, bc_pre=bc_pre, **common) + else: + selective_state_update_replayssm_state_and_output( + state_cached, x, dt, A, B, C, **common) + + torch.testing.assert_close( + out_cached[:batch], out_baseline[:batch], rtol=rtol, atol=atol) + if with_padding: + assert torch.equal( + out_cached[batch:], torch.full_like(out_cached[batch:], 42)) + + if bool(is_flush[:batch].all()): + torch.testing.assert_close( + state_cached[state_indices], state_baseline[state_indices], + rtol=rtol, atol=atol) + + write_pos = torch.where( + is_flush, torch.zeros_like(write_pos), write_pos + 1) + + assert torch.equal(state_cached[unused_states], state_before[unused_states]) + assert torch.equal(state_baseline[unused_states], state_before[unused_states]) diff --git a/tests/kernels/mamba/utils.py b/tests/kernels/mamba/utils.py index fb8a4b0a28ec..40d766bb3bf2 100644 --- a/tests/kernels/mamba/utils.py +++ b/tests/kernels/mamba/utils.py @@ -76,3 +76,289 @@ def selective_state_update_ref( if not has_heads: out = out.squeeze(1) return out + + +def selective_state_update_replayssm_state_and_output_ref( + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + z: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = False, + x_cache: torch.Tensor | None = None, + dt_cache: torch.Tensor | None = None, + B_cache: torch.Tensor | None = None, + write_pos: torch.Tensor | None = None, + max_cache_len: int = 16, +) -> torch.Tensor: + """Pure-PyTorch cached-dot reference for validation.""" + has_heads = state.dim() > 3 + if state.dim() == 3: + state = state.unsqueeze(1) + if x.dim() == 2: + x = x.unsqueeze(1) + if dt.dim() == 2: + dt = dt.unsqueeze(1) + if A.dim() == 2: + A = A.unsqueeze(0) + if B.dim() == 2: + B = B.unsqueeze(1) + if C.dim() == 2: + C = C.unsqueeze(1) + if D is not None and D.dim() == 1: + D = D.unsqueeze(0) + if z is not None and z.dim() == 2: + z = z.unsqueeze(1) + if dt_bias is not None and dt_bias.dim() == 1: + dt_bias = dt_bias.unsqueeze(0) + + batch, nheads, dim, dstate = state.shape + assert x.shape == (batch, nheads, dim) + assert dt.shape == x.shape + assert A.shape == (nheads, dim, dstate) + ngroups = B.shape[1] + assert nheads % ngroups == 0, "nheads must be divisible by ngroups" + assert B.shape == (batch, ngroups, dstate) + assert C.shape == B.shape + if D is not None: + assert D.shape == (nheads, dim) + if z is not None: + assert z.shape == x.shape + if dt_bias is not None: + assert dt_bias.shape == (nheads, dim) + + assert x_cache is not None + assert dt_cache is not None + assert B_cache is not None + assert x_cache.shape == (batch, nheads, max_cache_len, dim) + assert dt_cache.shape == (batch, nheads, max_cache_len) + assert B_cache.shape == (batch, ngroups, max_cache_len, dstate) + assert write_pos is not None + assert write_pos.shape == (batch,) and write_pos.dtype == torch.int32 + + ratio = nheads // ngroups + + dt_val = dt[:, :, 0].float() + if dt_bias is not None: + dt_val = dt_val + dt_bias[:, 0].float() + if dt_softplus: + dt_val = F.softplus(dt_val) + A_val = A[:, 0, 0].float() + C_heads = C.repeat_interleave(ratio, dim=1) + out = torch.empty(batch, nheads, dim, device=x.device, dtype=torch.float32) + + for b in range(batch): + cache_len = int(write_pos[b].item()) + is_flush = cache_len == max_cache_len - 1 + n_steps = cache_len + 1 + + dt_all = torch.zeros(nheads, n_steps, device=x.device, dtype=torch.float32) + if cache_len > 0: + dt_all[:, :cache_len] = dt_cache[b, :, :cache_len] + dt_all[:, cache_len] = dt_val[b] + + cumsum = torch.cumsum(dt_all, dim=-1) + total = cumsum[:, -1] + dA_cumsum = A_val[:, None] * cumsum + dA_total = A_val * total + total_decay = torch.exp(dA_total) + scale = dt_all * torch.exp(dA_total[:, None] - dA_cumsum) + + x_all = torch.zeros(nheads, dim, n_steps, device=x.device, dtype=x.dtype) + if cache_len > 0: + x_all[..., :cache_len] = x_cache[b, :, :cache_len, :].permute(0, 2, 1) + x_all[..., cache_len] = x[b] + + B_all = torch.zeros(ngroups, n_steps, dstate, device=B.device, dtype=B.dtype) + if cache_len > 0: + B_all[:, :cache_len, :] = B_cache[b, :, :cache_len, :] + B_all[:, cache_len, :] = B[b] + + B_heads = B_all.repeat_interleave(ratio, dim=0) + # Mirror the kernel: it scales B and rounds B*scale to the buffer dtype + # before the tl.dot reconstruction (the matmul runs on buffer-dtype + # operands). For bf16 buffers this rounds; for fp32 it is a no-op. + B_scaled = (B_heads.float() * scale[:, :, None]).to(B_heads.dtype) + delta = torch.einsum("hdk,hkn->hdn", x_all.float(), B_scaled.float()) + state_new = state[b].float() * total_decay[:, None, None] + delta + if is_flush: + state[b].copy_(state_new.to(state.dtype)) + else: + x_cache[b, :, cache_len, :] = x[b] + dt_cache[b, :, cache_len] = dt_val[b] + B_cache[b, :, cache_len, :] = B[b] + + out[b] = torch.einsum("hdn,hn->hd", state_new, C_heads[b].float()) + if D is not None: + out = out + (x.float() * D[None]).to(out.dtype) + if z is not None: + out = out * F.silu(z.float()) + out = out.to(x.dtype) + if not has_heads: + out = out.squeeze(1) + return out + + +def selective_state_update_replayssm_output_only_ref( + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + z: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = False, + x_cache: torch.Tensor | None = None, + dt_cache: torch.Tensor | None = None, + B_cache: torch.Tensor | None = None, + write_pos: torch.Tensor | None = None, + max_cache_len: int = 16, +) -> torch.Tensor: + """Pure-PyTorch cached-bc reference for validation.""" + has_heads = state.dim() > 3 + if state.dim() == 3: + state = state.unsqueeze(1) + if x.dim() == 2: + x = x.unsqueeze(1) + if dt.dim() == 2: + dt = dt.unsqueeze(1) + if A.dim() == 2: + A = A.unsqueeze(0) + if B.dim() == 2: + B = B.unsqueeze(1) + if C.dim() == 2: + C = C.unsqueeze(1) + if D is not None and D.dim() == 1: + D = D.unsqueeze(0) + if z is not None and z.dim() == 2: + z = z.unsqueeze(1) + if dt_bias is not None and dt_bias.dim() == 1: + dt_bias = dt_bias.unsqueeze(0) + + batch, nheads, dim, dstate = state.shape + assert x.shape == (batch, nheads, dim) + assert dt.shape == x.shape + assert A.shape == (nheads, dim, dstate) + ngroups = B.shape[1] + assert nheads % ngroups == 0, "nheads must be divisible by ngroups" + assert B.shape == (batch, ngroups, dstate) + assert C.shape == B.shape + + ratio = nheads // ngroups + + dt_val = dt[:, :, 0].float() + if dt_bias is not None: + dt_val = dt_val + dt_bias[:, 0].float() + if dt_softplus: + dt_val = F.softplus(dt_val) + A_val = A[:, 0, 0].float() + C_heads = C.repeat_interleave(ratio, dim=1) + out = torch.empty(batch, nheads, dim, device=x.device, dtype=torch.float32) + + assert x_cache is not None + assert dt_cache is not None + assert B_cache is not None + assert write_pos is not None + + for b in range(batch): + cache_len = int(write_pos[b].item()) + is_flush = cache_len == max_cache_len - 1 + n_steps = cache_len + 1 + + dt_all = torch.zeros(nheads, n_steps, device=x.device, dtype=torch.float32) + if cache_len > 0: + dt_all[:, :cache_len] = dt_cache[b, :, :cache_len] + dt_all[:, cache_len] = dt_val[b] + + cumsum = torch.cumsum(dt_all, dim=-1) + total = cumsum[:, -1] + dA_cumsum = A_val[:, None] * cumsum + dA_total = A_val * total + total_decay = torch.exp(dA_total) + scale = dt_all * torch.exp(dA_total[:, None] - dA_cumsum) + + x_all = torch.zeros(nheads, dim, n_steps, device=x.device, dtype=x.dtype) + if cache_len > 0: + x_all[..., :cache_len] = x_cache[b, :, :cache_len, :].permute(0, 2, 1) + x_all[..., cache_len] = x[b] + + B_all = torch.zeros(ngroups, n_steps, dstate, device=B.device, dtype=B.dtype) + if cache_len > 0: + B_all[:, :cache_len, :] = B_cache[b, :, :cache_len, :] + B_all[:, cache_len, :] = B[b] + + B_heads = B_all.repeat_interleave(ratio, dim=0) + C_heads_b = C_heads[b] + + if is_flush: + # Mirror the kernel's bf16 reconstruction (see the cached-dot ref). + B_scaled = (B_heads.float() * scale[:, :, None]).to(B_heads.dtype) + delta = torch.einsum("hdk,hkn->hdn", x_all.float(), B_scaled.float()) + state_new = state[b].float() * total_decay[:, None, None] + delta + state[b].copy_(state_new.to(state.dtype)) + out[b] = torch.einsum("hdn,hn->hd", state_new, C_heads_b.float()) + else: + checkpoint_out = torch.einsum( + "hdn,hn->hd", state[b].float(), C_heads_b.float() + ) + checkpoint_out = checkpoint_out * total_decay[:, None] + BC = torch.einsum("hkn,hn->hk", B_heads.float(), C_heads_b.float()) + cache_out = torch.einsum("hdk,hk->hd", x_all.float(), scale * BC) + out[b] = checkpoint_out + cache_out + x_cache[b, :, cache_len, :] = x[b] + dt_cache[b, :, cache_len] = dt_val[b] + B_cache[b, :, cache_len, :] = B[b] + + if D is not None: + out = out + (x.float() * D[None]).to(out.dtype) + if z is not None: + out = out * F.silu(z.float()) + out = out.to(x.dtype) + if not has_heads: + out = out.squeeze(1) + return out + + +def allocate_update_caches( + batch: int, + nheads: int, + ngroups: int, + dim: int, + dstate: int, + max_cache_len: int, + device: torch.device, + x_dtype: torch.dtype, + B_dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Allocate dense reference caches for standalone validation.""" + x_cache = torch.zeros( + batch, + nheads, + max_cache_len, + dim, + device=device, + dtype=x_dtype, + ) + dt_cache = torch.zeros( + batch, + nheads, + max_cache_len, + device=device, + dtype=torch.float32, + ) + B_cache = torch.zeros( + batch, + ngroups, + max_cache_len, + dstate, + device=device, + dtype=B_dtype, + ) + write_pos = torch.zeros(batch, dtype=torch.int32, device=device) + return x_cache, dt_cache, B_cache, write_pos diff --git a/tests/kernels/test_fused_recurrent_packed_decode.py b/tests/kernels/test_fused_recurrent_packed_decode.py index d63186bde118..c51142718885 100644 --- a/tests/kernels/test_fused_recurrent_packed_decode.py +++ b/tests/kernels/test_fused_recurrent_packed_decode.py @@ -41,11 +41,11 @@ def test_fused_recurrent_packed_decode_matches_reference( A_log = torch.randn((HV,), device=device, dtype=dtype) dt_bias = torch.randn((HV,), device=device, dtype=dtype) - # Continuous batching indices (include PAD_SLOT_ID=-1 cases). - ssm_state_indices = torch.arange(B, device=device, dtype=torch.int32) + # Continuous batching indices (slot 0 is the null block; include PAD_SLOT_ID=-1). + ssm_state_indices = torch.arange(1, B + 1, device=device, dtype=torch.int32) ssm_state_indices[-3:] = -1 - state0 = torch.randn((B, HV, V, K), device=device, dtype=dtype) + state0 = torch.randn((B + 1, HV, V, K), device=device, dtype=dtype) state_ref = state0.clone() state_packed = state0.clone() diff --git a/tests/kernels/test_replayssm_spec_decode_gdn.py b/tests/kernels/test_replayssm_spec_decode_gdn.py new file mode 100644 index 000000000000..9501cb6c2d8c --- /dev/null +++ b/tests/kernels/test_replayssm_spec_decode_gdn.py @@ -0,0 +1,663 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Speculative-decode correctness for the GDN (Gated DeltaNet) ReplaySSM kernel. + +A spec/verify step takes the checkpoint state + the committed ring buffer (the +per-step delta-rule ``d``/``k``/``g`` caches) + a window of ``max_spec_len = +1 + num_speculative_tokens`` draft tokens, and computes the recurrence OUTPUT at +each window position (causal). It does NOT write the state per draft -- only a +flush step folds the *committed* history into the checkpoint, so unaccepted +drafts can be rolled back. + +The history window is ``L = B + max_spec_len`` (block ``B`` = ``buffer_len``): the +physical circular buffer is ``next_pow2(L)`` and ``max_cache_len`` passed to the +kernel is the logical ``L`` (the history block ``BC = next_pow2(B)`` is derived +inside the kernel). + +GDN's standard decode kernel (``fused_recurrent_gated_delta_rule_replayssm``) and +the spec kernel SHARE the d/k/g cache layout (the spec one is circular via +``cache_base``; the standard one is linear). With ``cache_base=0`` they coincide, +so the oracle is built directly: + * build the committed history by running the standard kernel for ``wp`` tokens + (no flush, so the checkpoint stays fixed), + * spec: one verify call over the window reading those very caches, + * oracle: the standard kernel stepped one token at a time over the SAME window + from the post-history checkpoint+caches. +For the multi-step rollback the ground truth is the baseline +``fused_recurrent_gated_delta_rule_packed_decode`` over the accepted token stream. + +Precision: the spec kernel uses a chunked UT-transform whereas the standard / +packed kernels use the sequential recurrence, so even at fp32 they differ by +~1e-3 (an arithmetic-ordering difference, not a bug). bf16 adds the usual rounding +(rel ~1e-2), bounded per step and non-accumulating across the decode. +""" + +import pytest +import torch + +from vllm.model_executor.layers.fla.ops import ( + fused_recurrent_gated_delta_rule_packed_decode, + fused_recurrent_gated_delta_rule_replayssm, +) +from vllm.model_executor.layers.fla.ops.gdn_replayssm_spec_decode import ( + commit_gdn_replayssm_spec, + gdn_replayssm_spec_decode, + reset_gdn_replayssm_spec_cursors, +) +from vllm.utils.torch_utils import set_random_seed + +DEV = "cuda" +PAD_SLOT_ID = -1 + + +def _lbt(base_block: int, max_spec_len: int) -> tuple[int, int]: + """Logical flush threshold L = B + max_spec_len and physical pow2 buffer.""" + L = base_block + max_spec_len + buf = 1 << (L - 1).bit_length() + return L, buf + + +def _output_tol(both_fp32: bool) -> tuple[float, float]: + # Spec (chunked) vs standard/baseline (sequential): fp32 differs ~1.7e-3 + # by arithmetic ordering; bf16 adds rounding (~1e-2). atol is a small floor. + if both_fp32: + return 5e-3, 2e-3 + return 4e-2, 1e-2 + + +def _state_tol(both_fp32: bool) -> tuple[float, float]: + if both_fp32: + return 5e-3, 3e-3 + return 4e-2, 2e-2 + + +def _build_history( + *, + mqkv, + a, + b, + A_log, + dt_bias, + scale, + state, + d_cache, + k_cache, + g_cache, + slot: int, + wp: int, +): + """Run the standard decode for the first ``wp`` tokens into ``slot`` (no + flush, since wp <= B - max_spec_len < buf - 1), populating the shared caches + and leaving the checkpoint unchanged.""" + sbi = torch.tensor([slot], device=DEV, dtype=torch.int32) + for t in range(wp): + ot = torch.empty(1, 1, *state.shape[1:3], device=DEV, dtype=mqkv.dtype) + fused_recurrent_gated_delta_rule_replayssm( + mixed_qkv=mqkv[t : t + 1], + a=a[t : t + 1], + b=b[t : t + 1], + A_log=A_log, + dt_bias=dt_bias, + scale=scale, + initial_state=state, + d_cache=d_cache, + k_cache=k_cache, + g_cache=g_cache, + out=ot, + ssm_state_indices=sbi, + write_pos=torch.tensor([t], device=DEV, dtype=torch.int32), + use_qk_l2norm_in_kernel=True, + ) + + +def _standard_window_oracle( + *, + mqkv, + a, + b, + A_log, + dt_bias, + scale, + state, + d_cache, + k_cache, + g_cache, + slot: int, + wp: int, + spec_len: int, + HV: int, + V: int, +) -> torch.Tensor: + """Continue the standard decode over the window on CLONES of the post-history + state/caches; return the window outputs (spec_len, HV, V).""" + st = state.clone() + d_o, k_o, g_o = d_cache.clone(), k_cache.clone(), g_cache.clone() + sbi = torch.tensor([slot], device=DEV, dtype=torch.int32) + win = [] + for s in range(spec_len): + ot = torch.empty(1, 1, HV, V, device=DEV, dtype=mqkv.dtype) + fused_recurrent_gated_delta_rule_replayssm( + mixed_qkv=mqkv[wp + s : wp + s + 1], + a=a[wp + s : wp + s + 1], + b=b[wp + s : wp + s + 1], + A_log=A_log, + dt_bias=dt_bias, + scale=scale, + initial_state=st, + d_cache=d_o, + k_cache=k_o, + g_cache=g_o, + out=ot, + ssm_state_indices=sbi, + write_pos=torch.tensor([wp + s], device=DEV, dtype=torch.int32), + use_qk_l2norm_in_kernel=True, + ) + win.append(ot.reshape(1, HV, V).clone()) + return torch.cat(win, dim=0) + + +def _run_single_step( + *, + state_dtype, + act_dtype, + HQ, + HV, + K, + V, + buffer_len, # history block B + max_spec_len, + wp, + seed=0, + perturb=False, +): + set_random_seed(seed) + H = HQ + spec_len = max_spec_len + scale = K**-0.5 + qkv_dim = 2 * H * K + HV * V + L, buf = _lbt(buffer_len, max_spec_len) + both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32 + rtol, atol = _output_tol(both_fp32) + + num_slots = 2 + slot = 1 + A_log = torch.randn(HV, device=DEV, dtype=torch.float32) + dt_bias = torch.randn(HV, device=DEV, dtype=torch.float32) + + T_tot = wp + spec_len + mqkv = torch.randn(T_tot, qkv_dim, device=DEV, dtype=act_dtype) + a = torch.randn(T_tot, HV, device=DEV, dtype=act_dtype) + b = torch.randn(T_tot, HV, device=DEV, dtype=act_dtype) + S0 = torch.randn(num_slots, HV, V, K, device=DEV, dtype=state_dtype) * 0.1 + + d_cache = torch.zeros(num_slots, HV, buf, V, device=DEV, dtype=act_dtype) + k_cache = torch.zeros(num_slots, H, buf, K, device=DEV, dtype=act_dtype) + g_cache = torch.zeros(num_slots, HV, buf, device=DEV, dtype=torch.float32) + + state = S0.clone() + _build_history( + mqkv=mqkv, + a=a, + b=b, + A_log=A_log, + dt_bias=dt_bias, + scale=scale, + state=state, + d_cache=d_cache, + k_cache=k_cache, + g_cache=g_cache, + slot=slot, + wp=wp, + ) + # history build must not have touched the checkpoint (no flush). + torch.testing.assert_close(state, S0, rtol=0, atol=0) + + oracle = _standard_window_oracle( + mqkv=mqkv, + a=a, + b=b, + A_log=A_log, + dt_bias=dt_bias, + scale=scale, + state=state, + d_cache=d_cache, + k_cache=k_cache, + g_cache=g_cache, + slot=slot, + wp=wp, + spec_len=spec_len, + HV=HV, + V=V, + ) + + if perturb: + # teeth: corrupt one cached history key -> outputs must diverge. + k_cache[slot, 0, max(0, wp - 1), 0] += 5.0 + + state_spec = state.clone() + write_pos = torch.zeros(num_slots, dtype=torch.int32, device=DEV) + write_pos[slot] = wp + cache_base = torch.zeros(num_slots, dtype=torch.int32, device=DEV) + is_flush = torch.zeros(num_slots, dtype=torch.int8, device=DEV) + qsl = torch.tensor([0, spec_len], device=DEV, dtype=torch.int32) + sbi = torch.tensor([slot], device=DEV, dtype=torch.int32) + out_spec = torch.empty(spec_len, HV, V, device=DEV, dtype=act_dtype) + gdn_replayssm_spec_decode( + mixed_qkv=mqkv[wp:], + a=a[wp:], + b=b[wp:], + A_log=A_log, + dt_bias=dt_bias, + checkpoint_state=state_spec, + d_cache=d_cache, + k_cache=k_cache, + g_cache=g_cache, + out=out_spec, + query_start_loc=qsl, + ssm_state_indices=sbi, + write_pos=write_pos, + cache_base=cache_base, + is_flush=is_flush, + max_cache_len=L, + max_spec_len=max_spec_len, + scale=scale, + use_qk_l2norm_in_kernel=True, + ) + + torch.testing.assert_close(out_spec, oracle, rtol=rtol, atol=atol) + # non-flush verify must not touch the checkpoint. + torch.testing.assert_close(state_spec, S0, rtol=0, atol=0) + + +def _wp_set(base_block: int, max_spec_len: int) -> list[int]: + # Verify-path fills up to the max non-flush fill C = B - max_spec_len. + C = base_block - max_spec_len + return sorted({0, max(0, C // 2), max(0, C)}) + + +_PRECISIONS = [ + pytest.param((torch.float32, torch.float32), id="s32_a32"), + pytest.param((torch.float32, torch.bfloat16), id="s32_a16"), + pytest.param((torch.bfloat16, torch.bfloat16), id="s16_a16"), +] +# (num_q_heads, num_v_heads, head_k_dim, head_v_dim). The full sweep runs on the +# small shape; the production Qwen3.5 GDN shapes are exercised by +# test_spec_step_real_geometry. +_SMALL = pytest.param((2, 4, 64, 64), id="small") +_REAL_GEOMETRIES = [ + pytest.param((16, 32, 128, 128), id="qwen4b"), + pytest.param((16, 64, 128, 128), id="qwen122b"), +] +_BASE_BLOCKS = [16, 32] # history block B (replayssm_buffer_len) +_MAX_SPEC_LENS = [2, 4, 6, 8] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("precision", _PRECISIONS) +@pytest.mark.parametrize("geometry", [_SMALL]) +@pytest.mark.parametrize("base_block", _BASE_BLOCKS) +@pytest.mark.parametrize("max_spec_len", _MAX_SPEC_LENS) +def test_spec_step_matches_standard_decode( + precision, geometry, base_block, max_spec_len +): + state_dtype, act_dtype = precision + HQ, HV, K, V = geometry + for wp in _wp_set(base_block, max_spec_len): + _run_single_step( + state_dtype=state_dtype, + act_dtype=act_dtype, + HQ=HQ, + HV=HV, + K=K, + V=V, + buffer_len=base_block, + max_spec_len=max_spec_len, + wp=wp, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("precision", _PRECISIONS) +@pytest.mark.parametrize("geometry", _REAL_GEOMETRIES) +@pytest.mark.parametrize("max_spec_len", [2, 8]) +def test_spec_step_real_geometry(precision, geometry, max_spec_len): + # Production Qwen3.5 GDN shapes (4B v_heads=32 / 122B-A10B v_heads=64) at the + # deployable block B=16, both ends of the max_spec_len range. + state_dtype, act_dtype = precision + HQ, HV, K, V = geometry + for wp in _wp_set(16, max_spec_len): + _run_single_step( + state_dtype=state_dtype, + act_dtype=act_dtype, + HQ=HQ, + HV=HV, + K=K, + V=V, + buffer_len=16, + max_spec_len=max_spec_len, + wp=wp, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +def test_spec_step_teeth(): + _run_single_step( + state_dtype=torch.float32, + act_dtype=torch.float32, + HQ=2, + HV=4, + K=64, + V=64, + buffer_len=16, + max_spec_len=4, + wp=5, + ) + with pytest.raises(AssertionError): + _run_single_step( + state_dtype=torch.float32, + act_dtype=torch.float32, + HQ=2, + HV=4, + K=64, + V=64, + buffer_len=16, + max_spec_len=4, + wp=5, + perturb=True, + ) + + +def _run_rollback( + *, + state_dtype, + act_dtype, + HQ, + HV, + K, + V, + buffer_len, # history block B + max_spec_len, + num_steps=40, + seed=0, +): + """Drive the verify+commit loop from an empty buffer; accept a random k each + step. Baseline packed decode of the accepted stream is the output ground + truth; the committed checkpoint at each flush must match the baseline state at + the matching folded-token count (n_folded bookkeeping).""" + set_random_seed(seed) + H = HQ + scale = K**-0.5 + qkv_dim = 2 * H * K + HV * V + L, buf = _lbt(buffer_len, max_spec_len) + both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32 + o_rtol, o_atol = _output_tol(both_fp32) + s_rtol, s_atol = _state_tol(both_fp32) + + num_slots = 2 + slot = 1 + A_log = torch.randn(HV, device=DEV, dtype=torch.float32) + dt_bias = torch.randn(HV, device=DEV, dtype=torch.float32) + S0 = torch.randn(num_slots, HV, V, K, device=DEV, dtype=state_dtype) * 0.1 + state_spec = S0.clone() + state_base = S0.clone() # full pool; row in slot 1 + + d_cache = torch.zeros(num_slots, HV, buf, V, device=DEV, dtype=act_dtype) + k_cache = torch.zeros(num_slots, H, buf, K, device=DEV, dtype=act_dtype) + g_cache = torch.zeros(num_slots, HV, buf, device=DEV, dtype=torch.float32) + + write_pos = torch.zeros(num_slots, dtype=torch.int32, device=DEV) + cache_base = torch.zeros(num_slots, dtype=torch.int32, device=DEV) + is_flush = torch.zeros(num_slots, dtype=torch.int8, device=DEV) + sbi = torch.tensor([slot], device=DEV, dtype=torch.int32) + reset_gdn_replayssm_spec_cursors( + write_pos, + cache_base, + is_flush, + torch.ones(1, dtype=torch.int32, device=DEV), + sbi, + L, + max_spec_len, + ) + + n_folded = 0 + snapshots = {0: state_base[slot].clone()} + total_accepted = 0 + g = torch.Generator(device="cpu").manual_seed(seed + 1) + + for _ in range(num_steps): + spec_len = max_spec_len + mqkv = torch.randn(spec_len, qkv_dim, device=DEV, dtype=act_dtype) + a = torch.randn(spec_len, HV, device=DEV, dtype=act_dtype) + b = torch.randn(spec_len, HV, device=DEV, dtype=act_dtype) + qsl = torch.tensor([0, spec_len], device=DEV, dtype=torch.int32) + out_spec = torch.empty(spec_len, HV, V, device=DEV, dtype=act_dtype) + + wp_before = int(write_pos[slot].item()) + flush_before = int(is_flush[slot].item()) + gdn_replayssm_spec_decode( + mixed_qkv=mqkv, + a=a, + b=b, + A_log=A_log, + dt_bias=dt_bias, + checkpoint_state=state_spec, + d_cache=d_cache, + k_cache=k_cache, + g_cache=g_cache, + out=out_spec, + query_start_loc=qsl, + ssm_state_indices=sbi, + write_pos=write_pos, + cache_base=cache_base, + is_flush=is_flush, + max_cache_len=L, + max_spec_len=max_spec_len, + scale=scale, + use_qk_l2norm_in_kernel=True, + ) + if flush_before and wp_before > 0: + n_folded += wp_before + + k = int(torch.randint(1, spec_len + 1, (1,), generator=g).item()) + base_out = [] + for s in range(k): + ot = torch.empty(1, 1, HV, V, device=DEV, dtype=act_dtype) + fused_recurrent_gated_delta_rule_packed_decode( + mixed_qkv=mqkv[s : s + 1], + a=a[s : s + 1], + b=b[s : s + 1], + A_log=A_log, + dt_bias=dt_bias, + scale=scale, + initial_state=state_base, + out=ot, + ssm_state_indices=sbi, + use_qk_l2norm_in_kernel=True, + ) + base_out.append(ot.reshape(1, HV, V).clone()) + total_accepted += 1 + snapshots[total_accepted] = state_base[slot].clone() + base_out = torch.cat(base_out, dim=0) + + # (a) accepted-position outputs track the baseline decode. + torch.testing.assert_close(out_spec[:k], base_out, rtol=o_rtol, atol=o_atol) + + commit_gdn_replayssm_spec( + write_pos, + cache_base, + is_flush, + torch.tensor([k], device=DEV, dtype=torch.int32), + sbi, + L, + max_spec_len, + ) + + # (b) committed checkpoint == baseline state at the folded-token count. + if n_folded in snapshots: + torch.testing.assert_close( + state_spec[slot], snapshots[n_folded], rtol=s_rtol, atol=s_atol + ) + + assert n_folded > 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("precision", _PRECISIONS) +@pytest.mark.parametrize("base_block", _BASE_BLOCKS) +@pytest.mark.parametrize("max_spec_len", _MAX_SPEC_LENS) +def test_spec_rollback_tracks_baseline(precision, base_block, max_spec_len): + state_dtype, act_dtype = precision + _run_rollback( + state_dtype=state_dtype, + act_dtype=act_dtype, + HQ=2, + HV=4, + K=64, + V=64, + buffer_len=base_block, + max_spec_len=max_spec_len, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("precision", _PRECISIONS) +@pytest.mark.parametrize("with_padding", [False, True]) +def test_spec_continuous_batching(precision, with_padding): + """Sparse state slots via ssm_state_indices with PAD(-1) padding rows and a + DIFFERENT buffer fill per row, in one verify call. Each active row is checked + against its own standard-decode window oracle; padding rows are zeroed and + unused state slots are left untouched.""" + state_dtype, act_dtype = precision + set_random_seed(0) + HQ, HV, K, V = 2, 4, 64, 64 + H = HQ + scale = K**-0.5 + qkv_dim = 2 * H * K + HV * V + base_block = 16 + max_spec_len = 4 + spec_len = max_spec_len + L, buf = _lbt(base_block, max_spec_len) + both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32 + rtol, atol = _output_tol(both_fp32) + + batch = 3 + padding = 2 if with_padding else 0 + padded_batch = batch + padding + total_slots = 12 + C = base_block - max_spec_len + wps = [0, C // 2, C] + + A_log = torch.randn(HV, device=DEV, dtype=torch.float32) + dt_bias = torch.randn(HV, device=DEV, dtype=torch.float32) + state_indices = (torch.randperm(total_slots - 1, device=DEV)[:batch] + 1).to( + torch.int32 + ) + sbi = torch.cat( + [ + state_indices, + torch.full((padding,), PAD_SLOT_ID, dtype=torch.int32, device=DEV), + ] + ) + unused = torch.ones(total_slots, dtype=torch.bool, device=DEV) + unused[state_indices] = False + + S0 = torch.randn(total_slots, HV, V, K, device=DEV, dtype=state_dtype) * 0.1 + state_spec = S0.clone() + d_cache = torch.zeros(total_slots, HV, buf, V, device=DEV, dtype=act_dtype) + k_cache = torch.zeros(total_slots, H, buf, K, device=DEV, dtype=act_dtype) + g_cache = torch.zeros(total_slots, HV, buf, device=DEV, dtype=torch.float32) + + write_pos = torch.zeros(total_slots, dtype=torch.int32, device=DEV) + cache_base = torch.zeros(total_slots, dtype=torch.int32, device=DEV) + is_flush = torch.zeros(total_slots, dtype=torch.int8, device=DEV) + + oracles = [] + mqkv_pack = torch.zeros( + padded_batch * spec_len, qkv_dim, device=DEV, dtype=act_dtype + ) + a_pack = torch.zeros(padded_batch * spec_len, HV, device=DEV, dtype=act_dtype) + b_pack = torch.zeros(padded_batch * spec_len, HV, device=DEV, dtype=act_dtype) + for r in range(batch): + wp = wps[r] + slot = int(state_indices[r].item()) + T_tot = wp + spec_len + mqkv = torch.randn(T_tot, qkv_dim, device=DEV, dtype=act_dtype) + a = torch.randn(T_tot, HV, device=DEV, dtype=act_dtype) + b = torch.randn(T_tot, HV, device=DEV, dtype=act_dtype) + _build_history( + mqkv=mqkv, + a=a, + b=b, + A_log=A_log, + dt_bias=dt_bias, + scale=scale, + state=state_spec, + d_cache=d_cache, + k_cache=k_cache, + g_cache=g_cache, + slot=slot, + wp=wp, + ) + oracles.append( + _standard_window_oracle( + mqkv=mqkv, + a=a, + b=b, + A_log=A_log, + dt_bias=dt_bias, + scale=scale, + state=state_spec, + d_cache=d_cache, + k_cache=k_cache, + g_cache=g_cache, + slot=slot, + wp=wp, + spec_len=spec_len, + HV=HV, + V=V, + ) + ) + write_pos[slot] = wp + seg = slice(r * spec_len, (r + 1) * spec_len) + mqkv_pack[seg] = mqkv[wp:] + a_pack[seg] = a[wp:] + b_pack[seg] = b[wp:] + + state_before = state_spec.clone() + qsl = torch.arange( + 0, (padded_batch + 1) * spec_len, spec_len, device=DEV, dtype=torch.int32 + ) + out_spec = torch.full( + (padded_batch * spec_len, HV, V), 42.0, device=DEV, dtype=act_dtype + ) + gdn_replayssm_spec_decode( + mixed_qkv=mqkv_pack, + a=a_pack, + b=b_pack, + A_log=A_log, + dt_bias=dt_bias, + checkpoint_state=state_spec, + d_cache=d_cache, + k_cache=k_cache, + g_cache=g_cache, + out=out_spec, + query_start_loc=qsl, + ssm_state_indices=sbi, + write_pos=write_pos, + cache_base=cache_base, + is_flush=is_flush, + max_cache_len=L, + max_spec_len=max_spec_len, + scale=scale, + use_qk_l2norm_in_kernel=True, + ) + + for r in range(batch): + seg = slice(r * spec_len, (r + 1) * spec_len) + torch.testing.assert_close(out_spec[seg], oracles[r], rtol=rtol, atol=atol) + if with_padding: + pad = slice(batch * spec_len, padded_batch * spec_len) + assert torch.equal(out_spec[pad], torch.zeros_like(out_spec[pad])) + # non-flush verify leaves every state slot untouched. + torch.testing.assert_close(state_spec, state_before, rtol=0, atol=0) + assert torch.equal(state_spec[unused], S0[unused]) diff --git a/tests/kernels/test_replayssm_standard_decode_gdn.py b/tests/kernels/test_replayssm_standard_decode_gdn.py new file mode 100644 index 000000000000..ec69acaf1ef5 --- /dev/null +++ b/tests/kernels/test_replayssm_standard_decode_gdn.py @@ -0,0 +1,273 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Standard (autoregressive) decode correctness for the GDN ReplaySSM kernel. + +GDN ReplaySSM caches the per-step delta-rule inputs (the corrected value ``u``, +the key, and the gate) in a small ring buffer and reconstructs the recurrent +state on the fly, writing it back to HBM only when the buffer flushes. This +checks the cached decode kernel ``fused_recurrent_gated_delta_rule_replayssm`` +against the upstream baseline ``fused_recurrent_gated_delta_rule_packed_decode`` +over a multi-step decode, across the flush boundaries and with sparse +(continuous-batching) state allocation including padding rows. + +GDN uses only the state-reconstruction route (it reads the state at both k and +q, so there is no output-only variant). State precision and activation/buffer +precision are swept independently; Qwen3.5 defaults to fp32 SSM state, though +bf16 state is also supported. ``g_cache`` is always fp32. +""" + +import pytest +import torch + +from vllm.model_executor.layers.fla.ops import ( + fused_recurrent_gated_delta_rule_packed_decode, + fused_recurrent_gated_delta_rule_replayssm, +) +from vllm.utils.torch_utils import set_random_seed + + +def _output_tolerances(act_dtype: torch.dtype) -> tuple[float, float]: + # Anchored on the baseline packed-decode kernel; same regime as the existing + # GDN test. The output has no fp32 drift (tight 1e-4); keyed off act dtype. + if act_dtype == torch.float32: + return 1e-4, 1e-4 + return 1e-2, 2e-2 + + +def _state_tolerances(act_dtype: torch.dtype) -> tuple[float, float]: + # The reconstructed checkpoint state vs the baseline's sequential state + # differs a bit more than the output at fp32 (~8e-4, an fp32 reconstruction + # vs recurrence difference, not a precision bug) -- same looser regime the + # existing GDN test uses for state. + if act_dtype == torch.float32: + return 1e-3, 2e-3 + return 1e-2, 2e-2 + + +def _run_gdn_standard_decode( + *, + state_dtype: torch.dtype, + act_dtype: torch.dtype, + batch: int, + num_q_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + max_cache_len: int, + num_steps: int, + n_pad: int, + strided: bool = False, + seed: int = 0, +) -> None: + """Drive ``num_steps`` decode steps; check the ReplaySSM kernel matches the + baseline packed decode each step (and the checkpoint state at flushes). + State follows ``state_dtype``; activations/caches follow ``act_dtype``.""" + device = "cuda" + o_rtol, o_atol = _output_tolerances(act_dtype) + s_rtol, s_atol = _state_tolerances(act_dtype) + set_random_seed(seed) + scale = head_k_dim**-0.5 + qkv_dim = 2 * (num_q_heads * head_k_dim) + num_v_heads * head_v_dim + padded_batch = batch + n_pad + num_state_slots = batch + n_pad + 1 # +1 reserves a null block at slot 0 + + A_log = torch.randn(num_v_heads, device=device, dtype=act_dtype) + dt_bias = torch.randn(num_v_heads, device=device, dtype=act_dtype) + + # Sparse state allocation: active rows map to a random permutation of slots + # 1.., padding rows map to PAD_SLOT_ID=-1. + state_indices = ( + torch.randperm(num_state_slots - 1, device=device)[:batch] + 1 + ).to(torch.int32) + ssm_state_indices = torch.cat([ + state_indices, + torch.full((n_pad,), -1, dtype=torch.int32, device=device), + ]) + unused = torch.ones(num_state_slots, dtype=torch.bool, device=device) + unused[state_indices] = False + + state0 = torch.randn( + num_state_slots, num_v_heads, head_v_dim, head_k_dim, + device=device, dtype=state_dtype) + state_packed = state0.clone() + state_cached = state0.clone() + state_before = state0.clone() + + d_cache = torch.zeros( + num_state_slots, num_v_heads, max_cache_len, head_v_dim, + device=device, dtype=act_dtype) + k_cache = torch.zeros( + num_state_slots, num_q_heads, max_cache_len, head_k_dim, + device=device, dtype=act_dtype) + g_cache = torch.zeros( + num_state_slots, num_v_heads, max_cache_len, device=device, + dtype=torch.float32) + + for step in range(num_steps): + if strided: + proj = torch.randn( + padded_batch, qkv_dim + 64, device=device, dtype=act_dtype) + mixed_qkv = proj[:, :qkv_dim] + else: + mixed_qkv = torch.randn( + padded_batch, qkv_dim, device=device, dtype=act_dtype) + a = torch.randn(padded_batch, num_v_heads, device=device, dtype=act_dtype) + b = torch.randn(padded_batch, num_v_heads, device=device, dtype=act_dtype) + write_pos = torch.full( + (padded_batch,), step % max_cache_len, device=device, + dtype=torch.int32) + + out_packed = torch.empty( + padded_batch, 1, num_v_heads, head_v_dim, device=device, + dtype=act_dtype) + fused_recurrent_gated_delta_rule_packed_decode( + mixed_qkv=mixed_qkv, a=a, b=b, A_log=A_log, dt_bias=dt_bias, + scale=scale, initial_state=state_packed, out=out_packed, + ssm_state_indices=ssm_state_indices, use_qk_l2norm_in_kernel=True) + + out_cached = torch.empty( + padded_batch, 1, num_v_heads, head_v_dim, device=device, + dtype=act_dtype) + fused_recurrent_gated_delta_rule_replayssm( + mixed_qkv=mixed_qkv, a=a, b=b, A_log=A_log, dt_bias=dt_bias, + scale=scale, initial_state=state_cached, d_cache=d_cache, + k_cache=k_cache, g_cache=g_cache, out=out_cached, + ssm_state_indices=ssm_state_indices, write_pos=write_pos, + use_qk_l2norm_in_kernel=True) + + torch.testing.assert_close( + out_cached[:batch], out_packed[:batch], rtol=o_rtol, atol=o_atol) + if step % max_cache_len == max_cache_len - 1: + torch.testing.assert_close( + state_cached[state_indices], state_packed[state_indices], + rtol=s_rtol, atol=s_atol) + + # Padding / unused state slots are never written. + assert torch.equal(state_cached[unused], state_before[unused]) + + +# State/activation precisions. fp32 state is the default; bf16/fp16 are the +# reduced-footprint configs. fp16 appears as an activation dtype (sfp16_afp16, +# s32_afp16) and as a state dtype under bf16 activations (sfp16_a16): fp16 has a +# finer mantissa than bf16 at the same 2 bytes, so it is a more accurate state at +# no extra footprint. We skip fp16 state under fp32 activations. +_PRECISIONS = [ + pytest.param((torch.float32, torch.float32), id="s32_a32"), + pytest.param((torch.float32, torch.bfloat16), id="s32_a16"), + pytest.param((torch.bfloat16, torch.bfloat16), id="s16_a16"), + pytest.param((torch.float32, torch.float16), id="s32_afp16"), + pytest.param((torch.float16, torch.float16), id="sfp16_afp16"), + pytest.param((torch.float16, torch.bfloat16), id="sfp16_a16"), +] +_GEOMETRIES = [ + # (num_q_heads, num_v_heads, head_k_dim, head_v_dim) + pytest.param((2, 4, 64, 64), id="small"), + pytest.param((16, 32, 128, 128), id="qwen4b"), + pytest.param((16, 64, 128, 128), id="qwen122b"), +] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("precision", _PRECISIONS) +@pytest.mark.parametrize("geometry", _GEOMETRIES) +@pytest.mark.parametrize("max_cache_len", [4, 16]) +@pytest.mark.parametrize("strided", [False, True]) +def test_replayssm_standard_decode_gdn_matches_packed( + precision: tuple[torch.dtype, torch.dtype], + geometry: tuple[int, int, int, int], + max_cache_len: int, + strided: bool, +): + state_dtype, act_dtype = precision + num_q_heads, num_v_heads, head_k_dim, head_v_dim = geometry + _run_gdn_standard_decode( + state_dtype=state_dtype, + act_dtype=act_dtype, + batch=4, + num_q_heads=num_q_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + max_cache_len=max_cache_len, + num_steps=2 * max_cache_len + 1, + n_pad=2, + strided=strided, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("precision", _PRECISIONS) +def test_replayssm_standard_decode_gdn_per_row_write_pos( + precision: tuple[torch.dtype, torch.dtype], +): + # Each row is built up to a different prefix length, so when the batch runs + # together the rows sit at different ring positions. This validates that the + # kernel reads a per-row write_pos (independent cache cursors per sequence -- + # the continuous-batching case the main test's uniform write_pos misses). + state_dtype, act_dtype = precision + device = "cuda" + o_rtol, o_atol = _output_tolerances(act_dtype) + set_random_seed(1) + + batch = 4 + num_q_heads, num_v_heads, head_k_dim, head_v_dim = 2, 4, 64, 64 + max_cache_len = 4 + num_state_slots = 7 + scale = head_k_dim**-0.5 + qkv_dim = 2 * (num_q_heads * head_k_dim) + num_v_heads * head_v_dim + + A_log = torch.randn(num_v_heads, device=device, dtype=act_dtype) + dt_bias = torch.randn(num_v_heads, device=device, dtype=act_dtype) + ssm_state_indices = torch.tensor([4, 2, 6, 1], device=device, dtype=torch.int32) + prefix_lens = torch.tensor([0, 1, 2, 3], device=device, dtype=torch.int32) + + state0 = torch.randn( + num_state_slots, num_v_heads, head_v_dim, head_k_dim, + device=device, dtype=state_dtype) + state_packed = state0.clone() + state_cached = state0.clone() + d_cache = torch.zeros( + num_state_slots, num_v_heads, max_cache_len, head_v_dim, + device=device, dtype=act_dtype) + k_cache = torch.zeros( + num_state_slots, num_q_heads, max_cache_len, head_k_dim, + device=device, dtype=act_dtype) + g_cache = torch.zeros( + num_state_slots, num_v_heads, max_cache_len, device=device, + dtype=torch.float32) + + def decode_step(mixed_qkv, a, b, idx, write_pos): + out_p = torch.empty( + mixed_qkv.shape[0], 1, num_v_heads, head_v_dim, device=device, + dtype=act_dtype) + fused_recurrent_gated_delta_rule_packed_decode( + mixed_qkv=mixed_qkv, a=a, b=b, A_log=A_log, dt_bias=dt_bias, + scale=scale, initial_state=state_packed, out=out_p, + ssm_state_indices=idx, use_qk_l2norm_in_kernel=True) + out_c = torch.empty_like(out_p) + fused_recurrent_gated_delta_rule_replayssm( + mixed_qkv=mixed_qkv, a=a, b=b, A_log=A_log, dt_bias=dt_bias, + scale=scale, initial_state=state_cached, d_cache=d_cache, + k_cache=k_cache, g_cache=g_cache, out=out_c, + ssm_state_indices=idx, write_pos=write_pos, + use_qk_l2norm_in_kernel=True) + return out_p, out_c + + # Build each row up independently to its prefix length. + for row, prefix_len in enumerate(prefix_lens.tolist()): + idx = ssm_state_indices[row:row + 1] + for s in range(prefix_len): + mixed_qkv = torch.randn(1, qkv_dim, device=device, dtype=act_dtype) + a = torch.randn(1, num_v_heads, device=device, dtype=act_dtype) + b = torch.randn(1, num_v_heads, device=device, dtype=act_dtype) + decode_step( + mixed_qkv, a, b, idx, + torch.tensor([s], device=device, dtype=torch.int32)) + + # Run the whole batch with per-row write_pos = prefix_lens. + mixed_qkv = torch.randn(batch, qkv_dim, device=device, dtype=act_dtype) + a = torch.randn(batch, num_v_heads, device=device, dtype=act_dtype) + b = torch.randn(batch, num_v_heads, device=device, dtype=act_dtype) + out_packed, out_cached = decode_step( + mixed_qkv, a, b, ssm_state_indices, prefix_lens) + torch.testing.assert_close(out_cached, out_packed, rtol=o_rtol, atol=o_atol) diff --git a/tests/kernels/test_replayssm_teacher_decode_equivalence_gdn.py b/tests/kernels/test_replayssm_teacher_decode_equivalence_gdn.py new file mode 100644 index 000000000000..5167750757e8 --- /dev/null +++ b/tests/kernels/test_replayssm_teacher_decode_equivalence_gdn.py @@ -0,0 +1,207 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Prefill(chunk) == decode equivalence for the GDN ReplaySSM kernel. + +The GDN analog of the Mamba2 prefill-teacher: the chunked prefill kernel +(``chunk_gated_delta_rule``, the path the model runs at prefill) and +step-by-step decode must produce the same per-position outputs and final state +over a sequence. We build q/k/v/g/beta with ``fused_post_conv_prep`` (the same +shared prep the model's prefill uses, so the teacher's inputs match what the +decode kernels compute internally from ``mixed_qkv``), then feed: + + * the chunked prefill kernel (the teacher), + * the baseline step decode (``..._packed_decode``), + * the ReplaySSM step decode (ring buffer, ``..._replayssm``), + +and check the step decoders reproduce the teacher's per-position outputs and +final state. The chunked prefill kernel only supports bf16 *activations* +(q/k/v), so it is the (trusted) reference run with bf16 activations -- there is +no separate fp32 ground truth here (the fp32-activation decode path is covered +by the GDN standard-decode test). The recurrent *state* precision is still +swept: fp32 (the production config) and bf16, both with bf16 activations. + +Prefill (a chunked scan) and decode (a step recurrence) are different code +paths, so they differ numerically; we use chunked-scan tolerances (the same +regime as the Mamba2 prefill-teacher). + +``seqlen`` is a multiple of the buffer length, so the final step flushes and +ReplaySSM's stored checkpoint is the full final state. +""" + +import pytest +import torch + +from vllm.model_executor.layers.fla.ops import ( + chunk_gated_delta_rule, + fused_recurrent_gated_delta_rule_packed_decode, + fused_recurrent_gated_delta_rule_replayssm, +) +from vllm.model_executor.layers.fla.ops.fused_gdn_prefill_post_conv import ( + fused_post_conv_prep, +) +from vllm.utils.torch_utils import set_random_seed + + +def _output_tolerances(act_dtype: torch.dtype) -> tuple[float, float]: + # Chunked-scan-vs-step regime (the chunked prefill, not the decode, sets + # these), same as the Mamba2 prefill-teacher. Keyed off the activation dtype. + if act_dtype == torch.float32: + return 1e-2, 3e-2 + return 6e-2, 1e-1 + + +def _run_gdn_teacher_equivalence( + *, + state_dtype: torch.dtype, + act_dtype: torch.dtype, + batch: int, + num_q_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + seqlen: int, + max_cache_len: int, + seed: int = 0, +) -> None: + assert seqlen % max_cache_len == 0, "final step must be a flush" + device = "cuda" + H, HV, K, V = num_q_heads, num_v_heads, head_k_dim, head_v_dim + o_rtol, o_atol = _output_tolerances(act_dtype) + set_random_seed(seed) + scale = K**-0.5 + qkv_dim = 2 * (H * K) + HV * V + L = batch * seqlen + + # Step decoders use slot-indexed state (slot 0 is the null block). + num_state_slots = batch + 1 + ssm_state_indices = torch.arange( + 1, batch + 1, device=device, dtype=torch.int32) + + A_log = torch.randn(HV, device=device, dtype=act_dtype) + dt_bias = torch.randn(HV, device=device, dtype=act_dtype) + mixed_qkv_seq = torch.randn( + batch, seqlen, qkv_dim, device=device, dtype=act_dtype) + a_seq = torch.randn(batch, seqlen, HV, device=device, dtype=act_dtype) + b_seq = torch.randn(batch, seqlen, HV, device=device, dtype=act_dtype) + + def run_teacher(qkv_dtype: torch.dtype, st_dtype: torch.dtype): + # Same post-conv prep as the model's prefill (GQA + l2norm), so the + # teacher's q/k/v/g/beta match what the decode kernels derive from + # mixed_qkv. apply_l2norm=True here -> l2norm off in the chunk kernel. + q, k, v, g, beta = fused_post_conv_prep( + conv_output=mixed_qkv_seq.reshape(L, qkv_dim).to(qkv_dtype), + a=a_seq.reshape(L, HV).to(qkv_dtype), + b=b_seq.reshape(L, HV).to(qkv_dtype), + A_log=A_log.to(qkv_dtype), dt_bias=dt_bias.to(qkv_dtype), + num_k_heads=H, head_k_dim=K, head_v_dim=V, + apply_l2norm=True, output_g_exp=False) + q = q.reshape(batch, seqlen, *q.shape[1:]) + k = k.reshape(batch, seqlen, *k.shape[1:]) + v = v.reshape(batch, seqlen, *v.shape[1:]) + g = g.reshape(batch, seqlen, *g.shape[1:]) + beta = beta.reshape(batch, seqlen, *beta.shape[1:]) + state0 = torch.zeros(batch, HV, V, K, device=device, dtype=st_dtype) + out, final_state = chunk_gated_delta_rule( + q=q, k=k, v=v, g=g, beta=beta, scale=scale, initial_state=state0, + output_final_state=True, cu_seqlens=None, + use_qk_l2norm_in_kernel=False) + return out, final_state # (B, T, HV, V), (B, HV, V, K) + + # The chunked prefill kernel only supports bf16 activations, so it is the + # (trusted) reference at the test (bf16-activation) precision; no separate + # fp32 ground truth (fp32-activation decode is covered by the GDN + # standard-decode test). The recurrent state precision is still swept. + y_teacher, state_teacher = run_teacher(act_dtype, state_dtype) + + # Step decoders: baseline packed decode and ReplaySSM (ring buffer). + state_base = torch.zeros( + num_state_slots, HV, V, K, device=device, dtype=state_dtype) + state_dec = torch.zeros( + num_state_slots, HV, V, K, device=device, dtype=state_dtype) + d_cache = torch.zeros( + num_state_slots, HV, max_cache_len, V, device=device, dtype=act_dtype) + k_cache = torch.zeros( + num_state_slots, H, max_cache_len, K, device=device, dtype=act_dtype) + g_cache = torch.zeros( + num_state_slots, HV, max_cache_len, device=device, dtype=torch.float32) + + y_base = torch.empty(batch, seqlen, HV, V, device=device, dtype=act_dtype) + y_dec = torch.empty(batch, seqlen, HV, V, device=device, dtype=act_dtype) + for t in range(seqlen): + mixed_qkv = mixed_qkv_seq[:, t] + a = a_seq[:, t] + b = b_seq[:, t] + write_pos = torch.full( + (batch,), t % max_cache_len, device=device, dtype=torch.int32) + + out_b = torch.empty(batch, 1, HV, V, device=device, dtype=act_dtype) + fused_recurrent_gated_delta_rule_packed_decode( + mixed_qkv=mixed_qkv, a=a, b=b, A_log=A_log, dt_bias=dt_bias, + scale=scale, initial_state=state_base, out=out_b, + ssm_state_indices=ssm_state_indices, use_qk_l2norm_in_kernel=True) + y_base[:, t] = out_b[:, 0] + + out_d = torch.empty(batch, 1, HV, V, device=device, dtype=act_dtype) + fused_recurrent_gated_delta_rule_replayssm( + mixed_qkv=mixed_qkv, a=a, b=b, A_log=A_log, dt_bias=dt_bias, + scale=scale, initial_state=state_dec, d_cache=d_cache, + k_cache=k_cache, g_cache=g_cache, out=out_d, + ssm_state_indices=ssm_state_indices, write_pos=write_pos, + use_qk_l2norm_in_kernel=True) + y_dec[:, t] = out_d[:, 0] + + # Baseline and ReplaySSM step decode both reproduce the chunked prefill + # teacher's per-position outputs. + torch.testing.assert_close( + y_base.float(), y_teacher.float(), rtol=o_rtol, atol=o_atol) + torch.testing.assert_close( + y_dec.float(), y_teacher.float(), rtol=o_rtol, atol=o_atol) + # Final state (the last step is a flush, so state_dec is the full state). + # The step decoders' state is slot-indexed; gather active rows to match the + # teacher's dense per-sequence state. + torch.testing.assert_close( + state_base[ssm_state_indices].float(), state_teacher.float(), + rtol=o_rtol, atol=o_atol) + torch.testing.assert_close( + state_dec[ssm_state_indices].float(), state_teacher.float(), + rtol=o_rtol, atol=o_atol) + + +# The chunked prefill teacher is bf16-activation-only, so the activation dtype is +# bf16; the recurrent-state precision is still swept: fp32 (production), bf16, and +# fp16 (a finer-mantissa state than bf16 at the same 2 bytes). Fully-fp16 +# activations are out of scope here (the teacher cannot run them). +_PRECISIONS = [ + pytest.param((torch.float32, torch.bfloat16), id="s32_a16"), + pytest.param((torch.bfloat16, torch.bfloat16), id="s16_a16"), + pytest.param((torch.float16, torch.bfloat16), id="sfp16_a16"), +] +_GEOMETRIES = [ + # (num_q_heads, num_v_heads, head_k_dim, head_v_dim) + pytest.param((2, 4, 64, 64), id="small"), + pytest.param((16, 32, 128, 128), id="qwen4b"), +] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("precision", _PRECISIONS) +@pytest.mark.parametrize("geometry", _GEOMETRIES) +@pytest.mark.parametrize("max_cache_len", [4, 16]) +def test_replayssm_teacher_decode_equivalence_gdn( + precision: tuple[torch.dtype, torch.dtype], + geometry: tuple[int, int, int, int], + max_cache_len: int, +): + state_dtype, act_dtype = precision + num_q_heads, num_v_heads, head_k_dim, head_v_dim = geometry + _run_gdn_teacher_equivalence( + state_dtype=state_dtype, + act_dtype=act_dtype, + batch=4, + num_q_heads=num_q_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + seqlen=16, + max_cache_len=max_cache_len, + ) diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py new file mode 100644 index 000000000000..4f551f76376e --- /dev/null +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Engine-level parity: ReplaySSM standard decode vs the baseline SSM kernel.""" + +import pytest + +from ...models.utils import check_logprobs_close +from ...utils import large_gpu_mark + +# Mamba2 (Nemotron) and GDN (Qwen3.5) hybrids. +MODELS = [ + pytest.param( + "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16", marks=large_gpu_mark(min_gb=40) + ), + pytest.param("Qwen/Qwen3.5-4B", marks=large_gpu_mark(min_gb=40)), +] + +PROMPTS = [ + "The capital of France is", + "Once upon a time, in a small village,", +] + + +@pytest.mark.parametrize("model_name", MODELS) +def test_replayssm_decode_matches_baseline(vllm_runner, model_name): + # ReplaySSM reconstructs the state in different fp arithmetic, so greedy ids + # can diverge at a near-tie; compare logprobs, not exact ids. + common = dict( + max_model_len=1024, + trust_remote_code=True, + enable_prefix_caching=False, + mamba_cache_mode="none", + ) + with vllm_runner(model_name, **common) as llm: + baseline = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5) + with vllm_runner( + model_name, use_replayssm=True, replayssm_buffer_len=16, **common + ) as llm: + replay = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5) + + check_logprobs_close( + outputs_0_lst=baseline, + outputs_1_lst=replay, + name_0="baseline", + name_1="replayssm", + ) diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 2fb3358d55c3..c35a4d631176 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -36,6 +36,7 @@ ] MambaDType = Literal["auto", "float32", "float16", "bfloat16"] MambaCacheMode = Literal["all", "align", "none"] +ReplaySSMRoute = Literal["state_and_output", "output_only"] PrefixCachingHashAlgo = Literal["sha256", "sha256_cbor", "xxhash", "xxhash_cbor"] KVOffloadingBackend = Literal["native", "lmcache"] @@ -140,6 +141,27 @@ class CacheConfig: - "align": only cache the mamba state of the last token of each scheduler step and when the token is at position i * block_size. """ + replayssm_buffer_len: int = Field(default=16, gt=0) + """ReplaySSM history block B. Autoregressive decode (use_replayssm) flushes + the checkpoint every B steps. Speculative decode (use_replayssm_spec) keeps a + L = B + 1 + num_speculative_tokens history window (usable committed history + B - 1 - num_speculative_tokens) in a power-of-two next_pow2(L) ring buffer.""" + use_replayssm: bool = False + """Use the ReplaySSM Mamba2 decode kernel (cache recent SSM inputs instead + of writing the recurrent state back to HBM each step). Only supported for + autoregressive decode with mamba_cache_mode='none'.""" + replayssm_route: ReplaySSMRoute = "output_only" + """ReplaySSM compute route (only meaningful when use_replayssm is True): + - "output_only" (default): inner-product route. Computes the output from + the checkpoint state plus the cached inputs without materializing the + per-step state (the state is only built on flush steps). + - "state_and_output": outer-product route. Reconstructs the full SSM state + every step via tl.dot, then reads the output from it.""" + use_replayssm_spec: bool = False + """Use the ReplaySSM speculative-decode kernel (circular cache + early-flush) + for Mamba2 and GDN. Requires speculative decoding and mamba_cache_mode='none'; + reuses vLLM's causal_conv1d_update for the conv (hybrid). Mutually exclusive + with use_replayssm. replayssm_buffer_len must be >= 1 + num_speculative_tokens.""" # Will be set after profiling. num_gpu_blocks: int | None = field(default=None, init=False) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index de52e5228e99..55b7f5d6b416 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -38,7 +38,7 @@ from .kv_transfer import KVTransferConfig from .load import LoadConfig from .lora import LoRAConfig -from .mamba import MambaConfig +from .mamba import MambaBackendEnum, MambaConfig from .model import ModelConfig from .observability import ObservabilityConfig from .offload import OffloadConfig @@ -2211,6 +2211,76 @@ def validate_mamba_block_size(self) -> "VllmConfig": ) return self + @model_validator(mode="after") + def validate_mamba_cached_kernel(self) -> "VllmConfig": + if not self.cache_config.use_replayssm: + if self.cache_config.replayssm_route != "output_only": + raise ValueError( + "--replayssm-route is only meaningful when " + "--use-replayssm is enabled" + ) + return self + if self.cache_config.mamba_cache_mode != "none": + raise ValueError( + "--use-replayssm requires --mamba-cache-mode none" + ) + if self.num_speculative_tokens > 0: + raise ValueError( + "--use-replayssm does not support speculative decoding" + ) + if self.mamba_config.backend != MambaBackendEnum.TRITON: + raise ValueError( + "--use-replayssm requires --mamba-backend triton" + ) + if self.mamba_config.enable_stochastic_rounding: + raise ValueError( + "--use-replayssm does not support Mamba cache " + "stochastic rounding" + ) + return self + + @model_validator(mode="after") + def validate_mamba_cached_spec_kernel(self) -> "VllmConfig": + if not self.cache_config.use_replayssm_spec: + return self + # Inverted guard: the cached-SPEC kernel *requires* speculative decode + # (the opposite of --use-replayssm, which forbids it). + if self.cache_config.use_replayssm: + raise ValueError( + "--use-replayssm-spec is mutually exclusive with " + "--use-replayssm (different page shapes / decode paths)" + ) + if self.num_speculative_tokens <= 0: + raise ValueError( + "--use-replayssm-spec requires speculative decoding " + "(num_speculative_tokens > 0)" + ) + if self.cache_config.mamba_cache_mode != "none": + raise ValueError( + "--use-replayssm-spec requires --mamba-cache-mode none" + ) + if self.mamba_config.backend != MambaBackendEnum.TRITON: + raise ValueError( + "--use-replayssm-spec requires --mamba-backend triton" + ) + if self.mamba_config.enable_stochastic_rounding: + raise ValueError( + "--use-replayssm-spec does not support Mamba cache " + "stochastic rounding" + ) + max_spec_len = 1 + self.num_speculative_tokens + # replayssm_buffer_len is the history block B; the flush threshold is + # L = B + max_spec_len and the usable committed history is B - max_spec_len, + # so B must be at least max_spec_len. The physical ring buffer is the + # power-of-two next_pow2(L), so B itself need not be a power of two. + if self.cache_config.replayssm_buffer_len < max_spec_len: + raise ValueError( + "--use-replayssm-spec requires --replayssm-buffer-len " + f">= 1 + num_speculative_tokens ({max_spec_len}); " + f"got {self.cache_config.replayssm_buffer_len}" + ) + return self + _current_vllm_config: VllmConfig | None = None _current_prefix: str | None = None diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index efdd7696fdce..e3229ba02347 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -687,6 +687,10 @@ class EngineArgs: mamba_ssm_cache_dtype: MambaDType = CacheConfig.mamba_ssm_cache_dtype mamba_block_size: int | None = get_field(CacheConfig, "mamba_block_size") mamba_cache_mode: MambaCacheMode = CacheConfig.mamba_cache_mode + replayssm_buffer_len: int = CacheConfig.replayssm_buffer_len + use_replayssm: bool = CacheConfig.use_replayssm + replayssm_route: str = CacheConfig.replayssm_route + use_replayssm_spec: bool = CacheConfig.use_replayssm_spec mamba_backend: MambaBackendEnum = MambaBackendEnum.TRITON enable_mamba_cache_stochastic_rounding: bool = ( @@ -1187,6 +1191,20 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: cache_group.add_argument( "--mamba-cache-mode", **cache_kwargs["mamba_cache_mode"] ) + cache_group.add_argument( + "--replayssm-buffer-len", **cache_kwargs["replayssm_buffer_len"] + ) + cache_group.add_argument( + "--use-replayssm", **cache_kwargs["use_replayssm"] + ) + cache_group.add_argument( + "--replayssm-route", + **cache_kwargs["replayssm_route"], + ) + cache_group.add_argument( + "--use-replayssm-spec", + **cache_kwargs["use_replayssm_spec"], + ) cache_group.add_argument( "--kv-offloading-size", **cache_kwargs["kv_offloading_size"] ) @@ -1882,6 +1900,10 @@ def create_engine_config( mamba_ssm_cache_dtype=self.mamba_ssm_cache_dtype, mamba_block_size=self.mamba_block_size, mamba_cache_mode=self.mamba_cache_mode, + replayssm_buffer_len=self.replayssm_buffer_len, + use_replayssm=self.use_replayssm, + replayssm_route=self.replayssm_route, + use_replayssm_spec=self.use_replayssm_spec, kv_offloading_size=self.kv_offloading_size, kv_offloading_backend=self.kv_offloading_backend, ) diff --git a/vllm/model_executor/layers/fla/ops/__init__.py b/vllm/model_executor/layers/fla/ops/__init__.py index 1942d8980bc5..77ae2e9909c8 100644 --- a/vllm/model_executor/layers/fla/ops/__init__.py +++ b/vllm/model_executor/layers/fla/ops/__init__.py @@ -12,6 +12,9 @@ fused_recurrent_gated_delta_rule, fused_recurrent_gated_delta_rule_packed_decode, ) +from .fused_recurrent_replayssm import ( + fused_recurrent_gated_delta_rule_replayssm, +) from .fused_sigmoid_gating import fused_sigmoid_gating_delta_rule_update from .layernorm_guard import RMSNormGated @@ -20,6 +23,7 @@ "chunk_gated_delta_rule", "fused_recurrent_gated_delta_rule", "fused_recurrent_gated_delta_rule_packed_decode", + "fused_recurrent_gated_delta_rule_replayssm", "fused_post_conv_prep", "fused_sigmoid_gating_delta_rule_update", ] diff --git a/vllm/model_executor/layers/fla/ops/fused_recurrent_replayssm.py b/vllm/model_executor/layers/fla/ops/fused_recurrent_replayssm.py new file mode 100644 index 000000000000..97507d7faa5a --- /dev/null +++ b/vllm/model_executor/layers/fla/ops/fused_recurrent_replayssm.py @@ -0,0 +1,290 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# ruff: noqa: E501 + +from __future__ import annotations + +import torch + +from vllm.model_executor.layers.mamba.ops.replayssm_config import ( + get_replayssm_config, +) +from vllm.triton_utils import tl, triton + + +@triton.jit +def fused_recurrent_gated_delta_rule_replayssm_kernel( + mixed_qkv, a, b, A_log, dt_bias, o, h0, ht, + d_cache, k_cache, g_cache, ssm_state_indices, write_pos, scale, + stride_mixed_qkv_tok: tl.constexpr, + stride_a_tok: tl.constexpr, + stride_b_tok: tl.constexpr, + stride_init_state_token: tl.constexpr, + stride_final_state_token: tl.constexpr, + stride_indices_seq: tl.constexpr, + stride_d_slot: tl.constexpr, + stride_k_slot: tl.constexpr, + stride_g_slot: tl.constexpr, + H: tl.constexpr, HV: tl.constexpr, K: tl.constexpr, V: tl.constexpr, + BK: tl.constexpr, BV: tl.constexpr, BC: tl.constexpr, + NK: tl.constexpr, BKT: tl.constexpr, + MAX_CACHE_LEN: tl.constexpr, SOFTPLUS_THRESHOLD: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, +): + i_v = tl.program_id(0) + i_n = tl.program_id(1) + i_hv = tl.program_id(2) + i_h = i_hv // (HV // H) + + o_v = i_v * BV + tl.arange(0, BV) + o_c = tl.arange(0, BC) + mask_v = o_v < V + + # Resolve the physical state slot; zero the output and bail for padded rows. + state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq).to(tl.int64) + p_o = o + (i_n * HV + i_hv) * V + o_v + if state_idx <= 0: + tl.store(p_o, tl.zeros([BV], dtype=tl.float32).to(p_o.dtype.element_ty), mask=mask_v) + return + + # Per-row buffer cursor and flush flag; valid (committed) cache positions. + # vLLM: write_pos is per decode row (i_n), not per physical slot. + b_write_pos = tl.load(write_pos + i_n).to(tl.int64) + b_is_flush = b_write_pos == MAX_CACHE_LEN - 1 + cache_valid = o_c < b_write_pos + + # Gate for the current token: decay g, its exp alpha, and the beta mixing weight. + a_val = tl.load(a + i_n * stride_a_tok + i_hv).to(tl.float32) + b_val = tl.load(b + i_n * stride_b_tok + i_hv).to(tl.float32) + A_log_val = tl.load(A_log + i_hv).to(tl.float32) + dt_bias_val = tl.load(dt_bias + i_hv).to(tl.float32) + x = a_val + dt_bias_val + softplus_x = tl.where(x <= SOFTPLUS_THRESHOLD, tl.log(1.0 + tl.exp(x)), x) + g_val = -tl.exp(A_log_val) * softplus_x + alpha_val = tl.exp(g_val) + beta_val = tl.sigmoid(b_val).to(b.dtype.element_ty).to(tl.float32) + + # Replay decay over the committed cache, from the cached per-step gates g. + p_g_main = g_cache + state_idx * stride_g_slot + i_hv * MAX_CACHE_LEN + o_c + b_g_all = tl.load(p_g_main, mask=cache_valid, other=0.0).to(tl.float32) + b_g_prefix = tl.cumsum(b_g_all, axis=0) + b_g_total = tl.sum(b_g_all, axis=0) + b_replay_decay = tl.where(cache_valid, tl.exp(b_g_total - b_g_prefix), 0.0) + b_total_decay = tl.exp(b_g_total) + + # Cached delta-rule update vectors d (K-independent), scaled by the replay decay. + p_d_main = d_cache + state_idx * stride_d_slot + ((i_hv * MAX_CACHE_LEN + o_c[None, :]) * V + o_v[:, None]) + b_d_all = tl.load(p_d_main, mask=mask_v[:, None] & cache_valid[None, :], other=0).to(tl.float32) + b_d_scaled_tc = (b_d_all * b_replay_decay[None, :]).to(p_o.dtype.element_ty) # [BV, BC] + + # Current token value (for the delta-rule update). + v_off = (2 * H * K) + i_hv * V + o_v + b_v = tl.load(mixed_qkv + i_n * stride_mixed_qkv_tok + v_off, mask=mask_v, other=0).to(tl.float32) + + # Optional q/k L2 norm: full-vector reciprocal norms (computed, not kept). + if USE_QK_L2NORM_IN_KERNEL: + o_kf = tl.arange(0, BK) + mask_kf = o_kf < K + p_mix = mixed_qkv + i_n * stride_mixed_qkv_tok + qf = tl.load(p_mix + i_h * K + o_kf, mask=mask_kf, other=0).to(tl.float32) + kf = tl.load(p_mix + H * K + i_h * K + o_kf, mask=mask_kf, other=0).to(tl.float32) + q_rnorm = 1.0 / tl.sqrt(tl.sum(qf * qf) + 1e-6) + k_rnorm = 1.0 / tl.sqrt(tl.sum(kf * kf) + 1e-6) + else: + q_rnorm = 1.0 + k_rnorm = 1.0 + + # Reconstruct the state from the checkpoint + cached (d, k) in K tiles and read + # it with the current q and k. K-tiling avoids holding a full [BV, BK] tile. + # Also append the current key chunk to the ring cache (non-flush only). + b_state_q = tl.zeros([BV], dtype=tl.float32) + b_state_k = tl.zeros([BV], dtype=tl.float32) + cur_kq = tl.zeros([1], dtype=tl.float32) + write_k = (not b_is_flush) and (i_v == 0) and (i_hv == i_h * (HV // H)) + for kk in range(NK): + o_kt = kk * BKT + tl.arange(0, BKT) + mask_kt = o_kt < K + p_mix = mixed_qkv + i_n * stride_mixed_qkv_tok + q_c = tl.load(p_mix + i_h * K + o_kt, mask=mask_kt, other=0).to(tl.float32) * q_rnorm + k_c = tl.load(p_mix + H * K + i_h * K + o_kt, mask=mask_kt, other=0).to(tl.float32) * k_rnorm + q_cs = q_c * scale + cur_kq += tl.sum(k_c * q_cs) + + # Reconstruct this K tile of the state: S = total_decay * S_0 + d_scaled . k_cache. + p_h0_c = h0 + state_idx * stride_init_state_token + i_hv * V * K + o_v[:, None] * K + o_kt[None, :] + b_h0_c = tl.load(p_h0_c, mask=mask_v[:, None] & mask_kt[None, :], other=0).to(tl.float32) + p_k_c = k_cache + state_idx * stride_k_slot + ((i_h * MAX_CACHE_LEN + o_c[:, None]) * K + o_kt[None, :]) + b_k_all_c = tl.load(p_k_c, mask=cache_valid[:, None] & mask_kt[None, :], other=0).to(p_o.dtype.element_ty) + b_h_c = b_h0_c * b_total_decay + tl.dot(b_d_scaled_tc, b_k_all_c).to(tl.float32) # [BV, BKT] + + # Read the state with q and k (accumulated across K tiles). + b_state_q += tl.sum(b_h_c * q_cs[None, :], axis=1) + b_state_k += tl.sum(b_h_c * k_c[None, :], axis=1) + + if write_k: + p_cur_k = k_cache + state_idx * stride_k_slot + ((i_h * MAX_CACHE_LEN + b_write_pos) * K + o_kt) + tl.store(p_cur_k, k_c.to(p_o.dtype.element_ty), mask=mask_kt & (b_write_pos < MAX_CACHE_LEN)) + + # Current-token output: alpha*(S q) + d_cur * (k . q), with the new update + # vector d_cur = beta * (v - alpha*(S k)). + b_state_q *= alpha_val + b_state_k *= alpha_val + b_d_cur = beta_val * (b_v - b_state_k) + b_o = b_state_q + b_d_cur * tl.sum(cur_kq) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + + if b_is_flush: + # Flush: fold the current token into the checkpoint, S_t = alpha*S + d_t k_t^T, + # and persist it. Re-walk K chunks to rebuild S before applying the update. + for kk in range(NK): + o_kt = kk * BKT + tl.arange(0, BKT) + mask_kt = o_kt < K + p_mix = mixed_qkv + i_n * stride_mixed_qkv_tok + k_c = tl.load(p_mix + H * K + i_h * K + o_kt, mask=mask_kt, other=0).to(tl.float32) * k_rnorm + p_h0_c = h0 + state_idx * stride_init_state_token + i_hv * V * K + o_v[:, None] * K + o_kt[None, :] + b_h0_c = tl.load(p_h0_c, mask=mask_v[:, None] & mask_kt[None, :], other=0).to(tl.float32) + p_k_c = k_cache + state_idx * stride_k_slot + ((i_h * MAX_CACHE_LEN + o_c[:, None]) * K + o_kt[None, :]) + b_k_all_c = tl.load(p_k_c, mask=cache_valid[:, None] & mask_kt[None, :], other=0).to(p_o.dtype.element_ty) + b_h_c = b_h0_c * b_total_decay + tl.dot(b_d_scaled_tc, b_k_all_c).to(tl.float32) + b_h_new_c = alpha_val * b_h_c + b_d_cur[:, None] * k_c[None, :] + p_ht_c = ht + state_idx * stride_final_state_token + i_hv * V * K + o_v[:, None] * K + o_kt[None, :] + tl.store(p_ht_c, b_h_new_c.to(p_ht_c.dtype.element_ty), mask=mask_v[:, None] & mask_kt[None, :]) + else: + # Non-flush: append the current token's update vector d and gate g to the + # cache (the k chunks were already written inside the loop above). + p_cur_d = d_cache + state_idx * stride_d_slot + ((i_hv * MAX_CACHE_LEN + b_write_pos) * V + o_v) + tl.store(p_cur_d, b_d_cur.to(p_cur_d.dtype.element_ty), mask=mask_v & (b_write_pos < MAX_CACHE_LEN)) + if i_v == 0: + p_cur_g = g_cache + state_idx * stride_g_slot + i_hv * MAX_CACHE_LEN + b_write_pos + tl.store(p_cur_g, g_val, mask=b_write_pos < MAX_CACHE_LEN) + + +def fused_recurrent_gated_delta_rule_replayssm( + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + d_cache: torch.Tensor, + k_cache: torch.Tensor, + g_cache: torch.Tensor, + out: torch.Tensor, + ssm_state_indices: torch.Tensor, + write_pos: torch.Tensor, + use_qk_l2norm_in_kernel: bool = False, + block_v: int | None = None, + num_warps: int | None = None, + num_stages: int | None = None, + nk: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Cached GDN autoregressive decode (one new token per sequence). + + Same call surface as ``fused_recurrent_gated_delta_rule_packed_decode`` + plus the three ring caches (``d_cache``/``k_cache``/``g_cache``) and the + per-decode-row ``write_pos`` cursor. ``initial_state`` is both the + checkpoint read (h0) and the (flush-only) checkpoint write (ht), in place. + """ + if mixed_qkv.ndim != 2: + raise ValueError( + f"`mixed_qkv` must be a 2D tensor (got ndim={mixed_qkv.ndim})." + ) + if mixed_qkv.stride(-1) != 1: + raise ValueError("`mixed_qkv` must be contiguous in the last dim.") + if a.ndim != 2 or b.ndim != 2: + raise ValueError( + f"`a` and `b` must be 2D tensors (got a.ndim={a.ndim}, b.ndim={b.ndim})." + ) + if A_log.ndim != 1 or dt_bias.ndim != 1: + raise ValueError("`A_log`/`dt_bias` must be 1D tensors.") + if initial_state.ndim != 4: + raise ValueError( + f"`initial_state` must be a 4D tensor (got ndim={initial_state.ndim})." + ) + if not out.is_contiguous(): + raise ValueError("`out` must be contiguous.") + if write_pos.ndim != 1 or write_pos.dtype != torch.int32: + raise ValueError("`write_pos` must be a 1D int32 tensor.") + + B = mixed_qkv.shape[0] + num_state_slots, HV, V, K = initial_state.shape + qkv_dim = mixed_qkv.shape[1] + q_dim = (qkv_dim - HV * V) // 2 + if q_dim <= 0 or q_dim % K != 0: + raise ValueError( + f"Invalid packed `mixed_qkv` last dim={qkv_dim} for HV={HV}, V={V}, K={K}." + ) + H = q_dim // K + if H <= 0 or HV % H != 0: + raise ValueError(f"Invalid head config inferred from mixed_qkv: H={H}, HV={HV}.") + max_cache_len = d_cache.shape[2] + + # Launch config (block_v, num_warps, num_stages, nk) from the L-keyed config + # module; explicit kwargs override. Lets benchmarks/the config sweep pin it via + # override_replayssm_config("gdn_decode", ...). + cfg_bv, cfg_nw, cfg_ns, cfg_nk = get_replayssm_config( + "gdn_decode", L=max_cache_len + ) + if block_v is None: + block_v = cfg_bv + if num_warps is None: + num_warps = cfg_nw + if num_stages is None: + num_stages = cfg_ns + if nk is None: + nk = cfg_nk + + # Cache shape sanity (per state slot): d=(HV, L, V), k=(H, L, K), g=(HV, L). + if tuple(d_cache.shape[1:]) != (HV, max_cache_len, V): + raise ValueError( + f"`d_cache` per-slot shape must be {(HV, max_cache_len, V)} " + f"(got {tuple(d_cache.shape[1:])})." + ) + if tuple(k_cache.shape[1:]) != (H, max_cache_len, K): + raise ValueError( + f"`k_cache` per-slot shape must be {(H, max_cache_len, K)} " + f"(got {tuple(k_cache.shape[1:])})." + ) + if tuple(g_cache.shape[1:]) != (HV, max_cache_len): + raise ValueError( + f"`g_cache` per-slot shape must be {(HV, max_cache_len)} " + f"(got {tuple(g_cache.shape[1:])})." + ) + if g_cache.dtype != torch.float32: + raise ValueError(f"`g_cache` must be float32 (got {g_cache.dtype}).") + + BK = triton.next_power_of_2(K) + if triton.cdiv(K, BK) != 1: + raise ValueError(f"Cached decode kernel only supports NK_global=1 (got K={K}, BK={BK}).") + if BK % nk != 0: + raise ValueError(f"nk={nk} must divide BK={BK}.") + BKT = BK // nk + if BKT < 16: + raise ValueError(f"BKT={BKT} must be >=16 for tl.dot (nk={nk}, BK={BK}).") + # K-tiling keeps the per-program tile small enough that BV=64 (NV=1, half the + # grid -> fewer redundant cache/metadata loads) fits without register + # spilling. + BV = block_v if block_v is not None else min(triton.next_power_of_2(V), 64) + BC = max(16, triton.next_power_of_2(max_cache_len)) + + grid = (triton.cdiv(V, BV), B, HV) + fused_recurrent_gated_delta_rule_replayssm_kernel[grid]( + mixed_qkv=mixed_qkv, a=a, b=b, A_log=A_log, dt_bias=dt_bias, o=out, + h0=initial_state, ht=initial_state, + d_cache=d_cache, k_cache=k_cache, g_cache=g_cache, + ssm_state_indices=ssm_state_indices, write_pos=write_pos, scale=scale, + stride_mixed_qkv_tok=mixed_qkv.stride(0), + stride_a_tok=a.stride(0), stride_b_tok=b.stride(0), + stride_init_state_token=initial_state.stride(0), + stride_final_state_token=initial_state.stride(0), + stride_indices_seq=ssm_state_indices.stride(0), + stride_d_slot=d_cache.stride(0), + stride_k_slot=k_cache.stride(0), + stride_g_slot=g_cache.stride(0), + H=H, HV=HV, K=K, V=V, BK=BK, BV=BV, BC=BC, NK=nk, BKT=BKT, + MAX_CACHE_LEN=max_cache_len, SOFTPLUS_THRESHOLD=20.0, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + num_warps=num_warps, num_stages=num_stages, + ) + return out, initial_state diff --git a/vllm/model_executor/layers/fla/ops/gdn_replayssm_spec_decode.py b/vllm/model_executor/layers/fla/ops/gdn_replayssm_spec_decode.py new file mode 100644 index 000000000000..3bc3295e0efa --- /dev/null +++ b/vllm/model_executor/layers/fla/ops/gdn_replayssm_spec_decode.py @@ -0,0 +1,642 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# ruff: noqa: E501 + +from __future__ import annotations + +import torch + +from vllm.model_executor.layers.mamba.ops.replayssm_config import get_replayssm_config +from vllm.triton_utils import tl, triton + + +@triton.jit +def gdn_replayssm_spec_circular_kernel( + mixed_qkv, # [total_tokens, qkv_dim] packed, channel-last (q|k|v) + a, # [total_tokens, HV] + b, # [total_tokens, HV] + A_log, # [HV] fp32 + dt_bias, # [HV] fp32 + o, # [total_tokens, HV, V] preallocated output + h0, # [num_slots, HV, V, K] checkpoint state (== ht, in-place) + ht, # [num_slots, HV, V, K] + d_cache, # [num_slots, HV, L, V] + k_cache, # [num_slots, H, L, K] + g_cache, # [num_slots, HV, L] fp32 + query_start_loc, # [B+1] int packed cu_seqlens + ssm_state_indices, # [B] int physical block per request + write_pos, # [num_slots] int32 block-keyed + cache_base, # [num_slots] int32 block-keyed circular origin + is_flush_flags, # [num_slots] int8 block-keyed + scale, + stride_mqkv_t: tl.constexpr, # per-token stride of mixed_qkv + stride_a_t: tl.constexpr, + stride_b_t: tl.constexpr, + stride_o_t: tl.constexpr, # per-token stride of o (= HV*V) + stride_state_slot: tl.constexpr, + stride_d_slot: tl.constexpr, + stride_k_slot: tl.constexpr, + stride_g_slot: tl.constexpr, + stride_qsl: tl.constexpr, + stride_indices: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + BS: tl.constexpr, + BC: tl.constexpr, + NK: tl.constexpr, + BKT: tl.constexpr, + MAX_CACHE_LEN: tl.constexpr, + SOFTPLUS_THRESHOLD: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + IS_FLUSH: tl.constexpr, + NULL_BLOCK_ID: tl.constexpr, +): + i_v = tl.program_id(0) + i_n = tl.program_id(1) + i_hv = tl.program_id(2) + i_h = i_hv // (HV // H) + + o_v = i_v * BV + tl.arange(0, BV) + o_s = tl.arange(0, BS) + o_c = tl.arange(0, BC) + mask_v = o_v < V + + # --- per-request packed window --- + bos = tl.load(query_start_loc + i_n * stride_qsl).to(tl.int64) + eos = tl.load(query_start_loc + (i_n + 1) * stride_qsl).to(tl.int64) + spec_len = eos - bos # full window length + + state_idx = tl.load(ssm_state_indices + i_n * stride_indices).to(tl.int64) + + # output pointer (packed): token (bos + o_s), value-head i_hv, dim o_v + p_o = o + (bos + o_s[:, None]) * stride_o_t + i_hv * V + o_v[None, :] + + if IS_FLUSH: + if state_idx <= NULL_BLOCK_ID: + return + b_is_flush = tl.load(is_flush_flags + state_idx) != 0 + if not b_is_flush: + return + else: + if state_idx <= NULL_BLOCK_ID: + full_mask = (o_s < spec_len)[:, None] & mask_v[None, :] + tl.store( + p_o, + tl.zeros([BS, BV], dtype=tl.float32).to(p_o.dtype.element_ty), + mask=full_mask, + ) + return + b_is_flush = tl.load(is_flush_flags + state_idx) != 0 + if b_is_flush: + return + + b_write_pos = tl.load(write_pos + state_idx).to(tl.int64) + b_cache_base = tl.load(cache_base + state_idx).to(tl.int32) + + mask_s = o_s < spec_len + out_mask = mask_s[:, None] & mask_v[None, :] + + b_wp_i = b_write_pos.to(tl.int32) + cache_valid = o_c < b_write_pos + + # CIRCULAR physical slots (addresses only; masks/cumsums stay logical). + phys_c = (b_cache_base + o_c) & (MAX_CACHE_LEN - 1) # [BC] history + phys_spec = (b_cache_base + b_wp_i + o_s) & (MAX_CACHE_LEN - 1) # [BS] spec + + # ------------------------------------------------------------------ + # Block 0: gates / beta / local cumsum + committed-history replay decay. + # ------------------------------------------------------------------ + A_log_val = tl.load(A_log + i_hv).to(tl.float32) + dt_bias_val = tl.load(dt_bias + i_hv).to(tl.float32) + a_s = tl.load(a + (bos + o_s) * stride_a_t + i_hv, mask=mask_s, other=0.0).to( + tl.float32 + ) + b_s = tl.load(b + (bos + o_s) * stride_b_t + i_hv, mask=mask_s, other=0.0).to( + tl.float32 + ) + x = a_s + dt_bias_val + softplus_x = tl.where(x <= SOFTPLUS_THRESHOLD, tl.log(1.0 + tl.exp(x)), x) + g_s = tl.where(mask_s, -tl.exp(A_log_val) * softplus_x, 0.0) + beta_s = tl.where(mask_s, tl.sigmoid(b_s), 0.0) + G_s = tl.cumsum(g_s, axis=0) + expG_s = tl.exp(G_s) + + # committed-history replay decay from cached g (history loads -> phys_c) + p_g_main = g_cache + state_idx * stride_g_slot + i_hv * MAX_CACHE_LEN + phys_c + b_g_all = tl.load(p_g_main, mask=cache_valid, other=0.0).to(tl.float32) + b_g_prefix = tl.cumsum(b_g_all, axis=0) + b_g_total = tl.sum(b_g_all, axis=0) + b_replay_decay = tl.where(cache_valid, tl.exp(b_g_total - b_g_prefix), 0.0) + b_total_decay = tl.exp(b_g_total) + + p_d_main = d_cache + ( + state_idx * stride_d_slot + + (i_hv * MAX_CACHE_LEN + phys_c[None, :]) * V + + o_v[:, None] + ) + b_d_all = tl.load( + p_d_main, mask=mask_v[:, None] & cache_valid[None, :], other=0.0 + ).to(tl.float32) + b_d_scaled = (b_d_all * b_replay_decay[None, :]).to(mixed_qkv.dtype.element_ty) + + if USE_QK_L2NORM_IN_KERNEL: + qnorm_acc = tl.zeros([BS], dtype=tl.float32) + knorm_acc = tl.zeros([BS], dtype=tl.float32) + for kk in range(NK): + o_kt = kk * BKT + tl.arange(0, BKT) + mask_kt = o_kt < K + ld = mask_s[:, None] & mask_kt[None, :] + qn = tl.load( + mixed_qkv + + (bos + o_s[:, None]) * stride_mqkv_t + + i_h * K + + o_kt[None, :], + mask=ld, + other=0.0, + ).to(tl.float32) + knn = tl.load( + mixed_qkv + + (bos + o_s[:, None]) * stride_mqkv_t + + H * K + + i_h * K + + o_kt[None, :], + mask=ld, + other=0.0, + ).to(tl.float32) + qnorm_acc += tl.sum(qn * qn, axis=1) + knorm_acc += tl.sum(knn * knn, axis=1) + q_rnorm = tl.where(mask_s, 1.0 / tl.sqrt(qnorm_acc + 1e-6), 0.0) + k_rnorm = tl.where(mask_s, 1.0 / tl.sqrt(knorm_acc + 1e-6), 0.0) + else: + q_rnorm = tl.where(mask_s, 1.0, 0.0) + k_rnorm = tl.where(mask_s, 1.0, 0.0) + + # ------------------------------------------------------------------ + # K-Tiled Fused Projection and Intra-Spec Matrices (+ flush) + # ------------------------------------------------------------------ + hw_q = tl.zeros([BV, BS], dtype=tl.float32) + hw_k = tl.zeros([BV, BS], dtype=tl.float32) + if not IS_FLUSH: + scores_q = tl.zeros([BC, BS], dtype=tl.float32) + scores_k = tl.zeros([BC, BS], dtype=tl.float32) + kk_mat = tl.zeros([BS, BS], dtype=tl.float32) + kq_mat = tl.zeros([BS, BS], dtype=tl.float32) + + write_k = (i_v == 0) and (i_hv == i_h * (HV // H)) + + for kk in range(NK): + o_kt = kk * BKT + tl.arange(0, BKT) + mask_kt = o_kt < K + ld_s = mask_s[:, None] & mask_kt[None, :] + q_tile = tl.load( + mixed_qkv + (bos + o_s[:, None]) * stride_mqkv_t + i_h * K + o_kt[None, :], + mask=ld_s, + other=0.0, + ).to(tl.float32) + k_tile = tl.load( + mixed_qkv + + (bos + o_s[:, None]) * stride_mqkv_t + + H * K + + i_h * K + + o_kt[None, :], + mask=ld_s, + other=0.0, + ).to(tl.float32) + q_tile = (q_tile * (q_rnorm * scale)[:, None]).to(mixed_qkv.dtype.element_ty) + k_tile = (k_tile * k_rnorm[:, None]).to(mixed_qkv.dtype.element_ty) + + p_h0 = ( + h0 + + state_idx * stride_state_slot + + i_hv * V * K + + o_v[:, None] * K + + o_kt[None, :] + ) + sc_tile = tl.load( + p_h0, mask=mask_v[:, None] & mask_kt[None, :], other=0.0 + ).to(mixed_qkv.dtype.element_ty) + # cached-key history load -> phys_c + p_k = k_cache + ( + state_idx * stride_k_slot + + (i_h * MAX_CACHE_LEN + phys_c[:, None]) * K + + o_kt[None, :] + ) + khist_tile = tl.load( + p_k, mask=cache_valid[:, None] & mask_kt[None, :], other=0.0 + ).to(mixed_qkv.dtype.element_ty) + + qT = tl.trans(q_tile) + kT = tl.trans(k_tile) + kk_mat += tl.dot(k_tile, kT) + kq_mat += tl.dot(k_tile, qT) + + if IS_FLUSH: + sw_f = tl.dot(b_d_scaled, khist_tile, acc=b_total_decay * sc_tile.to(tl.float32)) + sw_tile = sw_f.to(mixed_qkv.dtype.element_ty) + hw_q += tl.dot(sw_tile, qT) + hw_k += tl.dot(sw_tile, kT) + p_ht = ( + ht + + state_idx * stride_state_slot + + i_hv * V * K + + o_v[:, None] * K + + o_kt[None, :] + ) + tl.store(p_ht, sw_tile, mask=mask_v[:, None] & mask_kt[None, :]) + else: + hw_q += tl.dot(sc_tile, qT) + hw_k += tl.dot(sc_tile, kT) + scores_q += tl.dot(khist_tile, qT) + scores_k += tl.dot(khist_tile, kT) + + if write_k: + # spec key store -> phys_spec (circular) + p_cur_k = k_cache + ( + state_idx * stride_k_slot + + (i_h * MAX_CACHE_LEN + phys_spec[:, None]) * K + + o_kt[None, :] + ) + tl.store( + p_cur_k, + k_tile, + mask=mask_s[:, None] + & mask_kt[None, :] + & ((b_write_pos + o_s[:, None]) < MAX_CACHE_LEN), + ) + + if not IS_FLUSH: + hw_q = b_total_decay * hw_q + tl.dot(b_d_scaled, scores_q.to(b_d_scaled.dtype)) + hw_k = b_total_decay * hw_k + tl.dot(b_d_scaled, scores_k.to(b_d_scaled.dtype)) + + # ------------------------------------------------------------------ + # strictly-lower A and T = (I + A)^{-1}. + # ------------------------------------------------------------------ + lower = (o_s[:, None] > o_s[None, :]) & mask_s[:, None] & mask_s[None, :] + diff_ij = G_s[:, None] - G_s[None, :] + A_mat = tl.where(lower, beta_s[:, None] * tl.exp(diff_ij) * kk_mat, 0.0) + + b_Ai = -A_mat + for ii in range(2, BS): + row = tl.sum(tl.where((o_s == ii)[:, None], -A_mat, 0.0), axis=0) + row = tl.where(o_s < ii, row, 0.0) + row = row + tl.sum(row[:, None] * b_Ai, axis=0) + b_Ai = tl.where((o_s == ii)[:, None], row, b_Ai) + T_mat = b_Ai + (o_s[:, None] == o_s[None, :]).to(tl.float32) + + # ------------------------------------------------------------------ + # R and D_spec = R @ T^T. + # ------------------------------------------------------------------ + p_v = ( + mixed_qkv + + (bos + o_s[None, :]) * stride_mqkv_t + + 2 * H * K + + i_hv * V + + o_v[:, None] + ) + v_tile = tl.load(p_v, mask=mask_v[:, None] & mask_s[None, :], other=0.0).to(tl.float32) + R_mat = beta_s[None, :] * (v_tile - expG_s[None, :] * hw_k) + D_spec = tl.zeros([BV, BS], dtype=tl.float32) + for j in tl.static_range(BS): + Rj = tl.sum(tl.where((o_s == j)[None, :], R_mat, 0.0), axis=1) + Tj = tl.sum(tl.where((o_s == j)[None, :], T_mat, 0.0), axis=1) + D_spec += Rj[:, None] * Tj[None, :] + + # ------------------------------------------------------------------ + # outputs. + # ------------------------------------------------------------------ + causalF = (o_s[:, None] <= o_s[None, :]) & mask_s[:, None] & mask_s[None, :] + diff_ji = G_s[None, :] - G_s[:, None] + F_mat = tl.where(causalF, tl.exp(diff_ji) * kq_mat, 0.0) + DF = tl.zeros([BV, BS], dtype=tl.float32) + for j in tl.static_range(BS): + Dj = tl.sum(tl.where((o_s == j)[None, :], D_spec, 0.0), axis=1) + Fj = tl.sum(tl.where((o_s == j)[:, None], F_mat, 0.0), axis=0) + DF += Dj[:, None] * Fj[None, :] + O_tile = expG_s[None, :] * hw_q + DF + + tl.store(p_o, tl.trans(O_tile).to(p_o.dtype.element_ty), mask=out_mask) + + # ------------------------------------------------------------------ + # write speculative d / g at circular positions (phys_spec). + # ------------------------------------------------------------------ + spec_pos = b_write_pos + o_s + spec_store_mask = mask_s & (spec_pos < MAX_CACHE_LEN) + p_cur_d = d_cache + ( + state_idx * stride_d_slot + + (i_hv * MAX_CACHE_LEN + phys_spec[None, :]) * V + + o_v[:, None] + ) + tl.store( + p_cur_d, + D_spec.to(p_cur_d.dtype.element_ty), + mask=mask_v[:, None] & spec_store_mask[None, :], + ) + if i_v == 0: + p_cur_g = g_cache + state_idx * stride_g_slot + i_hv * MAX_CACHE_LEN + phys_spec + tl.store(p_cur_g, g_s, mask=spec_store_mask) + + +@triton.jit +def _advance_gdn_spec_cursors_kernel( + write_pos_ptr, + cache_base_ptr, + is_flush_ptr, + num_accepted_ptr, + state_batch_indices_ptr, + n_rows, + stride_sbi: tl.constexpr, + stride_na: tl.constexpr, + MAX_CACHE_LEN: tl.constexpr, + MAX_SPEC_LEN: tl.constexpr, + CACHE_BUF_LEN: tl.constexpr, + BLOCK: tl.constexpr, + NULL_BLOCK_ID: tl.constexpr, +): + offs = tl.arange(0, BLOCK) + row_mask = offs < n_rows + blk = tl.load( + state_batch_indices_ptr + offs * stride_sbi, mask=row_mask, other=NULL_BLOCK_ID + ).to(tl.int64) + valid = row_mask & (blk > NULL_BLOCK_ID) + + write_pos = tl.load(write_pos_ptr + blk, mask=valid, other=0).to(tl.int32) + cache_base = tl.load(cache_base_ptr + blk, mask=valid, other=0).to(tl.int32) + is_flush_cur = tl.load(is_flush_ptr + blk, mask=valid, other=0).to(tl.int32) + num_acc = tl.load( + num_accepted_ptr + offs * stride_na, mask=valid, other=0 + ).to(tl.int32) + + total_commit = num_acc + flush_now = (total_commit > 0) & (is_flush_cur != 0) + + new_base = tl.where( + flush_now, (cache_base + write_pos) & (CACHE_BUF_LEN - 1), cache_base + ) + new_wp = tl.where(is_flush_cur != 0, total_commit, write_pos + total_commit).to( + tl.int32 + ) + # Early-flush one window early so every verify step satisfies + # write_pos + spec_len <= max_cache_len (the spec window never overflows). + next_is_flush = ((new_wp + 2 * MAX_SPEC_LEN) > MAX_CACHE_LEN).to(tl.int8) + + tl.store(write_pos_ptr + blk, new_wp, mask=valid) + tl.store(cache_base_ptr + blk, new_base, mask=valid) + tl.store(is_flush_ptr + blk, next_is_flush, mask=valid) + + +@triton.jit +def _reset_gdn_replayssm_spec_cursors_kernel( + write_pos_ptr, + cache_base_ptr, + is_flush_ptr, + do_reset_ptr, + state_batch_indices_ptr, + n_rows, + stride_sbi: tl.constexpr, + stride_reset: tl.constexpr, + INIT_FLUSH: tl.constexpr, + BLOCK: tl.constexpr, + NULL_BLOCK_ID: tl.constexpr, +): + offs = tl.arange(0, BLOCK) + row_mask = offs < n_rows + blk = tl.load( + state_batch_indices_ptr + offs * stride_sbi, mask=row_mask, other=NULL_BLOCK_ID + ).to(tl.int64) + do_reset = tl.load( + do_reset_ptr + offs * stride_reset, mask=row_mask, other=0 + ).to(tl.int32) + do = row_mask & (blk > NULL_BLOCK_ID) & (do_reset != 0) + + tl.store(write_pos_ptr + blk, tl.zeros_like(blk).to(tl.int32), mask=do) + tl.store(cache_base_ptr + blk, tl.zeros_like(blk).to(tl.int32), mask=do) + tl.store( + is_flush_ptr + blk, + tl.full([BLOCK], INIT_FLUSH, dtype=tl.int8), + mask=do, + ) + + +# --------------------------------------------------------------------------- +# Python wrappers. +# --------------------------------------------------------------------------- +def _launch_gdn_spec( + mixed_qkv, + a, + b, + A_log, + dt_bias, + out, + checkpoint_state, + d_cache, + k_cache, + g_cache, + query_start_loc, + ssm_state_indices, + write_pos, + cache_base, + is_flush, + scale, + max_cache_len, + max_spec_len, + use_qk_l2norm_in_kernel, + is_flush_kernel, + block_v, + num_warps, + num_stages, + nk, + null_block_id, +): + num_slots, HV, V, K = checkpoint_state.shape + qkv_dim = mixed_qkv.shape[1] + q_dim = (qkv_dim - HV * V) // 2 + H = q_dim // K + B = query_start_loc.shape[0] - 1 + # max_cache_len is the logical flush threshold L; the physical pow2 ring is + # d_cache.shape[2] = next_pow2(L) and wraps addresses / per-head strides. + buf = d_cache.shape[2] + assert buf & (buf - 1) == 0, "circular cache requires power-of-two buffer" + + BK = triton.next_power_of_2(K) + if triton.cdiv(K, BK) != 1: + raise ValueError(f"only NK_global=1 supported (K={K}, BK={BK}).") + if BK % nk != 0: + raise ValueError(f"nk={nk} must divide BK={BK}.") + BKT = BK // nk + if BKT < 16: + raise ValueError(f"BKT={BKT} must be >=16 for tl.dot.") + BV = block_v if block_v is not None else min(triton.next_power_of_2(V), 64) + BS = max(4, triton.next_power_of_2(max_spec_len)) + # History block decoupled from the physical buffer: with L = B + max_spec_len + # committed history never exceeds L - max_spec_len, so BC covers it while + # staying small (the L=B+T win). + BC = max(16, triton.next_power_of_2(max(1, max_cache_len - max_spec_len))) + + grid = (triton.cdiv(V, BV), B, HV) + gdn_replayssm_spec_circular_kernel[grid]( + mixed_qkv, + a, + b, + A_log, + dt_bias, + out, + checkpoint_state, + checkpoint_state, + d_cache, + k_cache, + g_cache, + query_start_loc, + ssm_state_indices, + write_pos, + cache_base, + is_flush, + scale, + mixed_qkv.stride(0), + a.stride(0), + b.stride(0), + out.stride(0), + checkpoint_state.stride(0), + d_cache.stride(0), + k_cache.stride(0), + g_cache.stride(0), + query_start_loc.stride(0), + ssm_state_indices.stride(0), + H=H, + HV=HV, + K=K, + V=V, + BK=BK, + BV=BV, + BS=BS, + BC=BC, + NK=nk, + BKT=BKT, + MAX_CACHE_LEN=buf, + SOFTPLUS_THRESHOLD=20.0, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + IS_FLUSH=is_flush_kernel, + NULL_BLOCK_ID=null_block_id, + num_warps=num_warps, + num_stages=num_stages, + ) + + +def gdn_replayssm_spec_decode( + mixed_qkv: torch.Tensor, # [total_tokens, qkv_dim] post-conv packed (q|k|v) + a: torch.Tensor, # [total_tokens, HV] + b: torch.Tensor, # [total_tokens, HV] + A_log: torch.Tensor, # [HV] fp32 + dt_bias: torch.Tensor, # [HV] fp32 + checkpoint_state: torch.Tensor, # [num_slots, HV, V, K] (in-place h0==ht) + d_cache: torch.Tensor, # [num_slots, HV, buf, V] + k_cache: torch.Tensor, # [num_slots, H, buf, K] + g_cache: torch.Tensor, # [num_slots, HV, buf] fp32 + out: torch.Tensor, # [total_tokens, HV, V] preallocated + query_start_loc: torch.Tensor, # [B+1] int + ssm_state_indices: torch.Tensor, # [B] int physical block per request + write_pos: torch.Tensor, # [num_slots] int32 block-keyed + cache_base: torch.Tensor, # [num_slots] int32 block-keyed + is_flush: torch.Tensor, # [num_slots] int8 block-keyed + max_cache_len: int, # logical flush threshold L = B + max_spec_len + max_spec_len: int, + scale: float | None = None, + use_qk_l2norm_in_kernel: bool = True, + null_block_id: int = 0, + launch_mode: str = "both", +): + """GDN cached speculative-decode on a CIRCULAR d/k/g cache (vLLM packed varlen). + + ``max_cache_len`` is the logical flush threshold L = B + max_spec_len; the + physical pow2 buffer is ``d_cache.shape[2]`` = ``next_pow2(L)`` and the history + block ``BC = next_pow2(L - max_spec_len)``. Two launches (verify + flush + ``IS_FLUSH`` specializations) with device-side per-row routing keep the step + CUDA-graph capturable. Cursors are advanced by ``commit_gdn_replayssm_spec``. + """ + if scale is None: + scale = checkpoint_state.shape[-1] ** -0.5 + if is_flush.dtype != torch.int8: + is_flush = is_flush.to(torch.int8) + vb, vw, vnk, vns = get_replayssm_config("gdn_spec_verify", max_spec_len=max_spec_len) + fb, fw, fnk, fns = get_replayssm_config("gdn_spec_flush", max_spec_len=max_spec_len) + + if launch_mode in ("both", "verify"): + _launch_gdn_spec( + mixed_qkv, a, b, A_log, dt_bias, out, checkpoint_state, + d_cache, k_cache, g_cache, query_start_loc, ssm_state_indices, + write_pos, cache_base, is_flush, scale, max_cache_len, max_spec_len, + use_qk_l2norm_in_kernel, False, vb, vw, vns, vnk, null_block_id, + ) + if launch_mode in ("both", "flush"): + _launch_gdn_spec( + mixed_qkv, a, b, A_log, dt_bias, out, checkpoint_state, + d_cache, k_cache, g_cache, query_start_loc, ssm_state_indices, + write_pos, cache_base, is_flush, scale, max_cache_len, max_spec_len, + use_qk_l2norm_in_kernel, True, fb, fw, fns, fnk, null_block_id, + ) + return out + + +def commit_gdn_replayssm_spec( + write_pos: torch.Tensor, + cache_base: torch.Tensor, + is_flush: torch.Tensor, + num_accepted: torch.Tensor, # [n_rows] int (already includes the bonus token) + state_batch_indices: torch.Tensor, # [n_rows] int physical block per row + max_cache_len: int, # logical flush threshold L + max_spec_len: int, + cache_buf_len: int | None = None, # physical pow2 buffer next_pow2(L) + null_block_id: int = 0, +): + """Advance the block-keyed cursors once per decode step (device-only).""" + if cache_buf_len is None: + cache_buf_len = triton.next_power_of_2(max_cache_len) + n_rows = state_batch_indices.shape[0] + BLOCK = triton.next_power_of_2(max(1, n_rows)) + _advance_gdn_spec_cursors_kernel[(1,)]( + write_pos, + cache_base, + is_flush, + num_accepted, + state_batch_indices, + n_rows, + stride_sbi=state_batch_indices.stride(0), + stride_na=num_accepted.stride(0), + MAX_CACHE_LEN=max_cache_len, + MAX_SPEC_LEN=max_spec_len, + CACHE_BUF_LEN=cache_buf_len, + BLOCK=BLOCK, + NULL_BLOCK_ID=null_block_id, + ) + + +def reset_gdn_replayssm_spec_cursors( + write_pos: torch.Tensor, + cache_base: torch.Tensor, + is_flush: torch.Tensor, + do_reset: torch.Tensor, # [n_rows] int/bool 1 for first-decode rows + state_batch_indices: torch.Tensor, # [n_rows] int + max_cache_len: int, + max_spec_len: int, + null_block_id: int = 0, +): + """Reset the cursors of first-decode rows (prefill->decode handoff).""" + n_rows = state_batch_indices.shape[0] + BLOCK = triton.next_power_of_2(max(1, n_rows)) + init_flush = 1 if 2 * max_spec_len > max_cache_len else 0 + _reset_gdn_replayssm_spec_cursors_kernel[(1,)]( + write_pos, + cache_base, + is_flush, + do_reset, + state_batch_indices, + n_rows, + stride_sbi=state_batch_indices.stride(0), + stride_reset=do_reset.stride(0), + INIT_FLUSH=init_flush, + BLOCK=BLOCK, + NULL_BLOCK_ID=null_block_id, + ) diff --git a/vllm/model_executor/layers/mamba/abstract.py b/vllm/model_executor/layers/mamba/abstract.py index 8bbb21d7bc90..83b4d3c47734 100644 --- a/vllm/model_executor/layers/mamba/abstract.py +++ b/vllm/model_executor/layers/mamba/abstract.py @@ -52,10 +52,16 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: page_size_padded=page_size_padded, mamba_type=self.mamba_type, mamba_cache_mode=vllm_config.cache_config.mamba_cache_mode, + # The cached GDN spec kernel never uses the per-token speculative + # checkpoint slots the baseline spec kernel writes. num_speculative_blocks=( - vllm_config.speculative_config.num_speculative_tokens - if vllm_config.speculative_config - else 0 + 0 + if vllm_config.cache_config.use_replayssm_spec + else ( + vllm_config.speculative_config.num_speculative_tokens + if vllm_config.speculative_config + else 0 + ) ), ) diff --git a/vllm/model_executor/layers/mamba/gdn/base.py b/vllm/model_executor/layers/mamba/gdn/base.py index 95ced2e0a4d3..ac8a62bfa30b 100644 --- a/vllm/model_executor/layers/mamba/gdn/base.py +++ b/vllm/model_executor/layers/mamba/gdn/base.py @@ -51,8 +51,16 @@ def mamba_type(self) -> MambaAttentionBackendEnum: return MambaAttentionBackendEnum.GDN_ATTN def get_state_dtype(self) -> tuple[torch.dtype, ...]: - return MambaStateDtypeCalculator.gated_delta_net_state_dtype( + if self.cache_config.use_replayssm_spec: + return MambaStateDtypeCalculator.gated_delta_net_spec_cached_state_dtype( + self.model_config.dtype, + self.cache_config.mamba_cache_dtype, + self.cache_config.mamba_ssm_cache_dtype, + self.cache_config.use_replayssm_spec, + ) + return MambaStateDtypeCalculator.gated_delta_net_cached_state_dtype( self.model_config.dtype, self.cache_config.mamba_cache_dtype, self.cache_config.mamba_ssm_cache_dtype, + self.cache_config.use_replayssm, ) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 06bfe5c5de27..02bbc0e5c39c 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -26,6 +26,7 @@ ) from vllm.model_executor.layers.fla.ops import ( fused_post_conv_prep, + fused_recurrent_gated_delta_rule_replayssm, fused_recurrent_gated_delta_rule_packed_decode, fused_sigmoid_gating_delta_rule_update, ) @@ -420,14 +421,28 @@ def forward_cutedsl( class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): def get_state_shape( self, - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...]]: - return MambaStateShapeCalculator.gated_delta_net_state_shape( + ) -> tuple[tuple[int, ...], ...]: + if self.cache_config.use_replayssm_spec: + return MambaStateShapeCalculator.gated_delta_net_spec_cached_state_shape( + self.tp_size, + self.num_k_heads, + self.num_v_heads, + self.head_k_dim, + self.head_v_dim, + self.conv_kernel_size, + self.cache_config.use_replayssm_spec, + self.cache_config.replayssm_buffer_len, + self.num_spec, + ) + return MambaStateShapeCalculator.gated_delta_net_cached_state_shape( self.tp_size, self.num_k_heads, self.num_v_heads, self.head_k_dim, self.head_v_dim, self.conv_kernel_size, + self.cache_config.use_replayssm, + self.cache_config.replayssm_buffer_len, self.num_spec, ) @@ -557,6 +572,18 @@ def __init__( self.enable_packed_recurrent_decode = ( envs.VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE ) + # Cached-decode kernel (reuses the engine-level mamba_* flags). When + # enabled, the paged GDN page grows to 5 tensors and the non-spec decode + # branch decodes through fused_recurrent_gated_delta_rule_replayssm. + self.use_cache_kernel = self.cache_config.use_replayssm + self.max_cache_len = self.cache_config.replayssm_buffer_len + # Cached-SPEC decode kernel (gdn_replayssm_spec_decode). When enabled, the + # GDN page grows to the same 5-tuple (fp32 checkpoint) and the spec verify + # path decodes through the circular cached kernel. + self.use_cache_spec_kernel = ( + self.cache_config.use_replayssm_spec + ) + self.max_spec_len = 1 + self.num_spec compilation_config = get_current_vllm_config().compilation_config if prefix in compilation_config.static_forward_context: @@ -1098,7 +1125,8 @@ def _warmup_prefill_kernels(self, qkv_or_qkvz: torch.Tensor, v_dim: int) -> None dtype = qkv_or_qkvz.dtype num_k_heads = self.num_k_heads // self.tp_size num_v_heads = self.num_v_heads // self.tp_size - _, state_dtype = self.get_state_dtype() + # get_state_dtype() is (conv, ssm[, d, k, g]); we only need the ssm dtype. + state_dtype = self.get_state_dtype()[1] # All kernels use BT = chunk_size, so a single pass with T = chunk_size # is sufficient to populate every autotuner cache. Mirror the real @@ -1283,12 +1311,24 @@ def _forward_core( attn_metadata = attn_metadata_raw[self.prefix] # type: ignore[index] assert isinstance(attn_metadata, GDNAttentionMetadata) - if ( - self.enable_packed_recurrent_decode - and attn_metadata.spec_sequence_masks is None + is_non_spec_decode = ( + attn_metadata.spec_sequence_masks is None and attn_metadata.num_prefills == 0 and attn_metadata.num_decodes > 0 - ): + ) + + # Cached decode kernel (own kernel; independent of the packed-recurrent + # env flag). + if self.use_cache_kernel and is_non_spec_decode: + return self._forward_core_decode_non_spec_cached( + mixed_qkv=mixed_qkv, + b=b, + a=a, + core_attn_out=core_attn_out, + attn_metadata=attn_metadata, + ) + + if self.enable_packed_recurrent_decode and is_non_spec_decode: return self._forward_core_decode_non_spec( mixed_qkv=mixed_qkv, b=b, @@ -1352,7 +1392,9 @@ def _forward_core( ], num_accepted_tokens=num_accepted_tokens, query_start_loc=spec_query_start_loc, - max_query_len=spec_state_indices_tensor.size(-1), + # Spec verify window = 1 + num_spec. Use the constant rather than + # the block-table width so the cached-spec path can request num_speculative_blocks=0 + max_query_len=self.max_spec_len, validate_data=False, ) @@ -1389,7 +1431,16 @@ def _forward_core( else: mixed_qkv_non_spec = None - query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec) + # The cached-spec kernel consumes the post-conv packed ``mixed_qkv_spec`` + # directly, so the split/contiguous rearrange into q/k/v (3 cat copies) is + # pure waste on that path -- skip it to remove GPU work (and the eager + # CPU-dispatch bubble) between the conv and the SSM kernel. + if spec_sequence_masks is not None and self.use_cache_spec_kernel: + query_spec, key_spec, value_spec = None, None, None + else: + query_spec, key_spec, value_spec = self.rearrange_mixed_qkv( + mixed_qkv_spec + ) # Split mixed non-spec-decode+prefill to process independently split_non_spec = ( @@ -1452,7 +1503,57 @@ def _forward_core( # 2. Recurrent attention # 2.1: Process the multi-query part - if spec_sequence_masks is not None: + if spec_sequence_masks is not None and self.use_cache_spec_kernel: + # Cached circular spec verify: reuse the post-conv packed + # ``mixed_qkv_spec`` (q|k|v) + raw ``a``/``b`` (read per-request via + # spec_query_start_loc, same as the baseline kernel). The d/k/g ring + # caches + fp32 checkpoint live in the grown 5-tuple page; cursors + # are block-keyed in the metadata. + from vllm.model_executor.layers.fla.ops.gdn_replayssm_spec_decode import ( + gdn_replayssm_spec_decode, + ) + + assert mixed_qkv_spec is not None + num_spec_decodes = attn_metadata.num_spec_decodes + d_cache = self_kv_cache[2] + k_cache = self_kv_cache[3] + g_cache = self_kv_cache[4] + total_spec = mixed_qkv_spec.shape[0] + cs_out = torch.empty( + total_spec, + self.num_v_heads // self.tp_size, + self.head_v_dim, + dtype=mixed_qkv_spec.dtype, + device=mixed_qkv_spec.device, + ) + gdn_replayssm_spec_decode( + mixed_qkv=mixed_qkv_spec, + a=a, + b=b, + A_log=self.A_log, + dt_bias=self.dt_bias, + checkpoint_state=ssm_state, + d_cache=d_cache, + k_cache=k_cache, + g_cache=g_cache, + out=cs_out, + query_start_loc=spec_query_start_loc[ # type: ignore[index] + : num_spec_decodes + 1 + ], + ssm_state_indices=spec_state_indices_tensor[ # type: ignore[index] + :num_spec_decodes, 0 + ], + write_pos=attn_metadata.spec_write_pos_d, + cache_base=attn_metadata.spec_cache_base_d, + is_flush=attn_metadata.spec_is_flush_d, + max_cache_len=self.max_cache_len + self.max_spec_len, + max_spec_len=self.max_spec_len, + scale=self.head_k_dim**-0.5, + use_qk_l2norm_in_kernel=True, + ) + core_attn_out_spec = cs_out.unsqueeze(0) + last_recurrent_state = None + elif spec_sequence_masks is not None: core_attn_out_spec, last_recurrent_state = ( fused_sigmoid_gating_delta_rule_update( A_log=self.A_log, @@ -1695,6 +1796,71 @@ def _forward_core_decode_non_spec( ) return + def _forward_core_decode_non_spec_cached( + self, + mixed_qkv: torch.Tensor, + b: torch.Tensor, + a: torch.Tensor, + core_attn_out: torch.Tensor, + attn_metadata: GDNAttentionMetadata, + ): + """ + Cached non-spec decode: amortizes the SSM-state HBM traffic by caching + the per-step d/k/g vectors in a ring buffer and reconstructing the + output from a checkpoint that is only rewritten every max_cache_len + steps. + """ + non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501 + write_pos_d = attn_metadata.write_pos_d + self_kv_cache = self.kv_cache + # conv_state must be (..., dim, width-1) for the conv kernels. + # DS layout stores it that way directly; SD layout needs a transpose. + conv_state = ( + self_kv_cache[0] + if is_conv_state_dim_first() + else self_kv_cache[0].transpose(-1, -2) + ) + ssm_state = self_kv_cache[1] + d_cache = self_kv_cache[2] + k_cache = self_kv_cache[3] + g_cache = self_kv_cache[4] + num_actual_tokens = attn_metadata.num_actual_tokens + + mixed_qkv = mixed_qkv[:num_actual_tokens] + b = b[:num_actual_tokens] + a = a[:num_actual_tokens] + + conv_weights = self.conv1d.weight.view( + self.conv1d.weight.size(0), self.conv1d.weight.size(2) + ) + mixed_qkv_non_spec = causal_conv1d_update( + mixed_qkv, + conv_state, + conv_weights, + self.conv1d.bias, + self.activation, + conv_state_indices=non_spec_state_indices_tensor[:num_actual_tokens], # type: ignore[index] + validate_data=False, + ) + out_buf = core_attn_out[:num_actual_tokens].unsqueeze(1) + fused_recurrent_gated_delta_rule_replayssm( + mixed_qkv=mixed_qkv_non_spec, + a=a, + b=b, + A_log=self.A_log, + dt_bias=self.dt_bias, + scale=self.head_k_dim**-0.5, + initial_state=ssm_state, + d_cache=d_cache, + k_cache=k_cache, + g_cache=g_cache, + out=out_buf, + ssm_state_indices=non_spec_state_indices_tensor[:num_actual_tokens], # type: ignore[index] + write_pos=write_pos_d, + use_qk_l2norm_in_kernel=True, + ) + return + def qwen_gdn_attention_core( qkv_or_qkvz: torch.Tensor, diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index a6524961ea92..cfdff5edc8d3 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -35,6 +35,15 @@ from vllm.model_executor.layers.mamba.ops.ssd_combined import ( mamba_chunk_scan_combined_varlen, ) +from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_state_and_output import ( # noqa: E501 + selective_state_update_replayssm_state_and_output, +) +from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_output_only import ( # noqa: E501 + selective_state_update_replayssm_output_only, +) +from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_spec import ( + selective_state_update_replayssm_spec, +) from vllm.model_executor.layers.mamba.ops.ssu_dispatch import selective_state_update from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.model_loader.weight_utils import ( @@ -494,14 +503,46 @@ def __init__( if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self - # The tuple is (conv_state, ssm_state) - self.kv_cache = (torch.tensor([]), torch.tensor([])) self.model_config = model_config self.cache_config = cache_config self.prefix = prefix + self.use_cache_kernel = ( + cache_config.use_replayssm if cache_config is not None else False + ) + self.max_cache_len = ( + cache_config.replayssm_buffer_len if cache_config is not None else 16 + ) + self.cached_kernel_variant = ( + cache_config.replayssm_route + if cache_config is not None + else "state_and_output" + ) + self.use_cache_spec_kernel = ( + cache_config.use_replayssm_spec + if cache_config is not None + else False + ) + if ( + self.use_cache_kernel or self.use_cache_spec_kernel + ) and self.num_heads % self.tp_size != 0: + raise ValueError( + "Mamba2 cached decode kernel requires tensor-parallel heads " + "to divide evenly" + ) + # The tuple is (conv_state, ssm_state); with state-and-output/bc decode enabled + # (conv_state, ssm_state, x_cache, dt_cache, B_cache); with cached-spec + # (hybrid) enabled (conv_state, ssm_state, post_conv_cache, dt_cache). + if self.use_cache_spec_kernel: + _n_state = 4 + elif self.use_cache_kernel: + _n_state = 5 + else: + _n_state = 2 + self.kv_cache = tuple(torch.tensor([]) for _ in range(_n_state)) self.num_spec = vllm_config.num_speculative_tokens + self.max_spec_len = 1 + self.num_spec if self.num_spec > 0: self.register_buffer( "_decode_state_offsets", @@ -591,7 +632,7 @@ def _warmup_ssd_kernels(self, projected_states: torch.Tensor) -> None: # Triton's autotuner includes tensor dtypes in its cache key, # so state_dtype must match what real inference uses. - _, ssm_state_dtype = self.get_state_dtype() + ssm_state_dtype = self.get_state_dtype()[1] # SSD kernel autotune keys depend on dtype and head dimensions, # not on sequence length or batch size, so a single shape suffices. @@ -702,6 +743,24 @@ def conv_ssm_forward( else self.kv_cache[0].transpose(-1, -2) ) ssm_state = self.kv_cache[1] + spec_post_conv_cache = spec_dt_cache = None + if self.use_cache_spec_kernel: + if len(self.kv_cache) != 4: + raise ValueError( + "Mamba2 cached-spec decode kernel requires four Mamba " + "state tensors (conv, ssm, post_conv_cache, dt_cache)" + ) + spec_post_conv_cache, spec_dt_cache = self.kv_cache[2:] + x_cache = dt_cache = B_cache = None + elif self.use_cache_kernel: + if len(self.kv_cache) != 5: + raise ValueError( + "Mamba2 cached decode kernel requires five Mamba " + "state tensors" + ) + x_cache, dt_cache, B_cache = self.kv_cache[2:] + else: + x_cache = dt_cache = B_cache = None has_initial_states_p = attn_metadata.has_initial_states_p prep_initial_states = attn_metadata.prep_initial_states chunk_size = attn_metadata.chunk_size @@ -1002,6 +1061,12 @@ def conv_ssm_forward( max_query_len=state_indices_tensor_d.size(-1), ) + # cached-spec (hybrid) feeds the full channel-last post-conv output + # ([num_decode_tokens, conv_dim]) straight to the SSM scatter (no + # split) and uses the raw (unexpanded) dt. Capture before reuse. + conv_out_spec = hidden_states_B_C_d + dt_d_raw = dt_d + hidden_states_d, B_d, C_d = self.split_hidden_states_B_C_fn( hidden_states_B_C_d ) @@ -1027,35 +1092,172 @@ def conv_ssm_forward( # - mamba_cache_params.ssm_state's slots will be selected # using state_indices_tensor_d # NOTE: final output is an in-place update of out tensor - selective_state_update( - ssm_state, - hidden_states_d, - dt_d, - A_d, - B_d, - C_d, - D_d, - dt_bias, - dt_softplus=True, - state_batch_indices=state_indices_tensor_d_input, - dst_state_batch_indices=state_indices_tensor_d_output, - out=preallocated_ssm_out_d.view(num_decode_tokens, -1, self.head_dim), - num_accepted_tokens=num_accepted_tokens, - cu_seqlens=query_start_loc_d, - is_blackwell=self.is_blackwell, + preallocated_ssm_out_d = preallocated_ssm_out_d.view( + num_decode_tokens, -1, self.head_dim ) + if ( + self.use_cache_spec_kernel + and attn_metadata.spec_write_pos_d is not None + ): + # Fires only on speculative-verify batches (spec cursors set in + # build()). Rare non-spec decode rows under the spec flag (e.g. a + # single-token prefill chunk replayed as decode) have no cursors + # and fall through to the baseline update below. + if is_mamba_cache_all: + raise ValueError( + "Mamba2 cached-spec decode kernel requires " + "mamba_cache_mode='none'" + ) + assert spec_post_conv_cache is not None + assert spec_dt_cache is not None + assert query_start_loc_d is not None + # Hybrid: causal_conv1d_update (above) already produced conv_out; + # feed it + raw dt + the prepared A/D/dt_bias to the circular + # scatter+scan. Cursors are block-keyed (advanced once per step + # by the commit in build()). out is the packed preallocated buf. + selective_state_update_replayssm_spec( + ssm_state, + spec_post_conv_cache, + spec_dt_cache, + conv_out_spec, + dt_d_raw, + A_d, + write_pos=attn_metadata.spec_write_pos_d, + post_conv_state_pos=attn_metadata.spec_post_origin_d, + is_flush=attn_metadata.spec_is_flush_d, + query_start_loc=query_start_loc_d, + state_batch_indices=state_indices_tensor_d[:, 0], + max_cache_len=self.max_cache_len + self.max_spec_len, + max_spec_len=self.max_spec_len, + d_inner=self.intermediate_size // self.tp_size, + ngroups=self.n_groups // self.tp_size, + dstate=self.ssm_state_size, + D=D_d, + z=None, + dt_bias=dt_bias, + dt_softplus=True, + out=preallocated_ssm_out_d, + bc_pre=attn_metadata.spec_bc_pre_scratch, + ) + elif self.use_cache_kernel: + if is_mamba_cache_all: + raise ValueError( + "Mamba2 cached decode kernel requires " + "mamba_cache_mode='none'" + ) + if num_accepted_tokens is not None or query_start_loc_d is not None: + raise ValueError( + "Mamba2 cached decode kernel does not support " + "speculative or varlen decode" + ) + if attn_metadata.write_pos_d is None: + raise ValueError( + "Mamba2 cached decode metadata is missing write_pos_d" + ) + if attn_metadata.is_flush_d is None: + raise ValueError( + "Mamba2 cached decode metadata is missing is_flush_d" + ) + assert x_cache is not None + assert dt_cache is not None + assert B_cache is not None + if self.cached_kernel_variant == "output_only": + if attn_metadata.bc_pre_scratch is None: + raise ValueError( + "Mamba2 output-only decode kernel requires " + "bc_pre_scratch in attention metadata" + ) + selective_state_update_replayssm_output_only( + ssm_state, + hidden_states_d, + dt_d, + A_d, + B_d, + C_d, + D_d, + dt_bias, + dt_softplus=True, + x_cache=x_cache, + dt_cache=dt_cache, + B_cache=B_cache, + bc_pre=attn_metadata.bc_pre_scratch, + write_pos=attn_metadata.write_pos_d, + is_flush=attn_metadata.is_flush_d, + max_cache_len=self.max_cache_len, + state_batch_indices=state_indices_tensor_d_input, + out=preallocated_ssm_out_d, + ) + else: + selective_state_update_replayssm_state_and_output( + ssm_state, + hidden_states_d, + dt_d, + A_d, + B_d, + C_d, + D_d, + dt_bias, + dt_softplus=True, + x_cache=x_cache, + dt_cache=dt_cache, + B_cache=B_cache, + write_pos=attn_metadata.write_pos_d, + is_flush=attn_metadata.is_flush_d, + max_cache_len=self.max_cache_len, + state_batch_indices=state_indices_tensor_d_input, + out=preallocated_ssm_out_d, + ) + else: + selective_state_update( + ssm_state, + hidden_states_d, + dt_d, + A_d, + B_d, + C_d, + D_d, + dt_bias, + dt_softplus=True, + state_batch_indices=state_indices_tensor_d_input, + dst_state_batch_indices=state_indices_tensor_d_output, + out=preallocated_ssm_out_d, + num_accepted_tokens=num_accepted_tokens, + cu_seqlens=query_start_loc_d, + is_blackwell=self.is_blackwell, + ) - def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]: + def get_state_dtype(self) -> tuple[torch.dtype, ...]: assert self.model_config is not None assert self.cache_config is not None - return MambaStateDtypeCalculator.mamba2_state_dtype( + if self.use_cache_spec_kernel: + return MambaStateDtypeCalculator.mamba2_spec_cached_state_dtype( + self.model_config.dtype, + self.cache_config.mamba_cache_dtype, + self.cache_config.mamba_ssm_cache_dtype, + use_replayssm_spec=self.use_cache_spec_kernel, + ) + return MambaStateDtypeCalculator.mamba2_cached_state_dtype( self.model_config.dtype, self.cache_config.mamba_cache_dtype, self.cache_config.mamba_ssm_cache_dtype, + use_replayssm=self.use_cache_kernel, ) - def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]: - return MambaStateShapeCalculator.mamba2_state_shape( + def get_state_shape(self) -> tuple[tuple[int, ...], ...]: + if self.use_cache_spec_kernel: + return MambaStateShapeCalculator.mamba2_spec_cached_state_shape( + intermediate_size=self.intermediate_size, + tp_world_size=get_tensor_model_parallel_world_size(), + n_groups=self.n_groups, + num_heads=self.num_heads, + head_dim=self.head_dim, + state_size=self.ssm_state_size, + conv_kernel=self.conv_kernel_size, + num_spec=self.num_spec, + use_replayssm_spec=self.use_cache_spec_kernel, + replayssm_buffer_len=self.max_cache_len, + ) + return MambaStateShapeCalculator.mamba2_cached_state_shape( intermediate_size=self.intermediate_size, tp_world_size=get_tensor_model_parallel_world_size(), n_groups=self.n_groups, @@ -1064,6 +1266,8 @@ def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]: state_size=self.ssm_state_size, conv_kernel=self.conv_kernel_size, num_spec=self.num_spec, + use_replayssm=self.use_cache_kernel, + replayssm_buffer_len=self.max_cache_len, ) @property diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 9e78b8222803..270d7c1e2898 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -80,6 +80,54 @@ def mamba2_state_dtype( model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype ) + @classmethod + def mamba2_cached_state_dtype( + cls, + model_dtype: ModelDType | torch.dtype, + mamba_cache_dtype: MambaDType, + mamba_ssm_cache_dtype: MambaDType, + use_replayssm: bool, + ) -> tuple[torch.dtype, ...]: + """Mamba2 state dtypes, extended for the state-and-output decode kernel. + + Returns the baseline ``(conv, ssm)`` dtypes when + ``use_replayssm`` is ``False``; otherwise appends the + state-and-output ring-buffer dtypes ``(x_cache, dt_cache, B_cache)`` = + ``(activation, fp32, activation)``. Must stay in sync with + ``MambaMixer2.get_state_dtype``. + """ + conv_dtype, ssm_dtype = cls.mamba2_state_dtype( + model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype + ) + if not use_replayssm: + return conv_dtype, ssm_dtype + activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype) + return conv_dtype, ssm_dtype, activation_dtype, torch.float32, activation_dtype + + @classmethod + def mamba2_spec_cached_state_dtype( + cls, + model_dtype: ModelDType | torch.dtype, + mamba_cache_dtype: MambaDType, + mamba_ssm_cache_dtype: MambaDType, + use_replayssm_spec: bool, + ) -> tuple[torch.dtype, ...]: + """Mamba2 state dtypes for the cached SPECULATIVE-decode (hybrid) kernel. + + Baseline ``(conv, ssm)`` when off; otherwise the hybrid 4-tuple + ``(conv, ssm_checkpoint, post_conv_cache, dt_cache)``. The checkpoint + and ``dt_cache`` are forced fp32 (the cached-spec reconstruction was + validated against an fp32 reference); ``post_conv_cache`` is activation + dtype. Must stay in sync with ``MambaMixer2.get_state_dtype``. + """ + conv_dtype, _ssm_dtype = cls.mamba2_state_dtype( + model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype + ) + if not use_replayssm_spec: + return conv_dtype, _ssm_dtype + activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype) + return conv_dtype, torch.float32, activation_dtype, torch.float32 + @classmethod def _mamba_state_dtype( cls, @@ -115,6 +163,56 @@ def gated_delta_net_state_dtype( model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype ) + @classmethod + def gated_delta_net_cached_state_dtype( + cls, + model_dtype: ModelDType | torch.dtype, + mamba_cache_dtype: MambaDType, + mamba_ssm_cache_dtype: MambaDType, + use_replayssm: bool, + ) -> tuple[torch.dtype, ...]: + """GDN state dtypes, extended for the cached decode kernel. + + Returns the baseline ``(conv, ssm)`` dtypes when + ``use_replayssm`` is ``False``; otherwise appends the ring + cache dtypes ``(d_cache, k_cache, g_cache)`` = + ``(activation, activation, float32)``. + """ + conv_dtype, ssm_dtype = cls._mamba_state_dtype( + model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype + ) + if not use_replayssm: + return conv_dtype, ssm_dtype + activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype) + return conv_dtype, ssm_dtype, activation_dtype, activation_dtype, torch.float32 + + @classmethod + def gated_delta_net_spec_cached_state_dtype( + cls, + model_dtype: ModelDType | torch.dtype, + mamba_cache_dtype: MambaDType, + mamba_ssm_cache_dtype: MambaDType, + use_replayssm_spec: bool, + ) -> tuple[torch.dtype, ...]: + """GDN state dtypes for the cached SPECULATIVE-decode kernel. + + Same ``d/k/g`` ring page as the non-spec cached path, but the ``ssm`` + checkpoint is forced to ``float32`` Returns the baseline ``(conv, ssm)`` when the flag is off. + """ + conv_dtype, ssm_dtype = cls._mamba_state_dtype( + model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype + ) + if not use_replayssm_spec: + return conv_dtype, ssm_dtype + activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype) + return ( + conv_dtype, + torch.float32, # fp32 checkpoint + activation_dtype, # d_cache + activation_dtype, # k_cache + torch.float32, # g_cache + ) + @classmethod def kda_state_dtype( cls, @@ -186,6 +284,110 @@ def mamba2_state_shape( temporal_state_shape = (divide(num_heads, tp_world_size), head_dim, state_size) return conv_state_shape, temporal_state_shape + @classmethod + def mamba2_cached_state_shape( + cls, + tp_world_size: int, + intermediate_size: int, + n_groups: int, + num_heads: int, + head_dim: int, + state_size: int, + conv_kernel: int, + use_replayssm: bool, + replayssm_buffer_len: int, + num_spec: int = 0, + ) -> tuple[tuple[int, ...], ...]: + """Mamba2 state shapes, extended for the state-and-output decode kernel. + + Returns the baseline ``(conv, ssm)`` shapes when + ``use_replayssm`` is ``False``; otherwise appends the + state-and-output ring-buffer shapes ``x_cache``/``dt_cache``/``B_cache``. + Group/head counts use the (un-extended) ``n_groups``/``num_heads`` + divided by ``tp_world_size``, matching ``MambaMixer2.get_state_shape``. + """ + conv_state_shape, temporal_state_shape = cls.mamba2_state_shape( + tp_world_size=tp_world_size, + intermediate_size=intermediate_size, + n_groups=n_groups, + num_heads=num_heads, + head_dim=head_dim, + state_size=state_size, + conv_kernel=conv_kernel, + num_spec=num_spec, + ) + if not use_replayssm: + return conv_state_shape, temporal_state_shape + + local_nheads = divide(num_heads, tp_world_size) + local_ngroups = divide(n_groups, tp_world_size) + x_cache_shape = (local_nheads, replayssm_buffer_len, head_dim) + dt_cache_shape = (local_nheads, replayssm_buffer_len) + B_cache_shape = (local_ngroups, replayssm_buffer_len, state_size) + return ( + conv_state_shape, + temporal_state_shape, + x_cache_shape, + dt_cache_shape, + B_cache_shape, + ) + + @classmethod + def mamba2_spec_cached_state_shape( + cls, + tp_world_size: int, + intermediate_size: int, + n_groups: int, + num_heads: int, + head_dim: int, + state_size: int, + conv_kernel: int, + use_replayssm_spec: bool, + replayssm_buffer_len: int, + num_spec: int = 0, + ) -> tuple[tuple[int, ...], ...]: + """Mamba2 state shapes for the cached SPECULATIVE-decode (hybrid) kernel. + + Baseline ``(conv, ssm)`` when off (conv keeps its spec sliding-window + size ``conv_kernel-1+num_spec`` -- the hybrid reuses + ``causal_conv1d_update``); otherwise appends the circular caches + ``post_conv_cache=(cache_buf_len, conv_dim_local)`` and + ``dt_cache=(local_nheads, cache_buf_len)``, where the L = B + max_spec_len + history window sizes ``cache_buf_len = next_pow2(replayssm_buffer_len + 1 + + num_spec)`` and ``conv_dim_local`` matches the post-conv x|B width (C is + not cached; read fresh from conv_out). Must stay in sync with + ``MambaMixer2.get_state_shape``. + """ + conv_state_shape, temporal_state_shape = cls.mamba2_state_shape( + tp_world_size=tp_world_size, + intermediate_size=intermediate_size, + n_groups=n_groups, + num_heads=num_heads, + head_dim=head_dim, + state_size=state_size, + conv_kernel=conv_kernel, + num_spec=num_spec, + ) + if not use_replayssm_spec: + return conv_state_shape, temporal_state_shape + n_groups_ext = n_groups + cls.extra_groups_for_head_shards( + n_groups, tp_world_size + ) + conv_dim_local = divide( + intermediate_size + n_groups_ext * state_size, tp_world_size + ) + # L = B + max_spec_len history window: physical pow2 buffer next_pow2(L). + cache_buf_len = 1 << (replayssm_buffer_len + num_spec).bit_length() + local_nheads = divide(num_heads, tp_world_size) + post_conv_cache_shape = (cache_buf_len, conv_dim_local) + dt_cache_shape = (local_nheads, cache_buf_len) + return ( + conv_state_shape, + temporal_state_shape, + post_conv_cache_shape, + dt_cache_shape, + ) + @classmethod def short_conv_state_shape( cls, @@ -233,6 +435,98 @@ def gated_delta_net_state_shape( ) return conv_state_shape, temporal_state_shape + @classmethod + def gated_delta_net_cached_state_shape( + cls, + tp_world_size: int, + num_k_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + conv_kernel_size: int, + use_replayssm: bool, + replayssm_buffer_len: int, + num_spec: int = 0, + ) -> tuple[tuple[int, ...], ...]: + """GDN state shapes, extended for the cached decode kernel. + + Returns the baseline ``(conv, ssm)`` shapes when + ``use_replayssm`` is ``False``; otherwise appends the cached + ring-buffer shapes ``d_cache``/``k_cache``/``g_cache``. Head counts use + the (un-extended) ``num_v_heads``/``num_k_heads`` divided by + ``tp_world_size``, matching ``gated_delta_net_state_shape``. + """ + conv_state_shape, temporal_state_shape = cls.gated_delta_net_state_shape( + tp_world_size, + num_k_heads, + num_v_heads, + head_k_dim, + head_v_dim, + conv_kernel_size, + num_spec, + ) + if not use_replayssm: + return conv_state_shape, temporal_state_shape + + local_v_heads = divide(num_v_heads, tp_world_size) + local_k_heads = divide(num_k_heads, tp_world_size) + d_cache_shape = (local_v_heads, replayssm_buffer_len, head_v_dim) + k_cache_shape = (local_k_heads, replayssm_buffer_len, head_k_dim) + g_cache_shape = (local_v_heads, replayssm_buffer_len) + return ( + conv_state_shape, + temporal_state_shape, + d_cache_shape, + k_cache_shape, + g_cache_shape, + ) + + @classmethod + def gated_delta_net_spec_cached_state_shape( + cls, + tp_world_size: int, + num_k_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + conv_kernel_size: int, + use_replayssm_spec: bool, + replayssm_buffer_len: int, + num_spec: int = 0, + ) -> tuple[tuple[int, ...], ...]: + """GDN state shapes for the cached SPECULATIVE-decode kernel. + + The circular ``d_cache``/``k_cache``/``g_cache`` use the L = B + max_spec_len + history window: a power-of-two buffer ``next_pow2(replayssm_buffer_len + 1 + + num_spec)``. Returns the baseline ``(conv, ssm)`` shapes when the flag is + off. The block-keyed cursors live in the GDN metadata builder, not the page. + """ + conv_state_shape, temporal_state_shape = cls.gated_delta_net_state_shape( + tp_world_size, + num_k_heads, + num_v_heads, + head_k_dim, + head_v_dim, + conv_kernel_size, + num_spec, + ) + if not use_replayssm_spec: + return conv_state_shape, temporal_state_shape + + cache_buf_len = 1 << (replayssm_buffer_len + num_spec).bit_length() + local_v_heads = divide(num_v_heads, tp_world_size) + local_k_heads = divide(num_k_heads, tp_world_size) + d_cache_shape = (local_v_heads, cache_buf_len, head_v_dim) + k_cache_shape = (local_k_heads, cache_buf_len, head_k_dim) + g_cache_shape = (local_v_heads, cache_buf_len) + return ( + conv_state_shape, + temporal_state_shape, + d_cache_shape, + k_cache_shape, + g_cache_shape, + ) + @classmethod def kda_state_shape( cls, diff --git a/vllm/model_executor/layers/mamba/ops/replayssm_config.py b/vllm/model_executor/layers/mamba/ops/replayssm_config.py new file mode 100644 index 000000000000..6d37106ed66b --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/replayssm_config.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Launch-config selection for the ReplaySSM decode kernels. + +Mirrors ``mamba_ssm.py``: a hard-coded heuristic per kernel, plus an +``override`` context manager for benchmarks/tests/config sweeps. Hardware is +auto-detected (Blackwell vs not) so call sites need not thread it through. + +Config tuples are tuned from the ablation bestcfg sweeps at batch >= 256 in the +deployed spec regime (spec window T >= 4); use ``override_replayssm_config`` to +pin a different config. +""" + +import functools +from contextlib import contextmanager + +from vllm.platforms import current_platform +from vllm.triton_utils import triton + + +@functools.cache +def _is_blackwell() -> bool: + try: + return current_platform.is_device_capability_family(100) + except Exception: + return False + + +# Per-kernel overrides keyed by the kernel name passed to get_replayssm_config. +_overrides: dict[str, tuple] = {} + + +@contextmanager +def override_replayssm_config(kernel: str, config: tuple): + """Pin ``kernel``'s launch config for the duration of the context.""" + prev = _overrides.get(kernel) + _overrides[kernel] = config + try: + yield + finally: + if prev is None: + _overrides.pop(kernel, None) + else: + _overrides[kernel] = prev + + +def _dstate_tile(dstate: int, tile: int) -> int: + return max(16, min(tile, triton.next_power_of_2(dstate))) + + +def _mamba2_spec_verify(dstate, base_block, max_spec_len, is_blackwell): + # (block_size_m, num_warps, dstate_tile, num_stages) + bsm = 64 if (is_blackwell or base_block <= 16) else 32 + return bsm, 2, _dstate_tile(dstate, 64), 2 + + +def _mamba2_spec_flush(dstate, base_block, max_spec_len, is_blackwell): + if base_block <= 16: + return 32, 2, _dstate_tile(dstate, 64), 2 + if is_blackwell: + return 64, 2, _dstate_tile(dstate, 128), 2 + return 32, 1, _dstate_tile(dstate, 128), 2 + + +def _gdn_spec(max_spec_len, is_blackwell): + # (block_v, num_warps, nk, num_stages); verify and flush share a config. + return 64, 1, (4 if max_spec_len >= 6 else 2), 2 + + +def _l_bucket(cache_len: int) -> int: + """Map an arbitrary buffer length to the nearest tuned bucket: + L<=8 -> 8, 8 16, L>16 -> 32. Default callers pass L=16.""" + if cache_len <= 8: + return 8 + if cache_len <= 16: + return 16 + return 32 + + +# Standard-decode launch configs keyed by L bucket, tuned for the FP32 production +# state (mamba_ssm_dtype=float32). bsm=64 is intentionally avoided here: it is the +# bf16/B300 optimum but catastrophic (register spill, ~0.3x) at fp32. These seed +# values preserve the previously shipped per-kernel constants; the config sweep +# (profiling_pr/config_sweep) refines them per bucket. +_OUTPUT_ONLY_BY_L = {8: (16, 1), 16: (16, 1), 32: (16, 1)} # (block_size_m, num_warps) +_STATE_AND_OUTPUT_BY_L = {8: (32, 1), 16: (32, 1), 32: (32, 1)} +# GDN standard decode: (block_v, num_warps, num_stages, nk). L-flat in the sweep. +_GDN_DECODE_BY_L = {8: (64, 1, 3, 2), 16: (64, 1, 3, 2), 32: (64, 1, 3, 2)} + + +def get_replayssm_config(kernel: str, **shape) -> tuple: + """Return the launch config for ``kernel`` (override > tuned default). + + kernel: one of "mamba2_spec_verify", "mamba2_spec_flush", + "mamba2_output_only", "mamba2_state_and_output", "gdn_decode", + "gdn_spec_verify", "gdn_spec_flush". ``shape`` carries the keying dims + (dstate / base_block / max_spec_len for spec; L for standard decode, + default 16), hardware is auto-detected. + """ + if kernel in _overrides: + return _overrides[kernel] + bw = _is_blackwell() + if kernel == "mamba2_spec_verify": + return _mamba2_spec_verify( + shape["dstate"], shape["base_block"], shape["max_spec_len"], bw + ) + if kernel == "mamba2_spec_flush": + return _mamba2_spec_flush( + shape["dstate"], shape["base_block"], shape["max_spec_len"], bw + ) + if kernel == "mamba2_output_only": + return _OUTPUT_ONLY_BY_L[_l_bucket(shape.get("L", 16))] + if kernel == "mamba2_state_and_output": + return _STATE_AND_OUTPUT_BY_L[_l_bucket(shape.get("L", 16))] + if kernel == "gdn_decode": + return _GDN_DECODE_BY_L[_l_bucket(shape.get("L", 16))] + if kernel in ("gdn_spec_verify", "gdn_spec_flush"): + return _gdn_spec(shape["max_spec_len"], bw) + raise ValueError(f"unknown ReplaySSM kernel config key: {kernel}") diff --git a/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py b/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py new file mode 100644 index 000000000000..14eacabe6719 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py @@ -0,0 +1,550 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# ruff: noqa: E501 + +import torch + +from vllm.model_executor.layers.mamba.ops.mamba_ssm import softplus +from vllm.model_executor.layers.mamba.ops.replayssm_config import get_replayssm_config +from vllm.triton_utils import tl, triton +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + + +@triton.heuristics( + { + "HAS_STATE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] + is not None + } +) +@triton.heuristics( + {"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])} +) +@triton.jit +def _replayssm_output_only_precompute_kernel( + B_ptr, + C_ptr, + B_cache_ptr, + write_pos_ptr, + is_flush_ptr, + bc_pre_ptr, + state_batch_indices_ptr, + null_block_id, + # Matrix dimensions + batch, + ngroups, + dstate, + # Input strides + stride_B_batch, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_group, + stride_C_dstate, + # Cache strides + stride_B_cache_batch, + stride_B_cache_group, + stride_B_cache_pos, + stride_B_cache_dstate, + stride_bc_pre_batch, + stride_bc_pre_group, + stride_bc_pre_pos, + stride_state_indices_batch, + stride_state_indices_T, + # Meta-parameters + MAX_CACHE_LEN: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + # heuristic-computed + BLOCK_SIZE_DSTATE: tl.constexpr, + HAS_STATE_BATCH_INDICES: tl.constexpr, +): + pid_b = tl.program_id(axis=0) + pid_g = tl.program_id(axis=1) + + # On flush steps the main kernel does not read bc_pre, so skip the work. + is_flush = tl.load(is_flush_ptr + pid_b) != 0 + if is_flush: + return + + if HAS_STATE_BATCH_INDICES: + state_batch_idx = tl.load( + state_batch_indices_ptr + + pid_b * stride_state_indices_batch + + 0 * stride_state_indices_T + ).to(tl.int64) + if state_batch_idx == null_block_id: + return + else: + state_batch_idx = pid_b + + offs_k = tl.arange(0, BLOCK_SIZE_K) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + + write_pos = tl.load(write_pos_ptr + pid_b).to(tl.int64) + + B_ptr += pid_b * stride_B_batch + pid_g * stride_B_group + C_ptr += pid_b * stride_C_batch + pid_g * stride_C_group + B_cache_ptr += ( + state_batch_idx * stride_B_cache_batch + pid_g * stride_B_cache_group + ) + bc_pre_ptr += pid_b * stride_bc_pre_batch + pid_g * stride_bc_pre_group + + B_cur = tl.load( + B_ptr + offs_n * stride_B_dstate, + mask=offs_n < dstate, + other=0.0, + ) + C = tl.load( + C_ptr + offs_n * stride_C_dstate, + mask=offs_n < dstate, + other=0.0, + ) + B_cache_ptrs = ( + B_cache_ptr + + offs_k[:, None] * stride_B_cache_pos + + offs_n[None, :] * stride_B_cache_dstate + ) + B_cache = tl.load( + B_cache_ptrs, + mask=(offs_k[:, None] < write_pos) & (offs_n[None, :] < dstate), + other=0.0, + ) + B_all = tl.where(offs_k[:, None] == write_pos, B_cur[None, :], B_cache) + bc = tl.sum(B_all.to(tl.float32) * C[None, :].to(tl.float32), axis=1) + + tl.store( + bc_pre_ptr + offs_k * stride_bc_pre_pos, + bc, + mask=(offs_k <= write_pos) & (offs_k < MAX_CACHE_LEN), + ) + + +@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) +@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) +@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) +@triton.heuristics( + { + "HAS_STATE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] + is not None + } +) +@triton.heuristics( + {"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])} +) +@triton.jit +def _replayssm_output_only_kernel( + # Pointers to matrices + state_ptr, + x_ptr, + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + x_cache_ptr, + dt_cache_ptr, + B_cache_ptr, + bc_pre_ptr, + write_pos_ptr, + is_flush_ptr, + state_batch_indices_ptr, + null_block_id, + # Matrix dimensions + batch, + nheads, + dim, + dstate, + nheads_ngroups_ratio, + # State strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # Input strides + stride_x_batch, + stride_x_head, + stride_x_dim, + stride_dt_batch, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + stride_B_batch, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_group, + stride_C_dstate, + stride_D_head, + stride_D_dim, + stride_z_batch, + stride_z_head, + stride_z_dim, + stride_out_batch, + stride_out_head, + stride_out_dim, + # Cache strides + stride_x_cache_batch, + stride_x_cache_head, + stride_x_cache_dim, + stride_x_cache_pos, + stride_dt_cache_batch, + stride_dt_cache_head, + stride_dt_cache_pos, + stride_B_cache_batch, + stride_B_cache_group, + stride_B_cache_pos, + stride_B_cache_dstate, + stride_bc_pre_batch, + stride_bc_pre_group, + stride_bc_pre_pos, + stride_state_indices_batch, + stride_state_indices_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + MAX_CACHE_LEN: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_K_CACHE: tl.constexpr, + BLOCK_SIZE_K_DOT: tl.constexpr, + # heuristic-computed + BLOCK_SIZE_DSTATE: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_STATE_BATCH_INDICES: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + pid_b = tl.program_id(axis=1) + pid_h = tl.program_id(axis=2) + + # Resolve the physical state slot for this decode row; skip padded rows. + if HAS_STATE_BATCH_INDICES: + state_batch_idx = tl.load(state_batch_indices_ptr + pid_b * stride_state_indices_batch + 0 * stride_state_indices_T).to(tl.int64) + if state_batch_idx == null_block_id: + return + else: + state_batch_idx = pid_b + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + + # Buffer cursor (number of cached tokens so far) and the flush flag. + write_pos = tl.load(write_pos_ptr + pid_b).to(tl.int64) + is_flush = tl.load(is_flush_ptr + pid_b) != 0 + + # Advance every pointer to this (row, head, group). + state_ptr += state_batch_idx * stride_state_batch + pid_h * stride_state_head + x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head + dt_ptr += pid_b * stride_dt_batch + pid_h * stride_dt_head + B_ptr += pid_b * stride_B_batch + (pid_h // nheads_ngroups_ratio) * stride_B_group + C_ptr += pid_b * stride_C_batch + (pid_h // nheads_ngroups_ratio) * stride_C_group + out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head + x_cache_ptr += state_batch_idx * stride_x_cache_batch + pid_h * stride_x_cache_head + dt_cache_ptr += state_batch_idx * stride_dt_cache_batch + pid_h * stride_dt_cache_head + B_cache_ptr += state_batch_idx * stride_B_cache_batch + (pid_h // nheads_ngroups_ratio) * stride_B_cache_group + bc_pre_ptr += pid_b * stride_bc_pre_batch + (pid_h // nheads_ngroups_ratio) * stride_bc_pre_group + + # Current-token dt (+ bias, softplus), scalar A, current x / C, checkpoint + # state S_0, and current-token B (shared by both routes below). + dt_cur = tl.load(dt_ptr).to(tl.float32) + if HAS_DT_BIAS: + dt_cur += tl.load(dt_bias_ptr + pid_h * stride_dt_bias_head).to(tl.float32) + if DT_SOFTPLUS: + dt_cur = tl.where(dt_cur <= 20.0, softplus(dt_cur), dt_cur) + A = tl.load(A_ptr + pid_h * stride_A_head).to(tl.float32) + x_cur = tl.load(x_ptr + offs_m * stride_x_dim, mask=offs_m < dim, other=0.0) + C = tl.load(C_ptr + offs_n * stride_C_dstate, mask=offs_n < dstate, other=0.0).to(tl.float32) + state_ptrs = state_ptr + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + state = tl.load(state_ptrs, mask=(offs_m[:, None] < dim) & (offs_n[None, :] < dstate), other=0.0) + B_cur = tl.load(B_ptr + offs_n * stride_B_dstate, mask=offs_n < dstate, other=0.0) + + if not is_flush: + # Output-only route: read y without materializing the state, using the + # precomputed k^T q products (`bc`): + # y = total_decay * (S_0 q) + sum_j s_j (k_j^T q) v_j. + # Then append the current token to the buffer. + offs_k_cache = tl.arange(0, BLOCK_SIZE_K_CACHE) + # dt over the window (history + current token), then the decay weights. + dt_all_cache = tl.load(dt_cache_ptr + offs_k_cache * stride_dt_cache_pos, mask=offs_k_cache < write_pos, other=0.0).to(tl.float32) + dt_all_cache = tl.where(offs_k_cache == write_pos, dt_cur, dt_all_cache) + dA_cumsum_cache = A * tl.cumsum(dt_all_cache, axis=0) + dA_total_cache = A * tl.sum(dt_all_cache, axis=0) + total_decay_cache = tl.exp(dA_total_cache) + scale_cache = dt_all_cache * tl.exp(dA_total_cache - dA_cumsum_cache) + scale_cache = tl.where(offs_k_cache <= write_pos, scale_cache, 0.0) + + # Gather buffered x over the window (history + current token). + x_all_cache_ptrs = x_cache_ptr + offs_m[:, None] * stride_x_cache_dim + offs_k_cache[None, :] * stride_x_cache_pos + x_all_cache = tl.load(x_all_cache_ptrs, mask=(offs_m[:, None] < dim) & (offs_k_cache[None, :] < write_pos), other=0.0) + x_all_cache = tl.where(offs_k_cache[None, :] == write_pos, x_cur[:, None], x_all_cache) + + # Decayed checkpoint readout plus the weighted sum of cached values. + checkpoint_out = tl.sum(state.to(tl.float32) * C[None, :], axis=1) * total_decay_cache + bc_cache = tl.load(bc_pre_ptr + offs_k_cache * stride_bc_pre_pos, mask=offs_k_cache <= write_pos, other=0.0) + cache_out = tl.sum(x_all_cache.to(tl.float32) * (scale_cache * bc_cache)[None, :], axis=1) + out = checkpoint_out + cache_out + + # Append the current token (x, dt, B) into the buffer at write_pos. + tl.store(x_cache_ptr + offs_m * stride_x_cache_dim + write_pos * stride_x_cache_pos, x_cur, mask=offs_m < dim) + if pid_m == 0: + tl.store(dt_cache_ptr + write_pos * stride_dt_cache_pos, dt_cur) + tl.store(B_cache_ptr + write_pos * stride_B_cache_pos + offs_n * stride_B_cache_dstate, B_cur, mask=offs_n < dstate) + else: + # Flush step: state route. Reconstruct the state from cached inputs, + # S_t = total_decay * S_0 + sum_j s_j (v_j k_j^T), persist it as the new + # checkpoint, then read y = S_t q. + offs_k_dot = tl.arange(0, BLOCK_SIZE_K_DOT) + dt_all_dot = tl.load(dt_cache_ptr + offs_k_dot * stride_dt_cache_pos, mask=offs_k_dot < write_pos, other=0.0).to(tl.float32) + dt_all_dot = tl.where(offs_k_dot == write_pos, dt_cur, dt_all_dot) + dA_cumsum_dot = A * tl.cumsum(dt_all_dot, axis=0) + dA_total_dot = A * tl.sum(dt_all_dot, axis=0) + total_decay_dot = tl.exp(dA_total_dot) + scale_dot = dt_all_dot * tl.exp(dA_total_dot - dA_cumsum_dot) + scale_dot = tl.where(offs_k_dot <= write_pos, scale_dot, 0.0) + + # Gather buffered x and B over the window (history + current token). + x_all_dot_ptrs = x_cache_ptr + offs_m[:, None] * stride_x_cache_dim + offs_k_dot[None, :] * stride_x_cache_pos + x_all_dot = tl.load(x_all_dot_ptrs, mask=(offs_m[:, None] < dim) & (offs_k_dot[None, :] < write_pos), other=0.0) + x_all_dot = tl.where(offs_k_dot[None, :] == write_pos, x_cur[:, None], x_all_dot) + B_all_dot_ptrs = B_cache_ptr + offs_k_dot[:, None] * stride_B_cache_pos + offs_n[None, :] * stride_B_cache_dstate + B_all_dot = tl.load(B_all_dot_ptrs, mask=(offs_k_dot[:, None] < write_pos) & (offs_n[None, :] < dstate), other=0.0) + B_all_dot = tl.where(offs_k_dot[:, None] == write_pos, B_cur[None, :], B_all_dot) + + # Reconstruct the state from cached inputs and store it as the checkpoint. + B_scaled = (B_all_dot.to(tl.float32) * scale_dot[:, None]).to(x_ptr.dtype.element_ty) + # tf32x3 keeps fp32 parity with the elementwise baseline (plain tf32 on + # fp32 inputs drifts ~1e-2); bf16/fp16 inputs are unaffected by this flag. + delta_state = tl.dot( + x_all_dot.to(x_ptr.dtype.element_ty), B_scaled, input_precision="tf32x3" + ) + state_new = state.to(tl.float32) * total_decay_dot + delta_state.to(tl.float32) + tl.store(state_ptrs, state_new.to(state.dtype), mask=(offs_m[:, None] < dim) & (offs_n[None, :] < dstate)) + out = tl.sum(state_new * C[None, :], axis=1) + + # Skip connection (D) and output gate (z). + if HAS_D: + D_ptr += pid_h * stride_D_head + D = tl.load(D_ptr + offs_m * stride_D_dim, mask=offs_m < dim, other=0.0).to(tl.float32) + out += x_cur.to(tl.float32) * D + if HAS_Z: + z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head + z = tl.load(z_ptr + offs_m * stride_z_dim, mask=offs_m < dim, other=0.0).to(tl.float32) + out *= z * tl.sigmoid(z) + + tl.store(out_ptr + offs_m * stride_out_dim, out, mask=offs_m < dim) + + +def selective_state_update_replayssm_output_only( + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + z: torch.Tensor | None = None, + dt_softplus: bool = False, + x_cache: torch.Tensor | None = None, + dt_cache: torch.Tensor | None = None, + B_cache: torch.Tensor | None = None, + bc_pre: torch.Tensor | None = None, + write_pos: torch.Tensor | None = None, + is_flush: torch.Tensor | None = None, + max_cache_len: int = 16, + state_batch_indices: torch.Tensor | None = None, + null_block_id: int = NULL_BLOCK_ID, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Cached-bc SSM update for vLLM's autoregressive Mamba2 decode path.""" + has_heads = state.dim() > 3 + if state.dim() == 3: + state = state.unsqueeze(1) + if x.dim() == 2: + x = x.unsqueeze(1) + if dt.dim() == 2: + dt = dt.unsqueeze(1) + if A.dim() == 2: + A = A.unsqueeze(0) + if B.dim() == 2: + B = B.unsqueeze(1) + if C.dim() == 2: + C = C.unsqueeze(1) + if D is not None and D.dim() == 1: + D = D.unsqueeze(0) + if z is not None and z.dim() == 2: + z = z.unsqueeze(1) + if dt_bias is not None and dt_bias.dim() == 1: + dt_bias = dt_bias.unsqueeze(0) + if out is not None and out.dim() == 2: + out = out.unsqueeze(1) + if state_batch_indices is not None and state_batch_indices.dim() == 1: + state_batch_indices = state_batch_indices.unsqueeze(1) + + _, nheads, dim, dstate = state.shape + batch = x.shape[0] + assert x.shape == (batch, nheads, dim) + assert dt.shape == x.shape + assert A.shape == (nheads, dim, dstate) + ngroups = B.shape[1] + assert nheads % ngroups == 0, "nheads must be divisible by ngroups" + assert B.shape == (batch, ngroups, dstate) + assert C.shape == B.shape + if D is not None: + assert D.shape == (nheads, dim) + if z is not None: + assert z.shape == x.shape + if dt_bias is not None: + assert dt_bias.shape == (nheads, dim) + assert out is not None and out.shape == x.shape + + assert A.stride(-1) == 0 and A.stride(-2) == 0, ( + "Cached kernel requires TIE_HDIM (A scalar per head)" + ) + assert dt.stride(-1) == 0, "Cached kernel requires TIE_HDIM (dt scalar per head)" + if dt_bias is not None: + assert dt_bias.stride(-1) == 0, ( + "Cached kernel requires TIE_HDIM (dt_bias scalar per head)" + ) + + assert x_cache is not None + assert dt_cache is not None + assert B_cache is not None + assert x_cache.shape[1:] == (nheads, max_cache_len, dim) + assert dt_cache.shape[1:] == (nheads, max_cache_len) + assert B_cache.shape[1:] == (ngroups, max_cache_len, dstate) + assert write_pos is not None and write_pos.shape[0] >= batch + assert write_pos.dtype == torch.int32 + assert is_flush is not None and is_flush.shape[0] >= batch + assert is_flush.dtype in (torch.bool, torch.int8) + assert bc_pre is not None + assert bc_pre.shape[0] >= batch and bc_pre.shape[1] >= ngroups + assert bc_pre.shape[2] == max_cache_len + assert bc_pre.dtype == torch.float32 + if state_batch_indices is not None: + assert state_batch_indices.shape[0] >= batch + assert state_batch_indices.shape[1] >= 1 + + block_size_k_cache = max(1, triton.next_power_of_2(max_cache_len)) + block_size_k_dot = max(16, block_size_k_cache) + block_size_m, num_warps = get_replayssm_config( + "mamba2_output_only", dstate=dstate, L=max_cache_len + ) + + grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads) + z_strides = (z.stride(0), z.stride(1), z.stride(2)) if z is not None else (0, 0, 0) + state_indices_strides = ( + (state_batch_indices.stride(0), state_batch_indices.stride(1)) + if state_batch_indices is not None + else (0, 0) + ) + + with torch.accelerator.device_index(x.device.index): + _replayssm_output_only_precompute_kernel[(batch, ngroups)]( + B, + C, + B_cache, + write_pos, + is_flush, + bc_pre, + state_batch_indices, + null_block_id, + batch, + ngroups, + dstate, + B.stride(0), + B.stride(1), + B.stride(2), + C.stride(0), + C.stride(1), + C.stride(2), + B_cache.stride(0), + B_cache.stride(1), + B_cache.stride(2), + B_cache.stride(3), + bc_pre.stride(0), + bc_pre.stride(1), + bc_pre.stride(2), + state_indices_strides[0], + state_indices_strides[1], + max_cache_len, + block_size_k_cache, + num_warps=2, + ) + _replayssm_output_only_kernel[grid]( + state, + x, + dt, + dt_bias, + A, + B, + C, + D, + z, + out, + x_cache, + dt_cache, + B_cache, + bc_pre, + write_pos, + is_flush, + state_batch_indices, + null_block_id, + batch, + nheads, + dim, + dstate, + nheads // ngroups, + state.stride(0), + state.stride(1), + state.stride(2), + state.stride(3), + x.stride(0), + x.stride(1), + x.stride(2), + dt.stride(0), + dt.stride(1), + dt_bias.stride(0) if dt_bias is not None else 0, + A.stride(0), + B.stride(0), + B.stride(1), + B.stride(2), + C.stride(0), + C.stride(1), + C.stride(2), + D.stride(0) if D is not None else 0, + D.stride(1) if D is not None else 0, + z_strides[0], + z_strides[1], + z_strides[2], + out.stride(0), + out.stride(1), + out.stride(2), + x_cache.stride(0), + x_cache.stride(1), + x_cache.stride(3), + x_cache.stride(2), + dt_cache.stride(0), + dt_cache.stride(1), + dt_cache.stride(2), + B_cache.stride(0), + B_cache.stride(1), + B_cache.stride(2), + B_cache.stride(3), + bc_pre.stride(0), + bc_pre.stride(1), + bc_pre.stride(2), + state_indices_strides[0], + state_indices_strides[1], + dt_softplus, + max_cache_len, + block_size_m, + block_size_k_cache, + block_size_k_dot, + num_warps=num_warps, + ) + + if not has_heads: + out = out.squeeze(1) + return out diff --git a/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_spec.py b/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_spec.py new file mode 100644 index 000000000000..8369d3f8a762 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_spec.py @@ -0,0 +1,742 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# ruff: noqa: E501 + +import torch + +from vllm.model_executor.layers.mamba.ops.mamba_ssm import softplus +from vllm.model_executor.layers.mamba.ops.replayssm_config import get_replayssm_config +from vllm.triton_utils import tl, triton +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + + +# ====================================================================== +# Fused scatter + precompute (grid: (batch, ngroups)) +# +# Scatters all conv_dim channels (x|B|C, partitioned by group) + dt of the +# fresh spec tokens into the circular post-conv / dt caches at +# ``(origin + write_pos + s) % buf``, and computes ``bc[k, s] = B_full[k] . C[s]`` +# over the window (history B from the cache + fresh spec B, no read-back). +# ====================================================================== +@triton.heuristics( + {"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])} +) +@triton.jit +def _fused_scatter_precompute_kernel( + conv_out_ptr, # (total_tokens, conv_dim) packed channel-last conv output + dt_spec_ptr, # (total_tokens, nheads) packed raw dt + post_conv_cache_ptr, # (num_blocks, buf, conv_dim) circular paged + dt_cache_ptr, # (num_blocks, nheads, buf) circular paged + write_pos_ptr, # (num_state_slots,) block-keyed + post_origin_ptr, # (num_state_slots,) block-keyed + bc_pre_ptr, # (max_bs, ngroups, max_cache_len, block_spec) dense per-row scratch + state_batch_indices_ptr, # (batch,) physical block per dense decode row + query_start_loc_ptr, # (batch + 1,) packed token offsets + null_block_id, + batch, + ngroups, + nheads, + dstate, + d_inner, + conv_dim, + max_cache_len, + stride_conv_out_tok, + stride_conv_out_c, + stride_dt_spec_tok, + stride_dt_spec_h, + stride_post_conv_cache_b, + stride_post_conv_cache_pos, + stride_post_conv_cache_c, + stride_dt_cache_b, + stride_dt_cache_h, + stride_dt_cache_pos, + stride_bc_pre_batch, + stride_bc_pre_group, + stride_bc_pre_pos, + stride_bc_pre_spec, + stride_state_indices_batch, + RATIO: tl.constexpr, + RATIO_P: tl.constexpr, + NCX: tl.constexpr, + BLOCK_CX: tl.constexpr, + CACHE_BUF_LEN: tl.constexpr, + BLOCK_SIZE_CACHE: tl.constexpr, + BLOCK_SIZE_SPEC: tl.constexpr, + BLOCK_HL: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, +): + pid_b = tl.program_id(0) + pid_g = tl.program_id(1) + state_batch_idx = tl.load( + state_batch_indices_ptr + pid_b * stride_state_indices_batch + ).to(tl.int64) + if state_batch_idx == null_block_id: + return + bos = tl.load(query_start_loc_ptr + pid_b).to(tl.int64) + eos = tl.load(query_start_loc_ptr + pid_b + 1).to(tl.int64) + spec_len = (eos - bos).to(tl.int32) + write_pos = tl.load(write_pos_ptr + state_batch_idx).to(tl.int32) + post_origin = tl.load(post_origin_ptr + state_batch_idx).to(tl.int32) + + offs_s = tl.arange(0, BLOCK_SIZE_SPEC) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + spec_valid = offs_s < spec_len + nmask = offs_n < dstate + phys_spec = (post_origin + write_pos + offs_s) & (CACHE_BUF_LEN - 1) + + b_c0 = d_inner + pid_g * dstate + c_c0 = d_inner + ngroups * dstate + pid_g * dstate + + src_base = conv_out_ptr + bos * stride_conv_out_tok + # fresh spec B / C [S, N] + B_spec = tl.load( + src_base + (b_c0 + offs_n[None, :]) * stride_conv_out_c + offs_s[:, None] * stride_conv_out_tok, + mask=spec_valid[:, None] & nmask[None, :], + other=0.0, + ) + C_spec = tl.load( + src_base + (c_c0 + offs_n[None, :]) * stride_conv_out_c + offs_s[:, None] * stride_conv_out_tok, + mask=spec_valid[:, None] & nmask[None, :], + other=0.0, + ) + cache_base = post_conv_cache_ptr + state_batch_idx * stride_post_conv_cache_b + # scatter B (C is not cached; read fresh from conv_out) + tl.store( + cache_base + phys_spec[:, None] * stride_post_conv_cache_pos + (b_c0 + offs_n[None, :]) * stride_post_conv_cache_c, + B_spec, + mask=spec_valid[:, None] & nmask[None, :], + ) + + # scatter x channels owned by this group: [g*RATIO_P, (g+1)*RATIO_P) + gx0 = pid_g * RATIO_P + for i in tl.static_range(NCX): + offs_cx = i * BLOCK_CX + tl.arange(0, BLOCK_CX) + cxm = offs_cx < RATIO_P + gx = gx0 + offs_cx + xv = tl.load( + src_base + + gx[None, :] * stride_conv_out_c + + offs_s[:, None] * stride_conv_out_tok, + mask=spec_valid[:, None] & cxm[None, :], + other=0.0, + ) + tl.store( + cache_base + + phys_spec[:, None] * stride_post_conv_cache_pos + + gx[None, :] * stride_post_conv_cache_c, + xv, + mask=spec_valid[:, None] & cxm[None, :], + ) + + # scatter dt for this group's heads + offs_hl = tl.arange(0, BLOCK_HL) + hlm = offs_hl < RATIO + gh = pid_g * RATIO + offs_hl + dt_base = dt_spec_ptr + bos * stride_dt_spec_tok + dtv = tl.load( + dt_base + offs_s[:, None] * stride_dt_spec_tok + gh[None, :] * stride_dt_spec_h, + mask=spec_valid[:, None] & hlm[None, :], + other=0.0, + ) + dtc_base = dt_cache_ptr + state_batch_idx * stride_dt_cache_b + tl.store( + dtc_base + + gh[None, :] * stride_dt_cache_h + + phys_spec[:, None] * stride_dt_cache_pos, + dtv, + mask=spec_valid[:, None] & hlm[None, :], + ) + + # bc: history B from cache + fresh spec B (no read-back of spec B) + offs_k = tl.arange(0, BLOCK_SIZE_CACHE) + hist_mask = offs_k < write_pos + cache_valid = (offs_k < max_cache_len) & (offs_k < (write_pos + spec_len)) + spec_tok = (offs_k >= write_pos) & (offs_k < (write_pos + spec_len)) + spec_off = offs_k - write_pos + phys_k = (post_origin + offs_k) & (CACHE_BUF_LEN - 1) + B_hist = tl.load( + cache_base + phys_k[:, None] * stride_post_conv_cache_pos + (b_c0 + offs_n[None, :]) * stride_post_conv_cache_c, + mask=hist_mask[:, None] & nmask[None, :], + other=0.0, + ) + B_specrows = tl.load( + src_base + (b_c0 + offs_n[None, :]) * stride_conv_out_c + spec_off[:, None] * stride_conv_out_tok, + mask=spec_tok[:, None] & nmask[None, :], + other=0.0, + ) + B_full = tl.where(spec_tok[:, None], B_specrows, B_hist) + B_full = tl.where(cache_valid[:, None], B_full.to(tl.float32), 0.0).to( + conv_out_ptr.dtype.element_ty + ) + bc = tl.dot( + B_full, + tl.trans(C_spec.to(conv_out_ptr.dtype.element_ty)), + input_precision="tf32x3", + ).to(tl.float32) + bc_ptrs = ( + bc_pre_ptr + + pid_b * stride_bc_pre_batch + + pid_g * stride_bc_pre_group + + offs_k[:, None] * stride_bc_pre_pos + + offs_s[None, :] * stride_bc_pre_spec + ) + tl.store( + bc_ptrs, + bc.to(bc_pre_ptr.dtype.element_ty), + mask=cache_valid[:, None] & spec_valid[None, :], + ) + + +# ====================================================================== +# Verify launch (non-flush rows): per-draft output from the fixed checkpoint +# S_0, no state write. dstate-tiled. Flush rows early-exit. +# ====================================================================== +@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) +@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) +@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.jit +def _replayssm_spec_nf_kernel( + state_ptr, x_cache_ptr, dt_cache_ptr, B_cache_ptr, C_src_ptr, bc_pre_ptr, + D_ptr, z_ptr, dt_bias_ptr, A_ptr, out_ptr, is_flush_flags_ptr, write_pos_ptr, + post_origin_ptr, state_batch_indices_ptr, query_start_loc_ptr, null_block_id, + batch, nheads, dim, dstate, max_cache_len, nheads_ngroups_ratio, + stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, + stride_x_cache_batch, stride_x_cache_head, stride_x_cache_dim, stride_x_cache_pos, + stride_dt_cache_batch, stride_dt_cache_head, stride_dt_cache_pos, + stride_B_cache_batch, stride_B_cache_group, stride_B_cache_dstate, stride_B_cache_pos, + stride_C_src_tok, stride_C_src_c, + stride_bc_pre_batch, stride_bc_pre_group, stride_bc_pre_pos, stride_bc_pre_spec, + stride_D_head, stride_D_dim, stride_z_tok, stride_z_head, stride_z_dim, + stride_dt_bias_head, stride_A_head, stride_out_tok, stride_out_head, stride_out_dim, + stride_state_indices_batch, + DT_SOFTPLUS: tl.constexpr, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_CACHE: tl.constexpr, + BLOCK_SIZE_SPEC: tl.constexpr, HAS_DT_BIAS: tl.constexpr, HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, CACHE_BUF_LEN: tl.constexpr, DSTATE_TILE: tl.constexpr, + NDS: tl.constexpr, BLOCK_SIZE_DSTATE: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + pid_b = tl.program_id(axis=1) + pid_h = tl.program_id(axis=2) + state_batch_idx = tl.load(state_batch_indices_ptr + pid_b * stride_state_indices_batch).to(tl.int64) + if state_batch_idx == null_block_id: + return + if tl.load(is_flush_flags_ptr + state_batch_idx) != 0: + return + bos = tl.load(query_start_loc_ptr + pid_b).to(tl.int64) + eos = tl.load(query_start_loc_ptr + pid_b + 1).to(tl.int64) + spec_len = (eos - bos).to(tl.int32) + write_pos = tl.load(write_pos_ptr + state_batch_idx).to(tl.int32) + post_origin = tl.load(post_origin_ptr + state_batch_idx).to(tl.int32) + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_k = tl.arange(0, BLOCK_SIZE_CACHE) + offs_s = tl.arange(0, BLOCK_SIZE_SPEC) + offs_nt = tl.arange(0, DSTATE_TILE) + spec_valid_mask = offs_s < spec_len + hist_mask = offs_k < write_pos + cache_valid_mask = (offs_k < max_cache_len) & (offs_k < (write_pos + spec_len)) + spec_token_mask = (offs_k >= write_pos) & (offs_k < (write_pos + spec_len)) + spec_cache_pos = write_pos + offs_s + spec_prefix_mask = spec_valid_mask[:, None] & spec_token_mask[None, :] & (offs_k[None, :] <= spec_cache_pos[:, None]) + phys_k = (post_origin + offs_k) & (CACHE_BUF_LEN - 1) + phys_spec = (post_origin + spec_cache_pos) & (CACHE_BUF_LEN - 1) + + state_ptr += state_batch_idx * stride_state_batch + pid_h * stride_state_head + x_cache_ptr += state_batch_idx * stride_x_cache_batch + pid_h * stride_x_cache_head + dt_cache_ptr += state_batch_idx * stride_dt_cache_batch + pid_h * stride_dt_cache_head + C_src_ptr += bos * stride_C_src_tok + (pid_h // nheads_ngroups_ratio) * dstate * stride_C_src_c + bc_pre_ptr += pid_b * stride_bc_pre_batch + (pid_h // nheads_ngroups_ratio) * stride_bc_pre_group + if HAS_D: + D_ptr += pid_h * stride_D_head + if HAS_Z: + z_ptr += bos * stride_z_tok + pid_h * stride_z_head + if HAS_DT_BIAS: + dt_bias_ptr += pid_h * stride_dt_bias_head + A_ptr += pid_h * stride_A_head + out_ptr += bos * stride_out_tok + pid_h * stride_out_head + A_val = tl.load(A_ptr).to(tl.float32) + dt_bias_val = tl.load(dt_bias_ptr).to(tl.float32) if HAS_DT_BIAS else 0.0 + + # dt over the window (+ bias / softplus), then the per-draft decay weights. + dt_blk = tl.load(dt_cache_ptr + phys_k * stride_dt_cache_pos, mask=cache_valid_mask, other=0.0).to(tl.float32) + dt_blk = tl.where(cache_valid_mask, dt_blk, 0.0) + if HAS_DT_BIAS: + dt_blk = tl.where(cache_valid_mask, dt_blk + dt_bias_val, 0.0) + if DT_SOFTPLUS: + dt_blk = tl.where(cache_valid_mask, tl.where(dt_blk <= 20.0, softplus(dt_blk), dt_blk), 0.0) + dt_cum = tl.cumsum(dt_blk, axis=0) + hist_total = tl.sum(tl.where(hist_mask, dt_blk, 0.0), axis=0) + spec_cum = tl.sum(tl.where(spec_prefix_mask, dt_blk[None, :], 0.0), axis=1) + spec_cum = tl.where(spec_valid_mask, spec_cum, 0.0) + spec_total = hist_total + spec_cum + checkpoint_decay = tl.where(spec_valid_mask, tl.exp(tl.minimum(A_val * spec_total, 0.0)), 0.0) + + # Causal weighted sum over cached values: spec_contrib = x_cache @ factor. + x_blk = tl.load(x_cache_ptr + phys_k[None, :] * stride_x_cache_pos + offs_m[:, None] * stride_x_cache_dim, mask=(offs_m[:, None] < dim) & cache_valid_mask[None, :], other=0.0) + x_ty = x_blk.to(x_cache_ptr.dtype.element_ty) + bc = tl.load(bc_pre_ptr + offs_k[:, None] * stride_bc_pre_pos + offs_s[None, :] * stride_bc_pre_spec, mask=cache_valid_mask[:, None] & spec_valid_mask[None, :], other=0.0).to(tl.float32) + spec_scale = dt_blk[:, None] * tl.exp(tl.minimum(A_val * (spec_total[None, :] - dt_cum[:, None]), 0.0)) + causal = spec_valid_mask[None, :] & cache_valid_mask[:, None] & (offs_k[:, None] <= spec_cache_pos[None, :]) + factor = tl.where(causal, bc * spec_scale, 0.0) + spec_contrib = tl.dot(x_ty, factor.to(x_cache_ptr.dtype.element_ty), input_precision="tf32x3").to(tl.float32) + + # Decayed checkpoint readout S_0 @ C, dstate-tiled. tf32x3 keeps fp32-act + # parity; bf16 act uses single-pass tf32 (the flag is a no-op on bf16 inputs). + checkpoint_out = tl.zeros([BLOCK_SIZE_M, BLOCK_SIZE_SPEC], dtype=tl.float32) + for i in tl.static_range(NDS): + offs_n = i * DSTATE_TILE + offs_nt + nmask = offs_n < dstate + st = tl.load(state_ptr + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate, mask=(offs_m[:, None] < dim) & nmask[None, :], other=0.0).to(tl.float32) + c_mask = spec_valid_mask[:, None] & nmask[None, :] + c_chunk = tl.load(C_src_ptr + offs_s[:, None] * stride_C_src_tok + offs_n[None, :] * stride_C_src_c, mask=c_mask, other=0.0).to(tl.float32) + if x_cache_ptr.dtype.element_ty == tl.float32: + checkpoint_out += tl.dot(st, tl.trans(c_chunk), input_precision="tf32x3").to(tl.float32) + else: + checkpoint_out += tl.dot(st, tl.trans(c_chunk), input_precision="tf32").to(tl.float32) + checkpoint_out *= checkpoint_decay[None, :] + out = tl.trans(checkpoint_out + spec_contrib) + + if HAS_D: + x_spec_sm = tl.load(x_cache_ptr + offs_m[None, :] * stride_x_cache_dim + phys_spec[:, None] * stride_x_cache_pos, mask=spec_valid_mask[:, None] & (offs_m[None, :] < dim), other=0.0).to(tl.float32) + D_val = tl.load(D_ptr + offs_m * stride_D_dim, mask=offs_m < dim, other=0.0).to(tl.float32) + out += x_spec_sm * D_val[None, :] + if HAS_Z: + z_val = tl.load(z_ptr + offs_s[:, None] * stride_z_tok + offs_m[None, :] * stride_z_dim, mask=spec_valid_mask[:, None] & (offs_m[None, :] < dim), other=0.0).to(tl.float32) + out *= z_val * tl.sigmoid(z_val) + out = tl.where(spec_valid_mask[:, None], out, 0.0) + tl.store(out_ptr + offs_s[:, None] * stride_out_tok + offs_m[None, :] * stride_out_dim, out, mask=spec_valid_mask[:, None] & (offs_m[None, :] < dim)) + + +# ====================================================================== +# Flush launch (flush rows): reconstruct the committed-history state S_1, store +# it as the checkpoint, and read the output via S_1 + the intra-spec window. +# Non-flush rows early-exit. dstate-tiled. +# ====================================================================== +@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) +@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) +@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.jit +def _replayssm_spec_fl_kernel( + state_ptr, x_cache_ptr, dt_cache_ptr, B_cache_ptr, C_src_ptr, bc_pre_ptr, + D_ptr, z_ptr, dt_bias_ptr, A_ptr, out_ptr, is_flush_flags_ptr, write_pos_ptr, + post_origin_ptr, state_batch_indices_ptr, query_start_loc_ptr, null_block_id, + batch, nheads, dim, dstate, max_cache_len, nheads_ngroups_ratio, + stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, + stride_x_cache_batch, stride_x_cache_head, stride_x_cache_dim, stride_x_cache_pos, + stride_dt_cache_batch, stride_dt_cache_head, stride_dt_cache_pos, + stride_B_cache_batch, stride_B_cache_group, stride_B_cache_dstate, stride_B_cache_pos, + stride_C_src_tok, stride_C_src_c, + stride_bc_pre_batch, stride_bc_pre_group, stride_bc_pre_pos, stride_bc_pre_spec, + stride_D_head, stride_D_dim, stride_z_tok, stride_z_head, stride_z_dim, + stride_dt_bias_head, stride_A_head, stride_out_tok, stride_out_head, stride_out_dim, + stride_state_indices_batch, + DT_SOFTPLUS: tl.constexpr, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_CACHE: tl.constexpr, + BLOCK_SIZE_SPEC: tl.constexpr, HAS_DT_BIAS: tl.constexpr, HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, CACHE_BUF_LEN: tl.constexpr, DSTATE_TILE: tl.constexpr, + NDS: tl.constexpr, BLOCK_SIZE_DSTATE: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + pid_b = tl.program_id(axis=1) + pid_h = tl.program_id(axis=2) + state_batch_idx = tl.load(state_batch_indices_ptr + pid_b * stride_state_indices_batch).to(tl.int64) + if state_batch_idx == null_block_id: + return + if tl.load(is_flush_flags_ptr + state_batch_idx) == 0: + return + bos = tl.load(query_start_loc_ptr + pid_b).to(tl.int64) + eos = tl.load(query_start_loc_ptr + pid_b + 1).to(tl.int64) + spec_len = (eos - bos).to(tl.int32) + write_pos = tl.load(write_pos_ptr + state_batch_idx).to(tl.int32) + post_origin = tl.load(post_origin_ptr + state_batch_idx).to(tl.int32) + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_k = tl.arange(0, BLOCK_SIZE_CACHE) + offs_s = tl.arange(0, BLOCK_SIZE_SPEC) + offs_nt = tl.arange(0, DSTATE_TILE) + spec_valid_mask = offs_s < spec_len + spec_cache_pos = write_pos + offs_s + phys_spec = (post_origin + spec_cache_pos) & (CACHE_BUF_LEN - 1) + hist_mask = offs_k < write_pos + phys_h = (post_origin + offs_k) & (CACHE_BUF_LEN - 1) + + state_ptr += state_batch_idx * stride_state_batch + pid_h * stride_state_head + x_cache_ptr += state_batch_idx * stride_x_cache_batch + pid_h * stride_x_cache_head + dt_cache_ptr += state_batch_idx * stride_dt_cache_batch + pid_h * stride_dt_cache_head + B_cache_ptr += state_batch_idx * stride_B_cache_batch + (pid_h // nheads_ngroups_ratio) * stride_B_cache_group + C_src_ptr += bos * stride_C_src_tok + (pid_h // nheads_ngroups_ratio) * dstate * stride_C_src_c + bc_pre_ptr += pid_b * stride_bc_pre_batch + (pid_h // nheads_ngroups_ratio) * stride_bc_pre_group + if HAS_D: + D_ptr += pid_h * stride_D_head + if HAS_Z: + z_ptr += bos * stride_z_tok + pid_h * stride_z_head + if HAS_DT_BIAS: + dt_bias_ptr += pid_h * stride_dt_bias_head + A_ptr += pid_h * stride_A_head + out_ptr += bos * stride_out_tok + pid_h * stride_out_head + A_val = tl.load(A_ptr).to(tl.float32) + dt_bias_val = tl.load(dt_bias_ptr).to(tl.float32) if HAS_DT_BIAS else 0.0 + + # History decay (for the S_1 reconstruction) and spec-prefix decay (output). + dt_h = tl.load(dt_cache_ptr + phys_h * stride_dt_cache_pos, mask=hist_mask, other=0.0).to(tl.float32) + dt_h = tl.where(hist_mask, dt_h, 0.0) + if HAS_DT_BIAS: + dt_h = tl.where(hist_mask, dt_h + dt_bias_val, 0.0) + if DT_SOFTPLUS: + dt_h = tl.where(hist_mask, tl.where(dt_h <= 20.0, softplus(dt_h), dt_h), 0.0) + hist_cum = tl.cumsum(dt_h, axis=0) + hist_total = tl.sum(dt_h, axis=0) + hist_decay = tl.exp(tl.minimum(A_val * hist_total, 0.0)) + hist_scale = tl.where(hist_mask, dt_h * tl.exp(tl.minimum(A_val * (hist_total - hist_cum), 0.0)), 0.0) + dt_s = tl.load(dt_cache_ptr + phys_spec * stride_dt_cache_pos, mask=spec_valid_mask, other=0.0).to(tl.float32) + dt_s = tl.where(spec_valid_mask, dt_s, 0.0) + if HAS_DT_BIAS: + dt_s = tl.where(spec_valid_mask, dt_s + dt_bias_val, 0.0) + if DT_SOFTPLUS: + dt_s = tl.where(spec_valid_mask, tl.where(dt_s <= 20.0, softplus(dt_s), dt_s), 0.0) + spec_cum = tl.cumsum(dt_s, axis=0) + spec_decay = tl.where(spec_valid_mask, tl.exp(tl.minimum(A_val * spec_cum, 0.0)), 0.0) + x_hist = tl.load(x_cache_ptr + phys_h[None, :] * stride_x_cache_pos + offs_m[:, None] * stride_x_cache_dim, mask=(offs_m[:, None] < dim) & hist_mask[None, :], other=0.0) + x_hist_ty = x_hist.to(x_cache_ptr.dtype.element_ty) + + # Reconstruct S_1 = S_0 * hist_decay + (x_hist @ scaled B_hist), store it, and + # accumulate the checkpoint readout S_1 @ C. dstate-tiled. + checkpoint_out = tl.zeros([BLOCK_SIZE_M, BLOCK_SIZE_SPEC], dtype=tl.float32) + for i in tl.static_range(NDS): + offs_n = i * DSTATE_TILE + offs_nt + nmask = offs_n < dstate + B_block = tl.load(B_cache_ptr + phys_h[:, None] * stride_B_cache_pos + offs_n[None, :] * stride_B_cache_dstate, mask=hist_mask[:, None] & nmask[None, :], other=0.0) + B_hist_scaled = (tl.where(hist_mask[:, None], B_block.to(tl.float32), 0.0) * hist_scale[:, None]).to(x_cache_ptr.dtype.element_ty) + delta = tl.dot(x_hist_ty, B_hist_scaled, input_precision="tf32x3").to(tl.float32) + st_ptrs = state_ptr + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + st = tl.load(st_ptrs, mask=(offs_m[:, None] < dim) & nmask[None, :], other=0.0) + S1 = st.to(tl.float32) * hist_decay + delta + if write_pos > 0: + tl.store(st_ptrs, S1.to(st.dtype), mask=(offs_m[:, None] < dim) & nmask[None, :]) + c_mask = spec_valid_mask[:, None] & nmask[None, :] + c_chunk = tl.load(C_src_ptr + offs_s[:, None] * stride_C_src_tok + offs_n[None, :] * stride_C_src_c, mask=c_mask, other=0.0).to(tl.float32) + if x_cache_ptr.dtype.element_ty == tl.float32: + checkpoint_out += tl.dot(S1, tl.trans(c_chunk), input_precision="tf32x3").to(tl.float32) + else: + checkpoint_out += tl.dot(S1, tl.trans(c_chunk), input_precision="tf32").to(tl.float32) + checkpoint_out *= spec_decay[None, :] + + # Intra-spec window contribution: intra = x_spec @ factor_intra (causal T x T). + bc_spec = tl.load(bc_pre_ptr + (write_pos + offs_k)[:, None] * stride_bc_pre_pos + offs_s[None, :] * stride_bc_pre_spec, mask=(offs_k[:, None] < spec_len) & spec_valid_mask[None, :], other=0.0).to(tl.float32) + dt_k = tl.load(dt_cache_ptr + ((post_origin + write_pos + offs_k) & (CACHE_BUF_LEN - 1)) * stride_dt_cache_pos, mask=offs_k < spec_len, other=0.0).to(tl.float32) + dt_k = tl.where(offs_k < spec_len, dt_k, 0.0) + if HAS_DT_BIAS: + dt_k = tl.where(offs_k < spec_len, dt_k + dt_bias_val, 0.0) + if DT_SOFTPLUS: + dt_k = tl.where(offs_k < spec_len, tl.where(dt_k <= 20.0, softplus(dt_k), dt_k), 0.0) + speccum_k = tl.cumsum(dt_k, axis=0) + causal = (offs_k[:, None] < spec_len) & spec_valid_mask[None, :] & (offs_k[:, None] <= offs_s[None, :]) + decay_ks = tl.exp(tl.minimum(A_val * (spec_cum[None, :] - speccum_k[:, None]), 0.0)) + factor_intra = tl.where(causal, bc_spec * dt_k[:, None] * decay_ks, 0.0) + x_src = tl.load(x_cache_ptr + ((post_origin + write_pos + offs_k)[None, :] & (CACHE_BUF_LEN - 1)) * stride_x_cache_pos + offs_m[:, None] * stride_x_cache_dim, mask=(offs_m[:, None] < dim) & (offs_k[None, :] < spec_len), other=0.0).to(x_cache_ptr.dtype.element_ty) + intra = tl.dot(x_src, factor_intra.to(x_cache_ptr.dtype.element_ty), input_precision="tf32x3").to(tl.float32) + out = tl.trans(checkpoint_out + intra) + + if HAS_D: + x_spec_sm = tl.load(x_cache_ptr + offs_m[None, :] * stride_x_cache_dim + phys_spec[:, None] * stride_x_cache_pos, mask=spec_valid_mask[:, None] & (offs_m[None, :] < dim), other=0.0).to(tl.float32) + D_val = tl.load(D_ptr + offs_m * stride_D_dim, mask=offs_m < dim, other=0.0).to(tl.float32) + out += x_spec_sm * D_val[None, :] + if HAS_Z: + z_val = tl.load(z_ptr + offs_s[:, None] * stride_z_tok + offs_m[None, :] * stride_z_dim, mask=spec_valid_mask[:, None] & (offs_m[None, :] < dim), other=0.0).to(tl.float32) + out *= z_val * tl.sigmoid(z_val) + out = tl.where(spec_valid_mask[:, None], out, 0.0) + tl.store(out_ptr + offs_s[:, None] * stride_out_tok + offs_m[None, :] * stride_out_dim, out, mask=spec_valid_mask[:, None] & (offs_m[None, :] < dim)) + + +@triton.jit +def _advance_write_pos_origin_kernel( + write_pos_ptr, + post_origin_ptr, + is_flush_ptr, + num_accepted_ptr, + state_batch_indices_ptr, + null_block_id, + batch, + stride_state_indices_batch, + MAX_CACHE_LEN: tl.constexpr, + MAX_SPEC_LEN: tl.constexpr, + CACHE_BUF_LEN: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + offs = tl.arange(0, BLOCK_SIZE) + row_mask = offs < batch + state_batch_idx = tl.load(state_batch_indices_ptr + offs * stride_state_indices_batch, mask=row_mask, other=null_block_id).to( + tl.int64 + ) + valid = row_mask & (state_batch_idx != null_block_id) + write_pos = tl.load( + write_pos_ptr + state_batch_idx, mask=valid, other=0 + ).to(tl.int32) + post_origin = tl.load( + post_origin_ptr + state_batch_idx, mask=valid, other=0 + ).to(tl.int32) + is_flush_cur = tl.load( + is_flush_ptr + state_batch_idx, mask=valid, other=0 + ).to(tl.int32) + num_accepted = tl.load(num_accepted_ptr + offs, mask=valid, other=0).to(tl.int32) + total_commit = tl.where(valid, num_accepted, 0).to(tl.int32) + flush_now = (total_commit > 0) & (is_flush_cur != 0) + new_origin = tl.where( + flush_now, (post_origin + write_pos) & (CACHE_BUF_LEN - 1), post_origin + ).to(tl.int32) + new_wp = tl.where( + total_commit <= 0, + write_pos, + tl.where(is_flush_cur != 0, total_commit, write_pos + total_commit), + ).to(tl.int32) + # Early-flush one window early so every verify step satisfies + # write_pos + spec_len <= max_cache_len (the spec window never overflows). + next_is_flush = ((new_wp + 2 * MAX_SPEC_LEN) > MAX_CACHE_LEN).to(tl.int8) + tl.store(post_origin_ptr + state_batch_idx, new_origin, mask=valid) + tl.store(write_pos_ptr + state_batch_idx, new_wp, mask=valid) + tl.store(is_flush_ptr + state_batch_idx, next_is_flush, mask=valid) + + +@triton.jit +def _reset_replayssm_spec_cursors_kernel( + write_pos_ptr, + post_origin_ptr, + is_flush_ptr, + first_decode_ptr, # (batch,) int8 mask + state_batch_indices_ptr, + null_block_id, + batch, + stride_state_indices_batch, + INIT_IS_FLUSH: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + offs = tl.arange(0, BLOCK_SIZE) + row_mask = offs < batch + state_batch_idx = tl.load(state_batch_indices_ptr + offs * stride_state_indices_batch, mask=row_mask, other=null_block_id).to( + tl.int64 + ) + first = tl.load(first_decode_ptr + offs, mask=row_mask, other=0).to(tl.int32) + do_reset = row_mask & (state_batch_idx != null_block_id) & (first != 0) + tl.store( + write_pos_ptr + state_batch_idx, + tl.zeros_like(state_batch_idx).to(tl.int32), + mask=do_reset, + ) + tl.store( + post_origin_ptr + state_batch_idx, + tl.zeros_like(state_batch_idx).to(tl.int32), + mask=do_reset, + ) + tl.store( + is_flush_ptr + state_batch_idx, + (tl.zeros_like(state_batch_idx) + INIT_IS_FLUSH).to(tl.int8), + mask=do_reset, + ) + + +def selective_state_update_replayssm_spec( + state_checkpoint: torch.Tensor, # (num_blocks, H, P, N) checkpoint (flush updates in place) + post_conv_cache: torch.Tensor, # (num_blocks, cache_buf_len, conv_dim) circular + dt_cache: torch.Tensor, # (num_blocks, H, cache_buf_len) circular + conv_out: torch.Tensor, # (total_tokens, conv_dim) packed channel-last post-conv + dt_spec: torch.Tensor, # (total_tokens, H) packed raw dt + A: torch.Tensor, # (H, P, N) TIE_HDIM (A.stride(-1)==A.stride(-2)==0) + write_pos: torch.Tensor, # (num_state_slots,) int32 block-keyed cursor + post_conv_state_pos: torch.Tensor, # (num_state_slots,) int32 circular origin + is_flush: torch.Tensor, # (num_state_slots,) int8 block-keyed flag + query_start_loc: torch.Tensor, # (batch + 1,) int32 packed offsets + state_batch_indices: torch.Tensor, # (batch,) int32 physical block per row + max_cache_len: int, # logical flush threshold L = B + max_spec_len + max_spec_len: int, + d_inner: int, + ngroups: int, + dstate: int, + D: torch.Tensor | None = None, + z: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = True, + out: torch.Tensor | None = None, + bc_pre: torch.Tensor | None = None, + null_block_id: int = NULL_BLOCK_ID, +) -> torch.Tensor: + """One Mamba2 speculative verify step on the paged CIRCULAR post-conv cache. + + Hybrid conv variant: ``conv_out`` is the post-conv output of vLLM's + ``causal_conv1d_update`` (packed channel-last ``[total_tokens, conv_dim]``). + ``max_cache_len`` is the logical flush threshold L = B + max_spec_len; the + physical pow2 buffer (= ``post_conv_cache.shape[1]`` = ``next_pow2(L)``) wraps + the ring, while the history/cache tile is ``next_pow2(L - max_spec_len)``. + Fuses scatter + bc-precompute in one ``(batch, ngroups)`` launch, then runs + two dedicated launches (verify + flush) with device-side row routing. Cursors + are block-keyed and advanced once per step by ``commit_replayssm_spec``. + """ + num_blocks, nheads, dim, n_state = state_checkpoint.shape + assert n_state == dstate + total_tokens, conv_dim = conv_out.shape + buf = post_conv_cache.shape[1] + cache_buf_len = buf + assert cache_buf_len & (cache_buf_len - 1) == 0, "cache_buf_len must be a power of two" + assert d_inner == nheads * dim + cache_conv_dim = d_inner + ngroups * dstate # x|B only (no C) + assert post_conv_cache.shape == (num_blocks, buf, cache_conv_dim) + assert dt_cache.shape == (num_blocks, nheads, buf) + assert dt_spec.shape == (total_tokens, nheads) + assert A.shape == (nheads, dim, dstate) and A.stride(-1) == 0 and A.stride(-2) == 0 + batch = state_batch_indices.shape[0] + assert query_start_loc.shape[0] == batch + 1 + + if out is None: + out = torch.empty(total_tokens, nheads, dim, device=conv_out.device, dtype=conv_out.dtype) + if total_tokens == 0: + return out + + L = max_cache_len + base_block = max(1, L - max_spec_len) + main_block = max(16, triton.next_power_of_2(base_block)) # history/cache tile + block_spec = max(1, triton.next_power_of_2(max_spec_len)) + pre_block = max(16, triton.next_power_of_2(L)) # precompute window tile + block_dstate = triton.next_power_of_2(dstate) + + bsm_v, nw_v, dt_v, ns_v = get_replayssm_config( + "mamba2_spec_verify", dstate=dstate, base_block=base_block, max_spec_len=max_spec_len + ) + bsm_f, nw_f, dt_f, ns_f = get_replayssm_config( + "mamba2_spec_flush", dstate=dstate, base_block=base_block, max_spec_len=max_spec_len + ) + dt_v = max(16, min(dt_v, block_dstate)); nds_v = triton.cdiv(block_dstate, dt_v) + dt_f = max(16, min(dt_f, block_dstate)); nds_f = triton.cdiv(block_dstate, dt_f) + + if bc_pre is None: + bc_pre = torch.empty(batch, ngroups, buf, block_spec, device=conv_out.device, dtype=conv_out.dtype) + sis = state_batch_indices.stride(0) + + # --- fused scatter + precompute (full-window bc over [0, L)) --- + ratio = nheads // ngroups + ratio_p = ratio * dim + BLOCK_CX = 256 + NCX = triton.cdiv(ratio_p, BLOCK_CX) + block_hl = max(1, triton.next_power_of_2(ratio)) + with torch.cuda.device(conv_out.device.index): + _fused_scatter_precompute_kernel[(batch, ngroups)]( + conv_out, dt_spec, post_conv_cache, dt_cache, write_pos, post_conv_state_pos, + bc_pre, state_batch_indices, query_start_loc, null_block_id, batch, ngroups, + nheads, dstate, d_inner, conv_dim, L, + conv_out.stride(0), conv_out.stride(1), dt_spec.stride(0), dt_spec.stride(1), + post_conv_cache.stride(0), post_conv_cache.stride(1), post_conv_cache.stride(2), + dt_cache.stride(0), dt_cache.stride(1), dt_cache.stride(2), + bc_pre.stride(0), bc_pre.stride(1), bc_pre.stride(2), bc_pre.stride(3), sis, + RATIO=ratio, RATIO_P=ratio_p, NCX=NCX, BLOCK_CX=BLOCK_CX, CACHE_BUF_LEN=cache_buf_len, + BLOCK_SIZE_CACHE=pre_block, BLOCK_SIZE_SPEC=block_spec, BLOCK_HL=block_hl, num_warps=4, + ) + + # views into the paged circular post-conv cache: x | B | C on the channel axis + x_view = post_conv_cache[:, :, :d_inner].view(num_blocks, buf, nheads, dim).permute(0, 2, 1, 3) + B_view = post_conv_cache[:, :, d_inner : d_inner + ngroups * dstate].view(num_blocks, buf, ngroups, dstate).permute(0, 2, 1, 3) + # C is not cached; the kernels read it fresh from this conv_out slice. + C_src = conv_out[:, d_inner + ngroups * dstate :] + z_strides = (z.stride(0), z.stride(1), z.stride(2)) if z is not None else (0, 0, 0) + + def _args(bsm): + grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads) + return grid, ( + state_checkpoint, x_view, dt_cache, B_view, C_src, bc_pre, D, z, dt_bias, A, + out, is_flush, write_pos, post_conv_state_pos, state_batch_indices, + query_start_loc, null_block_id, batch, nheads, dim, dstate, L, ratio, + state_checkpoint.stride(0), state_checkpoint.stride(1), state_checkpoint.stride(2), state_checkpoint.stride(3), + x_view.stride(0), x_view.stride(1), x_view.stride(3), x_view.stride(2), + dt_cache.stride(0), dt_cache.stride(1), dt_cache.stride(2), + B_view.stride(0), B_view.stride(1), B_view.stride(3), B_view.stride(2), + C_src.stride(0), C_src.stride(1), + bc_pre.stride(0), bc_pre.stride(1), bc_pre.stride(2), bc_pre.stride(3), + D.stride(0) if D is not None else 0, D.stride(1) if D is not None else 0, + z_strides[0], z_strides[1], z_strides[2], dt_bias.stride(0) if dt_bias is not None else 0, + A.stride(0), out.stride(0), out.stride(1), out.stride(2), sis, dt_softplus, bsm, + ) + + with torch.cuda.device(state_checkpoint.device.index): + grid, base = _args(bsm_v) + _replayssm_spec_nf_kernel[grid]( + *base, main_block, block_spec, + CACHE_BUF_LEN=cache_buf_len, DSTATE_TILE=dt_v, NDS=nds_v, num_warps=nw_v, num_stages=ns_v, + ) + grid, base = _args(bsm_f) + _replayssm_spec_fl_kernel[grid]( + *base, main_block, block_spec, + CACHE_BUF_LEN=cache_buf_len, DSTATE_TILE=dt_f, NDS=nds_f, num_warps=nw_f, num_stages=ns_f, + ) + return out + + +def commit_replayssm_spec( + write_pos: torch.Tensor, + post_conv_state_pos: torch.Tensor, + is_flush: torch.Tensor, + num_accepted_tokens: torch.Tensor, # (batch,) int32, INCLUDES bonus (min 1) + state_batch_indices: torch.Tensor, # (batch,) int32 + max_cache_len: int, # logical flush threshold L + max_spec_len: int, + cache_buf_len: int | None = None, # physical pow2 buffer next_pow2(L) + null_block_id: int = NULL_BLOCK_ID, +) -> None: + """CUDA-graph-safe block-keyed commit. Advances ``write_pos`` and the circular + origin (flush = O(1) bump) per decode row, and precomputes next-step + ``is_flush``. Maps vLLM ``num_accepted_tokens`` (incl. bonus) to the commit.""" + batch = state_batch_indices.shape[0] + if cache_buf_len is None: + cache_buf_len = max(1, triton.next_power_of_2(max_cache_len)) + BLOCK = max(1, triton.next_power_of_2(batch)) + with torch.cuda.device(write_pos.device.index): + _advance_write_pos_origin_kernel[(1,)]( + write_pos, + post_conv_state_pos, + is_flush, + num_accepted_tokens, + state_batch_indices, + null_block_id, + batch, + state_batch_indices.stride(0), + MAX_CACHE_LEN=max_cache_len, + MAX_SPEC_LEN=max_spec_len, + CACHE_BUF_LEN=cache_buf_len, + BLOCK_SIZE=BLOCK, + num_warps=1, + ) + + +def reset_replayssm_spec_cursors( + write_pos: torch.Tensor, + post_conv_state_pos: torch.Tensor, + is_flush: torch.Tensor, + first_decode_mask: torch.Tensor, # (batch,) int8 + state_batch_indices: torch.Tensor, # (batch,) int32 + max_cache_len: int, # logical flush threshold L + max_spec_len: int, + null_block_id: int = NULL_BLOCK_ID, +) -> None: + """Prefill->decode reset for first-decode rows (block-keyed). Seeds the + initial ``is_flush`` to match the steady-state early-flush cadence.""" + batch = state_batch_indices.shape[0] + BLOCK = max(1, triton.next_power_of_2(batch)) + init_is_flush = 1 if 2 * max_spec_len > max_cache_len else 0 + with torch.cuda.device(write_pos.device.index): + _reset_replayssm_spec_cursors_kernel[(1,)]( + write_pos, + post_conv_state_pos, + is_flush, + first_decode_mask, + state_batch_indices, + null_block_id, + batch, + state_batch_indices.stride(0), + INIT_IS_FLUSH=init_is_flush, + BLOCK_SIZE=BLOCK, + num_warps=1, + ) + + +__all__ = [ + "selective_state_update_replayssm_spec", + "commit_replayssm_spec", + "reset_replayssm_spec_cursors", +] diff --git a/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_state_and_output.py b/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_state_and_output.py new file mode 100644 index 000000000000..f38269199803 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_state_and_output.py @@ -0,0 +1,373 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# ruff: noqa: E501 + +import torch + +from vllm.model_executor.layers.mamba.ops.mamba_ssm import softplus +from vllm.model_executor.layers.mamba.ops.replayssm_config import get_replayssm_config +from vllm.triton_utils import tl, triton +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + + +@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) +@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) +@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) +@triton.heuristics( + { + "HAS_STATE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] + is not None + } +) +@triton.heuristics( + {"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])} +) +@triton.jit +def _replayssm_state_and_output_kernel( + # Pointers to matrices + state_ptr, + x_ptr, + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + x_cache_ptr, + dt_cache_ptr, + B_cache_ptr, + write_pos_ptr, + is_flush_ptr, + state_batch_indices_ptr, + null_block_id, + # Matrix dimensions + batch, + nheads, + dim, + dstate, + nheads_ngroups_ratio, + # State strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # Input strides + stride_x_batch, + stride_x_head, + stride_x_dim, + stride_dt_batch, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + stride_B_batch, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_group, + stride_C_dstate, + stride_D_head, + stride_D_dim, + stride_z_batch, + stride_z_head, + stride_z_dim, + stride_out_batch, + stride_out_head, + stride_out_dim, + # Cache strides + stride_x_cache_batch, + stride_x_cache_head, + stride_x_cache_dim, + stride_x_cache_pos, + stride_dt_cache_batch, + stride_dt_cache_head, + stride_dt_cache_pos, + stride_B_cache_batch, + stride_B_cache_group, + stride_B_cache_pos, + stride_B_cache_dstate, + stride_state_indices_batch, + stride_state_indices_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + MAX_CACHE_LEN: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + # heuristic-computed + BLOCK_SIZE_DSTATE: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_STATE_BATCH_INDICES: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + pid_b = tl.program_id(axis=1) + pid_h = tl.program_id(axis=2) + + # Resolve the physical state slot for this decode row; skip padded rows. + if HAS_STATE_BATCH_INDICES: + state_batch_idx = tl.load(state_batch_indices_ptr + pid_b * stride_state_indices_batch + 0 * stride_state_indices_T).to(tl.int64) + if state_batch_idx == null_block_id: + return + else: + state_batch_idx = pid_b + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + offs_k = tl.arange(0, BLOCK_SIZE_K) + + # Buffer cursor (number of cached tokens so far) and the flush flag. + write_pos = tl.load(write_pos_ptr + pid_b).to(tl.int64) + is_flush = tl.load(is_flush_ptr + pid_b) != 0 + + # Advance every pointer to this (row, head, group). + state_ptr += state_batch_idx * stride_state_batch + pid_h * stride_state_head + x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head + dt_ptr += pid_b * stride_dt_batch + pid_h * stride_dt_head + B_ptr += pid_b * stride_B_batch + (pid_h // nheads_ngroups_ratio) * stride_B_group + C_ptr += pid_b * stride_C_batch + (pid_h // nheads_ngroups_ratio) * stride_C_group + out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head + x_cache_ptr += state_batch_idx * stride_x_cache_batch + pid_h * stride_x_cache_head + dt_cache_ptr += state_batch_idx * stride_dt_cache_batch + pid_h * stride_dt_cache_head + B_cache_ptr += state_batch_idx * stride_B_cache_batch + (pid_h // nheads_ngroups_ratio) * stride_B_cache_group + + # Current-token dt (+ bias, softplus) and the per-head scalar A. + dt_cur = tl.load(dt_ptr).to(tl.float32) + if HAS_DT_BIAS: + dt_cur += tl.load(dt_bias_ptr + pid_h * stride_dt_bias_head).to(tl.float32) + if DT_SOFTPLUS: + dt_cur = tl.where(dt_cur <= 20.0, softplus(dt_cur), dt_cur) + A = tl.load(A_ptr + pid_h * stride_A_head).to(tl.float32) + + # dt over the window (cached history + current token at write_pos), then the + # decay weights: total decay exp(A*sum dt) and the per-position scale s_j. + dt_all = tl.load(dt_cache_ptr + offs_k * stride_dt_cache_pos, mask=offs_k < write_pos, other=0.0).to(tl.float32) + dt_all = tl.where(offs_k == write_pos, dt_cur, dt_all) + dA_cumsum = A * tl.cumsum(dt_all, axis=0) + dA_total = A * tl.sum(dt_all, axis=0) + total_decay = tl.exp(dA_total) + scale = dt_all * tl.exp(dA_total - dA_cumsum) + scale = tl.where(offs_k <= write_pos, scale, 0.0) + + # Current-token x / C, the checkpoint state S_0, and current-token B. + x_cur = tl.load(x_ptr + offs_m * stride_x_dim, mask=offs_m < dim, other=0.0) + C = tl.load(C_ptr + offs_n * stride_C_dstate, mask=offs_n < dstate, other=0.0).to(tl.float32) + state_ptrs = state_ptr + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + state = tl.load(state_ptrs, mask=(offs_m[:, None] < dim) & (offs_n[None, :] < dstate), other=0.0) + B_cur = tl.load(B_ptr + offs_n * stride_B_dstate, mask=offs_n < dstate, other=0.0) + + # Gather buffered x and B over the window (history + current token). + x_all_ptrs = x_cache_ptr + offs_m[:, None] * stride_x_cache_dim + offs_k[None, :] * stride_x_cache_pos + x_all = tl.load(x_all_ptrs, mask=(offs_m[:, None] < dim) & (offs_k[None, :] < write_pos), other=0.0) + x_all = tl.where(offs_k[None, :] == write_pos, x_cur[:, None], x_all) + B_all_ptrs = B_cache_ptr + offs_k[:, None] * stride_B_cache_pos + offs_n[None, :] * stride_B_cache_dstate + B_all = tl.load(B_all_ptrs, mask=(offs_k[:, None] < write_pos) & (offs_n[None, :] < dstate), other=0.0) + B_all = tl.where(offs_k[:, None] == write_pos, B_cur[None, :], B_all) + + # Reconstruct the state from cached inputs (outer-product / state route): + # S_t = total_decay * S_0 + sum_j s_j (v_j k_j^T). Store it back on a flush. + B_scaled = (B_all.to(tl.float32) * scale[:, None]).to(x_ptr.dtype.element_ty) + # tf32x3 keeps fp32 parity with the elementwise baseline (plain tf32 on + # fp32 inputs drifts ~1e-2); bf16/fp16 inputs are unaffected by this flag. + delta_state = tl.dot( + x_all.to(x_ptr.dtype.element_ty), B_scaled, input_precision="tf32x3" + ) + state_new = state.to(tl.float32) * total_decay + delta_state.to(tl.float32) + if is_flush: + tl.store(state_ptrs, state_new.to(state.dtype), mask=(offs_m[:, None] < dim) & (offs_n[None, :] < dstate)) + + # Read the output from the reconstructed state: y = S_t q. + out = tl.sum(state_new * C[None, :], axis=1) + + # Skip connection (D) and output gate (z). + if HAS_D: + D_ptr += pid_h * stride_D_head + D = tl.load(D_ptr + offs_m * stride_D_dim, mask=offs_m < dim, other=0.0).to(tl.float32) + out += x_cur.to(tl.float32) * D + if HAS_Z: + z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head + z = tl.load(z_ptr + offs_m * stride_z_dim, mask=offs_m < dim, other=0.0).to(tl.float32) + out *= z * tl.sigmoid(z) + + tl.store(out_ptr + offs_m * stride_out_dim, out, mask=offs_m < dim) + + # Non-flush step: append the current token (x, dt, B) into the buffer at write_pos. + if not is_flush: + tl.store(x_cache_ptr + offs_m * stride_x_cache_dim + write_pos * stride_x_cache_pos, x_cur, mask=offs_m < dim) + if pid_m == 0: + tl.store(dt_cache_ptr + write_pos * stride_dt_cache_pos, dt_cur) + tl.store(B_cache_ptr + write_pos * stride_B_cache_pos + offs_n * stride_B_cache_dstate, B_cur, mask=offs_n < dstate) + + +def selective_state_update_replayssm_state_and_output( + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + z: torch.Tensor | None = None, + dt_softplus: bool = False, + x_cache: torch.Tensor | None = None, + dt_cache: torch.Tensor | None = None, + B_cache: torch.Tensor | None = None, + write_pos: torch.Tensor | None = None, + is_flush: torch.Tensor | None = None, + max_cache_len: int = 16, + state_batch_indices: torch.Tensor | None = None, + null_block_id: int = NULL_BLOCK_ID, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Cached-dot SSM update for vLLM's autoregressive Mamba2 decode path.""" + has_heads = state.dim() > 3 + if state.dim() == 3: + state = state.unsqueeze(1) + if x.dim() == 2: + x = x.unsqueeze(1) + if dt.dim() == 2: + dt = dt.unsqueeze(1) + if A.dim() == 2: + A = A.unsqueeze(0) + if B.dim() == 2: + B = B.unsqueeze(1) + if C.dim() == 2: + C = C.unsqueeze(1) + if D is not None and D.dim() == 1: + D = D.unsqueeze(0) + if z is not None and z.dim() == 2: + z = z.unsqueeze(1) + if dt_bias is not None and dt_bias.dim() == 1: + dt_bias = dt_bias.unsqueeze(0) + if out is not None and out.dim() == 2: + out = out.unsqueeze(1) + if state_batch_indices is not None and state_batch_indices.dim() == 1: + state_batch_indices = state_batch_indices.unsqueeze(1) + + _, nheads, dim, dstate = state.shape + batch = x.shape[0] + assert x.shape == (batch, nheads, dim) + assert dt.shape == x.shape + assert A.shape == (nheads, dim, dstate) + ngroups = B.shape[1] + assert nheads % ngroups == 0, "nheads must be divisible by ngroups" + assert B.shape == (batch, ngroups, dstate) + assert C.shape == B.shape + if D is not None: + assert D.shape == (nheads, dim) + if z is not None: + assert z.shape == x.shape + if dt_bias is not None: + assert dt_bias.shape == (nheads, dim) + assert out is not None and out.shape == x.shape + + assert A.stride(-1) == 0 and A.stride(-2) == 0, ( + "Cached kernel requires TIE_HDIM (A scalar per head)" + ) + assert dt.stride(-1) == 0, "Cached kernel requires TIE_HDIM (dt scalar per head)" + if dt_bias is not None: + assert dt_bias.stride(-1) == 0, ( + "Cached kernel requires TIE_HDIM (dt_bias scalar per head)" + ) + + assert x_cache is not None + assert dt_cache is not None + assert B_cache is not None + assert x_cache.shape[1:] == (nheads, max_cache_len, dim) + assert dt_cache.shape[1:] == (nheads, max_cache_len) + assert B_cache.shape[1:] == (ngroups, max_cache_len, dstate) + assert write_pos is not None and write_pos.shape[0] >= batch + assert write_pos.dtype == torch.int32 + assert is_flush is not None and is_flush.shape[0] >= batch + assert is_flush.dtype in (torch.bool, torch.int8) + if state_batch_indices is not None: + assert state_batch_indices.shape[0] >= batch + assert state_batch_indices.shape[1] >= 1 + + block_size_k = max(16, triton.next_power_of_2(max_cache_len)) + block_size_m, num_warps = get_replayssm_config( + "mamba2_state_and_output", dstate=dstate, L=max_cache_len + ) + + grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads) + z_strides = (z.stride(0), z.stride(1), z.stride(2)) if z is not None else (0, 0, 0) + state_indices_strides = ( + (state_batch_indices.stride(0), state_batch_indices.stride(1)) + if state_batch_indices is not None + else (0, 0) + ) + + with torch.accelerator.device_index(x.device.index): + _replayssm_state_and_output_kernel[grid]( + state, + x, + dt, + dt_bias, + A, + B, + C, + D, + z, + out, + x_cache, + dt_cache, + B_cache, + write_pos, + is_flush, + state_batch_indices, + null_block_id, + batch, + nheads, + dim, + dstate, + nheads // ngroups, + state.stride(0), + state.stride(1), + state.stride(2), + state.stride(3), + x.stride(0), + x.stride(1), + x.stride(2), + dt.stride(0), + dt.stride(1), + dt_bias.stride(0) if dt_bias is not None else 0, + A.stride(0), + B.stride(0), + B.stride(1), + B.stride(2), + C.stride(0), + C.stride(1), + C.stride(2), + D.stride(0) if D is not None else 0, + D.stride(1) if D is not None else 0, + z_strides[0], + z_strides[1], + z_strides[2], + out.stride(0), + out.stride(1), + out.stride(2), + x_cache.stride(0), + x_cache.stride(1), + x_cache.stride(3), + x_cache.stride(2), + dt_cache.stride(0), + dt_cache.stride(1), + dt_cache.stride(2), + B_cache.stride(0), + B_cache.stride(1), + B_cache.stride(2), + B_cache.stride(3), + state_indices_strides[0], + state_indices_strides[1], + dt_softplus, + max_cache_len, + block_size_m, + block_size_k, + num_warps=num_warps, + ) + + if not has_heads: + out = out.squeeze(1) + return out diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index bd5cd358d8a4..ddf8e01fdaa3 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -814,18 +814,27 @@ class NemotronHForCausalLM( def get_mamba_state_dtype_from_config( cls, vllm_config: "VllmConfig", - ) -> tuple[torch.dtype, torch.dtype]: - return MambaStateDtypeCalculator.mamba2_state_dtype( + ) -> tuple[torch.dtype, ...]: + cache_config = vllm_config.cache_config + if cache_config.use_replayssm_spec: + return MambaStateDtypeCalculator.mamba2_spec_cached_state_dtype( + vllm_config.model_config.dtype, + cache_config.mamba_cache_dtype, + cache_config.mamba_ssm_cache_dtype, + use_replayssm_spec=cache_config.use_replayssm_spec, + ) + return MambaStateDtypeCalculator.mamba2_cached_state_dtype( vllm_config.model_config.dtype, - vllm_config.cache_config.mamba_cache_dtype, - vllm_config.cache_config.mamba_ssm_cache_dtype, + cache_config.mamba_cache_dtype, + cache_config.mamba_ssm_cache_dtype, + use_replayssm=cache_config.use_replayssm, ) @classmethod def get_mamba_state_shape_from_config( cls, vllm_config: "VllmConfig", - ) -> tuple[tuple[int, int], tuple[int, int, int]]: + ) -> tuple[tuple[int, ...], ...]: """Calculate shapes for Mamba's convolutional and state caches. Args: @@ -835,12 +844,28 @@ def get_mamba_state_shape_from_config( Tuple containing: - conv_state_shape: Shape for convolutional state cache - temporal_state_shape: Shape for state space model cache + - (when the state-and-output decode kernel is enabled) the + x_cache/dt_cache/B_cache ring-buffer shapes """ parallel_config = vllm_config.parallel_config + cache_config = vllm_config.cache_config hf_config = vllm_config.model_config.hf_config intermediate_size = hf_config.mamba_num_heads * hf_config.mamba_head_dim - return MambaStateShapeCalculator.mamba2_state_shape( + if cache_config.use_replayssm_spec: + return MambaStateShapeCalculator.mamba2_spec_cached_state_shape( + intermediate_size=intermediate_size, + tp_world_size=parallel_config.tensor_parallel_size, + n_groups=hf_config.n_groups, + num_heads=hf_config.mamba_num_heads, + head_dim=hf_config.mamba_head_dim, + state_size=hf_config.ssm_state_size, + conv_kernel=hf_config.conv_kernel, + num_spec=vllm_config.num_speculative_tokens, + use_replayssm_spec=cache_config.use_replayssm_spec, + replayssm_buffer_len=cache_config.replayssm_buffer_len, + ) + return MambaStateShapeCalculator.mamba2_cached_state_shape( intermediate_size=intermediate_size, tp_world_size=parallel_config.tensor_parallel_size, n_groups=hf_config.n_groups, @@ -849,6 +874,8 @@ def get_mamba_state_shape_from_config( state_size=hf_config.ssm_state_size, conv_kernel=hf_config.conv_kernel, num_spec=vllm_config.num_speculative_tokens, + use_replayssm=cache_config.use_replayssm, + replayssm_buffer_len=cache_config.replayssm_buffer_len, ) @classmethod diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index 2a7dc3e3621f..e809acac53bd 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -513,17 +513,25 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: def get_mamba_state_dtype_from_config( cls, vllm_config: "VllmConfig", - ) -> tuple[torch.dtype, torch.dtype]: - return MambaStateDtypeCalculator.gated_delta_net_state_dtype( + ) -> tuple[torch.dtype, ...]: + if vllm_config.cache_config.use_replayssm_spec: + return MambaStateDtypeCalculator.gated_delta_net_spec_cached_state_dtype( + vllm_config.model_config.dtype, + vllm_config.cache_config.mamba_cache_dtype, + vllm_config.cache_config.mamba_ssm_cache_dtype, + vllm_config.cache_config.use_replayssm_spec, + ) + return MambaStateDtypeCalculator.gated_delta_net_cached_state_dtype( vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype, vllm_config.cache_config.mamba_ssm_cache_dtype, + vllm_config.cache_config.use_replayssm, ) @classmethod def get_mamba_state_shape_from_config( cls, vllm_config: "VllmConfig" - ) -> tuple[tuple[int, int], tuple[int, int]]: + ) -> tuple[tuple[int, ...], ...]: parallel_config = vllm_config.parallel_config hf_config = vllm_config.model_config.hf_text_config tp_size = parallel_config.tensor_parallel_size @@ -532,13 +540,27 @@ def get_mamba_state_shape_from_config( if vllm_config.speculative_config else 0 ) - return MambaStateShapeCalculator.gated_delta_net_state_shape( + if vllm_config.cache_config.use_replayssm_spec: + return MambaStateShapeCalculator.gated_delta_net_spec_cached_state_shape( + tp_size, + hf_config.linear_num_key_heads, + hf_config.linear_num_value_heads, + hf_config.linear_key_head_dim, + hf_config.linear_value_head_dim, + hf_config.linear_conv_kernel_dim, + vllm_config.cache_config.use_replayssm_spec, + vllm_config.cache_config.replayssm_buffer_len, + num_spec, + ) + return MambaStateShapeCalculator.gated_delta_net_cached_state_shape( tp_size, hf_config.linear_num_key_heads, hf_config.linear_num_value_heads, hf_config.linear_key_head_dim, hf_config.linear_value_head_dim, hf_config.linear_conv_kernel_dim, + vllm_config.cache_config.use_replayssm, + vllm_config.cache_config.replayssm_buffer_len, num_spec, ) diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 22c6a3822870..b89e5baed199 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -448,6 +448,10 @@ class CommonAttentionMetadata: and for all rows outside async spec decode; optimistic for async-spec decode rows (assumes every draft was accepted). Not safe for kernels that need exact per-row context lengths on decode rows.""" + num_prompt_tokens_cpu: torch.Tensor | None = None + """(batch_size,) CPU prompt lengths, when available from the worker input + batch. Used by Mamba2 cached decode metadata to derive generated-token + counts without including prompt tokens.""" mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None """PrefixLM bidirectional ranges for multimodal tokens. Maps @@ -530,6 +534,7 @@ def unpadded( _num_computed_tokens_cpu=self._num_computed_tokens_cpu[:num_actual_reqs] if self._num_computed_tokens_cpu is not None else None, + num_prompt_tokens_cpu=maybe_slice_reqs(self.num_prompt_tokens_cpu), num_reqs=num_actual_reqs, num_actual_tokens=num_actual_tokens, max_query_len=self.max_query_len, @@ -547,6 +552,10 @@ def unpadded( dcp_local_seq_lens_cpu=maybe_slice_reqs(self.dcp_local_seq_lens_cpu), is_prefilling=maybe_slice_reqs(self.is_prefilling), rswa_prefix_lens=maybe_slice_reqs(self.rswa_prefix_lens), + seq_lens_cpu_upper_bound=maybe_slice_reqs(self.seq_lens_cpu_upper_bound), + positions=self.positions[:num_actual_tokens] + if self.positions is not None + else None, ) diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 340a304030ea..ad638d7723b0 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -65,6 +65,18 @@ class GDNAttentionMetadata: num_accepted_tokens: torch.Tensor | None = None # shape: [batch,] + # Per-decode-row ring write position for the cached decode kernel. + # shape: [num_decodes]; None unless use_replayssm is enabled. + write_pos_d: torch.Tensor | None = None + + # Cached-SPEC decode cursors: persistent, block-keyed (full (num_blocks,) + # fixed-address buffers indexed by spec_state_indices_tensor[:, 0]), advanced + # once per step by commit_gdn_replayssm_spec. None unless + # use_replayssm_spec is enabled. + spec_write_pos_d: torch.Tensor | None = None + spec_cache_base_d: torch.Tensor | None = None + spec_is_flush_d: torch.Tensor | None = None + # Pre-computed FLA chunk metadata (avoids GPU->CPU sync in prepare_chunk_indices) chunk_indices: torch.Tensor | None = None chunk_offsets: torch.Tensor | None = None @@ -165,6 +177,35 @@ def __init__( device=device, ) + # Cached decode kernel: persistent per-decode-row ring write position. + # write_pos is derived per request each step (decode_step % max_cache_len) + # so recycled paged blocks need no zero-init. + self.use_cached_kernel: bool = ( + vllm_config.cache_config.use_replayssm + ) + self.max_cache_len: int = vllm_config.cache_config.replayssm_buffer_len + if self.use_cached_kernel: + self.decode_write_pos_d: torch.Tensor = torch.empty( + (self.decode_cudagraph_max_bs,), + dtype=torch.int32, + device=device, + ) + + # Cached-SPEC decode: block-keyed cursors (sized num_gpu_blocks), + # allocated lazily on first build (num_gpu_blocks is unknown here), and + # advanced once per step by commit_gdn_replayssm_spec. + self.use_cache_spec_kernel: bool = ( + vllm_config.cache_config.use_replayssm_spec + ) + self.max_spec_len: int = 1 + self.num_spec + # L = B + max_spec_len history window; physical pow2 ring = next_pow2(L). + self.spec_flush_threshold = self.max_cache_len + self.max_spec_len + self.spec_cache_buf_len = 1 << (self.spec_flush_threshold - 1).bit_length() + self.cursor_device = device + self.spec_write_pos: torch.Tensor | None = None + self.spec_cache_base: torch.Tensor | None = None + self.spec_is_flush: torch.Tensor | None = None + def build( # type: ignore[override] self, common_prefix_len: int, @@ -410,6 +451,118 @@ def build( # type: ignore[override] f"num_decodes: {num_decodes}, num_spec_decodes: {num_spec_decodes}" ) + # Cached decode kernel: derive the per-request ring write position + # (write_pos = decode_step % max_cache_len). Only the non-spec decode + # path runs the cached kernel. + write_pos_d = None + spec_write_pos_d = None + spec_cache_base_d = None + spec_is_flush_d = None + if self.use_cached_kernel and spec_sequence_masks is None and num_decodes > 0: + num_prompt_tokens_cpu = m.num_prompt_tokens_cpu + num_computed_tokens_cpu = m._num_computed_tokens_cpu + if num_prompt_tokens_cpu is None or num_computed_tokens_cpu is None: + raise ValueError( + "use_replayssm requires CPU prompt and " + "computed-token counts to derive decode write positions" + ) + decode_steps_cpu = ( + num_computed_tokens_cpu[:num_decodes] + - num_prompt_tokens_cpu[:num_decodes] + ) + query_lens_cpu = ( + query_start_loc_cpu[1 : num_decodes + 1] + - query_start_loc_cpu[:num_decodes] + ) + valid_decode_rows = query_lens_cpu > 0 + if torch.any(decode_steps_cpu[valid_decode_rows] < 0).item(): + raise ValueError( + "use_replayssm requires decode-step counts that " + "exclude prompt tokens and start at zero" + ) + decode_steps_cpu = torch.where( + valid_decode_rows, + decode_steps_cpu, + torch.zeros_like(decode_steps_cpu), + ) + write_pos_cpu = torch.remainder(decode_steps_cpu, self.max_cache_len) + write_pos_d = async_tensor_h2d( + write_pos_cpu.to(torch.int32).tolist(), + dtype=torch.int32, + device=query_start_loc.device, + ) + + # Cached-SPEC decode: advance the block-keyed cursors once per step + # (commit-at-start, using the previous step's num_accepted), then reset + # first-decode rows. Runs on the UNPADDED spec tensors (block-keyed + # cursors are skipped for null blocks, so cudagraph padding is fine). + # The commit/reset run in build() (eager, not in the captured region); + # the cursors are full (num_gpu_blocks,) fixed-address buffers read by + # the captured verify kernel. + if self.use_cache_spec_kernel and num_spec_decodes > 0: + from vllm.model_executor.layers.fla.ops.gdn_replayssm_spec_decode import ( + commit_gdn_replayssm_spec, + reset_gdn_replayssm_spec_cursors, + ) + + assert spec_state_indices_tensor is not None + assert num_accepted_tokens is not None + if self.spec_write_pos is None: + n_blocks = self.vllm_config.cache_config.num_gpu_blocks + assert n_blocks is not None and n_blocks > 0, ( + "--use-replayssm-spec needs num_gpu_blocks at " + "build time to size the block-keyed cursor buffers" + ) + self.spec_write_pos = torch.zeros( + n_blocks, dtype=torch.int32, device=self.cursor_device + ) + self.spec_cache_base = torch.zeros( + n_blocks, dtype=torch.int32, device=self.cursor_device + ) + self.spec_is_flush = torch.zeros( + n_blocks, dtype=torch.int8, device=self.cursor_device + ) + sbi = spec_state_indices_tensor[:, 0] + commit_gdn_replayssm_spec( + self.spec_write_pos, + self.spec_cache_base, + self.spec_is_flush, + num_accepted_tokens.to(torch.int32), + sbi, + max_cache_len=self.spec_flush_threshold, + max_spec_len=self.max_spec_len, + cache_buf_len=self.spec_cache_buf_len, + ) + # prefill->decode reset for first-decode rows (cursors only; conv + # context lives in conv_state). A request's first spec verify has + # num_computed_tokens == num_prompt_tokens; that resets its (possibly + # recycled) block's cursors to write_pos=0. + num_prompt_tokens_cpu = m.num_prompt_tokens_cpu + if num_prompt_tokens_cpu is not None: + num_prompt_d = num_prompt_tokens_cpu.to( + context_lens_tensor.device, non_blocking=True + ) + first_decode_full = (context_lens_tensor == num_prompt_d).to( + torch.int8 + ) + spec_row_idx = ( + spec_sequence_masks_cpu.nonzero(as_tuple=True)[0] + .to(query_start_loc.device, non_blocking=True) + ) + first_decode_d = first_decode_full.index_select(0, spec_row_idx) + reset_gdn_replayssm_spec_cursors( + self.spec_write_pos, + self.spec_cache_base, + self.spec_is_flush, + first_decode_d, + sbi, + max_cache_len=self.spec_flush_threshold, + max_spec_len=self.max_spec_len, + ) + spec_write_pos_d = self.spec_write_pos + spec_cache_base_d = self.spec_cache_base + spec_is_flush_d = self.spec_is_flush + # Prepare per-request tensors for cudagraph. m.num_actual_tokens is # token-padded for FULL graph replay, but the GDN state/query/accepted # metadata below is indexed by request. @@ -482,6 +635,16 @@ def build( # type: ignore[override] non_spec_query_start_loc = self.non_spec_query_start_loc[: batch_size + 1] non_spec_query_start_loc[num_decodes + 1 :].fill_(non_spec_num_query_tokens) + if self.use_cached_kernel: + assert write_pos_d is not None + self.decode_write_pos_d[:num_decodes].copy_( + write_pos_d, non_blocking=True + ) + write_pos_d = self.decode_write_pos_d[:batch_size] + # Padded rows map to NULL_BLOCK_ID and hit the kernel's early + # return, so their write position is never read; zero is fine. + write_pos_d[num_decodes:].fill_(0) + attn_metadata = GDNAttentionMetadata( num_prefills=num_prefills, num_prefill_tokens=num_prefill_tokens, @@ -504,6 +667,10 @@ def build( # type: ignore[override] spec_token_indx=spec_token_indx, non_spec_token_indx=non_spec_token_indx, num_accepted_tokens=num_accepted_tokens, + write_pos_d=write_pos_d, + spec_write_pos_d=spec_write_pos_d, + spec_cache_base_d=spec_cache_base_d, + spec_is_flush_d=spec_is_flush_d, nums_dict=nums_dict, batch_ptr=batch_ptr, token_chunk_offset_ptr=token_chunk_offset_ptr, diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index 16e292e21d2f..204cb225f1c2 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -74,6 +74,21 @@ class BaseMambaAttentionMetadata: nums_dict: dict | None = None batch_ptr: torch.Tensor | None = None token_chunk_offset_ptr: torch.Tensor | None = None + write_pos_d: torch.Tensor | None = None + is_flush_d: torch.Tensor | None = None + # Shared per-step scratch for the output-only (output_only) variant: + # (decode_rows, ngroups, max_cache_len) fp32. None for the recurrent + # variant or when the cached kernel is disabled. + bc_pre_scratch: torch.Tensor | None = None + # cached-SPEC (hybrid) cursors: persistent, block-keyed (full (num_blocks,) + # buffers indexed by physical SSM block id), shared across all Mamba2 layers + # and advanced once per step by the commit. spec_bc_pre_scratch is the + # per-step (decode_rows, ngroups, max_cache_len, block_spec) fp32 scratch. + # All None unless the cached-spec kernel is enabled. + spec_write_pos_d: torch.Tensor | None = None + spec_post_origin_d: torch.Tensor | None = None + spec_is_flush_d: torch.Tensor | None = None + spec_bc_pre_scratch: torch.Tensor | None = None class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): @@ -98,6 +113,15 @@ def __init__( self.compilation_config = vllm_config.compilation_config self.num_spec_tokens: int = vllm_config.num_speculative_tokens self.use_spec_decode = self.num_spec_tokens > 0 + self.use_cached_kernel = vllm_config.cache_config.use_replayssm + self.max_cache_len = vllm_config.cache_config.replayssm_buffer_len + self.use_cache_spec_kernel = ( + vllm_config.cache_config.use_replayssm_spec + ) + self.max_spec_len = 1 + self.num_spec_tokens + # L = B + max_spec_len history window; physical pow2 ring = next_pow2(L). + self.spec_flush_threshold = self.max_cache_len + self.max_spec_len + self.spec_cache_buf_len = 1 << (self.spec_flush_threshold - 1).bit_length() assert isinstance(kv_cache_spec, MambaSpec) scheduler_config = vllm_config.scheduler_config @@ -158,6 +182,85 @@ def __init__( dtype=torch.int32, device=device, ) + if self.use_cached_kernel: + self.decode_write_pos_d: torch.Tensor = torch.empty( + (self.decode_cudagraph_max_bs,), + dtype=torch.int32, + device=device, + ) + self.decode_is_flush_d: torch.Tensor = torch.empty( + (self.decode_cudagraph_max_bs,), + dtype=torch.int8, + device=device, + ) + self.cached_kernel_variant: str = ( + vllm_config.cache_config.replayssm_route + ) + if self.cached_kernel_variant == "output_only": + # B_cache shape = (ngroups, max_cache_len, dstate). Index in + # MambaSpec.shapes is (conv_state, ssm_state, x_cache, + # dt_cache, B_cache) when the cached kernel is enabled. + if len(kv_cache_spec.shapes) < 5: + raise ValueError( + "output-only variant requires the 5-tensor Mamba2 " + "page (conv, ssm, x_cache, dt_cache, B_cache)" + ) + bc_ngroups = kv_cache_spec.shapes[4][0] + self.decode_bc_pre_scratch: torch.Tensor = torch.empty( + ( + self.decode_cudagraph_max_bs, + bc_ngroups, + self.max_cache_len, + ), + dtype=torch.float32, + device=device, + ) + else: + self.decode_bc_pre_scratch = None + else: + self.cached_kernel_variant = "state_and_output" + self.decode_bc_pre_scratch = None + + # cached-SPEC (hybrid): persistent block-keyed cursors are allocated + # lazily on first build() (they need num_gpu_blocks). The per-step CB + # scratch is fixed-address (CUDA-graph safe), sized here. ngroups is + # derived from the page shapes (conv, ssm, post_conv_cache, dt_cache). + self.spec_write_pos: torch.Tensor | None = None + self.spec_post_origin: torch.Tensor | None = None + self.spec_is_flush: torch.Tensor | None = None + self.decode_spec_bc_pre: torch.Tensor | None = None + if self.use_cache_spec_kernel: + if len(kv_cache_spec.shapes) < 4: + raise ValueError( + "cached-spec kernel requires the 4-tensor hybrid Mamba2 page " + "(conv, ssm, post_conv_cache, dt_cache)" + ) + local_nheads, head_dim, dstate = kv_cache_spec.shapes[1] + conv_dim_local = kv_cache_spec.shapes[2][1] + d_inner_local = local_nheads * head_dim + ngroups_local = (conv_dim_local - d_inner_local) // (2 * dstate) + block_spec = 1 << (max(1, self.max_spec_len) - 1).bit_length() + # This is a PER-STEP scratch consumed by the scatter on every decode + # step (eager AND cudagraph), indexed by pid_b in [0, num_decodes). + # It must therefore cover the max decode batch (max_num_seqs), NOT + # decode_cudagraph_max_bs -- the latter is 0 under enforce_eager + # (max_cudagraph_capture_size=0), which would make this scratch empty + # and the scatter write bc_pre[pid_b] out of bounds (IMA). Sizing by + # max_num_seqs is CUDA-graph safe: the captured [:num_decodes] slice + # shares the same (offset-0) base pointer and row strides. + spec_scratch_bs = max( + self.decode_cudagraph_max_bs, scheduler_config.max_num_seqs + ) + self.decode_spec_bc_pre = torch.empty( + ( + spec_scratch_bs, + ngroups_local, + self.spec_cache_buf_len, + block_spec, + ), + dtype=torch.float32, + device=device, + ) self._init_reorder_batch_threshold(1, self.use_spec_decode) if self.use_spec_decode: @@ -401,6 +504,11 @@ def _compute_common_metadata( has_prior_state = seq_lens_cpu > 1 prefill_to_decode = single_token_prefill_rows & has_prior_state if torch.any(prefill_to_decode).item(): + if self.use_cached_kernel and metadata.num_decodes > 0: + raise ValueError( + "--use-replayssm does not support single-token " + "prefill rows replayed through the decode path" + ) is_prefilling = is_prefilling.clone() is_prefilling[prefill_to_decode] = False common_attn_metadata = common_attn_metadata.replace( @@ -431,6 +539,8 @@ def _compute_common_metadata( # for causal_conv1d nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None + write_pos_d = None + is_flush_d = None if self.vllm_config.cache_config.mamba_cache_mode == "all": num_computed_tokens = common_attn_metadata.compute_num_computed_tokens() @@ -518,6 +628,137 @@ def _compute_common_metadata( num_reqs - num_prefills : num_reqs ] + if self.use_cached_kernel and num_decodes > 0: + num_prompt_tokens_cpu = common_attn_metadata.num_prompt_tokens_cpu + num_computed_tokens_cpu = common_attn_metadata._num_computed_tokens_cpu + if num_prompt_tokens_cpu is None or num_computed_tokens_cpu is None: + raise ValueError( + "--use-replayssm requires CPU prompt and " + "computed-token counts to derive decode write positions" + ) + decode_steps_cpu = ( + num_computed_tokens_cpu[:num_decodes] + - num_prompt_tokens_cpu[:num_decodes] + ) + query_lens_cpu = ( + common_attn_metadata.query_start_loc_cpu[1 : num_decodes + 1] + - common_attn_metadata.query_start_loc_cpu[:num_decodes] + ) + valid_decode_rows = query_lens_cpu > 0 + if torch.any(decode_steps_cpu[valid_decode_rows] < 0).item(): + raise ValueError( + "--use-replayssm requires decode-step counts " + "that exclude prompt tokens and start at zero" + ) + decode_steps_cpu = torch.where( + valid_decode_rows, + decode_steps_cpu, + torch.zeros_like(decode_steps_cpu), + ) + write_pos_cpu = torch.remainder(decode_steps_cpu, self.max_cache_len) + is_flush_cpu = (write_pos_cpu == self.max_cache_len - 1).to(torch.int8) + write_pos_d = async_tensor_h2d( + write_pos_cpu.to(torch.int32).tolist(), + dtype=torch.int32, + device=common_attn_metadata.query_start_loc.device, + ) + is_flush_d = async_tensor_h2d( + is_flush_cpu.tolist(), + dtype=torch.int8, + device=common_attn_metadata.query_start_loc.device, + ) + + bc_pre_scratch = None + if ( + self.use_cached_kernel + and self.cached_kernel_variant == "output_only" + and self.decode_bc_pre_scratch is not None + and num_decodes > 0 + ): + bc_pre_scratch = self.decode_bc_pre_scratch[:num_decodes] + + # cached-SPEC (hybrid): commit-at-start advances the persistent + # block-keyed cursors using the previous step's num_accepted_tokens, + # then first-decode rows are reset. The kernels read the full + # (num_gpu_blocks,) cursor buffers, indexed by physical SSM block id. + spec_write_pos_d = None + spec_post_origin_d = None + spec_is_flush_d = None + spec_bc_pre_scratch = None + if ( + self.use_cache_spec_kernel + and num_decodes > 0 + and self.use_spec_decode + and num_accepted_tokens is not None + ): + from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_spec import ( # noqa: E501 + commit_replayssm_spec, + reset_replayssm_spec_cursors, + ) + + cursor_device = common_attn_metadata.query_start_loc.device + if self.spec_write_pos is None: + n_blocks = self.vllm_config.cache_config.num_gpu_blocks + assert n_blocks is not None and n_blocks > 0, ( + "--use-replayssm-spec needs num_gpu_blocks at " + "build time to size the block-keyed cursor buffers" + ) + self.spec_write_pos = torch.zeros( + n_blocks, dtype=torch.int32, device=cursor_device + ) + self.spec_post_origin = torch.zeros( + n_blocks, dtype=torch.int32, device=cursor_device + ) + self.spec_is_flush = torch.zeros( + n_blocks, dtype=torch.int8, device=cursor_device + ) + sbi = state_indices_tensor_d[:, 0] + commit_replayssm_spec( + self.spec_write_pos, + self.spec_post_origin, + self.spec_is_flush, + num_accepted_tokens.to(torch.int32), + sbi, + max_cache_len=self.spec_flush_threshold, + max_spec_len=self.max_spec_len, + cache_buf_len=self.spec_cache_buf_len, + ) + # prefill->decode reset for first-decode rows (cursors only; no conv + # seed -- conv_state carries context). A request's first spec verify + # has num_computed_tokens == num_prompt_tokens; that resets its + # (possibly recycled) block's cursors to write_pos=0, and -- because + # the commit above runs BEFORE this reset -- also undoes any write_pos + # the first-decode commit advanced from the freshly-zeroed cursor. + # Derive the mask from the DEVICE-side compute_num_computed_tokens() + # (always populated), NOT _num_computed_tokens_cpu, which is None on + # the spec verify path -> the old guard silently skipped the reset, + # leaving recycled blocks with stale cursors and fresh blocks with a + # wrong first-decode write_pos (coherent-but-divergent output + + # acceptance drop). + num_prompt_tokens_cpu = common_attn_metadata.num_prompt_tokens_cpu + if num_prompt_tokens_cpu is not None: + ctx_lens = common_attn_metadata.compute_num_computed_tokens() + num_prompt_d = num_prompt_tokens_cpu.to( + ctx_lens.device, non_blocking=True + ) + first_decode_d = ( + ctx_lens[:num_decodes] == num_prompt_d[:num_decodes] + ).to(torch.int8) + reset_replayssm_spec_cursors( + self.spec_write_pos, + self.spec_post_origin, + self.spec_is_flush, + first_decode_d, + sbi, + max_cache_len=self.spec_flush_threshold, + max_spec_len=self.max_spec_len, + ) + spec_write_pos_d = self.spec_write_pos + spec_post_origin_d = self.spec_post_origin + spec_is_flush_d = self.spec_is_flush + if self.decode_spec_bc_pre is not None: + spec_bc_pre_scratch = self.decode_spec_bc_pre[:num_decodes] + metadata = self.metadata_cls( num_prefills=num_prefills, num_prefill_tokens=num_prefill_tokens, @@ -527,6 +768,13 @@ def _compute_common_metadata( has_initial_states_p=has_initial_states_p, state_indices_tensor_p=state_indices_tensor_p, state_indices_tensor_d=state_indices_tensor_d, + write_pos_d=write_pos_d, + is_flush_d=is_flush_d, + bc_pre_scratch=bc_pre_scratch, + spec_write_pos_d=spec_write_pos_d, + spec_post_origin_d=spec_post_origin_d, + spec_is_flush_d=spec_is_flush_d, + spec_bc_pre_scratch=spec_bc_pre_scratch, num_accepted_tokens=num_accepted_tokens, query_start_loc_d=query_start_loc_d, block_idx_last_scheduled_token=block_idx_last_scheduled_token, @@ -561,6 +809,17 @@ def _update_metadata_for_cudagraph_capture( block_idx_last_scheduled_token_prev_step = ( metadata.block_idx_last_scheduled_token_prev_step ) + write_pos_d = metadata.write_pos_d + is_flush_d = metadata.is_flush_d + bc_pre_scratch = metadata.bc_pre_scratch + # cached-spec cursors are full (num_blocks,) fixed-address buffers indexed + # by physical block id, so they need NO per-batch padding (padding rows + # carry NULL_BLOCK_ID in state_indices and are skipped by the kernels). + # Only the per-row CB scratch is re-sliced to the padded batch. + spec_write_pos_d = metadata.spec_write_pos_d + spec_post_origin_d = metadata.spec_post_origin_d + spec_is_flush_d = metadata.spec_is_flush_d + spec_bc_pre_scratch = metadata.spec_bc_pre_scratch if ( metadata.num_prefills == 0 and metadata.num_decodes <= self.decode_cudagraph_max_bs @@ -622,11 +881,44 @@ def _update_metadata_for_cudagraph_capture( ) block_idx_last_scheduled_token_prev_step[metadata.num_decodes :] = 0 + if self.use_cached_kernel: + assert write_pos_d is not None + assert is_flush_d is not None + self.decode_write_pos_d[: metadata.num_decodes].copy_( + write_pos_d[: metadata.num_decodes], + non_blocking=True, + ) + write_pos_d = self.decode_write_pos_d[:padded_bs] + write_pos_d[metadata.num_decodes :] = 0 + + self.decode_is_flush_d[: metadata.num_decodes].copy_( + is_flush_d[: metadata.num_decodes], + non_blocking=True, + ) + is_flush_d = self.decode_is_flush_d[:padded_bs] + is_flush_d[metadata.num_decodes :] = 0 + + if ( + self.cached_kernel_variant == "output_only" + and self.decode_bc_pre_scratch is not None + ): + bc_pre_scratch = self.decode_bc_pre_scratch[:padded_bs] + + if self.use_cache_spec_kernel and self.decode_spec_bc_pre is not None: + spec_bc_pre_scratch = self.decode_spec_bc_pre[:padded_bs] + return replace( metadata, state_indices_tensor_d=state_indices_tensor_d, query_start_loc_d=query_start_loc_d, num_accepted_tokens=num_accepted_tokens, + write_pos_d=write_pos_d, + is_flush_d=is_flush_d, + bc_pre_scratch=bc_pre_scratch, + spec_write_pos_d=spec_write_pos_d, + spec_post_origin_d=spec_post_origin_d, + spec_is_flush_d=spec_is_flush_d, + spec_bc_pre_scratch=spec_bc_pre_scratch, block_idx_last_scheduled_token=block_idx_last_scheduled_token, block_idx_last_computed_token=block_idx_last_computed_token, block_idx_last_scheduled_token_prev_step=( diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 53795951cb26..c19252084f16 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -2377,6 +2377,7 @@ def _get_block_table(kv_cache_gid: int): _seq_lens_cpu=seq_lens_cpu, _num_computed_tokens_cpu=num_computed_tokens_cpu, seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, + num_prompt_tokens_cpu=num_prompt_tokens_cpu, num_reqs=num_reqs_padded, num_actual_tokens=num_tokens_padded, max_query_len=max_query_len, diff --git a/vllm/v1/worker/ubatch_utils.py b/vllm/v1/worker/ubatch_utils.py index f4a76529023c..cc7cfb1490a2 100644 --- a/vllm/v1/worker/ubatch_utils.py +++ b/vllm/v1/worker/ubatch_utils.py @@ -193,6 +193,11 @@ def _make_metadata_with_slice( if attn_metadata._num_computed_tokens_cpu is not None else None ) + num_prompt_tokens_cpu = ( + attn_metadata.num_prompt_tokens_cpu[request_slice] + if attn_metadata.num_prompt_tokens_cpu is not None + else None + ) if splits_last_request: # NOTE: We use start_locs (the original query_start_loc_cpu) to calculate @@ -243,8 +248,12 @@ def _make_metadata_with_slice( block_table_tensor=block_table_tensor, slot_mapping=slot_mapping, seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, + num_prompt_tokens_cpu=num_prompt_tokens_cpu, _seq_lens_cpu=seq_lens_cpu, _num_computed_tokens_cpu=num_computed_tokens_cpu, + positions=attn_metadata.positions[token_slice] + if attn_metadata.positions is not None + else None, ) From c09b71e56b2de6bd48bf4dfa8a37e77e3f552950 Mon Sep 17 00:00:00 2001 From: Johnny-Liou Date: Tue, 7 Jul 2026 15:07:52 -0400 Subject: [PATCH 02/23] [Test] ReplaySSM: add tensor-parallel correctness tests for standard decode Signed-off-by: Johnny-Liou --- .../test_replayssm_standard_decode_mamba2.py | 111 ++++++++++++++++++ tests/v1/e2e/test_replayssm_decode.py | 32 +++-- 2 files changed, 134 insertions(+), 9 deletions(-) diff --git a/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py b/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py index 7d730fefb18c..26635024b38c 100644 --- a/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py +++ b/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py @@ -448,3 +448,114 @@ def test_replayssm_standard_decode_with_batch_indices( assert torch.equal(state_cached[unused_states], state_before[unused_states]) assert torch.equal(state_baseline[unused_states], state_before[unused_states]) + + +# Geometries with (nheads, ngroups) both divisible by the tp below. +_TP_GEOMETRIES = [ + pytest.param((8, 64, 64, 4), id="small"), + pytest.param((96, 80, 128, 8), id="nano4b"), +] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("precision", _PRECISIONS) +@pytest.mark.parametrize("geometry", _TP_GEOMETRIES) +@pytest.mark.parametrize("tp", [2]) +def test_replayssm_standard_decode_tp_head_shard_equivalence( + precision: tuple[torch.dtype, torch.dtype], + geometry: tuple[int, int, int, int], + tp: int, +): + """Tensor-parallel correctness at the kernel boundary (single GPU). + + The kernel has no cross-rank communication, so head sharding must be exactly + separable: one run over all heads must equal concatenating ``tp`` independent + per-rank runs on ``nheads // tp`` heads and ``ngroups // tp`` groups. This + guards the per-rank divisors and group->head mapping the state-shape wiring + relies on. The real TP1==TP2 engine check lives in the v1/e2e suite.""" + state_dtype, act_dtype = precision + nheads, headdim, dstate, ngroups = geometry + assert nheads % tp == 0 and ngroups % tp == 0 + device = "cuda" + both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32 + rtol, atol = _tolerances(torch.float32 if both_fp32 else torch.bfloat16) + set_random_seed(0) + + batch = 4 + max_cache_len = 4 + num_steps = 2 * max_cache_len + 1 + nh_s = nheads // tp + ng_s = ngroups // tp + + # Shards slice the tied params; never .contiguous() -- that would drop the + # stride-0 broadcast the kernel's TIE_HDIM asserts require. + A = _tied_A(nheads, headdim, dstate, device) + dt_bias = _tied_dt_bias(nheads, headdim, device) + D = torch.randn(nheads, headdim, device=device) + + state_full = torch.randn( + batch, nheads, headdim, dstate, dtype=state_dtype, device=device) + x_cache, dt_cache, B_cache, _ = allocate_update_caches( + batch, nheads, ngroups, headdim, dstate, max_cache_len, device, + act_dtype, act_dtype) + bc_pre = torch.empty( + batch, ngroups, max_cache_len, device=device, dtype=torch.float32) + + # Per-rank shards, each seeded from the matching head slice so all start equal. + shard_state = [] + shard_caches = [] + for r in range(tp): + h0 = r * nh_s + shard_state.append(state_full[:, h0 : h0 + nh_s].contiguous()) + xc, dtc, bc, _ = allocate_update_caches( + batch, nh_s, ng_s, headdim, dstate, max_cache_len, device, + act_dtype, act_dtype) + bcp = torch.empty( + batch, ng_s, max_cache_len, device=device, dtype=torch.float32) + shard_caches.append((xc, dtc, bc, bcp)) + + write_pos = torch.zeros(batch, dtype=torch.int32, device=device) + for _ in range(num_steps): + x = torch.randn(batch, nheads, headdim, device=device, dtype=act_dtype) + dt = _tied_dt(batch, nheads, headdim, device, act_dtype) + B = torch.randn(batch, ngroups, dstate, device=device, dtype=act_dtype) + C = torch.randn(batch, ngroups, dstate, device=device, dtype=act_dtype) + z = torch.randn_like(x) + is_flush = write_pos == max_cache_len - 1 + + out_full = torch.empty_like(x) + selective_state_update_replayssm_output_only( + state_full, x, dt, A, B, C, D=D, z=z, dt_bias=dt_bias, + dt_softplus=True, x_cache=x_cache, dt_cache=dt_cache, + B_cache=B_cache, bc_pre=bc_pre, write_pos=write_pos, + is_flush=is_flush, max_cache_len=max_cache_len, out=out_full) + + for r in range(tp): + h0 = r * nh_s + g0 = r * ng_s + xc, dtc, bc, bcp = shard_caches[r] + out_shard = torch.empty( + batch, nh_s, headdim, device=device, dtype=act_dtype) + selective_state_update_replayssm_output_only( + shard_state[r], + x[:, h0 : h0 + nh_s].contiguous(), + dt[:, h0 : h0 + nh_s], + A[h0 : h0 + nh_s], + B[:, g0 : g0 + ng_s].contiguous(), + C[:, g0 : g0 + ng_s].contiguous(), + D=D[h0 : h0 + nh_s].contiguous(), + z=z[:, h0 : h0 + nh_s].contiguous(), + dt_bias=dt_bias[h0 : h0 + nh_s], + dt_softplus=True, x_cache=xc, dt_cache=dtc, B_cache=bc, + bc_pre=bcp, write_pos=write_pos, is_flush=is_flush, + max_cache_len=max_cache_len, out=out_shard) + + torch.testing.assert_close( + out_full[:, h0 : h0 + nh_s], out_shard, rtol=rtol, atol=atol) + if bool(is_flush.any()): + torch.testing.assert_close( + state_full[:, h0 : h0 + nh_s][is_flush], + shard_state[r][is_flush], rtol=rtol, atol=atol) + + write_pos = torch.where( + is_flush, torch.zeros_like(write_pos), write_pos + 1) diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 4f551f76376e..1188845781ac 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -5,14 +5,14 @@ import pytest from ...models.utils import check_logprobs_close -from ...utils import large_gpu_mark +from ...utils import large_gpu_mark, multi_gpu_test # Mamba2 (Nemotron) and GDN (Qwen3.5) hybrids. +MAMBA2_MODEL = "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16" +GDN_MODEL = "Qwen/Qwen3.5-4B" MODELS = [ - pytest.param( - "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16", marks=large_gpu_mark(min_gb=40) - ), - pytest.param("Qwen/Qwen3.5-4B", marks=large_gpu_mark(min_gb=40)), + pytest.param(MAMBA2_MODEL, marks=large_gpu_mark(min_gb=40)), + pytest.param(GDN_MODEL, marks=large_gpu_mark(min_gb=40)), ] PROMPTS = [ @@ -21,15 +21,16 @@ ] -@pytest.mark.parametrize("model_name", MODELS) -def test_replayssm_decode_matches_baseline(vllm_runner, model_name): - # ReplaySSM reconstructs the state in different fp arithmetic, so greedy ids - # can diverge at a near-tie; compare logprobs, not exact ids. +def _check_replayssm_parity(vllm_runner, model_name, *, tensor_parallel_size=1): + # Compare logprobs, not greedy ids: ReplaySSM's fp arithmetic can flip a + # near-tie. Baseline and ReplaySSM run at the same TP, so TP numerics are + # common-mode and only ReplaySSM varies. common = dict( max_model_len=1024, trust_remote_code=True, enable_prefix_caching=False, mamba_cache_mode="none", + tensor_parallel_size=tensor_parallel_size, ) with vllm_runner(model_name, **common) as llm: baseline = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5) @@ -44,3 +45,16 @@ def test_replayssm_decode_matches_baseline(vllm_runner, model_name): name_0="baseline", name_1="replayssm", ) + + +@pytest.mark.parametrize("model_name", MODELS) +def test_replayssm_decode_matches_baseline(vllm_runner, model_name): + _check_replayssm_parity(vllm_runner, model_name) + + +@multi_gpu_test(num_gpus=2) +@pytest.mark.parametrize("model_name", [MAMBA2_MODEL, GDN_MODEL]) +def test_replayssm_decode_matches_baseline_tp2(vllm_runner, model_name): + # Tensor-parallel correctness: ReplaySSM's caches and checkpoint state are + # sharded per rank, so TP2 decode must still match the baseline at TP2. + _check_replayssm_parity(vllm_runner, model_name, tensor_parallel_size=2) From fcda26218239538ea9d562d194f0d4bcef0fc5f8 Mon Sep 17 00:00:00 2001 From: Johnny-Liou Date: Tue, 7 Jul 2026 16:09:16 -0400 Subject: [PATCH 03/23] [Kernel] ReplaySSM: decoupled dstate-tiled Mamba2 output_only decode kernel Signed-off-by: Johnny-Liou --- .../layers/mamba/ops/replayssm_config.py | 18 ++-- ...tive_state_update_replayssm_output_only.py | 90 +++++++++++++------ 2 files changed, 74 insertions(+), 34 deletions(-) diff --git a/vllm/model_executor/layers/mamba/ops/replayssm_config.py b/vllm/model_executor/layers/mamba/ops/replayssm_config.py index 6d37106ed66b..06801ebacaed 100644 --- a/vllm/model_executor/layers/mamba/ops/replayssm_config.py +++ b/vllm/model_executor/layers/mamba/ops/replayssm_config.py @@ -62,6 +62,14 @@ def _mamba2_spec_flush(dstate, base_block, max_spec_len, is_blackwell): return 32, 1, _dstate_tile(dstate, 128), 2 +def _mamba2_output_only(dstate, L, is_blackwell): + # (block_size_m, num_warps, nf_dstate_tile, fl_dstate_tile, num_stages); + # decoupled dstate tiling, serving-batch optimum, dtype-independent per device. + if is_blackwell: + return 64, 1, _dstate_tile(dstate, 32), _dstate_tile(dstate, 64), 2 + return 16, 1, _dstate_tile(dstate, 64), _dstate_tile(dstate, 128), 2 + + def _gdn_spec(max_spec_len, is_blackwell): # (block_v, num_warps, nk, num_stages); verify and flush share a config. return 64, 1, (4 if max_spec_len >= 6 else 2), 2 @@ -77,12 +85,8 @@ def _l_bucket(cache_len: int) -> int: return 32 -# Standard-decode launch configs keyed by L bucket, tuned for the FP32 production -# state (mamba_ssm_dtype=float32). bsm=64 is intentionally avoided here: it is the -# bf16/B300 optimum but catastrophic (register spill, ~0.3x) at fp32. These seed -# values preserve the previously shipped per-kernel constants; the config sweep -# (profiling_pr/config_sweep) refines them per bucket. -_OUTPUT_ONLY_BY_L = {8: (16, 1), 16: (16, 1), 32: (16, 1)} # (block_size_m, num_warps) +# state_and_output stays un-tiled (retained only for precision experiments); the +# output_only route uses the decoupled dstate-tiled config in _mamba2_output_only. _STATE_AND_OUTPUT_BY_L = {8: (32, 1), 16: (32, 1), 32: (32, 1)} # GDN standard decode: (block_v, num_warps, num_stages, nk). L-flat in the sweep. _GDN_DECODE_BY_L = {8: (64, 1, 3, 2), 16: (64, 1, 3, 2), 32: (64, 1, 3, 2)} @@ -109,7 +113,7 @@ def get_replayssm_config(kernel: str, **shape) -> tuple: shape["dstate"], shape["base_block"], shape["max_spec_len"], bw ) if kernel == "mamba2_output_only": - return _OUTPUT_ONLY_BY_L[_l_bucket(shape.get("L", 16))] + return _mamba2_output_only(shape["dstate"], shape.get("L", 16), bw) if kernel == "mamba2_state_and_output": return _STATE_AND_OUTPUT_BY_L[_l_bucket(shape.get("L", 16))] if kernel == "gdn_decode": diff --git a/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py b/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py index 14eacabe6719..1f2442f25a81 100644 --- a/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py +++ b/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py @@ -207,6 +207,10 @@ def _replayssm_output_only_kernel( BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K_CACHE: tl.constexpr, BLOCK_SIZE_K_DOT: tl.constexpr, + NF_DSTATE_TILE: tl.constexpr, + NF_NDS: tl.constexpr, + FL_DSTATE_TILE: tl.constexpr, + FL_NDS: tl.constexpr, # heuristic-computed BLOCK_SIZE_DSTATE: tl.constexpr, HAS_DT_BIAS: tl.constexpr, @@ -227,7 +231,7 @@ def _replayssm_output_only_kernel( state_batch_idx = pid_b offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + offs_n_full = tl.arange(0, BLOCK_SIZE_DSTATE) # Buffer cursor (number of cached tokens so far) and the flush flag. write_pos = tl.load(write_pos_ptr + pid_b).to(tl.int64) @@ -245,8 +249,8 @@ def _replayssm_output_only_kernel( B_cache_ptr += state_batch_idx * stride_B_cache_batch + (pid_h // nheads_ngroups_ratio) * stride_B_cache_group bc_pre_ptr += pid_b * stride_bc_pre_batch + (pid_h // nheads_ngroups_ratio) * stride_bc_pre_group - # Current-token dt (+ bias, softplus), scalar A, current x / C, checkpoint - # state S_0, and current-token B (shared by both routes below). + # Current-token dt (+ bias, softplus), scalar A, and current x. C, the + # checkpoint state S_0, and current-token B are read per dstate tile below. dt_cur = tl.load(dt_ptr).to(tl.float32) if HAS_DT_BIAS: dt_cur += tl.load(dt_bias_ptr + pid_h * stride_dt_bias_head).to(tl.float32) @@ -254,10 +258,6 @@ def _replayssm_output_only_kernel( dt_cur = tl.where(dt_cur <= 20.0, softplus(dt_cur), dt_cur) A = tl.load(A_ptr + pid_h * stride_A_head).to(tl.float32) x_cur = tl.load(x_ptr + offs_m * stride_x_dim, mask=offs_m < dim, other=0.0) - C = tl.load(C_ptr + offs_n * stride_C_dstate, mask=offs_n < dstate, other=0.0).to(tl.float32) - state_ptrs = state_ptr + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate - state = tl.load(state_ptrs, mask=(offs_m[:, None] < dim) & (offs_n[None, :] < dstate), other=0.0) - B_cur = tl.load(B_ptr + offs_n * stride_B_dstate, mask=offs_n < dstate, other=0.0) if not is_flush: # Output-only route: read y without materializing the state, using the @@ -279,8 +279,20 @@ def _replayssm_output_only_kernel( x_all_cache = tl.load(x_all_cache_ptrs, mask=(offs_m[:, None] < dim) & (offs_k_cache[None, :] < write_pos), other=0.0) x_all_cache = tl.where(offs_k_cache[None, :] == write_pos, x_cur[:, None], x_all_cache) - # Decayed checkpoint readout plus the weighted sum of cached values. - checkpoint_out = tl.sum(state.to(tl.float32) * C[None, :], axis=1) * total_decay_cache + # Decayed checkpoint readout sum_n S_0(m,n) q(n), streamed over NF dstate + # tiles so the (M, N) state slice is never held whole. + offs_nt = tl.arange(0, NF_DSTATE_TILE) + ck_acc = tl.zeros([BLOCK_SIZE_M], dtype=tl.float32) + for i in tl.static_range(NF_NDS): + offs_n = i * NF_DSTATE_TILE + offs_nt + nmask = offs_n < dstate + st = tl.load( + state_ptr + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate, + mask=(offs_m[:, None] < dim) & nmask[None, :], other=0.0, + ) + c_chunk = tl.load(C_ptr + offs_n * stride_C_dstate, mask=nmask, other=0.0).to(tl.float32) + ck_acc += tl.sum(st.to(tl.float32) * c_chunk[None, :], axis=1) + checkpoint_out = ck_acc * total_decay_cache bc_cache = tl.load(bc_pre_ptr + offs_k_cache * stride_bc_pre_pos, mask=offs_k_cache <= write_pos, other=0.0) cache_out = tl.sum(x_all_cache.to(tl.float32) * (scale_cache * bc_cache)[None, :], axis=1) out = checkpoint_out + cache_out @@ -288,12 +300,13 @@ def _replayssm_output_only_kernel( # Append the current token (x, dt, B) into the buffer at write_pos. tl.store(x_cache_ptr + offs_m * stride_x_cache_dim + write_pos * stride_x_cache_pos, x_cur, mask=offs_m < dim) if pid_m == 0: + B_cur = tl.load(B_ptr + offs_n_full * stride_B_dstate, mask=offs_n_full < dstate, other=0.0) tl.store(dt_cache_ptr + write_pos * stride_dt_cache_pos, dt_cur) - tl.store(B_cache_ptr + write_pos * stride_B_cache_pos + offs_n * stride_B_cache_dstate, B_cur, mask=offs_n < dstate) + tl.store(B_cache_ptr + write_pos * stride_B_cache_pos + offs_n_full * stride_B_cache_dstate, B_cur, mask=offs_n_full < dstate) else: # Flush step: state route. Reconstruct the state from cached inputs, # S_t = total_decay * S_0 + sum_j s_j (v_j k_j^T), persist it as the new - # checkpoint, then read y = S_t q. + # checkpoint, then read y = S_t q -- streamed over FL dstate tiles. offs_k_dot = tl.arange(0, BLOCK_SIZE_K_DOT) dt_all_dot = tl.load(dt_cache_ptr + offs_k_dot * stride_dt_cache_pos, mask=offs_k_dot < write_pos, other=0.0).to(tl.float32) dt_all_dot = tl.where(offs_k_dot == write_pos, dt_cur, dt_all_dot) @@ -303,24 +316,37 @@ def _replayssm_output_only_kernel( scale_dot = dt_all_dot * tl.exp(dA_total_dot - dA_cumsum_dot) scale_dot = tl.where(offs_k_dot <= write_pos, scale_dot, 0.0) - # Gather buffered x and B over the window (history + current token). + # Gather buffered x over the window (history + current token). x_all_dot_ptrs = x_cache_ptr + offs_m[:, None] * stride_x_cache_dim + offs_k_dot[None, :] * stride_x_cache_pos x_all_dot = tl.load(x_all_dot_ptrs, mask=(offs_m[:, None] < dim) & (offs_k_dot[None, :] < write_pos), other=0.0) x_all_dot = tl.where(offs_k_dot[None, :] == write_pos, x_cur[:, None], x_all_dot) - B_all_dot_ptrs = B_cache_ptr + offs_k_dot[:, None] * stride_B_cache_pos + offs_n[None, :] * stride_B_cache_dstate - B_all_dot = tl.load(B_all_dot_ptrs, mask=(offs_k_dot[:, None] < write_pos) & (offs_n[None, :] < dstate), other=0.0) - B_all_dot = tl.where(offs_k_dot[:, None] == write_pos, B_cur[None, :], B_all_dot) - - # Reconstruct the state from cached inputs and store it as the checkpoint. - B_scaled = (B_all_dot.to(tl.float32) * scale_dot[:, None]).to(x_ptr.dtype.element_ty) - # tf32x3 keeps fp32 parity with the elementwise baseline (plain tf32 on - # fp32 inputs drifts ~1e-2); bf16/fp16 inputs are unaffected by this flag. - delta_state = tl.dot( - x_all_dot.to(x_ptr.dtype.element_ty), B_scaled, input_precision="tf32x3" - ) - state_new = state.to(tl.float32) * total_decay_dot + delta_state.to(tl.float32) - tl.store(state_ptrs, state_new.to(state.dtype), mask=(offs_m[:, None] < dim) & (offs_n[None, :] < dstate)) - out = tl.sum(state_new * C[None, :], axis=1) + x_all_ty = x_all_dot.to(x_ptr.dtype.element_ty) + + # Distinct tile locals (_f) from the nf branch: differing tile widths + # would force a shape-mismatched merge at the if/else exit. + offs_nt_f = tl.arange(0, FL_DSTATE_TILE) + out = tl.zeros([BLOCK_SIZE_M], dtype=tl.float32) + for i in tl.static_range(FL_NDS): + offs_n_f = i * FL_DSTATE_TILE + offs_nt_f + nmask_f = offs_n_f < dstate + # Gather buffered B over the window (history + current token). + B_all_dot = tl.load( + B_cache_ptr + offs_k_dot[:, None] * stride_B_cache_pos + offs_n_f[None, :] * stride_B_cache_dstate, + mask=(offs_k_dot[:, None] < write_pos) & nmask_f[None, :], other=0.0, + ) + B_cur_tile = tl.load(B_ptr + offs_n_f * stride_B_dstate, mask=nmask_f, other=0.0) + B_all_dot = tl.where(offs_k_dot[:, None] == write_pos, B_cur_tile[None, :], B_all_dot) + # tf32x3 keeps fp32 parity with the elementwise baseline (plain tf32 on + # fp32 inputs drifts ~1e-2); bf16/fp16 inputs are unaffected by this flag. + B_scaled = (B_all_dot.to(tl.float32) * scale_dot[:, None]).to(x_ptr.dtype.element_ty) + delta_state = tl.dot(x_all_ty, B_scaled, input_precision="tf32x3") + state_ptrs = state_ptr + offs_m[:, None] * stride_state_dim + offs_n_f[None, :] * stride_state_dstate + state_mask = (offs_m[:, None] < dim) & nmask_f[None, :] + st_f = tl.load(state_ptrs, mask=state_mask, other=0.0) + state_new = st_f.to(tl.float32) * total_decay_dot + delta_state.to(tl.float32) + tl.store(state_ptrs, state_new.to(st_f.dtype), mask=state_mask) + c_chunk_f = tl.load(C_ptr + offs_n_f * stride_C_dstate, mask=nmask_f, other=0.0).to(tl.float32) + out += tl.sum(state_new * c_chunk_f[None, :], axis=1) # Skip connection (D) and output gate (z). if HAS_D: @@ -428,9 +454,14 @@ def selective_state_update_replayssm_output_only( block_size_k_cache = max(1, triton.next_power_of_2(max_cache_len)) block_size_k_dot = max(16, block_size_k_cache) - block_size_m, num_warps = get_replayssm_config( + block_size_m, num_warps, nf_tile, fl_tile, num_stages = get_replayssm_config( "mamba2_output_only", dstate=dstate, L=max_cache_len ) + bs_dstate = triton.next_power_of_2(dstate) + nf_dstate_tile = max(16, min(nf_tile, bs_dstate)) + nf_nds = triton.cdiv(bs_dstate, nf_dstate_tile) + fl_dstate_tile = max(16, min(fl_tile, bs_dstate)) + fl_nds = triton.cdiv(bs_dstate, fl_dstate_tile) grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads) z_strides = (z.stride(0), z.stride(1), z.stride(2)) if z is not None else (0, 0, 0) @@ -542,7 +573,12 @@ def selective_state_update_replayssm_output_only( block_size_m, block_size_k_cache, block_size_k_dot, + nf_dstate_tile, + nf_nds, + fl_dstate_tile, + fl_nds, num_warps=num_warps, + num_stages=num_stages, ) if not has_heads: From 5de8d976179ffab8e537c4362f00ee5f2f8b5234 Mon Sep 17 00:00:00 2001 From: Johnny-Liou Date: Wed, 8 Jul 2026 13:39:57 -0400 Subject: [PATCH 04/23] [Refactor] ReplaySSM: split cached state dtype/shape calculators into dedicated replayssm methods, drop flag params Signed-off-by: Johnny-Liou --- vllm/model_executor/layers/mamba/gdn/base.py | 12 +- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 31 +++-- .../layers/mamba/mamba_mixer2.py | 53 ++++---- .../layers/mamba/mamba_utils.py | 115 ++++++------------ vllm/model_executor/models/nemotron_h.py | 31 +++-- vllm/model_executor/models/qwen3_5.py | 30 +++-- 6 files changed, 140 insertions(+), 132 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn/base.py b/vllm/model_executor/layers/mamba/gdn/base.py index ac8a62bfa30b..3cb34a73475e 100644 --- a/vllm/model_executor/layers/mamba/gdn/base.py +++ b/vllm/model_executor/layers/mamba/gdn/base.py @@ -52,15 +52,19 @@ def mamba_type(self) -> MambaAttentionBackendEnum: def get_state_dtype(self) -> tuple[torch.dtype, ...]: if self.cache_config.use_replayssm_spec: - return MambaStateDtypeCalculator.gated_delta_net_spec_cached_state_dtype( + return MambaStateDtypeCalculator.gated_delta_net_replayssm_spec_state_dtype( self.model_config.dtype, self.cache_config.mamba_cache_dtype, self.cache_config.mamba_ssm_cache_dtype, - self.cache_config.use_replayssm_spec, ) - return MambaStateDtypeCalculator.gated_delta_net_cached_state_dtype( + elif self.cache_config.use_replayssm: + return MambaStateDtypeCalculator.gated_delta_net_replayssm_state_dtype( + self.model_config.dtype, + self.cache_config.mamba_cache_dtype, + self.cache_config.mamba_ssm_cache_dtype, + ) + return MambaStateDtypeCalculator.gated_delta_net_state_dtype( self.model_config.dtype, self.cache_config.mamba_cache_dtype, self.cache_config.mamba_ssm_cache_dtype, - self.cache_config.use_replayssm, ) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 02bbc0e5c39c..5858747985db 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -26,8 +26,8 @@ ) from vllm.model_executor.layers.fla.ops import ( fused_post_conv_prep, - fused_recurrent_gated_delta_rule_replayssm, fused_recurrent_gated_delta_rule_packed_decode, + fused_recurrent_gated_delta_rule_replayssm, fused_sigmoid_gating_delta_rule_update, ) from vllm.model_executor.layers.fla.ops.chunk import l2norm_fwd @@ -423,26 +423,34 @@ def get_state_shape( self, ) -> tuple[tuple[int, ...], ...]: if self.cache_config.use_replayssm_spec: - return MambaStateShapeCalculator.gated_delta_net_spec_cached_state_shape( + return MambaStateShapeCalculator.gated_delta_net_replayssm_spec_state_shape( self.tp_size, self.num_k_heads, self.num_v_heads, self.head_k_dim, self.head_v_dim, self.conv_kernel_size, - self.cache_config.use_replayssm_spec, self.cache_config.replayssm_buffer_len, self.num_spec, ) - return MambaStateShapeCalculator.gated_delta_net_cached_state_shape( + elif self.cache_config.use_replayssm: + return MambaStateShapeCalculator.gated_delta_net_replayssm_state_shape( + self.tp_size, + self.num_k_heads, + self.num_v_heads, + self.head_k_dim, + self.head_v_dim, + self.conv_kernel_size, + self.cache_config.replayssm_buffer_len, + self.num_spec, + ) + return MambaStateShapeCalculator.gated_delta_net_state_shape( self.tp_size, self.num_k_heads, self.num_v_heads, self.head_k_dim, self.head_v_dim, self.conv_kernel_size, - self.cache_config.use_replayssm, - self.cache_config.replayssm_buffer_len, self.num_spec, ) @@ -580,9 +588,7 @@ def __init__( # Cached-SPEC decode kernel (gdn_replayssm_spec_decode). When enabled, the # GDN page grows to the same 5-tuple (fp32 checkpoint) and the spec verify # path decodes through the circular cached kernel. - self.use_cache_spec_kernel = ( - self.cache_config.use_replayssm_spec - ) + self.use_cache_spec_kernel = self.cache_config.use_replayssm_spec self.max_spec_len = 1 + self.num_spec compilation_config = get_current_vllm_config().compilation_config @@ -1393,7 +1399,8 @@ def _forward_core( num_accepted_tokens=num_accepted_tokens, query_start_loc=spec_query_start_loc, # Spec verify window = 1 + num_spec. Use the constant rather than - # the block-table width so the cached-spec path can request num_speculative_blocks=0 + # the block-table width so the cached-spec path can request + # num_speculative_blocks=0 max_query_len=self.max_spec_len, validate_data=False, ) @@ -1438,9 +1445,7 @@ def _forward_core( if spec_sequence_masks is not None and self.use_cache_spec_kernel: query_spec, key_spec, value_spec = None, None, None else: - query_spec, key_spec, value_spec = self.rearrange_mixed_qkv( - mixed_qkv_spec - ) + query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec) # Split mixed non-spec-decode+prefill to process independently split_non_spec = ( diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index cfdff5edc8d3..b3bd303bdf12 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -32,18 +32,18 @@ causal_conv1d_update, ) from vllm.model_executor.layers.mamba.ops.layernorm_gated import rms_norm_gated -from vllm.model_executor.layers.mamba.ops.ssd_combined import ( - mamba_chunk_scan_combined_varlen, -) -from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_state_and_output import ( # noqa: E501 - selective_state_update_replayssm_state_and_output, -) from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_output_only import ( # noqa: E501 selective_state_update_replayssm_output_only, ) from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_spec import ( selective_state_update_replayssm_spec, ) +from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_state_and_output import ( # noqa: E501 + selective_state_update_replayssm_state_and_output, +) +from vllm.model_executor.layers.mamba.ops.ssd_combined import ( + mamba_chunk_scan_combined_varlen, +) from vllm.model_executor.layers.mamba.ops.ssu_dispatch import selective_state_update from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.model_loader.weight_utils import ( @@ -519,9 +519,7 @@ def __init__( else "state_and_output" ) self.use_cache_spec_kernel = ( - cache_config.use_replayssm_spec - if cache_config is not None - else False + cache_config.use_replayssm_spec if cache_config is not None else False ) if ( self.use_cache_kernel or self.use_cache_spec_kernel @@ -755,8 +753,7 @@ def conv_ssm_forward( elif self.use_cache_kernel: if len(self.kv_cache) != 5: raise ValueError( - "Mamba2 cached decode kernel requires five Mamba " - "state tensors" + "Mamba2 cached decode kernel requires five Mamba state tensors" ) x_cache, dt_cache, B_cache = self.kv_cache[2:] else: @@ -1142,8 +1139,7 @@ def conv_ssm_forward( elif self.use_cache_kernel: if is_mamba_cache_all: raise ValueError( - "Mamba2 cached decode kernel requires " - "mamba_cache_mode='none'" + "Mamba2 cached decode kernel requires mamba_cache_mode='none'" ) if num_accepted_tokens is not None or query_start_loc_d is not None: raise ValueError( @@ -1230,22 +1226,38 @@ def get_state_dtype(self) -> tuple[torch.dtype, ...]: assert self.model_config is not None assert self.cache_config is not None if self.use_cache_spec_kernel: - return MambaStateDtypeCalculator.mamba2_spec_cached_state_dtype( + return MambaStateDtypeCalculator.mamba2_replayssm_spec_state_dtype( + self.model_config.dtype, + self.cache_config.mamba_cache_dtype, + self.cache_config.mamba_ssm_cache_dtype, + ) + elif self.use_cache_kernel: + return MambaStateDtypeCalculator.mamba2_replayssm_state_dtype( self.model_config.dtype, self.cache_config.mamba_cache_dtype, self.cache_config.mamba_ssm_cache_dtype, - use_replayssm_spec=self.use_cache_spec_kernel, ) - return MambaStateDtypeCalculator.mamba2_cached_state_dtype( + return MambaStateDtypeCalculator.mamba2_state_dtype( self.model_config.dtype, self.cache_config.mamba_cache_dtype, self.cache_config.mamba_ssm_cache_dtype, - use_replayssm=self.use_cache_kernel, ) def get_state_shape(self) -> tuple[tuple[int, ...], ...]: if self.use_cache_spec_kernel: - return MambaStateShapeCalculator.mamba2_spec_cached_state_shape( + return MambaStateShapeCalculator.mamba2_replayssm_spec_state_shape( + intermediate_size=self.intermediate_size, + tp_world_size=get_tensor_model_parallel_world_size(), + n_groups=self.n_groups, + num_heads=self.num_heads, + head_dim=self.head_dim, + state_size=self.ssm_state_size, + conv_kernel=self.conv_kernel_size, + num_spec=self.num_spec, + replayssm_buffer_len=self.max_cache_len, + ) + elif self.use_cache_kernel: + return MambaStateShapeCalculator.mamba2_replayssm_state_shape( intermediate_size=self.intermediate_size, tp_world_size=get_tensor_model_parallel_world_size(), n_groups=self.n_groups, @@ -1254,10 +1266,9 @@ def get_state_shape(self) -> tuple[tuple[int, ...], ...]: state_size=self.ssm_state_size, conv_kernel=self.conv_kernel_size, num_spec=self.num_spec, - use_replayssm_spec=self.use_cache_spec_kernel, replayssm_buffer_len=self.max_cache_len, ) - return MambaStateShapeCalculator.mamba2_cached_state_shape( + return MambaStateShapeCalculator.mamba2_state_shape( intermediate_size=self.intermediate_size, tp_world_size=get_tensor_model_parallel_world_size(), n_groups=self.n_groups, @@ -1266,8 +1277,6 @@ def get_state_shape(self) -> tuple[tuple[int, ...], ...]: state_size=self.ssm_state_size, conv_kernel=self.conv_kernel_size, num_spec=self.num_spec, - use_replayssm=self.use_cache_kernel, - replayssm_buffer_len=self.max_cache_len, ) @property diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 270d7c1e2898..6ca56fd74a24 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -81,50 +81,41 @@ def mamba2_state_dtype( ) @classmethod - def mamba2_cached_state_dtype( + def mamba2_replayssm_state_dtype( cls, model_dtype: ModelDType | torch.dtype, mamba_cache_dtype: MambaDType, mamba_ssm_cache_dtype: MambaDType, - use_replayssm: bool, ) -> tuple[torch.dtype, ...]: - """Mamba2 state dtypes, extended for the state-and-output decode kernel. - - Returns the baseline ``(conv, ssm)`` dtypes when - ``use_replayssm`` is ``False``; otherwise appends the - state-and-output ring-buffer dtypes ``(x_cache, dt_cache, B_cache)`` = - ``(activation, fp32, activation)``. Must stay in sync with - ``MambaMixer2.get_state_dtype``. + """Mamba2 ReplaySSM state dtypes: baseline ``(conv, ssm)`` plus the + ring-buffer dtypes ``(x_cache, dt_cache, B_cache)`` = + ``(activation, fp32, activation)``. Call only when use_replayssm is on; + must stay in sync with ``MambaMixer2.get_state_dtype``. """ conv_dtype, ssm_dtype = cls.mamba2_state_dtype( model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype ) - if not use_replayssm: - return conv_dtype, ssm_dtype activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype) return conv_dtype, ssm_dtype, activation_dtype, torch.float32, activation_dtype @classmethod - def mamba2_spec_cached_state_dtype( + def mamba2_replayssm_spec_state_dtype( cls, model_dtype: ModelDType | torch.dtype, mamba_cache_dtype: MambaDType, mamba_ssm_cache_dtype: MambaDType, - use_replayssm_spec: bool, ) -> tuple[torch.dtype, ...]: - """Mamba2 state dtypes for the cached SPECULATIVE-decode (hybrid) kernel. - - Baseline ``(conv, ssm)`` when off; otherwise the hybrid 4-tuple + """Mamba2 ReplaySSM state dtypes for the SPECULATIVE-decode (hybrid) + kernel: the hybrid 4-tuple ``(conv, ssm_checkpoint, post_conv_cache, dt_cache)``. The checkpoint and ``dt_cache`` are forced fp32 (the cached-spec reconstruction was validated against an fp32 reference); ``post_conv_cache`` is activation - dtype. Must stay in sync with ``MambaMixer2.get_state_dtype``. + dtype. Call only when use_replayssm_spec is on; must stay in sync with + ``MambaMixer2.get_state_dtype``. """ conv_dtype, _ssm_dtype = cls.mamba2_state_dtype( model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype ) - if not use_replayssm_spec: - return conv_dtype, _ssm_dtype activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype) return conv_dtype, torch.float32, activation_dtype, torch.float32 @@ -164,46 +155,39 @@ def gated_delta_net_state_dtype( ) @classmethod - def gated_delta_net_cached_state_dtype( + def gated_delta_net_replayssm_state_dtype( cls, model_dtype: ModelDType | torch.dtype, mamba_cache_dtype: MambaDType, mamba_ssm_cache_dtype: MambaDType, - use_replayssm: bool, ) -> tuple[torch.dtype, ...]: - """GDN state dtypes, extended for the cached decode kernel. - - Returns the baseline ``(conv, ssm)`` dtypes when - ``use_replayssm`` is ``False``; otherwise appends the ring + """GDN ReplaySSM state dtypes: baseline ``(conv, ssm)`` plus the ring cache dtypes ``(d_cache, k_cache, g_cache)`` = - ``(activation, activation, float32)``. + ``(activation, activation, float32)``. Call only when use_replayssm is + on. """ conv_dtype, ssm_dtype = cls._mamba_state_dtype( model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype ) - if not use_replayssm: - return conv_dtype, ssm_dtype activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype) return conv_dtype, ssm_dtype, activation_dtype, activation_dtype, torch.float32 @classmethod - def gated_delta_net_spec_cached_state_dtype( + def gated_delta_net_replayssm_spec_state_dtype( cls, model_dtype: ModelDType | torch.dtype, mamba_cache_dtype: MambaDType, mamba_ssm_cache_dtype: MambaDType, - use_replayssm_spec: bool, ) -> tuple[torch.dtype, ...]: - """GDN state dtypes for the cached SPECULATIVE-decode kernel. + """GDN ReplaySSM state dtypes for the SPECULATIVE-decode kernel. - Same ``d/k/g`` ring page as the non-spec cached path, but the ``ssm`` - checkpoint is forced to ``float32`` Returns the baseline ``(conv, ssm)`` when the flag is off. + Same ``d/k/g`` ring page as the non-spec ReplaySSM path, but the ``ssm`` + checkpoint is forced to ``float32``. Call only when use_replayssm_spec + is on. """ conv_dtype, ssm_dtype = cls._mamba_state_dtype( model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype ) - if not use_replayssm_spec: - return conv_dtype, ssm_dtype activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype) return ( conv_dtype, @@ -285,7 +269,7 @@ def mamba2_state_shape( return conv_state_shape, temporal_state_shape @classmethod - def mamba2_cached_state_shape( + def mamba2_replayssm_state_shape( cls, tp_world_size: int, intermediate_size: int, @@ -294,17 +278,15 @@ def mamba2_cached_state_shape( head_dim: int, state_size: int, conv_kernel: int, - use_replayssm: bool, replayssm_buffer_len: int, num_spec: int = 0, ) -> tuple[tuple[int, ...], ...]: - """Mamba2 state shapes, extended for the state-and-output decode kernel. - - Returns the baseline ``(conv, ssm)`` shapes when - ``use_replayssm`` is ``False``; otherwise appends the - state-and-output ring-buffer shapes ``x_cache``/``dt_cache``/``B_cache``. - Group/head counts use the (un-extended) ``n_groups``/``num_heads`` - divided by ``tp_world_size``, matching ``MambaMixer2.get_state_shape``. + """Mamba2 ReplaySSM state shapes: baseline ``(conv, ssm)`` plus the + ring-buffer shapes ``x_cache``/``dt_cache``/``B_cache``. Delegates to + ``mamba2_state_shape`` for ``(conv, ssm)`` so the ring buffers keep the + un-extended ``n_groups`` (that method extends n_groups only in its own + scope). Call only when use_replayssm is on; must stay in sync with + ``MambaMixer2.get_state_shape``. """ conv_state_shape, temporal_state_shape = cls.mamba2_state_shape( tp_world_size=tp_world_size, @@ -316,9 +298,6 @@ def mamba2_cached_state_shape( conv_kernel=conv_kernel, num_spec=num_spec, ) - if not use_replayssm: - return conv_state_shape, temporal_state_shape - local_nheads = divide(num_heads, tp_world_size) local_ngroups = divide(n_groups, tp_world_size) x_cache_shape = (local_nheads, replayssm_buffer_len, head_dim) @@ -333,7 +312,7 @@ def mamba2_cached_state_shape( ) @classmethod - def mamba2_spec_cached_state_shape( + def mamba2_replayssm_spec_state_shape( cls, tp_world_size: int, intermediate_size: int, @@ -342,21 +321,19 @@ def mamba2_spec_cached_state_shape( head_dim: int, state_size: int, conv_kernel: int, - use_replayssm_spec: bool, replayssm_buffer_len: int, num_spec: int = 0, ) -> tuple[tuple[int, ...], ...]: - """Mamba2 state shapes for the cached SPECULATIVE-decode (hybrid) kernel. - - Baseline ``(conv, ssm)`` when off (conv keeps its spec sliding-window + """Mamba2 ReplaySSM state shapes for the SPECULATIVE-decode (hybrid) + kernel: baseline ``(conv, ssm)`` (conv keeps its spec sliding-window size ``conv_kernel-1+num_spec`` -- the hybrid reuses - ``causal_conv1d_update``); otherwise appends the circular caches + ``causal_conv1d_update``) plus the circular caches ``post_conv_cache=(cache_buf_len, conv_dim_local)`` and ``dt_cache=(local_nheads, cache_buf_len)``, where the L = B + max_spec_len history window sizes ``cache_buf_len = next_pow2(replayssm_buffer_len + 1 + num_spec)`` and ``conv_dim_local`` matches the post-conv x|B width (C is - not cached; read fresh from conv_out). Must stay in sync with - ``MambaMixer2.get_state_shape``. + not cached; read fresh from conv_out). Call only when use_replayssm_spec + is on; must stay in sync with ``MambaMixer2.get_state_shape``. """ conv_state_shape, temporal_state_shape = cls.mamba2_state_shape( tp_world_size=tp_world_size, @@ -368,8 +345,6 @@ def mamba2_spec_cached_state_shape( conv_kernel=conv_kernel, num_spec=num_spec, ) - if not use_replayssm_spec: - return conv_state_shape, temporal_state_shape n_groups_ext = n_groups + cls.extra_groups_for_head_shards( n_groups, tp_world_size ) @@ -436,7 +411,7 @@ def gated_delta_net_state_shape( return conv_state_shape, temporal_state_shape @classmethod - def gated_delta_net_cached_state_shape( + def gated_delta_net_replayssm_state_shape( cls, tp_world_size: int, num_k_heads: int, @@ -444,17 +419,14 @@ def gated_delta_net_cached_state_shape( head_k_dim: int, head_v_dim: int, conv_kernel_size: int, - use_replayssm: bool, replayssm_buffer_len: int, num_spec: int = 0, ) -> tuple[tuple[int, ...], ...]: - """GDN state shapes, extended for the cached decode kernel. - - Returns the baseline ``(conv, ssm)`` shapes when - ``use_replayssm`` is ``False``; otherwise appends the cached + """GDN ReplaySSM state shapes: baseline ``(conv, ssm)`` plus the cached ring-buffer shapes ``d_cache``/``k_cache``/``g_cache``. Head counts use the (un-extended) ``num_v_heads``/``num_k_heads`` divided by - ``tp_world_size``, matching ``gated_delta_net_state_shape``. + ``tp_world_size``, matching ``gated_delta_net_state_shape``. Call only + when use_replayssm is on. """ conv_state_shape, temporal_state_shape = cls.gated_delta_net_state_shape( tp_world_size, @@ -465,9 +437,6 @@ def gated_delta_net_cached_state_shape( conv_kernel_size, num_spec, ) - if not use_replayssm: - return conv_state_shape, temporal_state_shape - local_v_heads = divide(num_v_heads, tp_world_size) local_k_heads = divide(num_k_heads, tp_world_size) d_cache_shape = (local_v_heads, replayssm_buffer_len, head_v_dim) @@ -482,7 +451,7 @@ def gated_delta_net_cached_state_shape( ) @classmethod - def gated_delta_net_spec_cached_state_shape( + def gated_delta_net_replayssm_spec_state_shape( cls, tp_world_size: int, num_k_heads: int, @@ -490,16 +459,15 @@ def gated_delta_net_spec_cached_state_shape( head_k_dim: int, head_v_dim: int, conv_kernel_size: int, - use_replayssm_spec: bool, replayssm_buffer_len: int, num_spec: int = 0, ) -> tuple[tuple[int, ...], ...]: - """GDN state shapes for the cached SPECULATIVE-decode kernel. + """GDN ReplaySSM state shapes for the SPECULATIVE-decode kernel. The circular ``d_cache``/``k_cache``/``g_cache`` use the L = B + max_spec_len history window: a power-of-two buffer ``next_pow2(replayssm_buffer_len + 1 + - num_spec)``. Returns the baseline ``(conv, ssm)`` shapes when the flag is - off. The block-keyed cursors live in the GDN metadata builder, not the page. + num_spec)``. Call only when use_replayssm_spec is on. The block-keyed + cursors live in the GDN metadata builder, not the page. """ conv_state_shape, temporal_state_shape = cls.gated_delta_net_state_shape( tp_world_size, @@ -510,9 +478,6 @@ def gated_delta_net_spec_cached_state_shape( conv_kernel_size, num_spec, ) - if not use_replayssm_spec: - return conv_state_shape, temporal_state_shape - cache_buf_len = 1 << (replayssm_buffer_len + num_spec).bit_length() local_v_heads = divide(num_v_heads, tp_world_size) local_k_heads = divide(num_k_heads, tp_world_size) diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index ddf8e01fdaa3..d09add9333ef 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -817,17 +817,21 @@ def get_mamba_state_dtype_from_config( ) -> tuple[torch.dtype, ...]: cache_config = vllm_config.cache_config if cache_config.use_replayssm_spec: - return MambaStateDtypeCalculator.mamba2_spec_cached_state_dtype( + return MambaStateDtypeCalculator.mamba2_replayssm_spec_state_dtype( vllm_config.model_config.dtype, cache_config.mamba_cache_dtype, cache_config.mamba_ssm_cache_dtype, - use_replayssm_spec=cache_config.use_replayssm_spec, ) - return MambaStateDtypeCalculator.mamba2_cached_state_dtype( + elif cache_config.use_replayssm: + return MambaStateDtypeCalculator.mamba2_replayssm_state_dtype( + vllm_config.model_config.dtype, + cache_config.mamba_cache_dtype, + cache_config.mamba_ssm_cache_dtype, + ) + return MambaStateDtypeCalculator.mamba2_state_dtype( vllm_config.model_config.dtype, cache_config.mamba_cache_dtype, cache_config.mamba_ssm_cache_dtype, - use_replayssm=cache_config.use_replayssm, ) @classmethod @@ -853,7 +857,19 @@ def get_mamba_state_shape_from_config( intermediate_size = hf_config.mamba_num_heads * hf_config.mamba_head_dim if cache_config.use_replayssm_spec: - return MambaStateShapeCalculator.mamba2_spec_cached_state_shape( + return MambaStateShapeCalculator.mamba2_replayssm_spec_state_shape( + intermediate_size=intermediate_size, + tp_world_size=parallel_config.tensor_parallel_size, + n_groups=hf_config.n_groups, + num_heads=hf_config.mamba_num_heads, + head_dim=hf_config.mamba_head_dim, + state_size=hf_config.ssm_state_size, + conv_kernel=hf_config.conv_kernel, + num_spec=vllm_config.num_speculative_tokens, + replayssm_buffer_len=cache_config.replayssm_buffer_len, + ) + elif cache_config.use_replayssm: + return MambaStateShapeCalculator.mamba2_replayssm_state_shape( intermediate_size=intermediate_size, tp_world_size=parallel_config.tensor_parallel_size, n_groups=hf_config.n_groups, @@ -862,10 +878,9 @@ def get_mamba_state_shape_from_config( state_size=hf_config.ssm_state_size, conv_kernel=hf_config.conv_kernel, num_spec=vllm_config.num_speculative_tokens, - use_replayssm_spec=cache_config.use_replayssm_spec, replayssm_buffer_len=cache_config.replayssm_buffer_len, ) - return MambaStateShapeCalculator.mamba2_cached_state_shape( + return MambaStateShapeCalculator.mamba2_state_shape( intermediate_size=intermediate_size, tp_world_size=parallel_config.tensor_parallel_size, n_groups=hf_config.n_groups, @@ -874,8 +889,6 @@ def get_mamba_state_shape_from_config( state_size=hf_config.ssm_state_size, conv_kernel=hf_config.conv_kernel, num_spec=vllm_config.num_speculative_tokens, - use_replayssm=cache_config.use_replayssm, - replayssm_buffer_len=cache_config.replayssm_buffer_len, ) @classmethod diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index e809acac53bd..82101b0ebdca 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -515,17 +515,21 @@ def get_mamba_state_dtype_from_config( vllm_config: "VllmConfig", ) -> tuple[torch.dtype, ...]: if vllm_config.cache_config.use_replayssm_spec: - return MambaStateDtypeCalculator.gated_delta_net_spec_cached_state_dtype( + return MambaStateDtypeCalculator.gated_delta_net_replayssm_spec_state_dtype( vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype, vllm_config.cache_config.mamba_ssm_cache_dtype, - vllm_config.cache_config.use_replayssm_spec, ) - return MambaStateDtypeCalculator.gated_delta_net_cached_state_dtype( + elif vllm_config.cache_config.use_replayssm: + return MambaStateDtypeCalculator.gated_delta_net_replayssm_state_dtype( + vllm_config.model_config.dtype, + vllm_config.cache_config.mamba_cache_dtype, + vllm_config.cache_config.mamba_ssm_cache_dtype, + ) + return MambaStateDtypeCalculator.gated_delta_net_state_dtype( vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype, vllm_config.cache_config.mamba_ssm_cache_dtype, - vllm_config.cache_config.use_replayssm, ) @classmethod @@ -541,26 +545,34 @@ def get_mamba_state_shape_from_config( else 0 ) if vllm_config.cache_config.use_replayssm_spec: - return MambaStateShapeCalculator.gated_delta_net_spec_cached_state_shape( + return MambaStateShapeCalculator.gated_delta_net_replayssm_spec_state_shape( + tp_size, + hf_config.linear_num_key_heads, + hf_config.linear_num_value_heads, + hf_config.linear_key_head_dim, + hf_config.linear_value_head_dim, + hf_config.linear_conv_kernel_dim, + vllm_config.cache_config.replayssm_buffer_len, + num_spec, + ) + elif vllm_config.cache_config.use_replayssm: + return MambaStateShapeCalculator.gated_delta_net_replayssm_state_shape( tp_size, hf_config.linear_num_key_heads, hf_config.linear_num_value_heads, hf_config.linear_key_head_dim, hf_config.linear_value_head_dim, hf_config.linear_conv_kernel_dim, - vllm_config.cache_config.use_replayssm_spec, vllm_config.cache_config.replayssm_buffer_len, num_spec, ) - return MambaStateShapeCalculator.gated_delta_net_cached_state_shape( + return MambaStateShapeCalculator.gated_delta_net_state_shape( tp_size, hf_config.linear_num_key_heads, hf_config.linear_num_value_heads, hf_config.linear_key_head_dim, hf_config.linear_value_head_dim, hf_config.linear_conv_kernel_dim, - vllm_config.cache_config.use_replayssm, - vllm_config.cache_config.replayssm_buffer_len, num_spec, ) From 69eacf68b4f1a44c35300d1ae5ba96f8e9389a50 Mon Sep 17 00:00:00 2001 From: Johnny-Liou Date: Wed, 8 Jul 2026 13:40:25 -0400 Subject: [PATCH 05/23] [Bugfix] ReplaySSM: handle single-token prefill-as-decode rows as one-token flushes instead of raising Signed-off-by: Johnny-Liou --- vllm/v1/attention/backends/mamba_attn.py | 33 +++++++++++------------- 1 file changed, 15 insertions(+), 18 deletions(-) diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index 204cb225f1c2..9d24ac2a5a7b 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -115,9 +115,7 @@ def __init__( self.use_spec_decode = self.num_spec_tokens > 0 self.use_cached_kernel = vllm_config.cache_config.use_replayssm self.max_cache_len = vllm_config.cache_config.replayssm_buffer_len - self.use_cache_spec_kernel = ( - vllm_config.cache_config.use_replayssm_spec - ) + self.use_cache_spec_kernel = vllm_config.cache_config.use_replayssm_spec self.max_spec_len = 1 + self.num_spec_tokens # L = B + max_spec_len history window; physical pow2 ring = next_pow2(L). self.spec_flush_threshold = self.max_cache_len + self.max_spec_len @@ -193,9 +191,7 @@ def __init__( dtype=torch.int8, device=device, ) - self.cached_kernel_variant: str = ( - vllm_config.cache_config.replayssm_route - ) + self.cached_kernel_variant: str = vllm_config.cache_config.replayssm_route if self.cached_kernel_variant == "output_only": # B_cache shape = (ngroups, max_cache_len, dstate). Index in # MambaSpec.shapes is (conv_state, ssm_state, x_cache, @@ -504,11 +500,8 @@ def _compute_common_metadata( has_prior_state = seq_lens_cpu > 1 prefill_to_decode = single_token_prefill_rows & has_prior_state if torch.any(prefill_to_decode).item(): - if self.use_cached_kernel and metadata.num_decodes > 0: - raise ValueError( - "--use-replayssm does not support single-token " - "prefill rows replayed through the decode path" - ) + # ReplaySSM handles these rows as single-token flushes (see the + # write-position derivation below), same as the baseline decode path. is_prefilling = is_prefilling.clone() is_prefilling[prefill_to_decode] = False common_attn_metadata = common_attn_metadata.replace( @@ -645,18 +638,22 @@ def _compute_common_metadata( - common_attn_metadata.query_start_loc_cpu[:num_decodes] ) valid_decode_rows = query_lens_cpu > 0 - if torch.any(decode_steps_cpu[valid_decode_rows] < 0).item(): - raise ValueError( - "--use-replayssm requires decode-step counts " - "that exclude prompt tokens and start at zero" - ) + # A single-token prefill row replayed as decode (query_len==1 with + # prior state) has decode_steps < 0; force it to a one-token flush + # (write_pos=0, is_flush=1). The flush branch reads an empty history + # window, so it applies exactly one recurrence step off the checkpoint + # -- identical to the baseline decode kernel for that row. The split + # (treat_short_extends_as_decodes=False) admits only such rows here. + leftover_prompt = valid_decode_rows & (decode_steps_cpu < 0) decode_steps_cpu = torch.where( - valid_decode_rows, + valid_decode_rows & ~leftover_prompt, decode_steps_cpu, torch.zeros_like(decode_steps_cpu), ) write_pos_cpu = torch.remainder(decode_steps_cpu, self.max_cache_len) - is_flush_cpu = (write_pos_cpu == self.max_cache_len - 1).to(torch.int8) + is_flush_cpu = ( + (write_pos_cpu == self.max_cache_len - 1) | leftover_prompt + ).to(torch.int8) write_pos_d = async_tensor_h2d( write_pos_cpu.to(torch.int32).tolist(), dtype=torch.int32, From 40c8795c4f23ed027f9b1772a74ae2846f346d33 Mon Sep 17 00:00:00 2001 From: Johnny-Liou Date: Sun, 12 Jul 2026 19:36:29 -0400 Subject: [PATCH 06/23] [Bugfix] ReplaySSM GDN: route mixed prefill+decode batches to cached kernel Signed-off-by: Johnny-Liou --- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 67 +++++++++++++------ 1 file changed, 48 insertions(+), 19 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 5858747985db..aea091f02b70 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -1584,25 +1584,54 @@ def _forward_core( # 2.2: Process non-spec-decode part if split_non_spec: - query_decode, key_decode, value_decode = self.rearrange_mixed_qkv( - mixed_qkv_non_spec[:num_decode_tokens] # type: ignore[index] - ) - core_attn_out_decode, _ = fused_sigmoid_gating_delta_rule_update( - A_log=self.A_log, - a=a[:num_decode_tokens], - b=b[:num_decode_tokens], - dt_bias=self.dt_bias, - q=query_decode, - k=key_decode, - v=value_decode, - initial_state=ssm_state, - inplace_final_state=True, - cu_seqlens=non_spec_query_start_loc[ # type: ignore[index] - : attn_metadata.num_decodes + 1 - ], - ssm_state_indices=non_spec_state_indices_tensor, - use_qk_l2norm_in_kernel=True, - ) + if self.use_cache_kernel: + out_decode = torch.empty( + num_decode_tokens, + 1, + self.num_v_heads // self.tp_size, + self.head_v_dim, + dtype=mixed_qkv_non_spec.dtype, # type: ignore[union-attr] + device=mixed_qkv_non_spec.device, # type: ignore[union-attr] + ) + fused_recurrent_gated_delta_rule_replayssm( + mixed_qkv=mixed_qkv_non_spec[:num_decode_tokens].contiguous(), # type: ignore[index] + a=a[:num_decode_tokens], + b=b[:num_decode_tokens], + A_log=self.A_log, + dt_bias=self.dt_bias, + scale=self.head_k_dim**-0.5, + initial_state=ssm_state, + d_cache=self_kv_cache[2], + k_cache=self_kv_cache[3], + g_cache=self_kv_cache[4], + out=out_decode, + ssm_state_indices=non_spec_state_indices_tensor[ # type: ignore[index] + : attn_metadata.num_decodes + ], + write_pos=attn_metadata.write_pos_d, + use_qk_l2norm_in_kernel=True, + ) + core_attn_out_decode = out_decode.transpose(0, 1) + else: + query_decode, key_decode, value_decode = self.rearrange_mixed_qkv( + mixed_qkv_non_spec[:num_decode_tokens] # type: ignore[index] + ) + core_attn_out_decode, _ = fused_sigmoid_gating_delta_rule_update( + A_log=self.A_log, + a=a[:num_decode_tokens], + b=b[:num_decode_tokens], + dt_bias=self.dt_bias, + q=query_decode, + k=key_decode, + v=value_decode, + initial_state=ssm_state, + inplace_final_state=True, + cu_seqlens=non_spec_query_start_loc[ # type: ignore[index] + : attn_metadata.num_decodes + 1 + ], + ssm_state_indices=non_spec_state_indices_tensor, + use_qk_l2norm_in_kernel=True, + ) else: core_attn_out_decode = None From 883180edf3c5f892bdb3bee450be153e0bd12de0 Mon Sep 17 00:00:00 2001 From: Johnny-Liou Date: Sun, 12 Jul 2026 22:05:01 -0400 Subject: [PATCH 07/23] [Bugfix] ReplaySSM Mamba2: size bc_pre scratch to max_num_seqs Signed-off-by: Johnny-Liou --- vllm/v1/attention/backends/mamba_attn.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index 9d24ac2a5a7b..8fae389aadef 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -202,9 +202,12 @@ def __init__( "page (conv, ssm, x_cache, dt_cache, B_cache)" ) bc_ngroups = kv_cache_spec.shapes[4][0] + bc_scratch_bs = max( + self.decode_cudagraph_max_bs, scheduler_config.max_num_seqs + ) self.decode_bc_pre_scratch: torch.Tensor = torch.empty( ( - self.decode_cudagraph_max_bs, + bc_scratch_bs, bc_ngroups, self.max_cache_len, ), From 8f4ab99e62962f036bd59390376749d54a02e6fd Mon Sep 17 00:00:00 2001 From: Johnny-Liou Date: Mon, 13 Jul 2026 12:53:05 -0400 Subject: [PATCH 08/23] [Misc] Apply ruff-format to fused_recurrent_replayssm.py Signed-off-by: Johnny-Liou --- .../fla/ops/fused_recurrent_replayssm.py | 219 ++++++++++++++---- 1 file changed, 174 insertions(+), 45 deletions(-) diff --git a/vllm/model_executor/layers/fla/ops/fused_recurrent_replayssm.py b/vllm/model_executor/layers/fla/ops/fused_recurrent_replayssm.py index 97507d7faa5a..3ec867b535c1 100644 --- a/vllm/model_executor/layers/fla/ops/fused_recurrent_replayssm.py +++ b/vllm/model_executor/layers/fla/ops/fused_recurrent_replayssm.py @@ -14,8 +14,20 @@ @triton.jit def fused_recurrent_gated_delta_rule_replayssm_kernel( - mixed_qkv, a, b, A_log, dt_bias, o, h0, ht, - d_cache, k_cache, g_cache, ssm_state_indices, write_pos, scale, + mixed_qkv, + a, + b, + A_log, + dt_bias, + o, + h0, + ht, + d_cache, + k_cache, + g_cache, + ssm_state_indices, + write_pos, + scale, stride_mixed_qkv_tok: tl.constexpr, stride_a_tok: tl.constexpr, stride_b_tok: tl.constexpr, @@ -25,10 +37,17 @@ def fused_recurrent_gated_delta_rule_replayssm_kernel( stride_d_slot: tl.constexpr, stride_k_slot: tl.constexpr, stride_g_slot: tl.constexpr, - H: tl.constexpr, HV: tl.constexpr, K: tl.constexpr, V: tl.constexpr, - BK: tl.constexpr, BV: tl.constexpr, BC: tl.constexpr, - NK: tl.constexpr, BKT: tl.constexpr, - MAX_CACHE_LEN: tl.constexpr, SOFTPLUS_THRESHOLD: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + BC: tl.constexpr, + NK: tl.constexpr, + BKT: tl.constexpr, + MAX_CACHE_LEN: tl.constexpr, + SOFTPLUS_THRESHOLD: tl.constexpr, USE_QK_L2NORM_IN_KERNEL: tl.constexpr, ): i_v = tl.program_id(0) @@ -44,7 +63,9 @@ def fused_recurrent_gated_delta_rule_replayssm_kernel( state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq).to(tl.int64) p_o = o + (i_n * HV + i_hv) * V + o_v if state_idx <= 0: - tl.store(p_o, tl.zeros([BV], dtype=tl.float32).to(p_o.dtype.element_ty), mask=mask_v) + tl.store( + p_o, tl.zeros([BV], dtype=tl.float32).to(p_o.dtype.element_ty), mask=mask_v + ) return # Per-row buffer cursor and flush flag; valid (committed) cache positions. @@ -73,13 +94,23 @@ def fused_recurrent_gated_delta_rule_replayssm_kernel( b_total_decay = tl.exp(b_g_total) # Cached delta-rule update vectors d (K-independent), scaled by the replay decay. - p_d_main = d_cache + state_idx * stride_d_slot + ((i_hv * MAX_CACHE_LEN + o_c[None, :]) * V + o_v[:, None]) - b_d_all = tl.load(p_d_main, mask=mask_v[:, None] & cache_valid[None, :], other=0).to(tl.float32) - b_d_scaled_tc = (b_d_all * b_replay_decay[None, :]).to(p_o.dtype.element_ty) # [BV, BC] + p_d_main = ( + d_cache + + state_idx * stride_d_slot + + ((i_hv * MAX_CACHE_LEN + o_c[None, :]) * V + o_v[:, None]) + ) + b_d_all = tl.load( + p_d_main, mask=mask_v[:, None] & cache_valid[None, :], other=0 + ).to(tl.float32) + b_d_scaled_tc = (b_d_all * b_replay_decay[None, :]).to( + p_o.dtype.element_ty + ) # [BV, BC] # Current token value (for the delta-rule update). v_off = (2 * H * K) + i_hv * V + o_v - b_v = tl.load(mixed_qkv + i_n * stride_mixed_qkv_tok + v_off, mask=mask_v, other=0).to(tl.float32) + b_v = tl.load( + mixed_qkv + i_n * stride_mixed_qkv_tok + v_off, mask=mask_v, other=0 + ).to(tl.float32) # Optional q/k L2 norm: full-vector reciprocal norms (computed, not kept). if USE_QK_L2NORM_IN_KERNEL: @@ -87,7 +118,9 @@ def fused_recurrent_gated_delta_rule_replayssm_kernel( mask_kf = o_kf < K p_mix = mixed_qkv + i_n * stride_mixed_qkv_tok qf = tl.load(p_mix + i_h * K + o_kf, mask=mask_kf, other=0).to(tl.float32) - kf = tl.load(p_mix + H * K + i_h * K + o_kf, mask=mask_kf, other=0).to(tl.float32) + kf = tl.load(p_mix + H * K + i_h * K + o_kf, mask=mask_kf, other=0).to( + tl.float32 + ) q_rnorm = 1.0 / tl.sqrt(tl.sum(qf * qf) + 1e-6) k_rnorm = 1.0 / tl.sqrt(tl.sum(kf * kf) + 1e-6) else: @@ -105,25 +138,57 @@ def fused_recurrent_gated_delta_rule_replayssm_kernel( o_kt = kk * BKT + tl.arange(0, BKT) mask_kt = o_kt < K p_mix = mixed_qkv + i_n * stride_mixed_qkv_tok - q_c = tl.load(p_mix + i_h * K + o_kt, mask=mask_kt, other=0).to(tl.float32) * q_rnorm - k_c = tl.load(p_mix + H * K + i_h * K + o_kt, mask=mask_kt, other=0).to(tl.float32) * k_rnorm + q_c = ( + tl.load(p_mix + i_h * K + o_kt, mask=mask_kt, other=0).to(tl.float32) + * q_rnorm + ) + k_c = ( + tl.load(p_mix + H * K + i_h * K + o_kt, mask=mask_kt, other=0).to( + tl.float32 + ) + * k_rnorm + ) q_cs = q_c * scale cur_kq += tl.sum(k_c * q_cs) # Reconstruct this K tile of the state: S = total_decay * S_0 + d_scaled . k_cache. - p_h0_c = h0 + state_idx * stride_init_state_token + i_hv * V * K + o_v[:, None] * K + o_kt[None, :] - b_h0_c = tl.load(p_h0_c, mask=mask_v[:, None] & mask_kt[None, :], other=0).to(tl.float32) - p_k_c = k_cache + state_idx * stride_k_slot + ((i_h * MAX_CACHE_LEN + o_c[:, None]) * K + o_kt[None, :]) - b_k_all_c = tl.load(p_k_c, mask=cache_valid[:, None] & mask_kt[None, :], other=0).to(p_o.dtype.element_ty) - b_h_c = b_h0_c * b_total_decay + tl.dot(b_d_scaled_tc, b_k_all_c).to(tl.float32) # [BV, BKT] + p_h0_c = ( + h0 + + state_idx * stride_init_state_token + + i_hv * V * K + + o_v[:, None] * K + + o_kt[None, :] + ) + b_h0_c = tl.load(p_h0_c, mask=mask_v[:, None] & mask_kt[None, :], other=0).to( + tl.float32 + ) + p_k_c = ( + k_cache + + state_idx * stride_k_slot + + ((i_h * MAX_CACHE_LEN + o_c[:, None]) * K + o_kt[None, :]) + ) + b_k_all_c = tl.load( + p_k_c, mask=cache_valid[:, None] & mask_kt[None, :], other=0 + ).to(p_o.dtype.element_ty) + b_h_c = b_h0_c * b_total_decay + tl.dot(b_d_scaled_tc, b_k_all_c).to( + tl.float32 + ) # [BV, BKT] # Read the state with q and k (accumulated across K tiles). b_state_q += tl.sum(b_h_c * q_cs[None, :], axis=1) b_state_k += tl.sum(b_h_c * k_c[None, :], axis=1) if write_k: - p_cur_k = k_cache + state_idx * stride_k_slot + ((i_h * MAX_CACHE_LEN + b_write_pos) * K + o_kt) - tl.store(p_cur_k, k_c.to(p_o.dtype.element_ty), mask=mask_kt & (b_write_pos < MAX_CACHE_LEN)) + p_cur_k = ( + k_cache + + state_idx * stride_k_slot + + ((i_h * MAX_CACHE_LEN + b_write_pos) * K + o_kt) + ) + tl.store( + p_cur_k, + k_c.to(p_o.dtype.element_ty), + mask=mask_kt & (b_write_pos < MAX_CACHE_LEN), + ) # Current-token output: alpha*(S q) + d_cur * (k . q), with the new update # vector d_cur = beta * (v - alpha*(S k)). @@ -140,22 +205,63 @@ def fused_recurrent_gated_delta_rule_replayssm_kernel( o_kt = kk * BKT + tl.arange(0, BKT) mask_kt = o_kt < K p_mix = mixed_qkv + i_n * stride_mixed_qkv_tok - k_c = tl.load(p_mix + H * K + i_h * K + o_kt, mask=mask_kt, other=0).to(tl.float32) * k_rnorm - p_h0_c = h0 + state_idx * stride_init_state_token + i_hv * V * K + o_v[:, None] * K + o_kt[None, :] - b_h0_c = tl.load(p_h0_c, mask=mask_v[:, None] & mask_kt[None, :], other=0).to(tl.float32) - p_k_c = k_cache + state_idx * stride_k_slot + ((i_h * MAX_CACHE_LEN + o_c[:, None]) * K + o_kt[None, :]) - b_k_all_c = tl.load(p_k_c, mask=cache_valid[:, None] & mask_kt[None, :], other=0).to(p_o.dtype.element_ty) - b_h_c = b_h0_c * b_total_decay + tl.dot(b_d_scaled_tc, b_k_all_c).to(tl.float32) + k_c = ( + tl.load(p_mix + H * K + i_h * K + o_kt, mask=mask_kt, other=0).to( + tl.float32 + ) + * k_rnorm + ) + p_h0_c = ( + h0 + + state_idx * stride_init_state_token + + i_hv * V * K + + o_v[:, None] * K + + o_kt[None, :] + ) + b_h0_c = tl.load( + p_h0_c, mask=mask_v[:, None] & mask_kt[None, :], other=0 + ).to(tl.float32) + p_k_c = ( + k_cache + + state_idx * stride_k_slot + + ((i_h * MAX_CACHE_LEN + o_c[:, None]) * K + o_kt[None, :]) + ) + b_k_all_c = tl.load( + p_k_c, mask=cache_valid[:, None] & mask_kt[None, :], other=0 + ).to(p_o.dtype.element_ty) + b_h_c = b_h0_c * b_total_decay + tl.dot(b_d_scaled_tc, b_k_all_c).to( + tl.float32 + ) b_h_new_c = alpha_val * b_h_c + b_d_cur[:, None] * k_c[None, :] - p_ht_c = ht + state_idx * stride_final_state_token + i_hv * V * K + o_v[:, None] * K + o_kt[None, :] - tl.store(p_ht_c, b_h_new_c.to(p_ht_c.dtype.element_ty), mask=mask_v[:, None] & mask_kt[None, :]) + p_ht_c = ( + ht + + state_idx * stride_final_state_token + + i_hv * V * K + + o_v[:, None] * K + + o_kt[None, :] + ) + tl.store( + p_ht_c, + b_h_new_c.to(p_ht_c.dtype.element_ty), + mask=mask_v[:, None] & mask_kt[None, :], + ) else: # Non-flush: append the current token's update vector d and gate g to the # cache (the k chunks were already written inside the loop above). - p_cur_d = d_cache + state_idx * stride_d_slot + ((i_hv * MAX_CACHE_LEN + b_write_pos) * V + o_v) - tl.store(p_cur_d, b_d_cur.to(p_cur_d.dtype.element_ty), mask=mask_v & (b_write_pos < MAX_CACHE_LEN)) + p_cur_d = ( + d_cache + + state_idx * stride_d_slot + + ((i_hv * MAX_CACHE_LEN + b_write_pos) * V + o_v) + ) + tl.store( + p_cur_d, + b_d_cur.to(p_cur_d.dtype.element_ty), + mask=mask_v & (b_write_pos < MAX_CACHE_LEN), + ) if i_v == 0: - p_cur_g = g_cache + state_idx * stride_g_slot + i_hv * MAX_CACHE_LEN + b_write_pos + p_cur_g = ( + g_cache + state_idx * stride_g_slot + i_hv * MAX_CACHE_LEN + b_write_pos + ) tl.store(p_cur_g, g_val, mask=b_write_pos < MAX_CACHE_LEN) @@ -217,15 +323,15 @@ def fused_recurrent_gated_delta_rule_replayssm( ) H = q_dim // K if H <= 0 or HV % H != 0: - raise ValueError(f"Invalid head config inferred from mixed_qkv: H={H}, HV={HV}.") + raise ValueError( + f"Invalid head config inferred from mixed_qkv: H={H}, HV={HV}." + ) max_cache_len = d_cache.shape[2] # Launch config (block_v, num_warps, num_stages, nk) from the L-keyed config # module; explicit kwargs override. Lets benchmarks/the config sweep pin it via # override_replayssm_config("gdn_decode", ...). - cfg_bv, cfg_nw, cfg_ns, cfg_nk = get_replayssm_config( - "gdn_decode", L=max_cache_len - ) + cfg_bv, cfg_nw, cfg_ns, cfg_nk = get_replayssm_config("gdn_decode", L=max_cache_len) if block_v is None: block_v = cfg_bv if num_warps is None: @@ -256,7 +362,9 @@ def fused_recurrent_gated_delta_rule_replayssm( BK = triton.next_power_of_2(K) if triton.cdiv(K, BK) != 1: - raise ValueError(f"Cached decode kernel only supports NK_global=1 (got K={K}, BK={BK}).") + raise ValueError( + f"Cached decode kernel only supports NK_global=1 (got K={K}, BK={BK})." + ) if BK % nk != 0: raise ValueError(f"nk={nk} must divide BK={BK}.") BKT = BK // nk @@ -270,21 +378,42 @@ def fused_recurrent_gated_delta_rule_replayssm( grid = (triton.cdiv(V, BV), B, HV) fused_recurrent_gated_delta_rule_replayssm_kernel[grid]( - mixed_qkv=mixed_qkv, a=a, b=b, A_log=A_log, dt_bias=dt_bias, o=out, - h0=initial_state, ht=initial_state, - d_cache=d_cache, k_cache=k_cache, g_cache=g_cache, - ssm_state_indices=ssm_state_indices, write_pos=write_pos, scale=scale, + mixed_qkv=mixed_qkv, + a=a, + b=b, + A_log=A_log, + dt_bias=dt_bias, + o=out, + h0=initial_state, + ht=initial_state, + d_cache=d_cache, + k_cache=k_cache, + g_cache=g_cache, + ssm_state_indices=ssm_state_indices, + write_pos=write_pos, + scale=scale, stride_mixed_qkv_tok=mixed_qkv.stride(0), - stride_a_tok=a.stride(0), stride_b_tok=b.stride(0), + stride_a_tok=a.stride(0), + stride_b_tok=b.stride(0), stride_init_state_token=initial_state.stride(0), stride_final_state_token=initial_state.stride(0), stride_indices_seq=ssm_state_indices.stride(0), stride_d_slot=d_cache.stride(0), stride_k_slot=k_cache.stride(0), stride_g_slot=g_cache.stride(0), - H=H, HV=HV, K=K, V=V, BK=BK, BV=BV, BC=BC, NK=nk, BKT=BKT, - MAX_CACHE_LEN=max_cache_len, SOFTPLUS_THRESHOLD=20.0, + H=H, + HV=HV, + K=K, + V=V, + BK=BK, + BV=BV, + BC=BC, + NK=nk, + BKT=BKT, + MAX_CACHE_LEN=max_cache_len, + SOFTPLUS_THRESHOLD=20.0, USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, - num_warps=num_warps, num_stages=num_stages, + num_warps=num_warps, + num_stages=num_stages, ) return out, initial_state From ead7f5dc9fe7437ed2c474c8794ea49befb01354 Mon Sep 17 00:00:00 2001 From: Johnny-Liou Date: Mon, 13 Jul 2026 12:57:21 -0400 Subject: [PATCH 09/23] [Kernel] ReplaySSM GDN: switch from bf16 to fp16 d/k caches for finer state reconstruction Signed-off-by: Johnny-Liou --- .../layers/fla/ops/fused_recurrent_replayssm.py | 8 ++++---- vllm/model_executor/layers/mamba/mamba_utils.py | 11 +++++++---- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/vllm/model_executor/layers/fla/ops/fused_recurrent_replayssm.py b/vllm/model_executor/layers/fla/ops/fused_recurrent_replayssm.py index 3ec867b535c1..526c7c6bbb67 100644 --- a/vllm/model_executor/layers/fla/ops/fused_recurrent_replayssm.py +++ b/vllm/model_executor/layers/fla/ops/fused_recurrent_replayssm.py @@ -103,7 +103,7 @@ def fused_recurrent_gated_delta_rule_replayssm_kernel( p_d_main, mask=mask_v[:, None] & cache_valid[None, :], other=0 ).to(tl.float32) b_d_scaled_tc = (b_d_all * b_replay_decay[None, :]).to( - p_o.dtype.element_ty + p_d_main.dtype.element_ty ) # [BV, BC] # Current token value (for the delta-rule update). @@ -169,7 +169,7 @@ def fused_recurrent_gated_delta_rule_replayssm_kernel( ) b_k_all_c = tl.load( p_k_c, mask=cache_valid[:, None] & mask_kt[None, :], other=0 - ).to(p_o.dtype.element_ty) + ).to(p_k_c.dtype.element_ty) b_h_c = b_h0_c * b_total_decay + tl.dot(b_d_scaled_tc, b_k_all_c).to( tl.float32 ) # [BV, BKT] @@ -186,7 +186,7 @@ def fused_recurrent_gated_delta_rule_replayssm_kernel( ) tl.store( p_cur_k, - k_c.to(p_o.dtype.element_ty), + k_c.to(p_cur_k.dtype.element_ty), mask=mask_kt & (b_write_pos < MAX_CACHE_LEN), ) @@ -228,7 +228,7 @@ def fused_recurrent_gated_delta_rule_replayssm_kernel( ) b_k_all_c = tl.load( p_k_c, mask=cache_valid[:, None] & mask_kt[None, :], other=0 - ).to(p_o.dtype.element_ty) + ).to(p_k_c.dtype.element_ty) b_h_c = b_h0_c * b_total_decay + tl.dot(b_d_scaled_tc, b_k_all_c).to( tl.float32 ) diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 6ca56fd74a24..34396b80f5f9 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -162,15 +162,18 @@ def gated_delta_net_replayssm_state_dtype( mamba_ssm_cache_dtype: MambaDType, ) -> tuple[torch.dtype, ...]: """GDN ReplaySSM state dtypes: baseline ``(conv, ssm)`` plus the ring - cache dtypes ``(d_cache, k_cache, g_cache)`` = - ``(activation, activation, float32)``. Call only when use_replayssm is - on. + cache dtypes ``(d_cache, k_cache, g_cache)``. The ``d``/``k`` input + caches use fp16 for bf16 activations; ``g_cache`` is float32. Call only + when use_replayssm is on. """ conv_dtype, ssm_dtype = cls._mamba_state_dtype( model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype ) activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype) - return conv_dtype, ssm_dtype, activation_dtype, activation_dtype, torch.float32 + cache_dtype = ( + torch.float16 if activation_dtype == torch.bfloat16 else activation_dtype + ) + return conv_dtype, ssm_dtype, cache_dtype, cache_dtype, torch.float32 @classmethod def gated_delta_net_replayssm_spec_state_dtype( From ba634f611116bc9256ab4247e29347b430327e56 Mon Sep 17 00:00:00 2001 From: Johnny-Liou Date: Mon, 13 Jul 2026 20:41:33 -0400 Subject: [PATCH 10/23] [Kernel] ReplaySSM GDN: hardware-aware standard-decode launch config Signed-off-by: Johnny-Liou --- .../layers/mamba/ops/replayssm_config.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/mamba/ops/replayssm_config.py b/vllm/model_executor/layers/mamba/ops/replayssm_config.py index 06801ebacaed..541e764840e5 100644 --- a/vllm/model_executor/layers/mamba/ops/replayssm_config.py +++ b/vllm/model_executor/layers/mamba/ops/replayssm_config.py @@ -88,8 +88,12 @@ def _l_bucket(cache_len: int) -> int: # state_and_output stays un-tiled (retained only for precision experiments); the # output_only route uses the decoupled dstate-tiled config in _mamba2_output_only. _STATE_AND_OUTPUT_BY_L = {8: (32, 1), 16: (32, 1), 32: (32, 1)} -# GDN standard decode: (block_v, num_warps, num_stages, nk). L-flat in the sweep. -_GDN_DECODE_BY_L = {8: (64, 1, 3, 2), 16: (64, 1, 3, 2), 32: (64, 1, 3, 2)} + + +def _gdn_decode(L, is_blackwell): + if is_blackwell: + return 128, 1, 3, 4 + return 64, 1, 3, 2 def get_replayssm_config(kernel: str, **shape) -> tuple: @@ -117,7 +121,7 @@ def get_replayssm_config(kernel: str, **shape) -> tuple: if kernel == "mamba2_state_and_output": return _STATE_AND_OUTPUT_BY_L[_l_bucket(shape.get("L", 16))] if kernel == "gdn_decode": - return _GDN_DECODE_BY_L[_l_bucket(shape.get("L", 16))] + return _gdn_decode(shape.get("L", 16), bw) if kernel in ("gdn_spec_verify", "gdn_spec_flush"): return _gdn_spec(shape["max_spec_len"], bw) raise ValueError(f"unknown ReplaySSM kernel config key: {kernel}") From 6475acbabe0bc650fb679fb0d41da383dd764cd0 Mon Sep 17 00:00:00 2001 From: Johnny-Liou Date: Mon, 13 Jul 2026 20:53:22 -0400 Subject: [PATCH 11/23] [Bugfix] ReplaySSM GDN: route single-token prefill-as-decode rows to prefill instead of raising Signed-off-by: Johnny-Liou --- vllm/v1/attention/backends/gdn_attn.py | 24 +++++++++++------------- 1 file changed, 11 insertions(+), 13 deletions(-) diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index ad638d7723b0..b736b5f2c21a 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -180,9 +180,7 @@ def __init__( # Cached decode kernel: persistent per-decode-row ring write position. # write_pos is derived per request each step (decode_step % max_cache_len) # so recycled paged blocks need no zero-init. - self.use_cached_kernel: bool = ( - vllm_config.cache_config.use_replayssm - ) + self.use_cached_kernel: bool = vllm_config.cache_config.use_replayssm self.max_cache_len: int = vllm_config.cache_config.replayssm_buffer_len if self.use_cached_kernel: self.decode_write_pos_d: torch.Tensor = torch.empty( @@ -194,9 +192,7 @@ def __init__( # Cached-SPEC decode: block-keyed cursors (sized num_gpu_blocks), # allocated lazily on first build (num_gpu_blocks is unknown here), and # advanced once per step by commit_gdn_replayssm_spec. - self.use_cache_spec_kernel: bool = ( - vllm_config.cache_config.use_replayssm_spec - ) + self.use_cache_spec_kernel: bool = vllm_config.cache_config.use_replayssm_spec self.max_spec_len: int = 1 + self.num_spec # L = B + max_spec_len history window; physical pow2 ring = next_pow2(L). self.spec_flush_threshold = self.max_cache_len + self.max_spec_len @@ -250,8 +246,13 @@ def build( # type: ignore[override] ) if spec_sequence_masks is None: + # ReplaySSM routes single-token prefill-as-decode rows to prefill. num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( - split_decodes_and_prefills(m, decode_threshold=1) + split_decodes_and_prefills( + m, + decode_threshold=1, + treat_short_extends_as_decodes=not self.use_cached_kernel, + ) ) num_spec_decode_tokens = 0 spec_token_indx = None @@ -542,12 +543,9 @@ def build( # type: ignore[override] num_prompt_d = num_prompt_tokens_cpu.to( context_lens_tensor.device, non_blocking=True ) - first_decode_full = (context_lens_tensor == num_prompt_d).to( - torch.int8 - ) - spec_row_idx = ( - spec_sequence_masks_cpu.nonzero(as_tuple=True)[0] - .to(query_start_loc.device, non_blocking=True) + first_decode_full = (context_lens_tensor == num_prompt_d).to(torch.int8) + spec_row_idx = spec_sequence_masks_cpu.nonzero(as_tuple=True)[0].to( + query_start_loc.device, non_blocking=True ) first_decode_d = first_decode_full.index_select(0, spec_row_idx) reset_gdn_replayssm_spec_cursors( From 15a31c8ee4b43ea6828fa7a5532d864273504473 Mon Sep 17 00:00:00 2001 From: Johnny-Liou Date: Thu, 16 Jul 2026 13:31:42 -0400 Subject: [PATCH 12/23] [Bugfix] ReplaySSM Mamba2 spec: derive bc_pre scratch groups from the x|B page width Signed-off-by: Johnny-Liou --- .../test_replayssm_spec_decode_mamba2.py | 146 +++- .../selective_state_update_replayssm_spec.py | 699 ++++++++++++++---- vllm/v1/attention/backends/mamba_attn.py | 8 +- 3 files changed, 705 insertions(+), 148 deletions(-) diff --git a/tests/kernels/mamba/test_replayssm_spec_decode_mamba2.py b/tests/kernels/mamba/test_replayssm_spec_decode_mamba2.py index 85e3e954db92..a6a2d8806898 100644 --- a/tests/kernels/mamba/test_replayssm_spec_decode_mamba2.py +++ b/tests/kernels/mamba/test_replayssm_spec_decode_mamba2.py @@ -250,7 +250,7 @@ def _run_single_step( conv_out = _pack_window_conv_out(x[wp:], B[wp:], C[wp:], d_inner, G, N, act_dtype) dt_spec = dt[wp:].float() - z_spec = z[wp:] if has_z else None + z_spec = z[wp:] if z is not None else None write_pos = torch.zeros(num_blocks, dtype=torch.int32, device=DEV) write_pos[1] = wp post_origin = torch.zeros(num_blocks, dtype=torch.int32, device=DEV) @@ -500,7 +500,7 @@ def _run_rollback( base_out = [] for s in range(k): - ot = torch.empty(1, H, P, device=DEV, dtype=act_dtype) + out_t = torch.empty(1, H, P, device=DEV, dtype=act_dtype) selective_state_update( state_base, x[s : s + 1], @@ -509,13 +509,13 @@ def _run_rollback( Bw[s : s + 1], Cw[s : s + 1], D=D, - z=zw[s : s + 1] if has_z else None, + z=zw[s : s + 1] if zw is not None else None, dt_bias=dt_bias[:, None].expand(H, P), dt_softplus=True, state_batch_indices=sbi, - out=ot, + out=out_t, ) - base_out.append(ot.clone()) + base_out.append(out_t.clone()) total_accepted += 1 snapshots[total_accepted] = state_base[1].clone() base_out = torch.cat(base_out, dim=0) @@ -707,3 +707,139 @@ def test_spec_continuous_batching(precision, with_padding): # non-flush verify leaves every state slot untouched. assert torch.equal(state_spec[unused], S0[unused]) torch.testing.assert_close(state_spec, S0, rtol=0, atol=0) + + +_TP_SHARD_GEOMETRIES = [ + pytest.param((32, 64, 128, 2), id="super120b_tp4"), + pytest.param((16, 64, 128, 1), id="super120b_tp8"), +] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize( + "precision", + [ + pytest.param((torch.float32, torch.float32), id="s32_a32"), + pytest.param((torch.float32, torch.bfloat16), id="s32_a16"), + ], +) +@pytest.mark.parametrize("geometry", _TP_SHARD_GEOMETRIES) +def test_spec_external_bc_pre_builder_shape(precision, geometry): + """Multi-row verify with an externally allocated bc_pre scratch, as the + engine provides it (fp32, oversized batch dim sliced to the batch), at + the TP4/TP8 per-rank shards of Super-120B. Each row's outputs must match + its standard-decode oracle; the checkpoint stays untouched.""" + state_dtype, act_dtype = precision + set_random_seed(0) + H, P, N, G = geometry + d_inner = H * P + base_block = 16 + max_spec_len = 4 + spec_len = max_spec_len + L, buf = _lbt(base_block, max_spec_len) + both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32 + rtol, atol = _tolerances(both_fp32) + + batch = 3 + total_slots = 8 + C = base_block - max_spec_len + wps = [C, C // 2, C] # per-row non-flush fills incl. the tight edge + + A = _tied_A(H, P, N) + dt_bias = torch.rand(H, device=DEV) - 4.0 + D = torch.randn(H, P, device=DEV) + sbi = torch.arange(1, 1 + batch, device=DEV, dtype=torch.int32) + + S0 = torch.randn(total_slots, H, P, N, device=DEV, dtype=state_dtype) * 0.1 + state_spec = S0.clone() + post_conv_cache = torch.zeros( + total_slots, buf, d_inner + G * N, device=DEV, dtype=act_dtype + ) + dt_cache = torch.zeros(total_slots, H, buf, device=DEV, dtype=torch.float32) + write_pos = torch.zeros(total_slots, dtype=torch.int32, device=DEV) + post_origin = torch.zeros(total_slots, dtype=torch.int32, device=DEV) + is_flush = torch.zeros(total_slots, dtype=torch.int8, device=DEV) + + oracles = [] + conv_dim = d_inner + 2 * G * N + conv_out = torch.zeros(batch * spec_len, conv_dim, device=DEV, dtype=act_dtype) + dt_spec = torch.zeros(batch * spec_len, H, device=DEV, dtype=torch.float32) + for r in range(batch): + wp = wps[r] + slot = int(sbi[r].item()) + T_tot = wp + spec_len + x = torch.randn(T_tot, H, P, device=DEV, dtype=act_dtype) + dt = torch.randn(T_tot, H, device=DEV, dtype=act_dtype) + Bv = torch.randn(T_tot, G, N, device=DEV, dtype=act_dtype) + Cv = torch.randn(T_tot, G, N, device=DEV, dtype=act_dtype) + oracles.append( + _standard_window_oracle( + S0_slot=S0[slot], + x_all=x, + dt_all=dt, + B_all=Bv, + C_all=Cv, + z_all=None, + A=A, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + wp=wp, + spec_len=spec_len, + buffer_len=buf, + act_dtype=act_dtype, + ) + ) + _scatter_packed_history( + post_conv_cache, dt_cache, slot, x[:wp], Bv[:wp], dt[:wp], d_inner, G, N + ) + write_pos[slot] = wp + seg = slice(r * spec_len, (r + 1) * spec_len) + conv_out[seg] = _pack_window_conv_out( + x[wp:], Bv[wp:], Cv[wp:], d_inner, G, N, act_dtype + ) + dt_spec[seg] = dt[wp:].float() + + # Engine-style external scratch: fp32, oversized batch dim sliced to the + # batch, NaN-filled so an unwritten-but-read cell fails loudly. + block_spec = 1 << (max_spec_len - 1).bit_length() + scratch = torch.full( + (batch + 5, G, buf, block_spec), + float("nan"), + device=DEV, + dtype=torch.float32, + ) + + qsl = torch.arange( + 0, (batch + 1) * spec_len, spec_len, device=DEV, dtype=torch.int32 + ) + out_spec = torch.empty(batch * spec_len, H, P, device=DEV, dtype=act_dtype) + selective_state_update_replayssm_spec( + state_spec, + post_conv_cache, + dt_cache, + conv_out, + dt_spec, + A, + write_pos=write_pos, + post_conv_state_pos=post_origin, + is_flush=is_flush, + query_start_loc=qsl, + state_batch_indices=sbi, + max_cache_len=L, + max_spec_len=max_spec_len, + d_inner=d_inner, + ngroups=G, + dstate=N, + D=D, + z=None, + dt_bias=dt_bias, + dt_softplus=True, + out=out_spec, + bc_pre=scratch[:batch], + ) + + for r in range(batch): + seg = slice(r * spec_len, (r + 1) * spec_len) + torch.testing.assert_close(out_spec[seg], oracles[r], rtol=rtol, atol=atol) + torch.testing.assert_close(state_spec, S0, rtol=0, atol=0) diff --git a/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_spec.py b/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_spec.py index 8369d3f8a762..c23780cd053d 100644 --- a/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_spec.py +++ b/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_spec.py @@ -90,19 +90,25 @@ def _fused_scatter_precompute_kernel( src_base = conv_out_ptr + bos * stride_conv_out_tok # fresh spec B / C [S, N] B_spec = tl.load( - src_base + (b_c0 + offs_n[None, :]) * stride_conv_out_c + offs_s[:, None] * stride_conv_out_tok, + src_base + + (b_c0 + offs_n[None, :]) * stride_conv_out_c + + offs_s[:, None] * stride_conv_out_tok, mask=spec_valid[:, None] & nmask[None, :], other=0.0, ) C_spec = tl.load( - src_base + (c_c0 + offs_n[None, :]) * stride_conv_out_c + offs_s[:, None] * stride_conv_out_tok, + src_base + + (c_c0 + offs_n[None, :]) * stride_conv_out_c + + offs_s[:, None] * stride_conv_out_tok, mask=spec_valid[:, None] & nmask[None, :], other=0.0, ) cache_base = post_conv_cache_ptr + state_batch_idx * stride_post_conv_cache_b # scatter B (C is not cached; read fresh from conv_out) tl.store( - cache_base + phys_spec[:, None] * stride_post_conv_cache_pos + (b_c0 + offs_n[None, :]) * stride_post_conv_cache_c, + cache_base + + phys_spec[:, None] * stride_post_conv_cache_pos + + (b_c0 + offs_n[None, :]) * stride_post_conv_cache_c, B_spec, mask=spec_valid[:, None] & nmask[None, :], ) @@ -155,12 +161,16 @@ def _fused_scatter_precompute_kernel( spec_off = offs_k - write_pos phys_k = (post_origin + offs_k) & (CACHE_BUF_LEN - 1) B_hist = tl.load( - cache_base + phys_k[:, None] * stride_post_conv_cache_pos + (b_c0 + offs_n[None, :]) * stride_post_conv_cache_c, + cache_base + + phys_k[:, None] * stride_post_conv_cache_pos + + (b_c0 + offs_n[None, :]) * stride_post_conv_cache_c, mask=hist_mask[:, None] & nmask[None, :], other=0.0, ) B_specrows = tl.load( - src_base + (b_c0 + offs_n[None, :]) * stride_conv_out_c + spec_off[:, None] * stride_conv_out_tok, + src_base + + (b_c0 + offs_n[None, :]) * stride_conv_out_c + + spec_off[:, None] * stride_conv_out_tok, mask=spec_tok[:, None] & nmask[None, :], other=0.0, ) @@ -194,31 +204,84 @@ def _fused_scatter_precompute_kernel( @triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) @triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) @triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) -@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics( + {"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])} +) @triton.jit def _replayssm_spec_nf_kernel( - state_ptr, x_cache_ptr, dt_cache_ptr, B_cache_ptr, C_src_ptr, bc_pre_ptr, - D_ptr, z_ptr, dt_bias_ptr, A_ptr, out_ptr, is_flush_flags_ptr, write_pos_ptr, - post_origin_ptr, state_batch_indices_ptr, query_start_loc_ptr, null_block_id, - batch, nheads, dim, dstate, max_cache_len, nheads_ngroups_ratio, - stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, - stride_x_cache_batch, stride_x_cache_head, stride_x_cache_dim, stride_x_cache_pos, - stride_dt_cache_batch, stride_dt_cache_head, stride_dt_cache_pos, - stride_B_cache_batch, stride_B_cache_group, stride_B_cache_dstate, stride_B_cache_pos, - stride_C_src_tok, stride_C_src_c, - stride_bc_pre_batch, stride_bc_pre_group, stride_bc_pre_pos, stride_bc_pre_spec, - stride_D_head, stride_D_dim, stride_z_tok, stride_z_head, stride_z_dim, - stride_dt_bias_head, stride_A_head, stride_out_tok, stride_out_head, stride_out_dim, + state_ptr, + x_cache_ptr, + dt_cache_ptr, + B_cache_ptr, + C_src_ptr, + bc_pre_ptr, + D_ptr, + z_ptr, + dt_bias_ptr, + A_ptr, + out_ptr, + is_flush_flags_ptr, + write_pos_ptr, + post_origin_ptr, + state_batch_indices_ptr, + query_start_loc_ptr, + null_block_id, + batch, + nheads, + dim, + dstate, + max_cache_len, + nheads_ngroups_ratio, + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + stride_x_cache_batch, + stride_x_cache_head, + stride_x_cache_dim, + stride_x_cache_pos, + stride_dt_cache_batch, + stride_dt_cache_head, + stride_dt_cache_pos, + stride_B_cache_batch, + stride_B_cache_group, + stride_B_cache_dstate, + stride_B_cache_pos, + stride_C_src_tok, + stride_C_src_c, + stride_bc_pre_batch, + stride_bc_pre_group, + stride_bc_pre_pos, + stride_bc_pre_spec, + stride_D_head, + stride_D_dim, + stride_z_tok, + stride_z_head, + stride_z_dim, + stride_dt_bias_head, + stride_A_head, + stride_out_tok, + stride_out_head, + stride_out_dim, stride_state_indices_batch, - DT_SOFTPLUS: tl.constexpr, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_CACHE: tl.constexpr, - BLOCK_SIZE_SPEC: tl.constexpr, HAS_DT_BIAS: tl.constexpr, HAS_D: tl.constexpr, - HAS_Z: tl.constexpr, CACHE_BUF_LEN: tl.constexpr, DSTATE_TILE: tl.constexpr, - NDS: tl.constexpr, BLOCK_SIZE_DSTATE: tl.constexpr, + DT_SOFTPLUS: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_CACHE: tl.constexpr, + BLOCK_SIZE_SPEC: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + CACHE_BUF_LEN: tl.constexpr, + DSTATE_TILE: tl.constexpr, + NDS: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, ): pid_m = tl.program_id(axis=0) pid_b = tl.program_id(axis=1) pid_h = tl.program_id(axis=2) - state_batch_idx = tl.load(state_batch_indices_ptr + pid_b * stride_state_indices_batch).to(tl.int64) + state_batch_idx = tl.load( + state_batch_indices_ptr + pid_b * stride_state_indices_batch + ).to(tl.int64) if state_batch_idx == null_block_id: return if tl.load(is_flush_flags_ptr + state_batch_idx) != 0: @@ -238,15 +301,27 @@ def _replayssm_spec_nf_kernel( cache_valid_mask = (offs_k < max_cache_len) & (offs_k < (write_pos + spec_len)) spec_token_mask = (offs_k >= write_pos) & (offs_k < (write_pos + spec_len)) spec_cache_pos = write_pos + offs_s - spec_prefix_mask = spec_valid_mask[:, None] & spec_token_mask[None, :] & (offs_k[None, :] <= spec_cache_pos[:, None]) + spec_prefix_mask = ( + spec_valid_mask[:, None] + & spec_token_mask[None, :] + & (offs_k[None, :] <= spec_cache_pos[:, None]) + ) phys_k = (post_origin + offs_k) & (CACHE_BUF_LEN - 1) phys_spec = (post_origin + spec_cache_pos) & (CACHE_BUF_LEN - 1) state_ptr += state_batch_idx * stride_state_batch + pid_h * stride_state_head x_cache_ptr += state_batch_idx * stride_x_cache_batch + pid_h * stride_x_cache_head - dt_cache_ptr += state_batch_idx * stride_dt_cache_batch + pid_h * stride_dt_cache_head - C_src_ptr += bos * stride_C_src_tok + (pid_h // nheads_ngroups_ratio) * dstate * stride_C_src_c - bc_pre_ptr += pid_b * stride_bc_pre_batch + (pid_h // nheads_ngroups_ratio) * stride_bc_pre_group + dt_cache_ptr += ( + state_batch_idx * stride_dt_cache_batch + pid_h * stride_dt_cache_head + ) + C_src_ptr += ( + bos * stride_C_src_tok + + (pid_h // nheads_ngroups_ratio) * dstate * stride_C_src_c + ) + bc_pre_ptr += ( + pid_b * stride_bc_pre_batch + + (pid_h // nheads_ngroups_ratio) * stride_bc_pre_group + ) if HAS_D: D_ptr += pid_h * stride_D_head if HAS_Z: @@ -259,27 +334,53 @@ def _replayssm_spec_nf_kernel( dt_bias_val = tl.load(dt_bias_ptr).to(tl.float32) if HAS_DT_BIAS else 0.0 # dt over the window (+ bias / softplus), then the per-draft decay weights. - dt_blk = tl.load(dt_cache_ptr + phys_k * stride_dt_cache_pos, mask=cache_valid_mask, other=0.0).to(tl.float32) + dt_blk = tl.load( + dt_cache_ptr + phys_k * stride_dt_cache_pos, mask=cache_valid_mask, other=0.0 + ).to(tl.float32) dt_blk = tl.where(cache_valid_mask, dt_blk, 0.0) if HAS_DT_BIAS: dt_blk = tl.where(cache_valid_mask, dt_blk + dt_bias_val, 0.0) if DT_SOFTPLUS: - dt_blk = tl.where(cache_valid_mask, tl.where(dt_blk <= 20.0, softplus(dt_blk), dt_blk), 0.0) + dt_blk = tl.where( + cache_valid_mask, tl.where(dt_blk <= 20.0, softplus(dt_blk), dt_blk), 0.0 + ) dt_cum = tl.cumsum(dt_blk, axis=0) hist_total = tl.sum(tl.where(hist_mask, dt_blk, 0.0), axis=0) spec_cum = tl.sum(tl.where(spec_prefix_mask, dt_blk[None, :], 0.0), axis=1) spec_cum = tl.where(spec_valid_mask, spec_cum, 0.0) spec_total = hist_total + spec_cum - checkpoint_decay = tl.where(spec_valid_mask, tl.exp(tl.minimum(A_val * spec_total, 0.0)), 0.0) + checkpoint_decay = tl.where( + spec_valid_mask, tl.exp(tl.minimum(A_val * spec_total, 0.0)), 0.0 + ) # Causal weighted sum over cached values: spec_contrib = x_cache @ factor. - x_blk = tl.load(x_cache_ptr + phys_k[None, :] * stride_x_cache_pos + offs_m[:, None] * stride_x_cache_dim, mask=(offs_m[:, None] < dim) & cache_valid_mask[None, :], other=0.0) + x_blk = tl.load( + x_cache_ptr + + phys_k[None, :] * stride_x_cache_pos + + offs_m[:, None] * stride_x_cache_dim, + mask=(offs_m[:, None] < dim) & cache_valid_mask[None, :], + other=0.0, + ) x_ty = x_blk.to(x_cache_ptr.dtype.element_ty) - bc = tl.load(bc_pre_ptr + offs_k[:, None] * stride_bc_pre_pos + offs_s[None, :] * stride_bc_pre_spec, mask=cache_valid_mask[:, None] & spec_valid_mask[None, :], other=0.0).to(tl.float32) - spec_scale = dt_blk[:, None] * tl.exp(tl.minimum(A_val * (spec_total[None, :] - dt_cum[:, None]), 0.0)) - causal = spec_valid_mask[None, :] & cache_valid_mask[:, None] & (offs_k[:, None] <= spec_cache_pos[None, :]) + bc = tl.load( + bc_pre_ptr + + offs_k[:, None] * stride_bc_pre_pos + + offs_s[None, :] * stride_bc_pre_spec, + mask=cache_valid_mask[:, None] & spec_valid_mask[None, :], + other=0.0, + ).to(tl.float32) + spec_scale = dt_blk[:, None] * tl.exp( + tl.minimum(A_val * (spec_total[None, :] - dt_cum[:, None]), 0.0) + ) + causal = ( + spec_valid_mask[None, :] + & cache_valid_mask[:, None] + & (offs_k[:, None] <= spec_cache_pos[None, :]) + ) factor = tl.where(causal, bc * spec_scale, 0.0) - spec_contrib = tl.dot(x_ty, factor.to(x_cache_ptr.dtype.element_ty), input_precision="tf32x3").to(tl.float32) + spec_contrib = tl.dot( + x_ty, factor.to(x_cache_ptr.dtype.element_ty), input_precision="tf32x3" + ).to(tl.float32) # Decayed checkpoint readout S_0 @ C, dstate-tiled. tf32x3 keeps fp32-act # parity; bf16 act uses single-pass tf32 (the flag is a no-op on bf16 inputs). @@ -287,25 +388,57 @@ def _replayssm_spec_nf_kernel( for i in tl.static_range(NDS): offs_n = i * DSTATE_TILE + offs_nt nmask = offs_n < dstate - st = tl.load(state_ptr + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate, mask=(offs_m[:, None] < dim) & nmask[None, :], other=0.0).to(tl.float32) + st = tl.load( + state_ptr + + offs_m[:, None] * stride_state_dim + + offs_n[None, :] * stride_state_dstate, + mask=(offs_m[:, None] < dim) & nmask[None, :], + other=0.0, + ).to(tl.float32) c_mask = spec_valid_mask[:, None] & nmask[None, :] - c_chunk = tl.load(C_src_ptr + offs_s[:, None] * stride_C_src_tok + offs_n[None, :] * stride_C_src_c, mask=c_mask, other=0.0).to(tl.float32) + c_chunk = tl.load( + C_src_ptr + + offs_s[:, None] * stride_C_src_tok + + offs_n[None, :] * stride_C_src_c, + mask=c_mask, + other=0.0, + ).to(tl.float32) if x_cache_ptr.dtype.element_ty == tl.float32: - checkpoint_out += tl.dot(st, tl.trans(c_chunk), input_precision="tf32x3").to(tl.float32) + checkpoint_out += tl.dot( + st, tl.trans(c_chunk), input_precision="tf32x3" + ).to(tl.float32) else: - checkpoint_out += tl.dot(st, tl.trans(c_chunk), input_precision="tf32").to(tl.float32) + checkpoint_out += tl.dot(st, tl.trans(c_chunk), input_precision="tf32").to( + tl.float32 + ) checkpoint_out *= checkpoint_decay[None, :] out = tl.trans(checkpoint_out + spec_contrib) if HAS_D: - x_spec_sm = tl.load(x_cache_ptr + offs_m[None, :] * stride_x_cache_dim + phys_spec[:, None] * stride_x_cache_pos, mask=spec_valid_mask[:, None] & (offs_m[None, :] < dim), other=0.0).to(tl.float32) - D_val = tl.load(D_ptr + offs_m * stride_D_dim, mask=offs_m < dim, other=0.0).to(tl.float32) + x_spec_sm = tl.load( + x_cache_ptr + + offs_m[None, :] * stride_x_cache_dim + + phys_spec[:, None] * stride_x_cache_pos, + mask=spec_valid_mask[:, None] & (offs_m[None, :] < dim), + other=0.0, + ).to(tl.float32) + D_val = tl.load(D_ptr + offs_m * stride_D_dim, mask=offs_m < dim, other=0.0).to( + tl.float32 + ) out += x_spec_sm * D_val[None, :] if HAS_Z: - z_val = tl.load(z_ptr + offs_s[:, None] * stride_z_tok + offs_m[None, :] * stride_z_dim, mask=spec_valid_mask[:, None] & (offs_m[None, :] < dim), other=0.0).to(tl.float32) + z_val = tl.load( + z_ptr + offs_s[:, None] * stride_z_tok + offs_m[None, :] * stride_z_dim, + mask=spec_valid_mask[:, None] & (offs_m[None, :] < dim), + other=0.0, + ).to(tl.float32) out *= z_val * tl.sigmoid(z_val) out = tl.where(spec_valid_mask[:, None], out, 0.0) - tl.store(out_ptr + offs_s[:, None] * stride_out_tok + offs_m[None, :] * stride_out_dim, out, mask=spec_valid_mask[:, None] & (offs_m[None, :] < dim)) + tl.store( + out_ptr + offs_s[:, None] * stride_out_tok + offs_m[None, :] * stride_out_dim, + out, + mask=spec_valid_mask[:, None] & (offs_m[None, :] < dim), + ) # ====================================================================== @@ -316,31 +449,84 @@ def _replayssm_spec_nf_kernel( @triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) @triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) @triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) -@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics( + {"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])} +) @triton.jit def _replayssm_spec_fl_kernel( - state_ptr, x_cache_ptr, dt_cache_ptr, B_cache_ptr, C_src_ptr, bc_pre_ptr, - D_ptr, z_ptr, dt_bias_ptr, A_ptr, out_ptr, is_flush_flags_ptr, write_pos_ptr, - post_origin_ptr, state_batch_indices_ptr, query_start_loc_ptr, null_block_id, - batch, nheads, dim, dstate, max_cache_len, nheads_ngroups_ratio, - stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, - stride_x_cache_batch, stride_x_cache_head, stride_x_cache_dim, stride_x_cache_pos, - stride_dt_cache_batch, stride_dt_cache_head, stride_dt_cache_pos, - stride_B_cache_batch, stride_B_cache_group, stride_B_cache_dstate, stride_B_cache_pos, - stride_C_src_tok, stride_C_src_c, - stride_bc_pre_batch, stride_bc_pre_group, stride_bc_pre_pos, stride_bc_pre_spec, - stride_D_head, stride_D_dim, stride_z_tok, stride_z_head, stride_z_dim, - stride_dt_bias_head, stride_A_head, stride_out_tok, stride_out_head, stride_out_dim, + state_ptr, + x_cache_ptr, + dt_cache_ptr, + B_cache_ptr, + C_src_ptr, + bc_pre_ptr, + D_ptr, + z_ptr, + dt_bias_ptr, + A_ptr, + out_ptr, + is_flush_flags_ptr, + write_pos_ptr, + post_origin_ptr, + state_batch_indices_ptr, + query_start_loc_ptr, + null_block_id, + batch, + nheads, + dim, + dstate, + max_cache_len, + nheads_ngroups_ratio, + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + stride_x_cache_batch, + stride_x_cache_head, + stride_x_cache_dim, + stride_x_cache_pos, + stride_dt_cache_batch, + stride_dt_cache_head, + stride_dt_cache_pos, + stride_B_cache_batch, + stride_B_cache_group, + stride_B_cache_dstate, + stride_B_cache_pos, + stride_C_src_tok, + stride_C_src_c, + stride_bc_pre_batch, + stride_bc_pre_group, + stride_bc_pre_pos, + stride_bc_pre_spec, + stride_D_head, + stride_D_dim, + stride_z_tok, + stride_z_head, + stride_z_dim, + stride_dt_bias_head, + stride_A_head, + stride_out_tok, + stride_out_head, + stride_out_dim, stride_state_indices_batch, - DT_SOFTPLUS: tl.constexpr, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_CACHE: tl.constexpr, - BLOCK_SIZE_SPEC: tl.constexpr, HAS_DT_BIAS: tl.constexpr, HAS_D: tl.constexpr, - HAS_Z: tl.constexpr, CACHE_BUF_LEN: tl.constexpr, DSTATE_TILE: tl.constexpr, - NDS: tl.constexpr, BLOCK_SIZE_DSTATE: tl.constexpr, + DT_SOFTPLUS: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_CACHE: tl.constexpr, + BLOCK_SIZE_SPEC: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + CACHE_BUF_LEN: tl.constexpr, + DSTATE_TILE: tl.constexpr, + NDS: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, ): pid_m = tl.program_id(axis=0) pid_b = tl.program_id(axis=1) pid_h = tl.program_id(axis=2) - state_batch_idx = tl.load(state_batch_indices_ptr + pid_b * stride_state_indices_batch).to(tl.int64) + state_batch_idx = tl.load( + state_batch_indices_ptr + pid_b * stride_state_indices_batch + ).to(tl.int64) if state_batch_idx == null_block_id: return if tl.load(is_flush_flags_ptr + state_batch_idx) == 0: @@ -363,10 +549,21 @@ def _replayssm_spec_fl_kernel( state_ptr += state_batch_idx * stride_state_batch + pid_h * stride_state_head x_cache_ptr += state_batch_idx * stride_x_cache_batch + pid_h * stride_x_cache_head - dt_cache_ptr += state_batch_idx * stride_dt_cache_batch + pid_h * stride_dt_cache_head - B_cache_ptr += state_batch_idx * stride_B_cache_batch + (pid_h // nheads_ngroups_ratio) * stride_B_cache_group - C_src_ptr += bos * stride_C_src_tok + (pid_h // nheads_ngroups_ratio) * dstate * stride_C_src_c - bc_pre_ptr += pid_b * stride_bc_pre_batch + (pid_h // nheads_ngroups_ratio) * stride_bc_pre_group + dt_cache_ptr += ( + state_batch_idx * stride_dt_cache_batch + pid_h * stride_dt_cache_head + ) + B_cache_ptr += ( + state_batch_idx * stride_B_cache_batch + + (pid_h // nheads_ngroups_ratio) * stride_B_cache_group + ) + C_src_ptr += ( + bos * stride_C_src_tok + + (pid_h // nheads_ngroups_ratio) * dstate * stride_C_src_c + ) + bc_pre_ptr += ( + pid_b * stride_bc_pre_batch + + (pid_h // nheads_ngroups_ratio) * stride_bc_pre_group + ) if HAS_D: D_ptr += pid_h * stride_D_head if HAS_Z: @@ -379,7 +576,9 @@ def _replayssm_spec_fl_kernel( dt_bias_val = tl.load(dt_bias_ptr).to(tl.float32) if HAS_DT_BIAS else 0.0 # History decay (for the S_1 reconstruction) and spec-prefix decay (output). - dt_h = tl.load(dt_cache_ptr + phys_h * stride_dt_cache_pos, mask=hist_mask, other=0.0).to(tl.float32) + dt_h = tl.load( + dt_cache_ptr + phys_h * stride_dt_cache_pos, mask=hist_mask, other=0.0 + ).to(tl.float32) dt_h = tl.where(hist_mask, dt_h, 0.0) if HAS_DT_BIAS: dt_h = tl.where(hist_mask, dt_h + dt_bias_val, 0.0) @@ -388,16 +587,30 @@ def _replayssm_spec_fl_kernel( hist_cum = tl.cumsum(dt_h, axis=0) hist_total = tl.sum(dt_h, axis=0) hist_decay = tl.exp(tl.minimum(A_val * hist_total, 0.0)) - hist_scale = tl.where(hist_mask, dt_h * tl.exp(tl.minimum(A_val * (hist_total - hist_cum), 0.0)), 0.0) - dt_s = tl.load(dt_cache_ptr + phys_spec * stride_dt_cache_pos, mask=spec_valid_mask, other=0.0).to(tl.float32) + hist_scale = tl.where( + hist_mask, dt_h * tl.exp(tl.minimum(A_val * (hist_total - hist_cum), 0.0)), 0.0 + ) + dt_s = tl.load( + dt_cache_ptr + phys_spec * stride_dt_cache_pos, mask=spec_valid_mask, other=0.0 + ).to(tl.float32) dt_s = tl.where(spec_valid_mask, dt_s, 0.0) if HAS_DT_BIAS: dt_s = tl.where(spec_valid_mask, dt_s + dt_bias_val, 0.0) if DT_SOFTPLUS: - dt_s = tl.where(spec_valid_mask, tl.where(dt_s <= 20.0, softplus(dt_s), dt_s), 0.0) + dt_s = tl.where( + spec_valid_mask, tl.where(dt_s <= 20.0, softplus(dt_s), dt_s), 0.0 + ) spec_cum = tl.cumsum(dt_s, axis=0) - spec_decay = tl.where(spec_valid_mask, tl.exp(tl.minimum(A_val * spec_cum, 0.0)), 0.0) - x_hist = tl.load(x_cache_ptr + phys_h[None, :] * stride_x_cache_pos + offs_m[:, None] * stride_x_cache_dim, mask=(offs_m[:, None] < dim) & hist_mask[None, :], other=0.0) + spec_decay = tl.where( + spec_valid_mask, tl.exp(tl.minimum(A_val * spec_cum, 0.0)), 0.0 + ) + x_hist = tl.load( + x_cache_ptr + + phys_h[None, :] * stride_x_cache_pos + + offs_m[:, None] * stride_x_cache_dim, + mask=(offs_m[:, None] < dim) & hist_mask[None, :], + other=0.0, + ) x_hist_ty = x_hist.to(x_cache_ptr.dtype.element_ty) # Reconstruct S_1 = S_0 * hist_decay + (x_hist @ scaled B_hist), store it, and @@ -406,47 +619,117 @@ def _replayssm_spec_fl_kernel( for i in tl.static_range(NDS): offs_n = i * DSTATE_TILE + offs_nt nmask = offs_n < dstate - B_block = tl.load(B_cache_ptr + phys_h[:, None] * stride_B_cache_pos + offs_n[None, :] * stride_B_cache_dstate, mask=hist_mask[:, None] & nmask[None, :], other=0.0) - B_hist_scaled = (tl.where(hist_mask[:, None], B_block.to(tl.float32), 0.0) * hist_scale[:, None]).to(x_cache_ptr.dtype.element_ty) - delta = tl.dot(x_hist_ty, B_hist_scaled, input_precision="tf32x3").to(tl.float32) - st_ptrs = state_ptr + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + B_block = tl.load( + B_cache_ptr + + phys_h[:, None] * stride_B_cache_pos + + offs_n[None, :] * stride_B_cache_dstate, + mask=hist_mask[:, None] & nmask[None, :], + other=0.0, + ) + B_hist_scaled = ( + tl.where(hist_mask[:, None], B_block.to(tl.float32), 0.0) + * hist_scale[:, None] + ).to(x_cache_ptr.dtype.element_ty) + delta = tl.dot(x_hist_ty, B_hist_scaled, input_precision="tf32x3").to( + tl.float32 + ) + st_ptrs = ( + state_ptr + + offs_m[:, None] * stride_state_dim + + offs_n[None, :] * stride_state_dstate + ) st = tl.load(st_ptrs, mask=(offs_m[:, None] < dim) & nmask[None, :], other=0.0) S1 = st.to(tl.float32) * hist_decay + delta if write_pos > 0: - tl.store(st_ptrs, S1.to(st.dtype), mask=(offs_m[:, None] < dim) & nmask[None, :]) + tl.store( + st_ptrs, S1.to(st.dtype), mask=(offs_m[:, None] < dim) & nmask[None, :] + ) c_mask = spec_valid_mask[:, None] & nmask[None, :] - c_chunk = tl.load(C_src_ptr + offs_s[:, None] * stride_C_src_tok + offs_n[None, :] * stride_C_src_c, mask=c_mask, other=0.0).to(tl.float32) + c_chunk = tl.load( + C_src_ptr + + offs_s[:, None] * stride_C_src_tok + + offs_n[None, :] * stride_C_src_c, + mask=c_mask, + other=0.0, + ).to(tl.float32) if x_cache_ptr.dtype.element_ty == tl.float32: - checkpoint_out += tl.dot(S1, tl.trans(c_chunk), input_precision="tf32x3").to(tl.float32) + checkpoint_out += tl.dot( + S1, tl.trans(c_chunk), input_precision="tf32x3" + ).to(tl.float32) else: - checkpoint_out += tl.dot(S1, tl.trans(c_chunk), input_precision="tf32").to(tl.float32) + checkpoint_out += tl.dot(S1, tl.trans(c_chunk), input_precision="tf32").to( + tl.float32 + ) checkpoint_out *= spec_decay[None, :] # Intra-spec window contribution: intra = x_spec @ factor_intra (causal T x T). - bc_spec = tl.load(bc_pre_ptr + (write_pos + offs_k)[:, None] * stride_bc_pre_pos + offs_s[None, :] * stride_bc_pre_spec, mask=(offs_k[:, None] < spec_len) & spec_valid_mask[None, :], other=0.0).to(tl.float32) - dt_k = tl.load(dt_cache_ptr + ((post_origin + write_pos + offs_k) & (CACHE_BUF_LEN - 1)) * stride_dt_cache_pos, mask=offs_k < spec_len, other=0.0).to(tl.float32) + bc_spec = tl.load( + bc_pre_ptr + + (write_pos + offs_k)[:, None] * stride_bc_pre_pos + + offs_s[None, :] * stride_bc_pre_spec, + mask=(offs_k[:, None] < spec_len) & spec_valid_mask[None, :], + other=0.0, + ).to(tl.float32) + dt_k = tl.load( + dt_cache_ptr + + ((post_origin + write_pos + offs_k) & (CACHE_BUF_LEN - 1)) + * stride_dt_cache_pos, + mask=offs_k < spec_len, + other=0.0, + ).to(tl.float32) dt_k = tl.where(offs_k < spec_len, dt_k, 0.0) if HAS_DT_BIAS: dt_k = tl.where(offs_k < spec_len, dt_k + dt_bias_val, 0.0) if DT_SOFTPLUS: - dt_k = tl.where(offs_k < spec_len, tl.where(dt_k <= 20.0, softplus(dt_k), dt_k), 0.0) + dt_k = tl.where( + offs_k < spec_len, tl.where(dt_k <= 20.0, softplus(dt_k), dt_k), 0.0 + ) speccum_k = tl.cumsum(dt_k, axis=0) - causal = (offs_k[:, None] < spec_len) & spec_valid_mask[None, :] & (offs_k[:, None] <= offs_s[None, :]) + causal = ( + (offs_k[:, None] < spec_len) + & spec_valid_mask[None, :] + & (offs_k[:, None] <= offs_s[None, :]) + ) decay_ks = tl.exp(tl.minimum(A_val * (spec_cum[None, :] - speccum_k[:, None]), 0.0)) factor_intra = tl.where(causal, bc_spec * dt_k[:, None] * decay_ks, 0.0) - x_src = tl.load(x_cache_ptr + ((post_origin + write_pos + offs_k)[None, :] & (CACHE_BUF_LEN - 1)) * stride_x_cache_pos + offs_m[:, None] * stride_x_cache_dim, mask=(offs_m[:, None] < dim) & (offs_k[None, :] < spec_len), other=0.0).to(x_cache_ptr.dtype.element_ty) - intra = tl.dot(x_src, factor_intra.to(x_cache_ptr.dtype.element_ty), input_precision="tf32x3").to(tl.float32) + x_src = tl.load( + x_cache_ptr + + ((post_origin + write_pos + offs_k)[None, :] & (CACHE_BUF_LEN - 1)) + * stride_x_cache_pos + + offs_m[:, None] * stride_x_cache_dim, + mask=(offs_m[:, None] < dim) & (offs_k[None, :] < spec_len), + other=0.0, + ).to(x_cache_ptr.dtype.element_ty) + intra = tl.dot( + x_src, factor_intra.to(x_cache_ptr.dtype.element_ty), input_precision="tf32x3" + ).to(tl.float32) out = tl.trans(checkpoint_out + intra) if HAS_D: - x_spec_sm = tl.load(x_cache_ptr + offs_m[None, :] * stride_x_cache_dim + phys_spec[:, None] * stride_x_cache_pos, mask=spec_valid_mask[:, None] & (offs_m[None, :] < dim), other=0.0).to(tl.float32) - D_val = tl.load(D_ptr + offs_m * stride_D_dim, mask=offs_m < dim, other=0.0).to(tl.float32) + x_spec_sm = tl.load( + x_cache_ptr + + offs_m[None, :] * stride_x_cache_dim + + phys_spec[:, None] * stride_x_cache_pos, + mask=spec_valid_mask[:, None] & (offs_m[None, :] < dim), + other=0.0, + ).to(tl.float32) + D_val = tl.load(D_ptr + offs_m * stride_D_dim, mask=offs_m < dim, other=0.0).to( + tl.float32 + ) out += x_spec_sm * D_val[None, :] if HAS_Z: - z_val = tl.load(z_ptr + offs_s[:, None] * stride_z_tok + offs_m[None, :] * stride_z_dim, mask=spec_valid_mask[:, None] & (offs_m[None, :] < dim), other=0.0).to(tl.float32) + z_val = tl.load( + z_ptr + offs_s[:, None] * stride_z_tok + offs_m[None, :] * stride_z_dim, + mask=spec_valid_mask[:, None] & (offs_m[None, :] < dim), + other=0.0, + ).to(tl.float32) out *= z_val * tl.sigmoid(z_val) out = tl.where(spec_valid_mask[:, None], out, 0.0) - tl.store(out_ptr + offs_s[:, None] * stride_out_tok + offs_m[None, :] * stride_out_dim, out, mask=spec_valid_mask[:, None] & (offs_m[None, :] < dim)) + tl.store( + out_ptr + offs_s[:, None] * stride_out_tok + offs_m[None, :] * stride_out_dim, + out, + mask=spec_valid_mask[:, None] & (offs_m[None, :] < dim), + ) @triton.jit @@ -466,19 +749,21 @@ def _advance_write_pos_origin_kernel( ): offs = tl.arange(0, BLOCK_SIZE) row_mask = offs < batch - state_batch_idx = tl.load(state_batch_indices_ptr + offs * stride_state_indices_batch, mask=row_mask, other=null_block_id).to( - tl.int64 - ) + state_batch_idx = tl.load( + state_batch_indices_ptr + offs * stride_state_indices_batch, + mask=row_mask, + other=null_block_id, + ).to(tl.int64) valid = row_mask & (state_batch_idx != null_block_id) - write_pos = tl.load( - write_pos_ptr + state_batch_idx, mask=valid, other=0 - ).to(tl.int32) - post_origin = tl.load( - post_origin_ptr + state_batch_idx, mask=valid, other=0 - ).to(tl.int32) - is_flush_cur = tl.load( - is_flush_ptr + state_batch_idx, mask=valid, other=0 - ).to(tl.int32) + write_pos = tl.load(write_pos_ptr + state_batch_idx, mask=valid, other=0).to( + tl.int32 + ) + post_origin = tl.load(post_origin_ptr + state_batch_idx, mask=valid, other=0).to( + tl.int32 + ) + is_flush_cur = tl.load(is_flush_ptr + state_batch_idx, mask=valid, other=0).to( + tl.int32 + ) num_accepted = tl.load(num_accepted_ptr + offs, mask=valid, other=0).to(tl.int32) total_commit = tl.where(valid, num_accepted, 0).to(tl.int32) flush_now = (total_commit > 0) & (is_flush_cur != 0) @@ -513,9 +798,11 @@ def _reset_replayssm_spec_cursors_kernel( ): offs = tl.arange(0, BLOCK_SIZE) row_mask = offs < batch - state_batch_idx = tl.load(state_batch_indices_ptr + offs * stride_state_indices_batch, mask=row_mask, other=null_block_id).to( - tl.int64 - ) + state_batch_idx = tl.load( + state_batch_indices_ptr + offs * stride_state_indices_batch, + mask=row_mask, + other=null_block_id, + ).to(tl.int64) first = tl.load(first_decode_ptr + offs, mask=row_mask, other=0).to(tl.int32) do_reset = row_mask & (state_batch_idx != null_block_id) & (first != 0) tl.store( @@ -576,7 +863,9 @@ def selective_state_update_replayssm_spec( total_tokens, conv_dim = conv_out.shape buf = post_conv_cache.shape[1] cache_buf_len = buf - assert cache_buf_len & (cache_buf_len - 1) == 0, "cache_buf_len must be a power of two" + assert cache_buf_len & (cache_buf_len - 1) == 0, ( + "cache_buf_len must be a power of two" + ) assert d_inner == nheads * dim cache_conv_dim = d_inner + ngroups * dstate # x|B only (no C) assert post_conv_cache.shape == (num_blocks, buf, cache_conv_dim) @@ -587,7 +876,9 @@ def selective_state_update_replayssm_spec( assert query_start_loc.shape[0] == batch + 1 if out is None: - out = torch.empty(total_tokens, nheads, dim, device=conv_out.device, dtype=conv_out.dtype) + out = torch.empty( + total_tokens, nheads, dim, device=conv_out.device, dtype=conv_out.dtype + ) if total_tokens == 0: return out @@ -599,16 +890,43 @@ def selective_state_update_replayssm_spec( block_dstate = triton.next_power_of_2(dstate) bsm_v, nw_v, dt_v, ns_v = get_replayssm_config( - "mamba2_spec_verify", dstate=dstate, base_block=base_block, max_spec_len=max_spec_len + "mamba2_spec_verify", + dstate=dstate, + base_block=base_block, + max_spec_len=max_spec_len, ) bsm_f, nw_f, dt_f, ns_f = get_replayssm_config( - "mamba2_spec_flush", dstate=dstate, base_block=base_block, max_spec_len=max_spec_len + "mamba2_spec_flush", + dstate=dstate, + base_block=base_block, + max_spec_len=max_spec_len, ) - dt_v = max(16, min(dt_v, block_dstate)); nds_v = triton.cdiv(block_dstate, dt_v) - dt_f = max(16, min(dt_f, block_dstate)); nds_f = triton.cdiv(block_dstate, dt_f) + dt_v = max(16, min(dt_v, block_dstate)) + nds_v = triton.cdiv(block_dstate, dt_v) + dt_f = max(16, min(dt_f, block_dstate)) + nds_f = triton.cdiv(block_dstate, dt_f) if bc_pre is None: - bc_pre = torch.empty(batch, ngroups, buf, block_spec, device=conv_out.device, dtype=conv_out.dtype) + bc_pre = torch.empty( + batch, + ngroups, + buf, + block_spec, + device=conv_out.device, + dtype=conv_out.dtype, + ) + else: + # The kernels index bc_pre with its raw strides on a (batch, ngroups) + # grid; an under-sized group dim aliases across rows (races / OOB). + assert ( + bc_pre.shape[0] >= batch + and bc_pre.shape[1] == ngroups + and bc_pre.shape[2] >= L + and bc_pre.shape[3] >= block_spec + ), ( + f"bc_pre shape {tuple(bc_pre.shape)} incompatible with " + f"(batch={batch}, ngroups={ngroups}, L={L}, block_spec={block_spec})" + ) sis = state_batch_indices.stride(0) # --- fused scatter + precompute (full-window bc over [0, L)) --- @@ -617,22 +935,62 @@ def selective_state_update_replayssm_spec( BLOCK_CX = 256 NCX = triton.cdiv(ratio_p, BLOCK_CX) block_hl = max(1, triton.next_power_of_2(ratio)) - with torch.cuda.device(conv_out.device.index): + with torch.accelerator.device_index(conv_out.device.index): _fused_scatter_precompute_kernel[(batch, ngroups)]( - conv_out, dt_spec, post_conv_cache, dt_cache, write_pos, post_conv_state_pos, - bc_pre, state_batch_indices, query_start_loc, null_block_id, batch, ngroups, - nheads, dstate, d_inner, conv_dim, L, - conv_out.stride(0), conv_out.stride(1), dt_spec.stride(0), dt_spec.stride(1), - post_conv_cache.stride(0), post_conv_cache.stride(1), post_conv_cache.stride(2), - dt_cache.stride(0), dt_cache.stride(1), dt_cache.stride(2), - bc_pre.stride(0), bc_pre.stride(1), bc_pre.stride(2), bc_pre.stride(3), sis, - RATIO=ratio, RATIO_P=ratio_p, NCX=NCX, BLOCK_CX=BLOCK_CX, CACHE_BUF_LEN=cache_buf_len, - BLOCK_SIZE_CACHE=pre_block, BLOCK_SIZE_SPEC=block_spec, BLOCK_HL=block_hl, num_warps=4, + conv_out, + dt_spec, + post_conv_cache, + dt_cache, + write_pos, + post_conv_state_pos, + bc_pre, + state_batch_indices, + query_start_loc, + null_block_id, + batch, + ngroups, + nheads, + dstate, + d_inner, + conv_dim, + L, + conv_out.stride(0), + conv_out.stride(1), + dt_spec.stride(0), + dt_spec.stride(1), + post_conv_cache.stride(0), + post_conv_cache.stride(1), + post_conv_cache.stride(2), + dt_cache.stride(0), + dt_cache.stride(1), + dt_cache.stride(2), + bc_pre.stride(0), + bc_pre.stride(1), + bc_pre.stride(2), + bc_pre.stride(3), + sis, + RATIO=ratio, + RATIO_P=ratio_p, + NCX=NCX, + BLOCK_CX=BLOCK_CX, + CACHE_BUF_LEN=cache_buf_len, + BLOCK_SIZE_CACHE=pre_block, + BLOCK_SIZE_SPEC=block_spec, + BLOCK_HL=block_hl, + num_warps=4, ) # views into the paged circular post-conv cache: x | B | C on the channel axis - x_view = post_conv_cache[:, :, :d_inner].view(num_blocks, buf, nheads, dim).permute(0, 2, 1, 3) - B_view = post_conv_cache[:, :, d_inner : d_inner + ngroups * dstate].view(num_blocks, buf, ngroups, dstate).permute(0, 2, 1, 3) + x_view = ( + post_conv_cache[:, :, :d_inner] + .view(num_blocks, buf, nheads, dim) + .permute(0, 2, 1, 3) + ) + B_view = ( + post_conv_cache[:, :, d_inner : d_inner + ngroups * dstate] + .view(num_blocks, buf, ngroups, dstate) + .permute(0, 2, 1, 3) + ) # C is not cached; the kernels read it fresh from this conv_out slice. C_src = conv_out[:, d_inner + ngroups * dstate :] z_strides = (z.stride(0), z.stride(1), z.stride(2)) if z is not None else (0, 0, 0) @@ -640,30 +998,87 @@ def selective_state_update_replayssm_spec( def _args(bsm): grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads) return grid, ( - state_checkpoint, x_view, dt_cache, B_view, C_src, bc_pre, D, z, dt_bias, A, - out, is_flush, write_pos, post_conv_state_pos, state_batch_indices, - query_start_loc, null_block_id, batch, nheads, dim, dstate, L, ratio, - state_checkpoint.stride(0), state_checkpoint.stride(1), state_checkpoint.stride(2), state_checkpoint.stride(3), - x_view.stride(0), x_view.stride(1), x_view.stride(3), x_view.stride(2), - dt_cache.stride(0), dt_cache.stride(1), dt_cache.stride(2), - B_view.stride(0), B_view.stride(1), B_view.stride(3), B_view.stride(2), - C_src.stride(0), C_src.stride(1), - bc_pre.stride(0), bc_pre.stride(1), bc_pre.stride(2), bc_pre.stride(3), - D.stride(0) if D is not None else 0, D.stride(1) if D is not None else 0, - z_strides[0], z_strides[1], z_strides[2], dt_bias.stride(0) if dt_bias is not None else 0, - A.stride(0), out.stride(0), out.stride(1), out.stride(2), sis, dt_softplus, bsm, + state_checkpoint, + x_view, + dt_cache, + B_view, + C_src, + bc_pre, + D, + z, + dt_bias, + A, + out, + is_flush, + write_pos, + post_conv_state_pos, + state_batch_indices, + query_start_loc, + null_block_id, + batch, + nheads, + dim, + dstate, + L, + ratio, + state_checkpoint.stride(0), + state_checkpoint.stride(1), + state_checkpoint.stride(2), + state_checkpoint.stride(3), + x_view.stride(0), + x_view.stride(1), + x_view.stride(3), + x_view.stride(2), + dt_cache.stride(0), + dt_cache.stride(1), + dt_cache.stride(2), + B_view.stride(0), + B_view.stride(1), + B_view.stride(3), + B_view.stride(2), + C_src.stride(0), + C_src.stride(1), + bc_pre.stride(0), + bc_pre.stride(1), + bc_pre.stride(2), + bc_pre.stride(3), + D.stride(0) if D is not None else 0, + D.stride(1) if D is not None else 0, + z_strides[0], + z_strides[1], + z_strides[2], + dt_bias.stride(0) if dt_bias is not None else 0, + A.stride(0), + out.stride(0), + out.stride(1), + out.stride(2), + sis, + dt_softplus, + bsm, ) - with torch.cuda.device(state_checkpoint.device.index): + with torch.accelerator.device_index(state_checkpoint.device.index): grid, base = _args(bsm_v) _replayssm_spec_nf_kernel[grid]( - *base, main_block, block_spec, - CACHE_BUF_LEN=cache_buf_len, DSTATE_TILE=dt_v, NDS=nds_v, num_warps=nw_v, num_stages=ns_v, + *base, + main_block, + block_spec, + CACHE_BUF_LEN=cache_buf_len, + DSTATE_TILE=dt_v, + NDS=nds_v, + num_warps=nw_v, + num_stages=ns_v, ) grid, base = _args(bsm_f) _replayssm_spec_fl_kernel[grid]( - *base, main_block, block_spec, - CACHE_BUF_LEN=cache_buf_len, DSTATE_TILE=dt_f, NDS=nds_f, num_warps=nw_f, num_stages=ns_f, + *base, + main_block, + block_spec, + CACHE_BUF_LEN=cache_buf_len, + DSTATE_TILE=dt_f, + NDS=nds_f, + num_warps=nw_f, + num_stages=ns_f, ) return out @@ -686,7 +1101,7 @@ def commit_replayssm_spec( if cache_buf_len is None: cache_buf_len = max(1, triton.next_power_of_2(max_cache_len)) BLOCK = max(1, triton.next_power_of_2(batch)) - with torch.cuda.device(write_pos.device.index): + with torch.accelerator.device_index(write_pos.device.index): _advance_write_pos_origin_kernel[(1,)]( write_pos, post_conv_state_pos, @@ -719,7 +1134,7 @@ def reset_replayssm_spec_cursors( batch = state_batch_indices.shape[0] BLOCK = max(1, triton.next_power_of_2(batch)) init_is_flush = 1 if 2 * max_spec_len > max_cache_len else 0 - with torch.cuda.device(write_pos.device.index): + with torch.accelerator.device_index(write_pos.device.index): _reset_replayssm_spec_cursors_kernel[(1,)]( write_pos, post_conv_state_pos, diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index 8fae389aadef..a105248505bd 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -237,7 +237,13 @@ def __init__( local_nheads, head_dim, dstate = kv_cache_spec.shapes[1] conv_dim_local = kv_cache_spec.shapes[2][1] d_inner_local = local_nheads * head_dim - ngroups_local = (conv_dim_local - d_inner_local) // (2 * dstate) + # The spec post_conv page caches x|B only (C is read fresh from + # conv_out), so the width beyond d_inner is ngroups_local * dstate. + ngroups_local = (conv_dim_local - d_inner_local) // dstate + assert ngroups_local >= 1, ( + f"invalid ngroups_local={ngroups_local} derived from spec page " + f"width {conv_dim_local} (d_inner {d_inner_local}, dstate {dstate})" + ) block_spec = 1 << (max(1, self.max_spec_len) - 1).bit_length() # This is a PER-STEP scratch consumed by the scatter on every decode # step (eager AND cudagraph), indexed by pid_b in [0, num_decodes). From 784042467188235674b6ad83f636049162ce898f Mon Sep 17 00:00:00 2001 From: Johnny-Liou Date: Thu, 16 Jul 2026 13:32:18 -0400 Subject: [PATCH 13/23] [Bugfix] ReplaySSM Mamba2 spec: pass the spec window length as conv max_query_len Signed-off-by: Johnny-Liou --- vllm/model_executor/layers/mamba/mamba_mixer2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index b3bd303bdf12..bfa9908d5930 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -1055,7 +1055,7 @@ def conv_ssm_forward( initial_state_idx=block_idx_last_computed_token_d, num_accepted_tokens=num_accepted_tokens, query_start_loc=query_start_loc_d, - max_query_len=state_indices_tensor_d.size(-1), + max_query_len=self.max_spec_len, ) # cached-spec (hybrid) feeds the full channel-last post-conv output From 5629acf28272822362a11c98fd3de990d0d94680 Mon Sep 17 00:00:00 2001 From: Johnny-Liou Date: Thu, 16 Jul 2026 13:36:56 -0400 Subject: [PATCH 14/23] [Bugfix] ReplaySSM Mamba2 spec: keep cursor metadata on draft-less steps Signed-off-by: Johnny-Liou --- vllm/v1/worker/gpu_model_runner.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index c19252084f16..9ca4ea9f6a67 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -2441,7 +2441,14 @@ def _build_attn_group_metadata( ) extra_attn_metadata_args = {} - if use_spec_decode and isinstance( + # Under ReplaySSM-spec, draft-less steps must still reach the + # spec cursor machinery: + needs_mamba_spec_args = use_spec_decode or ( + self.speculative_config is not None + and self.cache_config.use_replayssm_spec + and isinstance(builder, Mamba2AttentionMetadataBuilder) + ) + if needs_mamba_spec_args and isinstance( builder, (Mamba2AttentionMetadataBuilder, GDNAttentionMetadataBuilder) ): assert ubid is None, "UBatching not supported with GDN yet" From 5d9dfc0c4b222051e083447ab6998c9ec277e792 Mon Sep 17 00:00:00 2001 From: danielafrimi <45691845+danielafrimi@users.noreply.github.com> Date: Tue, 7 Jul 2026 14:45:13 +0300 Subject: [PATCH 15/23] [BugFix] Fix ModelOpt mixed-precision quantization for sparse `quantized_layers` configs. (#47318) Signed-off-by: Daniel Afrimi Signed-off-by: (cherry picked from commit 0a2965b1b35f34fde29c9a653696bac0ef1e499d) --- tests/quantization/test_modelopt.py | 43 +++++++++++++++++++ .../layers/quantization/modelopt.py | 33 +++++++++----- 2 files changed, 66 insertions(+), 10 deletions(-) diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index 32450231487e..002c0e058c84 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -133,6 +133,49 @@ def test_modelopt_mixed_precision_quantizes_parallel_lm_head(): assert isinstance(method, ModelOptNvFp4LinearMethod) +def test_modelopt_mixed_precision_infers_fused_gate_up_projection(): + from vllm.model_executor.layers.linear import LinearBase + + config = _mixed_precision_config( + { + "model.layers.0.mlp.gate_proj": {"quant_algo": "NVFP4"}, + "model.layers.0.mlp.up_proj": {"quant_algo": "NVFP4"}, + } + ) + + fake_layer = MagicMock(spec=LinearBase) + with patch( + "vllm.model_executor.layers.quantization.modelopt.init_nvfp4_linear_kernel" + ): + method = config.get_quant_method(fake_layer, "model.layers.0.mlp.gate_up_proj") + + assert isinstance(method, ModelOptNvFp4LinearMethod) + + +@pytest.mark.parametrize( + ("quantized_prefix", "missing_prefix"), + [ + ("model.layers.0.mlp.gate_proj", "model.layers.0.mlp.down_proj"), + ("model.layers.0.self_attn.o_proj", "model.layers.0.self_attn.qkv_proj"), + ], +) +def test_modelopt_mixed_precision_does_not_infer_missing_sibling_linear( + quantized_prefix, missing_prefix +): + from vllm.model_executor.layers.linear import LinearBase + + config = _mixed_precision_config( + { + quantized_prefix: {"quant_algo": "NVFP4"}, + } + ) + + fake_layer = MagicMock(spec=LinearBase) + method = config.get_quant_method(fake_layer, missing_prefix) + + assert isinstance(method, UnquantizedLinearMethod) + + def test_vocab_parallel_embedding_weight_loader_accepts_scalar_scale(): holder = Mock() scale = torch.nn.Parameter(torch.empty(1)) diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 8fa1cb4d5445..a122e2e9be4d 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -2450,16 +2450,29 @@ def _resolve_quant_algo(self, prefix: str) -> str | None: if key.startswith(parent_dot): return info["quant_algo"].upper() - # 4. Parent-prefix fallback for fused projections (qkv_proj, gate_up_proj). - for candidate in self._quantized_layer_prefix_candidates(prefix): - parent_dot = candidate.rsplit(".", 1)[0] + "." - algos = { - info["quant_algo"].upper() - for key, info in self.quantized_layers.items() - if key.startswith(parent_dot) and "." not in key[len(parent_dot) :] - } - if len(algos) == 1: - return algos.pop() + # 4. Parent-prefix fallback for fused projections whose config lists + # shard names instead of vLLM's packed module name. + fused_projection_shards = { + "qkv_proj": ("q_proj", "k_proj", "v_proj"), + "gate_up_proj": ("gate_proj", "up_proj"), + } + shard_names = fused_projection_shards.get(proj_name) + if shard_names is not None: + for candidate in self._quantized_layer_prefix_candidates(prefix): + parent_dot = candidate.rsplit(".", 1)[0] + "." + shard_algos: set[str] = set() + for shard_name in shard_names: + shard_prefix = f"{parent_dot}{shard_name}" + if shard_prefix in self.quantized_layers: + algo = self.quantized_layers[shard_prefix]["quant_algo"].upper() + shard_algos.add(algo) + if len(shard_algos) == 1: + return shard_algos.pop() + if len(shard_algos) > 1: + raise ValueError( + f"Mixed quant_algo within fused layer {prefix}: " + f"{shard_algos}. All shards must use the same quantization." + ) return None From 6167578c240860773aeca736dbfa19b6ad839113 Mon Sep 17 00:00:00 2001 From: Johnny-Liou Date: Mon, 20 Jul 2026 14:31:18 -0400 Subject: [PATCH 16/23] [Bugfix] ReplaySSM GDN spec: keep SSM state at fp32 through verify/flush Signed-off-by: Johnny-Liou --- .../fla/ops/gdn_replayssm_spec_decode.py | 142 +++++++++++++----- .../layers/mamba/mamba_utils.py | 13 +- 2 files changed, 113 insertions(+), 42 deletions(-) diff --git a/vllm/model_executor/layers/fla/ops/gdn_replayssm_spec_decode.py b/vllm/model_executor/layers/fla/ops/gdn_replayssm_spec_decode.py index 3bc3295e0efa..d32d88895984 100644 --- a/vllm/model_executor/layers/fla/ops/gdn_replayssm_spec_decode.py +++ b/vllm/model_executor/layers/fla/ops/gdn_replayssm_spec_decode.py @@ -54,6 +54,7 @@ def gdn_replayssm_spec_circular_kernel( USE_QK_L2NORM_IN_KERNEL: tl.constexpr, IS_FLUSH: tl.constexpr, NULL_BLOCK_ID: tl.constexpr, + DOT_PRECISION: tl.constexpr, ): i_v = tl.program_id(0) i_n = tl.program_id(1) @@ -141,7 +142,7 @@ def gdn_replayssm_spec_circular_kernel( b_d_all = tl.load( p_d_main, mask=mask_v[:, None] & cache_valid[None, :], other=0.0 ).to(tl.float32) - b_d_scaled = (b_d_all * b_replay_decay[None, :]).to(mixed_qkv.dtype.element_ty) + b_d_scaled = b_d_all * b_replay_decay[None, :] if USE_QK_L2NORM_IN_KERNEL: qnorm_acc = tl.zeros([BS], dtype=tl.float32) @@ -206,8 +207,8 @@ def gdn_replayssm_spec_circular_kernel( mask=ld_s, other=0.0, ).to(tl.float32) - q_tile = (q_tile * (q_rnorm * scale)[:, None]).to(mixed_qkv.dtype.element_ty) - k_tile = (k_tile * k_rnorm[:, None]).to(mixed_qkv.dtype.element_ty) + q_tile = q_tile * (q_rnorm * scale)[:, None] + k_tile = k_tile * k_rnorm[:, None] p_h0 = ( h0 @@ -216,9 +217,9 @@ def gdn_replayssm_spec_circular_kernel( + o_v[:, None] * K + o_kt[None, :] ) - sc_tile = tl.load( - p_h0, mask=mask_v[:, None] & mask_kt[None, :], other=0.0 - ).to(mixed_qkv.dtype.element_ty) + sc_tile = tl.load(p_h0, mask=mask_v[:, None] & mask_kt[None, :], other=0.0).to( + tl.float32 + ) # cached-key history load -> phys_c p_k = k_cache + ( state_idx * stride_k_slot @@ -227,18 +228,22 @@ def gdn_replayssm_spec_circular_kernel( ) khist_tile = tl.load( p_k, mask=cache_valid[:, None] & mask_kt[None, :], other=0.0 - ).to(mixed_qkv.dtype.element_ty) + ).to(tl.float32) qT = tl.trans(q_tile) kT = tl.trans(k_tile) - kk_mat += tl.dot(k_tile, kT) - kq_mat += tl.dot(k_tile, qT) + kk_mat += tl.dot(k_tile, kT, input_precision=DOT_PRECISION) + kq_mat += tl.dot(k_tile, qT, input_precision=DOT_PRECISION) if IS_FLUSH: - sw_f = tl.dot(b_d_scaled, khist_tile, acc=b_total_decay * sc_tile.to(tl.float32)) - sw_tile = sw_f.to(mixed_qkv.dtype.element_ty) - hw_q += tl.dot(sw_tile, qT) - hw_k += tl.dot(sw_tile, kT) + sw_f = tl.dot( + b_d_scaled, + khist_tile, + acc=b_total_decay * sc_tile, + input_precision=DOT_PRECISION, + ) + hw_q += tl.dot(sw_f, qT, input_precision=DOT_PRECISION) + hw_k += tl.dot(sw_f, kT, input_precision=DOT_PRECISION) p_ht = ( ht + state_idx * stride_state_slot @@ -246,12 +251,17 @@ def gdn_replayssm_spec_circular_kernel( + o_v[:, None] * K + o_kt[None, :] ) - tl.store(p_ht, sw_tile, mask=mask_v[:, None] & mask_kt[None, :]) + # fp32 store to the fp32 checkpoint page + tl.store( + p_ht, + sw_f.to(p_ht.dtype.element_ty), + mask=mask_v[:, None] & mask_kt[None, :], + ) else: - hw_q += tl.dot(sc_tile, qT) - hw_k += tl.dot(sc_tile, kT) - scores_q += tl.dot(khist_tile, qT) - scores_k += tl.dot(khist_tile, kT) + hw_q += tl.dot(sc_tile, qT, input_precision=DOT_PRECISION) + hw_k += tl.dot(sc_tile, kT, input_precision=DOT_PRECISION) + scores_q += tl.dot(khist_tile, qT, input_precision=DOT_PRECISION) + scores_k += tl.dot(khist_tile, kT, input_precision=DOT_PRECISION) if write_k: # spec key store -> phys_spec (circular) @@ -262,15 +272,19 @@ def gdn_replayssm_spec_circular_kernel( ) tl.store( p_cur_k, - k_tile, + k_tile.to(p_cur_k.dtype.element_ty), mask=mask_s[:, None] & mask_kt[None, :] & ((b_write_pos + o_s[:, None]) < MAX_CACHE_LEN), ) if not IS_FLUSH: - hw_q = b_total_decay * hw_q + tl.dot(b_d_scaled, scores_q.to(b_d_scaled.dtype)) - hw_k = b_total_decay * hw_k + tl.dot(b_d_scaled, scores_k.to(b_d_scaled.dtype)) + hw_q = b_total_decay * hw_q + tl.dot( + b_d_scaled, scores_q, input_precision=DOT_PRECISION + ) + hw_k = b_total_decay * hw_k + tl.dot( + b_d_scaled, scores_k, input_precision=DOT_PRECISION + ) # ------------------------------------------------------------------ # strictly-lower A and T = (I + A)^{-1}. @@ -297,7 +311,9 @@ def gdn_replayssm_spec_circular_kernel( + i_hv * V + o_v[:, None] ) - v_tile = tl.load(p_v, mask=mask_v[:, None] & mask_s[None, :], other=0.0).to(tl.float32) + v_tile = tl.load(p_v, mask=mask_v[:, None] & mask_s[None, :], other=0.0).to( + tl.float32 + ) R_mat = beta_s[None, :] * (v_tile - expG_s[None, :] * hw_k) D_spec = tl.zeros([BV, BS], dtype=tl.float32) for j in tl.static_range(BS): @@ -366,9 +382,9 @@ def _advance_gdn_spec_cursors_kernel( write_pos = tl.load(write_pos_ptr + blk, mask=valid, other=0).to(tl.int32) cache_base = tl.load(cache_base_ptr + blk, mask=valid, other=0).to(tl.int32) is_flush_cur = tl.load(is_flush_ptr + blk, mask=valid, other=0).to(tl.int32) - num_acc = tl.load( - num_accepted_ptr + offs * stride_na, mask=valid, other=0 - ).to(tl.int32) + num_acc = tl.load(num_accepted_ptr + offs * stride_na, mask=valid, other=0).to( + tl.int32 + ) total_commit = num_acc flush_now = (total_commit > 0) & (is_flush_cur != 0) @@ -407,9 +423,9 @@ def _reset_gdn_replayssm_spec_cursors_kernel( blk = tl.load( state_batch_indices_ptr + offs * stride_sbi, mask=row_mask, other=NULL_BLOCK_ID ).to(tl.int64) - do_reset = tl.load( - do_reset_ptr + offs * stride_reset, mask=row_mask, other=0 - ).to(tl.int32) + do_reset = tl.load(do_reset_ptr + offs * stride_reset, mask=row_mask, other=0).to( + tl.int32 + ) do = row_mask & (blk > NULL_BLOCK_ID) & (do_reset != 0) tl.store(write_pos_ptr + blk, tl.zeros_like(blk).to(tl.int32), mask=do) @@ -450,6 +466,7 @@ def _launch_gdn_spec( num_stages, nk, null_block_id, + dot_precision, ): num_slots, HV, V, K = checkpoint_state.shape qkv_dim = mixed_qkv.shape[1] @@ -520,6 +537,7 @@ def _launch_gdn_spec( USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, IS_FLUSH=is_flush_kernel, NULL_BLOCK_ID=null_block_id, + DOT_PRECISION=dot_precision, num_warps=num_warps, num_stages=num_stages, ) @@ -547,6 +565,8 @@ def gdn_replayssm_spec_decode( use_qk_l2norm_in_kernel: bool = True, null_block_id: int = 0, launch_mode: str = "both", + dot_precision: str = "tf32", + dot_precision_flush: str | None = None, ): """GDN cached speculative-decode on a CIRCULAR d/k/g cache (vLLM packed varlen). @@ -560,22 +580,70 @@ def gdn_replayssm_spec_decode( scale = checkpoint_state.shape[-1] ** -0.5 if is_flush.dtype != torch.int8: is_flush = is_flush.to(torch.int8) - vb, vw, vnk, vns = get_replayssm_config("gdn_spec_verify", max_spec_len=max_spec_len) + if dot_precision_flush is None: + dot_precision_flush = dot_precision + vb, vw, vnk, vns = get_replayssm_config( + "gdn_spec_verify", max_spec_len=max_spec_len + ) fb, fw, fnk, fns = get_replayssm_config("gdn_spec_flush", max_spec_len=max_spec_len) if launch_mode in ("both", "verify"): _launch_gdn_spec( - mixed_qkv, a, b, A_log, dt_bias, out, checkpoint_state, - d_cache, k_cache, g_cache, query_start_loc, ssm_state_indices, - write_pos, cache_base, is_flush, scale, max_cache_len, max_spec_len, - use_qk_l2norm_in_kernel, False, vb, vw, vns, vnk, null_block_id, + mixed_qkv, + a, + b, + A_log, + dt_bias, + out, + checkpoint_state, + d_cache, + k_cache, + g_cache, + query_start_loc, + ssm_state_indices, + write_pos, + cache_base, + is_flush, + scale, + max_cache_len, + max_spec_len, + use_qk_l2norm_in_kernel, + False, + vb, + vw, + vns, + vnk, + null_block_id, + dot_precision, ) if launch_mode in ("both", "flush"): _launch_gdn_spec( - mixed_qkv, a, b, A_log, dt_bias, out, checkpoint_state, - d_cache, k_cache, g_cache, query_start_loc, ssm_state_indices, - write_pos, cache_base, is_flush, scale, max_cache_len, max_spec_len, - use_qk_l2norm_in_kernel, True, fb, fw, fns, fnk, null_block_id, + mixed_qkv, + a, + b, + A_log, + dt_bias, + out, + checkpoint_state, + d_cache, + k_cache, + g_cache, + query_start_loc, + ssm_state_indices, + write_pos, + cache_base, + is_flush, + scale, + max_cache_len, + max_spec_len, + use_qk_l2norm_in_kernel, + True, + fb, + fw, + fns, + fnk, + null_block_id, + dot_precision_flush, ) return out diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 34396b80f5f9..f1f7eb22d84f 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -184,19 +184,22 @@ def gated_delta_net_replayssm_spec_state_dtype( ) -> tuple[torch.dtype, ...]: """GDN ReplaySSM state dtypes for the SPECULATIVE-decode kernel. - Same ``d/k/g`` ring page as the non-spec ReplaySSM path, but the ``ssm`` - checkpoint is forced to ``float32``. Call only when use_replayssm_spec - is on. + The ``ssm`` checkpoint is forced to ``float32``; the ``d``/``k`` ring + caches use fp16 for bf16 activations (same rule as the non-spec path). + Call only when use_replayssm_spec is on. """ conv_dtype, ssm_dtype = cls._mamba_state_dtype( model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype ) activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype) + cache_dtype = ( + torch.float16 if activation_dtype == torch.bfloat16 else activation_dtype + ) return ( conv_dtype, torch.float32, # fp32 checkpoint - activation_dtype, # d_cache - activation_dtype, # k_cache + cache_dtype, # d_cache + cache_dtype, # k_cache torch.float32, # g_cache ) From ec6945c484dd338e2668771d77d95c3d2b083690 Mon Sep 17 00:00:00 2001 From: Johnny-Liou Date: Mon, 20 Jul 2026 14:32:07 -0400 Subject: [PATCH 17/23] [Bugfix] ReplaySSM GDN spec: route every post-prefill decode row through the spec kernel Signed-off-by: Johnny-Liou --- vllm/v1/attention/backends/gdn_attn.py | 30 +++++++++++++++++++++++++- vllm/v1/worker/gpu_model_runner.py | 5 ++++- 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index b736b5f2c21a..543e64de5f3e 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -224,7 +224,33 @@ def build( # type: ignore[override] ) spec_sequence_masks_cpu: torch.Tensor | None = None - if ( + if self.use_cache_spec_kernel and num_accepted_tokens is not None: + # ReplaySSM spec: every post-prefill row must run through the spec + # kernel (a draft-less row is a T=1 window). The baseline decode / + # prefill paths read the checkpoint page, which lags the committed + # ring history, so routing any decode row there corrupts the state. + # num_decode_draft_tokens_cpu cannot drive this mask: it is stale + # on draft-less steps and -1 for decode rows whose drafts were + # dropped. + is_prefilling_cpu = m.is_prefilling + assert is_prefilling_cpu is not None + query_lens_cpu_all = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] + spec_sequence_masks_cpu = ( + ~is_prefilling_cpu[: query_lens_cpu_all.shape[0]] + ) & (query_lens_cpu_all > 0) + num_spec_decodes = int(spec_sequence_masks_cpu.sum().item()) + if num_spec_decodes == 0: + spec_sequence_masks = None + spec_sequence_masks_cpu = None + else: + assert ( + int(query_lens_cpu_all[spec_sequence_masks_cpu].max().item()) + <= self.num_spec + 1 + ), "ReplaySSM-spec decode row wider than the spec window" + spec_sequence_masks = async_tensor_h2d( + spec_sequence_masks_cpu, device=query_start_loc.device + ) + elif ( not self.use_spec_decode or num_decode_draft_tokens_cpu is None or num_decode_draft_tokens_cpu[num_decode_draft_tokens_cpu >= 0] @@ -508,6 +534,8 @@ def build( # type: ignore[override] assert spec_state_indices_tensor is not None assert num_accepted_tokens is not None + # non-None whenever num_spec_decodes > 0 (set together above) + assert spec_sequence_masks_cpu is not None if self.spec_write_pos is None: n_blocks = self.vllm_config.cache_config.num_gpu_blocks assert n_blocks is not None and n_blocks > 0, ( diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 9ca4ea9f6a67..1b0e1800b22b 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -2446,7 +2446,10 @@ def _build_attn_group_metadata( needs_mamba_spec_args = use_spec_decode or ( self.speculative_config is not None and self.cache_config.use_replayssm_spec - and isinstance(builder, Mamba2AttentionMetadataBuilder) + and isinstance( + builder, + (Mamba2AttentionMetadataBuilder, GDNAttentionMetadataBuilder), + ) ) if needs_mamba_spec_args and isinstance( builder, (Mamba2AttentionMetadataBuilder, GDNAttentionMetadataBuilder) From 66270e82b77a6ecbc69e2f21a310672926f66f55 Mon Sep 17 00:00:00 2001 From: Johnny-Liou Date: Mon, 20 Jul 2026 15:59:54 -0400 Subject: [PATCH 18/23] [Bench] ReplaySSM spec e2e: size cudagraph captures to spec-window multiples Signed-off-by: Johnny-Liou --- benchmarks/replayssm/e2e_spec_decode_throughput.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/benchmarks/replayssm/e2e_spec_decode_throughput.py b/benchmarks/replayssm/e2e_spec_decode_throughput.py index c5a8f309af6c..fce383061d75 100644 --- a/benchmarks/replayssm/e2e_spec_decode_throughput.py +++ b/benchmarks/replayssm/e2e_spec_decode_throughput.py @@ -114,8 +114,16 @@ def run_worker(args): # Avoid the FlashInfer GDN-prefill cutlass-DSL JIT stall on Blackwell # (same default as the decode benchmark). additional_config={"gdn_prefill_backend": "triton"}, + # Ladder of spec_window multiples: covers the baseline spec path's + # <= max_concurrency sub-batches when bs > cap, and fixes the + # bs=1,T=6 "cudagraph size must be a multiple of T" abort. compilation_config={ - "max_cudagraph_capture_size": max(8, args.batch_size * spec_window) + "cudagraph_capture_sizes": sorted({ + min(r, args.batch_size) * spec_window + for r in (1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128, + 192, 256, 384, 512) + }), + "max_cudagraph_capture_size": args.batch_size * spec_window, }, ) if args.disable_flashinfer_autotune: From 99dcd7600195c3a7ffaa23eb782591d04df1dc65 Mon Sep 17 00:00:00 2001 From: Johnny-Liou Date: Mon, 20 Jul 2026 16:05:07 -0400 Subject: [PATCH 19/23] [Perf] ReplaySSM GDN spec: tune verify/flush launch configs on Blackwell Signed-off-by: Johnny-Liou --- .../fla/ops/gdn_replayssm_spec_decode.py | 10 ++++- .../layers/mamba/ops/replayssm_config.py | 42 +++++++++++++++++-- 2 files changed, 46 insertions(+), 6 deletions(-) diff --git a/vllm/model_executor/layers/fla/ops/gdn_replayssm_spec_decode.py b/vllm/model_executor/layers/fla/ops/gdn_replayssm_spec_decode.py index d32d88895984..f937be49245f 100644 --- a/vllm/model_executor/layers/fla/ops/gdn_replayssm_spec_decode.py +++ b/vllm/model_executor/layers/fla/ops/gdn_replayssm_spec_decode.py @@ -583,9 +583,15 @@ def gdn_replayssm_spec_decode( if dot_precision_flush is None: dot_precision_flush = dot_precision vb, vw, vnk, vns = get_replayssm_config( - "gdn_spec_verify", max_spec_len=max_spec_len + "gdn_spec_verify", + max_spec_len=max_spec_len, + head_k_dim=checkpoint_state.shape[-1], + ) + fb, fw, fnk, fns = get_replayssm_config( + "gdn_spec_flush", + max_spec_len=max_spec_len, + head_k_dim=checkpoint_state.shape[-1], ) - fb, fw, fnk, fns = get_replayssm_config("gdn_spec_flush", max_spec_len=max_spec_len) if launch_mode in ("both", "verify"): _launch_gdn_spec( diff --git a/vllm/model_executor/layers/mamba/ops/replayssm_config.py b/vllm/model_executor/layers/mamba/ops/replayssm_config.py index 541e764840e5..e6a933158a16 100644 --- a/vllm/model_executor/layers/mamba/ops/replayssm_config.py +++ b/vllm/model_executor/layers/mamba/ops/replayssm_config.py @@ -70,11 +70,43 @@ def _mamba2_output_only(dstate, L, is_blackwell): return 16, 1, _dstate_tile(dstate, 64), _dstate_tile(dstate, 128), 2 -def _gdn_spec(max_spec_len, is_blackwell): - # (block_v, num_warps, nk, num_stages); verify and flush share a config. +# (block_v, num_warps, nk, num_stages); verify and flush are tuned +# independently per spec window on Blackwell. Non-Blackwell and untuned T +# fall back to the shared default. +_GDN_SPEC_VERIFY_BLACKWELL = { + 2: (128, 2, 4, 3), + 4: (128, 2, 4, 3), + 6: (128, 2, 8, 4), + 8: (128, 2, 8, 4), +} +_GDN_SPEC_FLUSH_BLACKWELL = { + 2: (64, 1, 8, 4), + 4: (64, 1, 8, 4), + 6: (64, 1, 8, 2), + 8: (64, 1, 8, 2), +} + + +def _gdn_spec_default(max_spec_len): return 64, 1, (4 if max_spec_len >= 6 else 2), 2 +def _gdn_spec_verify(max_spec_len, head_k_dim, is_blackwell): + if is_blackwell and head_k_dim == 128: + return _GDN_SPEC_VERIFY_BLACKWELL.get( + max_spec_len, _gdn_spec_default(max_spec_len) + ) + return _gdn_spec_default(max_spec_len) + + +def _gdn_spec_flush(max_spec_len, head_k_dim, is_blackwell): + if is_blackwell and head_k_dim == 128: + return _GDN_SPEC_FLUSH_BLACKWELL.get( + max_spec_len, _gdn_spec_default(max_spec_len) + ) + return _gdn_spec_default(max_spec_len) + + def _l_bucket(cache_len: int) -> int: """Map an arbitrary buffer length to the nearest tuned bucket: L<=8 -> 8, 8 16, L>16 -> 32. Default callers pass L=16.""" @@ -122,6 +154,8 @@ def get_replayssm_config(kernel: str, **shape) -> tuple: return _STATE_AND_OUTPUT_BY_L[_l_bucket(shape.get("L", 16))] if kernel == "gdn_decode": return _gdn_decode(shape.get("L", 16), bw) - if kernel in ("gdn_spec_verify", "gdn_spec_flush"): - return _gdn_spec(shape["max_spec_len"], bw) + if kernel == "gdn_spec_verify": + return _gdn_spec_verify(shape["max_spec_len"], shape["head_k_dim"], bw) + if kernel == "gdn_spec_flush": + return _gdn_spec_flush(shape["max_spec_len"], shape["head_k_dim"], bw) raise ValueError(f"unknown ReplaySSM kernel config key: {kernel}") From 7e917cdcd57f65e4a693547397a2a0dc26651217 Mon Sep 17 00:00:00 2001 From: Amey Naik <212485788+ameynaik-hub@users.noreply.github.com> Date: Fri, 24 Jul 2026 15:34:54 -0700 Subject: [PATCH 20/23] GDN ucache CuTeDSL spec-decode backend: core + fp16 state/cache defaults MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Opt-in `flashinfer_ucache` GDN cached-spec backend (single fused CuTeDSL verify+flush kernel) alongside PR #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 --- .../layers/fla/ops/gdn_ucache_spec.py | 308 ++++++++++++++++++ vllm/model_executor/layers/mamba/gdn/base.py | 2 + .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 72 +++- .../layers/mamba/mamba_utils.py | 202 +++++++++++- vllm/model_executor/models/config.py | 18 +- vllm/model_executor/models/qwen3_5.py | 1 + vllm/v1/attention/backends/gdn_attn.py | 173 +++++++--- 7 files changed, 728 insertions(+), 48 deletions(-) create mode 100644 vllm/model_executor/layers/fla/ops/gdn_ucache_spec.py diff --git a/vllm/model_executor/layers/fla/ops/gdn_ucache_spec.py b/vllm/model_executor/layers/fla/ops/gdn_ucache_spec.py new file mode 100644 index 000000000000..60542c16a270 --- /dev/null +++ b/vllm/model_executor/layers/fla/ops/gdn_ucache_spec.py @@ -0,0 +1,308 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Adapter for the FlashInfer CuTeDSL GDN "ucache" spec-decode kernel. + +Replaces the Triton ``gdn_replayssm_spec_decode`` verify+flush pair with the +single fused kernel ``gated_delta_rule_mtp_ucache_flush`` (ReplaySSM Alg. 8): +per-request device-side routing runs the verify path for rows with +``hist_len < flush_min`` and additionally folds the ring into the +checkpoint (and restarts the ring at slots ``[0, T)``) for rows at or past the +threshold. One kernel per layer per step; CUDA-graph capturable. + +Ring page layout on this backend (allocated via ``ring_slots=16``): + page[1] = checkpoint [blocks, HV, V, K] fp16 default (GDN_UCACHE_STATE_DTYPE) + page[2] = u_cache [blocks, HV, 16, V] fp16 default (GDN_UCACHE_RING_DTYPE) + page[3] = k_cache [blocks, HK, 16, K] fp16 default (GDN_UCACHE_RING_DTYPE) + page[4] = g_cache [blocks, HV, 16] f32 (abs cumulative log-decay) + +Cursor model: ONE block-keyed persistent ``hist_len`` buffer, committed +eagerly in the metadata builder (outside any captured region) by +``commit_gdn_ucache_hist`` with the previous step's acceptance: +``new = (old >= flush_min ? 0 : old) + accepted`` — the ring restart of a +flush step is folded into the commit, so ``hist_len`` is strictly read-only +inside the captured forward (the kernel wrapper is called with +``restart_hist_on_flush=False``; see the kernel-repo flag). The builder +gathers the block-keyed values into fixed-address request-keyed buffers that +the captured kernel reads. +""" + +import importlib.util +import os +from typing import Any + +import torch + +from vllm.logger import init_logger +from vllm.triton_utils import tl, triton + +logger = init_logger(__name__) + +# The kernel's hardcoded ring depth; replayssm_buffer_len must equal this. +UCACHE_W_RING = 16 + +_KMOD: Any = None + +# Reserved null page id (mirrors vllm/v1/attention/backends/utils.py). +NULL_BLOCK_ID = 0 + +# Padded-row sentinel: the kernel retires the whole CTA at entry for rows +# with state index < 0 (pad-skip), so padded rows cost ~nothing instead of a +# full T-step verify. Requires kernel commit 455b0f6+ (_exit_cta_if_neg); +# rows with index 0 still run the legal P=0 verify against null page 0. +UCACHE_PAD_ROW_ID = -1 + + +def ucache_flush_min(max_spec_len: int) -> int: + """Lazy flush threshold: flush when [P, P+T) would overflow W_RING. + + Identical cadence to the Triton cursors' early-flush predicate + ``(write_pos + 2*max_spec_len) > logical_L`` for buffer_len == W_RING. + """ + return UCACHE_W_RING - max_spec_len + 1 + + +def load_ucache_kernel_module(strided_qkv: bool = True): + """Load the flush-kernel module exactly once. + + ``SGLANG_GDN_WY_STRIDED_QKV`` is read by the module at import time, so it + must be set *before* the import executes. ``VLLM_GDN_UCACHE_MODULE`` (an + absolute path to gdn_decode_bf16_wy_ucache_flush.py) loads the file + directly via importlib — the module is import-self-contained, so this + avoids putting a whole flashinfer fork on sys.path and shadowing the + installed flashinfer package. + """ + global _KMOD + if _KMOD is not None: + return _KMOD + os.environ.setdefault( + "SGLANG_GDN_WY_STRIDED_QKV", "1" if strided_qkv else "0" + ) + # Default kernel dtypes: fp16 SSM-state checkpoint + fp16 u/k rings with + # bf16 IO (inputs). Read by the module at import time; setdefault keeps + # explicit user overrides (set both to bf16 for bf16 state/rings). The + # vLLM-side pool allocation defaults match (see + # gated_delta_net_replayssm_spec_state_dtype). + os.environ.setdefault("GDN_UCACHE_STATE_DTYPE", "fp16") + os.environ.setdefault("GDN_UCACHE_RING_DTYPE", "fp16") + path = os.environ.get("VLLM_GDN_UCACHE_MODULE") + if path: + spec = importlib.util.spec_from_file_location( + "gdn_ucache_flush_kernel", path + ) + assert spec is not None and spec.loader is not None, ( + f"cannot load ucache kernel module from {path!r}" + ) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + else: + from flashinfer.gdn_kernels import ( # type: ignore[import-not-found] + gdn_decode_bf16_wy_ucache_flush as mod, + ) + assert mod.W_RING == UCACHE_W_RING, ( + f"ucache kernel W_RING={mod.W_RING} != expected {UCACHE_W_RING}" + ) + logger.info_once( + "GDN spec backend flashinfer_ucache: loaded kernel module from %s " + "(strided_qkv=%s)", + path or "flashinfer.gdn_kernels (PYTHONPATH)", + strided_qkv, + ) + _KMOD = mod + return mod + + +_PAD_BUFS: dict = {} + + +def _pad_scratch(key, shape, dtype, device, fill=0): + buf = _PAD_BUFS.get(key) + if buf is None or buf.shape != torch.Size(shape): + buf = torch.full(shape, fill, dtype=dtype, device=device) + _PAD_BUFS[key] = buf + return buf + + +def gdn_ucache_spec_verify( + *, + mixed_qkv_spec: torch.Tensor, # [total_spec, 2*HK*K + HV*V] packed q|k|v + a: torch.Tensor, # [num_tokens, HV] (chunk view; rows [0:total_spec] used) + b: torch.Tensor, # [num_tokens, HV] + A_log: torch.Tensor, + dt_bias: torch.Tensor, + ssm_state: torch.Tensor, # [blocks, HV, V, K] bf16 checkpoint pool + u_cache: torch.Tensor, # [blocks, HV, 16, V] bf16 + k_cache: torch.Tensor, # [blocks, HK, 16, K] bf16 + g_cache: torch.Tensor, # [blocks, HV, 16] f32 + hist_len: torch.Tensor, # [B] int32, request-keyed (gathered by builder) + state_indices: torch.Tensor, # [B] int32 physical block per request + num_spec_decodes: int, + max_spec_len: int, + num_k_heads: int, # per-rank HK + head_k_dim: int, + head_v_dim: int, + scale: float, + strided_qkv: bool = True, + pad_to: int | None = None, +) -> torch.Tensor: + """Run the fused verify+flush kernel; returns [total_spec, HV, V] view. + + ``pad_to``: pad the request batch to this fixed size with null-page rows. + The strided-mode JIT cache key includes B, so without padding every new + batch size seen on an eager (mixed prefill+decode ramp) step costs a + full CuTeDSL compile (~60s) during batch-admission ramps. One bucket => + one cubin. Pad rows: hist=0, state index = -1 (pad-skip: their CTAs + retire at kernel entry; outputs discarded). + """ + mod = load_ucache_kernel_module(strided_qkv) + B, T = num_spec_decodes, max_spec_len + total_spec = mixed_qkv_spec.shape[0] + assert total_spec == B * T, ( + f"flashinfer_ucache spec backend requires uniform {T}-token verify " + f"rows; got total_spec={total_spec}, num_spec_decodes={B}. " + "Non-uniform spec batches cannot fall back per-step (incompatible " + "persistent ring formats)." + ) + HK, K, V = num_k_heads, head_k_dim, head_v_dim + HV = (mixed_qkv_spec.shape[1] - 2 * HK * K) // V + + # Whole-fp16 kernel mode (GDN_UCACHE_IO_DTYPE=fp16): cast the PACKED qkv + # slice ONCE, before slicing — the q/k/v sub-views then share one token + # stride, keeping the wrapper's strided static-descriptor mode engaged + # (REQUIRED: block-strided vLLM state pools reject the contiguous-dynamic + # path). bf16 -> fp16 is exact in range. Output is cast back below. + _io = getattr(mod, "IO_TORCH", torch.bfloat16) + _out_dtype = mixed_qkv_spec.dtype + if _io is not mixed_qkv_spec.dtype: + mixed_qkv_spec = mixed_qkv_spec.to(_io) + + real_total = total_spec + if pad_to is not None and B < pad_to: + dev = mixed_qkv_spec.device + qkv_dim = mixed_qkv_spec.shape[1] + packed = _pad_scratch( + ("qkv", dev), (pad_to * T, qkv_dim), mixed_qkv_spec.dtype, dev + ) + packed[:total_spec].copy_(mixed_qkv_spec) + a_buf = _pad_scratch(("a", dev), (pad_to * T, HV), a.dtype, dev) + b_buf = _pad_scratch(("b", dev), (pad_to * T, HV), b.dtype, dev) + a_buf[:total_spec].copy_(a[:total_spec]) + b_buf[:total_spec].copy_(b[:total_spec]) + mixed_qkv_spec, a, b = packed, a_buf, b_buf + # Stage hist/idx into fixed pad_to-length scratch (pad rows: + # hist=0, state index=-1 pad-skip sentinel). + hist_buf = _pad_scratch(("hist", dev), (pad_to,), torch.int32, dev) + hist_buf[:B].copy_(hist_len) + hist_buf[B:].fill_(0) + hist_len = hist_buf + idx_buf = _pad_scratch( + ("idx", dev), (pad_to,), state_indices.dtype, dev, + fill=UCACHE_PAD_ROW_ID + ) + idx_buf[:B].copy_(state_indices) + idx_buf[B:].fill_(UCACHE_PAD_ROW_ID) + state_indices = idx_buf + B = pad_to + total_spec = B * T + + qkv = mixed_qkv_spec.view(B, T, -1) + # Last-dim slices of one packed row share token stride -> the wrapper's + # opt-in strided path reads them zero-copy (SGLANG_GDN_WY_STRIDED_QKV=1). + q = qkv[..., : HK * K].unflatten(-1, (HK, K)) + k = qkv[..., HK * K : 2 * HK * K].unflatten(-1, (HK, K)) + v = qkv[..., 2 * HK * K :].unflatten(-1, (HV, V)) + # a/b are chunk() views (token stride 2*HV); rows [0:total_spec] is the + # exact window the Triton spec kernel reads (parity). Materialize + # contiguous copies for the kernel wrappers. + a_spec = a[:total_spec].contiguous().reshape(B, T, HV) + b_spec = b[:total_spec].contiguous().reshape(B, T, HV) + + if _io is not a_spec.dtype: + a_spec = a_spec.to(_io) + b_spec = b_spec.to(_io) + + out = mod.gated_delta_rule_mtp_ucache_flush( + # nn.Parameters carry requires_grad=True, which DLPack refuses to + # export; detach() shares storage. + A_log=A_log.detach(), + a=a_spec, + dt_bias=dt_bias.detach(), + q=q, + k=k, + v=v, + b=b_spec, + initial_state_source=ssm_state, + initial_state_indices=state_indices, + k_cache=k_cache, + u_cache=u_cache, + g_cache=g_cache, + hist_len=hist_len, + scale=scale, + use_qk_l2norm_in_kernel=True, + flush_min=ucache_flush_min(T), + restart_hist_on_flush=False, # builder-owned restart (see module doc) + ) + if out.dtype is not _out_dtype: + out = out.to(_out_dtype) + return out.reshape(B * T, HV, V)[:real_total] + + +@triton.jit +def _commit_gdn_ucache_hist_kernel( + hist_ptr, # [num_gpu_blocks] int32, block-keyed + num_accepted_ptr, # [n_rows] int32 (previous step's acceptance) + sbi_ptr, # base ptr of spec_state_indices_tensor[:, 0] + first_decode_ptr, # [n_rows] int8 (dummy when HAS_RESET == False) + n_rows, + sbi_stride, + FLUSH_MIN: tl.constexpr, + HAS_RESET: tl.constexpr, + NULL_BLOCK: tl.constexpr, + BLOCK: tl.constexpr, +): + offs = tl.arange(0, BLOCK) + m = offs < n_rows + blk = tl.load(sbi_ptr + offs * sbi_stride, mask=m, other=NULL_BLOCK).to( + tl.int64 + ) + valid = m & (blk > NULL_BLOCK) + old = tl.load(hist_ptr + blk, mask=valid, other=0).to(tl.int32) + acc = tl.load(num_accepted_ptr + offs, mask=valid, other=0).to(tl.int32) + # A row whose previous verify launched with hist >= FLUSH_MIN flushed + # in-kernel (every layer): its ring restarted at [0, T), so committed + # history restarts at 0 before adding the accepted count. + new = tl.where(old >= FLUSH_MIN, 0, old) + acc + if HAS_RESET: + fd = tl.load(first_decode_ptr + offs, mask=valid, other=0).to(tl.int32) + new = tl.where(fd != 0, 0, new) # prefill->decode / block-recycle + tl.store(hist_ptr + blk, new, mask=valid) + + +def commit_gdn_ucache_hist( + hist_len: torch.Tensor, # [num_gpu_blocks] int32, block-keyed + num_accepted_tokens: torch.Tensor, # [n_rows] int32 + state_indices: torch.Tensor, # [n_rows] view of block ids (may be strided) + first_decode: torch.Tensor | None, # [n_rows] int8 or None + *, + flush_min: int, +) -> None: + """Eager (outside-capture) hist commit; mirrors commit_gdn_replayssm_spec. + + Invariant afterwards: hist <= (flush_min - 1) + T == W_RING, the kernel's + legal [0, W_RING] range. + """ + n_rows = state_indices.shape[0] + if n_rows == 0: + return + BLOCK = max(triton.next_power_of_2(n_rows), 16) + _commit_gdn_ucache_hist_kernel[(1,)]( + hist_len, + num_accepted_tokens, + state_indices, + first_decode if first_decode is not None else hist_len, + n_rows, + state_indices.stride(0), + FLUSH_MIN=flush_min, + HAS_RESET=first_decode is not None, + NULL_BLOCK=NULL_BLOCK_ID, + BLOCK=BLOCK, + ) diff --git a/vllm/model_executor/layers/mamba/gdn/base.py b/vllm/model_executor/layers/mamba/gdn/base.py index 3cb34a73475e..cd0ac0b03e3f 100644 --- a/vllm/model_executor/layers/mamba/gdn/base.py +++ b/vllm/model_executor/layers/mamba/gdn/base.py @@ -36,6 +36,7 @@ def __init__( self.hidden_size = config.hidden_size self.activation = config.hidden_act self.layer_norm_epsilon = config.rms_norm_eps + self.vllm_config = vllm_config self.model_config = vllm_config.model_config self.cache_config = vllm_config.cache_config self.quant_config = vllm_config.quant_config @@ -56,6 +57,7 @@ def get_state_dtype(self) -> tuple[torch.dtype, ...]: self.model_config.dtype, self.cache_config.mamba_cache_dtype, self.cache_config.mamba_ssm_cache_dtype, + vllm_config=self.vllm_config, ) elif self.cache_config.use_replayssm: return MambaStateDtypeCalculator.gated_delta_net_replayssm_state_dtype( diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index aea091f02b70..2abf3f23397e 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -42,7 +42,9 @@ from vllm.model_executor.layers.mamba.mamba_mixer2 import mamba_v2_sharded_weight_loader from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateShapeCalculator, + gdn_spec_ucache_strided, is_conv_state_dim_first, + resolve_gdn_spec_backend, ) from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( causal_conv1d_fn, @@ -432,6 +434,10 @@ def get_state_shape( self.conv_kernel_size, self.cache_config.replayssm_buffer_len, self.num_spec, + # ucache kernel: fixed 16-slot linear ring (u/k/g), no pow2. + ring_slots=16 + if self.gdn_spec_backend == "flashinfer_ucache" + else None, ) elif self.cache_config.use_replayssm: return MambaStateShapeCalculator.gated_delta_net_replayssm_state_shape( @@ -590,6 +596,30 @@ def __init__( # path decodes through the circular cached kernel. self.use_cache_spec_kernel = self.cache_config.use_replayssm_spec self.max_spec_len = 1 + self.num_spec + # Cached-SPEC kernel backend: "triton" (gdn_replayssm_spec_decode) or + # "flashinfer_ucache" (CuTeDSL fused verify+flush; 16-slot u/k/g ring, + # bf16 checkpoint, hist_len cursors owned by the metadata builder). + vllm_config_for_backend = get_current_vllm_config() + self.gdn_spec_backend = ( + resolve_gdn_spec_backend(vllm_config_for_backend) + if self.use_cache_spec_kernel + else "triton" + ) + self.gdn_spec_ucache_strided = gdn_spec_ucache_strided( + vllm_config_for_backend + ) + # Pad eager (mixed-batch ramp) calls to one fixed request-batch size: + # the strided-mode JIT key includes B, so unpadded ramps compile a + # fresh cubin per batch size seen (~60s each, measured). + self.gdn_spec_ucache_pad_b = ( + vllm_config_for_backend.scheduler_config.max_num_seqs + ) + if self.gdn_spec_backend != "triton": + logger.info_once( + "GDN cached-SPEC decode backend: %s (strided_qkv=%s)", + self.gdn_spec_backend, + self.gdn_spec_ucache_strided, + ) compilation_config = get_current_vllm_config().compilation_config if prefix in compilation_config.static_forward_context: @@ -1508,7 +1538,47 @@ def _forward_core( # 2. Recurrent attention # 2.1: Process the multi-query part - if spec_sequence_masks is not None and self.use_cache_spec_kernel: + if ( + spec_sequence_masks is not None + and self.use_cache_spec_kernel + and self.gdn_spec_backend == "flashinfer_ucache" + ): + # CuTeDSL fused verify+flush kernel: single launch per layer-step, + # device-side per-row flush routing. Page tuple: [2]=u_cache, + # [3]=k_cache, [4]=g_cache (16-slot linear ring); bf16 checkpoint + # folded in-place on flush. hist_len cursors are committed eagerly + # by the metadata builder and read-only here. + from vllm.model_executor.layers.fla.ops.gdn_ucache_spec import ( + gdn_ucache_spec_verify, + ) + + assert mixed_qkv_spec is not None + assert attn_metadata.spec_hist_len_d is not None + assert attn_metadata.spec_state_indices_col0_d is not None + n = attn_metadata.num_spec_decodes + core_attn_out_spec = gdn_ucache_spec_verify( + mixed_qkv_spec=mixed_qkv_spec, + a=a, + b=b, + A_log=self.A_log, + dt_bias=self.dt_bias, + ssm_state=ssm_state, + u_cache=self_kv_cache[2], + k_cache=self_kv_cache[3], + g_cache=self_kv_cache[4], + hist_len=attn_metadata.spec_hist_len_d[:n], + state_indices=attn_metadata.spec_state_indices_col0_d[:n], + num_spec_decodes=n, + max_spec_len=self.max_spec_len, + num_k_heads=self.num_k_heads // self.tp_size, + head_k_dim=self.head_k_dim, + head_v_dim=self.head_v_dim, + scale=self.head_k_dim**-0.5, + strided_qkv=self.gdn_spec_ucache_strided, + pad_to=self.gdn_spec_ucache_pad_b, + ).unsqueeze(0) + last_recurrent_state = None + elif spec_sequence_masks is not None and self.use_cache_spec_kernel: # Cached circular spec verify: reuse the post-conv packed # ``mixed_qkv_spec`` (q|k|v) + raw ``a``/``b`` (read per-request via # spec_query_start_loc, same as the baseline kernel). The d/k/g ring diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index f1f7eb22d84f..a597cf850e12 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import functools +import os from collections.abc import Callable from dataclasses import dataclass from typing import Literal, TypeAlias @@ -181,12 +182,18 @@ def gated_delta_net_replayssm_spec_state_dtype( model_dtype: ModelDType | torch.dtype, mamba_cache_dtype: MambaDType, mamba_ssm_cache_dtype: MambaDType, + vllm_config=None, ) -> tuple[torch.dtype, ...]: """GDN ReplaySSM state dtypes for the SPECULATIVE-decode kernel. - The ``ssm`` checkpoint is forced to ``float32``; the ``d``/``k`` ring - caches use fp16 for bf16 activations (same rule as the non-spec path). - Call only when use_replayssm_spec is on. + The ``ssm`` checkpoint is forced to ``float32`` unless an explicit + ``mamba_ssm_cache_dtype`` overrides it (``auto`` keeps the upstream + fp32 default). The ``d``/``k`` ring caches use fp16 for bf16 + activations (same rule as the non-spec path) — except on the + flashinfer_ucache spec backend, whose rings must match the CuTeDSL + kernel's compiled IO dtype (bf16 default / fp16 with + GDN_UCACHE_IO_DTYPE=fp16); pass ``vllm_config`` so the backend can be + resolved. Call only when use_replayssm_spec is on. """ conv_dtype, ssm_dtype = cls._mamba_state_dtype( model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype @@ -195,11 +202,63 @@ def gated_delta_net_replayssm_spec_state_dtype( cache_dtype = ( torch.float16 if activation_dtype == torch.bfloat16 else activation_dtype ) + # Explicit --mamba-ssm-cache-dtype overrides the ckpt dtype; "auto" + # keeps the upstream force-fp32 default (_mamba_state_dtype would + # map "auto" to the model dtype, NOT fp32). On the flashinfer_ucache + # backend "auto" instead resolves to the kernel's state dtype + # (default fp16) below. + ckpt_dtype = torch.float32 if mamba_ssm_cache_dtype == "auto" else ssm_dtype + # Backend-aware ring default: the ucache kernel reads the u/k rings at + # its compiled IO dtype, so allocation must match it (the fp16-ring + # upstream rule only applies to the dtype-agnostic Triton kernel). + default_ring_dtype = cache_dtype + if vllm_config is not None: + try: + backend = resolve_gdn_spec_backend(vllm_config) + except Exception: + backend = "triton" + if backend == "flashinfer_ucache": + # ucache defaults: fp16 SSM-state checkpoint + fp16 u/k rings + # with bf16 input IO. The adapter setdefaults the kernel + # module's GDN_UCACHE_STATE/RING_DTYPE envs to the same + # values, so pool allocation and the compiled kernel dtypes + # agree with no flags. Set the envs to bf16 for bf16 mode. + _st_env = os.environ.get( + "GDN_UCACHE_STATE_DTYPE", "fp16" + ).lower() + if mamba_ssm_cache_dtype == "auto": + ckpt_dtype = ( + torch.bfloat16 + if _st_env in ("bf16", "bfloat16") + else torch.float16 + ) + _ring_env_kernel = os.environ.get( + "GDN_UCACHE_RING_DTYPE", "fp16" + ).lower() + default_ring_dtype = ( + torch.bfloat16 + if _ring_env_kernel in ("bf16", "bfloat16") + else torch.float16 + ) + # Ring-dtype override (default: see above): + # VLLM_REPLAYSSM_RING_DTYPE=fp16 -> fp16 u/k rings (ucache whole-fp16 + # mode: pair with GDN_UCACHE_IO_DTYPE=fp16 so the cute kernel + # compiles fp16 IO; the adapter casts activations at the call). + # VLLM_REPLAYSSM_RING_DTYPE=bf16|fp32 -> bf16 / fp32 u/k rings. + _ring_env = os.environ.get("VLLM_REPLAYSSM_RING_DTYPE", "").lower() + if _ring_env in ("fp16", "float16", "half"): + ring_dtype = torch.float16 + elif _ring_env in ("bf16", "bfloat16"): + ring_dtype = torch.bfloat16 + elif _ring_env in ("fp32", "float32"): + ring_dtype = torch.float32 + else: + ring_dtype = default_ring_dtype return ( conv_dtype, - torch.float32, # fp32 checkpoint - cache_dtype, # d_cache - cache_dtype, # k_cache + ckpt_dtype, + ring_dtype, # d_cache + ring_dtype, # k_cache torch.float32, # g_cache ) @@ -467,6 +526,7 @@ def gated_delta_net_replayssm_spec_state_shape( conv_kernel_size: int, replayssm_buffer_len: int, num_spec: int = 0, + ring_slots: int | None = None, ) -> tuple[tuple[int, ...], ...]: """GDN ReplaySSM state shapes for the SPECULATIVE-decode kernel. @@ -474,6 +534,11 @@ def gated_delta_net_replayssm_spec_state_shape( history window: a power-of-two buffer ``next_pow2(replayssm_buffer_len + 1 + num_spec)``. Call only when use_replayssm_spec is on. The block-keyed cursors live in the GDN metadata builder, not the page. + + ``ring_slots`` overrides the physical ring depth for backends with a + fixed linear (non-circular) ring — e.g. the flashinfer_ucache kernel + uses exactly W_RING=16 slots (page[2] becomes its u_cache, page[3] its + k_cache, page[4] its g_cache; dtypes unchanged). """ conv_state_shape, temporal_state_shape = cls.gated_delta_net_state_shape( tp_world_size, @@ -484,7 +549,11 @@ def gated_delta_net_replayssm_spec_state_shape( conv_kernel_size, num_spec, ) - cache_buf_len = 1 << (replayssm_buffer_len + num_spec).bit_length() + cache_buf_len = ( + ring_slots + if ring_slots is not None + else 1 << (replayssm_buffer_len + num_spec).bit_length() + ) local_v_heads = divide(num_v_heads, tp_world_size) local_k_heads = divide(num_k_heads, tp_world_size) d_cache_shape = (local_v_heads, cache_buf_len, head_v_dim) @@ -620,3 +689,122 @@ def gated_delta_net_state_copy_func(cls): @classmethod def kda_state_copy_func(cls): return (get_conv_copy_spec, get_temporal_copy_spec) + + +def _ucache_kernel_available() -> tuple[bool, str]: + """Init-time check that the ucache CuTeDSL kernel module is loadable. + + Mirrors load_ucache_kernel_module's resolution order without importing + (and JIT-compiling) the module: an explicit VLLM_GDN_UCACHE_MODULE path + must exist, else flashinfer.gdn_kernels must provide the module. + """ + path = os.environ.get("VLLM_GDN_UCACHE_MODULE") + if path: + return ( + os.path.isfile(path), + f"VLLM_GDN_UCACHE_MODULE points to a missing file: {path!r}", + ) + import importlib.util + + try: + found = ( + importlib.util.find_spec( + "flashinfer.gdn_kernels.gdn_decode_bf16_wy_ucache_flush" + ) + is not None + ) + except ModuleNotFoundError: + found = False + return ( + found, + "ucache CuTeDSL kernel module not found: set VLLM_GDN_UCACHE_MODULE=" + "/abs/path/to/gdn_decode_bf16_wy_ucache_flush.py or install a " + "FlashInfer build that provides " + "flashinfer.gdn_kernels.gdn_decode_bf16_wy_ucache_flush", + ) + + +def resolve_gdn_spec_backend(vllm_config) -> str: + """Resolve the GDN cached-SPEC decode backend. + + Returns "triton" (default, PR #47576 gdn_replayssm_spec_decode) or + "flashinfer_ucache" (CuTeDSL gated_delta_rule_mtp_ucache_flush). Selected + via additional_config["gdn_spec_backend"]; constraint violations raise at + init (fail loudly rather than silently falling back to triton). + """ + additional_config = vllm_config.additional_config + requested = ( + str(additional_config.get("gdn_spec_backend", "triton")).strip().lower() + if isinstance(additional_config, dict) + else "triton" + ) + if requested in ("triton", "auto", ""): + return "triton" + if requested != "flashinfer_ucache": + raise ValueError(f"unknown gdn_spec_backend={requested!r}") + + from vllm.model_executor.layers.fla.ops.gdn_ucache_spec import UCACHE_W_RING + from vllm.platforms import current_platform + + cache_config = vllm_config.cache_config + spec_config = vllm_config.speculative_config + hf_config = vllm_config.model_config.hf_text_config + max_spec_len = 1 + ( + spec_config.num_speculative_tokens if spec_config is not None else 0 + ) + ssm_dtype = get_kv_cache_torch_dtype( + cache_config.mamba_ssm_cache_dtype, vllm_config.model_config.dtype + ) + checks = [ + (cache_config.use_replayssm_spec, "requires --use-replayssm-spec"), + ( + not cache_config.use_replayssm, + "incompatible with non-spec --use-replayssm (its Triton-format " + "ring shares the same page tuple)", + ), + ( + cache_config.replayssm_buffer_len == UCACHE_W_RING, + f"requires --replayssm-buffer-len {UCACHE_W_RING} (kernel W_RING)", + ), + ( + max_spec_len in (4, 8), + f"verify window T={max_spec_len} unsupported (native T in {{4,8}})", + ), + ( + getattr(hf_config, "linear_key_head_dim", None) == 128 + and getattr(hf_config, "linear_value_head_dim", None) == 128, + "requires linear key/value head dims == 128", + ), + ( + cache_config.mamba_ssm_cache_dtype == "auto" + or ssm_dtype in (torch.bfloat16, torch.float16), + "requires --mamba-ssm-cache-dtype auto (resolves to the kernel " + "state dtype, default fp16), bfloat16, or float16; an explicit " + "dtype must match the kernel module's GDN_UCACHE_STATE_DTYPE " + "(its wrapper asserts the pool dtype on first call)", + ), + _ucache_kernel_available(), + ( + current_platform.is_cuda() + and current_platform.get_device_capability().major >= 9, + "requires SM90+", + ), + ] + for ok, msg in checks: + if not ok: + raise ValueError(f"gdn_spec_backend=flashinfer_ucache: {msg}") + return "flashinfer_ucache" + + +def gdn_spec_ucache_strided(vllm_config) -> bool: + """Whether the ucache kernel uses the zero-copy strided q/k/v path. + + Default True. Set additional_config["gdn_spec_ucache_strided"]=false for + enforce-eager debugging: the strided JIT cache key includes (B, pool), so + arbitrary eager batch sizes would compile per size; the non-strided path + is batch-dynamic (one cubin) at the cost of q/k/v .contiguous() copies. + """ + additional_config = vllm_config.additional_config + if isinstance(additional_config, dict): + return bool(additional_config.get("gdn_spec_ucache_strided", True)) + return True diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index cb01a12efdf7..b650149df546 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -753,7 +753,23 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: hf_text_config = vllm_config.model_config.hf_text_config mamba_ssm_dtype = getattr(hf_text_config, "mamba_ssm_dtype", None) if cache_config.mamba_ssm_cache_dtype == "auto": - if mamba_ssm_dtype is not None: + # flashinfer_ucache GDN spec backend: 'auto' resolves to the + # CuTeDSL kernel's state dtype (fp16 unless + # GDN_UCACHE_STATE_DTYPE=bf16) rather than the HF-config value — + # the ucache kernel does not read an fp32 checkpoint. + import os + + _ac = vllm_config.additional_config + if ( + isinstance(_ac, dict) + and str(_ac.get("gdn_spec_backend", "")).strip().lower() + == "flashinfer_ucache" + ): + _st = os.environ.get("GDN_UCACHE_STATE_DTYPE", "fp16").lower() + cache_config.mamba_ssm_cache_dtype = ( + "bfloat16" if _st in ("bf16", "bfloat16") else "float16" + ) + elif mamba_ssm_dtype is not None: cache_config.mamba_ssm_cache_dtype = mamba_ssm_dtype elif ( mamba_ssm_dtype is not None diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index 82101b0ebdca..82372af4a9a7 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -519,6 +519,7 @@ def get_mamba_state_dtype_from_config( vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype, vllm_config.cache_config.mamba_ssm_cache_dtype, + vllm_config=vllm_config, ) elif vllm_config.cache_config.use_replayssm: return MambaStateDtypeCalculator.gated_delta_net_replayssm_state_dtype( diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 543e64de5f3e..ad017b8d9d38 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -77,6 +77,12 @@ class GDNAttentionMetadata: spec_cache_base_d: torch.Tensor | None = None spec_is_flush_d: torch.Tensor | None = None + # flashinfer_ucache spec backend: request-keyed ring fill levels (gathered + # each build from the block-keyed master) + contiguous col-0 state indices + # for the CuTeDSL verify+flush kernel. None on the triton backend. + spec_hist_len_d: torch.Tensor | None = None # shape: [batch,] int32 + spec_state_indices_col0_d: torch.Tensor | None = None # shape: [batch,] int32 + # Pre-computed FLA chunk metadata (avoids GPU->CPU sync in prepare_chunk_indices) chunk_indices: torch.Tensor | None = None chunk_offsets: torch.Tensor | None = None @@ -202,6 +208,38 @@ def __init__( self.spec_cache_base: torch.Tensor | None = None self.spec_is_flush: torch.Tensor | None = None + # flashinfer_ucache spec backend: one block-keyed hist_len master + # (lazy, sized num_gpu_blocks) committed eagerly per step, plus + # fixed-address request-keyed buffers the captured kernel reads. + from vllm.model_executor.layers.mamba.mamba_utils import ( + resolve_gdn_spec_backend, + ) + + self.gdn_spec_backend = ( + resolve_gdn_spec_backend(vllm_config) + if self.use_cache_spec_kernel + else "triton" + ) + self.spec_hist_len: torch.Tensor | None = None + if self.gdn_spec_backend == "flashinfer_ucache": + from vllm.model_executor.layers.fla.ops.gdn_ucache_spec import ( + UCACHE_PAD_ROW_ID, + ucache_flush_min, + ) + + self.ucache_flush_min = ucache_flush_min(self.max_spec_len) + self.ucache_pad_row_id = UCACHE_PAD_ROW_ID + max_reqs = max( + vllm_config.scheduler_config.max_num_seqs, + self.decode_cudagraph_max_bs, + ) + self.spec_hist_len_gathered: torch.Tensor = torch.zeros( + (max_reqs,), dtype=torch.int32, device=device + ) + self.spec_state_indices_col0: torch.Tensor = torch.full( + (max_reqs,), UCACHE_PAD_ROW_ID, dtype=torch.int32, device=device + ) + def build( # type: ignore[override] self, common_prefix_len: int, @@ -485,6 +523,8 @@ def build( # type: ignore[override] spec_write_pos_d = None spec_cache_base_d = None spec_is_flush_d = None + spec_hist_len_d = None + spec_state_indices_col0_d = None if self.use_cached_kernel and spec_sequence_masks is None and num_decodes > 0: num_prompt_tokens_cpu = m.num_prompt_tokens_cpu num_computed_tokens_cpu = m._num_computed_tokens_cpu @@ -527,45 +567,16 @@ def build( # type: ignore[override] # the cursors are full (num_gpu_blocks,) fixed-address buffers read by # the captured verify kernel. if self.use_cache_spec_kernel and num_spec_decodes > 0: - from vllm.model_executor.layers.fla.ops.gdn_replayssm_spec_decode import ( - commit_gdn_replayssm_spec, - reset_gdn_replayssm_spec_cursors, - ) - assert spec_state_indices_tensor is not None assert num_accepted_tokens is not None # non-None whenever num_spec_decodes > 0 (set together above) assert spec_sequence_masks_cpu is not None - if self.spec_write_pos is None: - n_blocks = self.vllm_config.cache_config.num_gpu_blocks - assert n_blocks is not None and n_blocks > 0, ( - "--use-replayssm-spec needs num_gpu_blocks at " - "build time to size the block-keyed cursor buffers" - ) - self.spec_write_pos = torch.zeros( - n_blocks, dtype=torch.int32, device=self.cursor_device - ) - self.spec_cache_base = torch.zeros( - n_blocks, dtype=torch.int32, device=self.cursor_device - ) - self.spec_is_flush = torch.zeros( - n_blocks, dtype=torch.int8, device=self.cursor_device - ) sbi = spec_state_indices_tensor[:, 0] - commit_gdn_replayssm_spec( - self.spec_write_pos, - self.spec_cache_base, - self.spec_is_flush, - num_accepted_tokens.to(torch.int32), - sbi, - max_cache_len=self.spec_flush_threshold, - max_spec_len=self.max_spec_len, - cache_buf_len=self.spec_cache_buf_len, - ) - # prefill->decode reset for first-decode rows (cursors only; conv - # context lives in conv_state). A request's first spec verify has - # num_computed_tokens == num_prompt_tokens; that resets its (possibly - # recycled) block's cursors to write_pos=0. + # prefill->decode reset predicate for first-decode rows (cursors + # only; conv context lives in conv_state). A request's first spec + # verify has num_computed_tokens == num_prompt_tokens; that resets + # its (possibly recycled) block's cursors to write_pos=0. + first_decode_d = None num_prompt_tokens_cpu = m.num_prompt_tokens_cpu if num_prompt_tokens_cpu is not None: num_prompt_d = num_prompt_tokens_cpu.to( @@ -576,18 +587,86 @@ def build( # type: ignore[override] query_start_loc.device, non_blocking=True ) first_decode_d = first_decode_full.index_select(0, spec_row_idx) - reset_gdn_replayssm_spec_cursors( + + if self.gdn_spec_backend == "flashinfer_ucache": + from vllm.model_executor.layers.fla.ops.gdn_ucache_spec import ( + commit_gdn_ucache_hist, + ) + + if self.spec_hist_len is None: + n_blocks = self.vllm_config.cache_config.num_gpu_blocks + assert n_blocks is not None and n_blocks > 0, ( + "--use-replayssm-spec needs num_gpu_blocks at " + "build time to size the block-keyed hist buffer" + ) + self.spec_hist_len = torch.zeros( + n_blocks, dtype=torch.int32, device=self.cursor_device + ) + # Commit-at-start with the previous step's acceptance; the + # flush-step ring restart is folded in (hist >= flush_min -> 0 + # before adding accepted). Eager, outside any captured region. + commit_gdn_ucache_hist( + self.spec_hist_len, + num_accepted_tokens.to(torch.int32), + sbi, + first_decode_d, + flush_min=self.ucache_flush_min, + ) + # Gather block-keyed -> request-keyed fixed-address buffers + # (the captured kernel reads these; padded rows filled below). + torch.index_select( + self.spec_hist_len, + 0, + sbi.to(torch.int64), + out=self.spec_hist_len_gathered[:num_spec_decodes], + ) + self.spec_state_indices_col0[:num_spec_decodes].copy_(sbi) + spec_hist_len_d = self.spec_hist_len_gathered + spec_state_indices_col0_d = self.spec_state_indices_col0 + else: + from vllm.model_executor.layers.fla.ops.gdn_replayssm_spec_decode import ( + commit_gdn_replayssm_spec, + reset_gdn_replayssm_spec_cursors, + ) + + if self.spec_write_pos is None: + n_blocks = self.vllm_config.cache_config.num_gpu_blocks + assert n_blocks is not None and n_blocks > 0, ( + "--use-replayssm-spec needs num_gpu_blocks at " + "build time to size the block-keyed cursor buffers" + ) + self.spec_write_pos = torch.zeros( + n_blocks, dtype=torch.int32, device=self.cursor_device + ) + self.spec_cache_base = torch.zeros( + n_blocks, dtype=torch.int32, device=self.cursor_device + ) + self.spec_is_flush = torch.zeros( + n_blocks, dtype=torch.int8, device=self.cursor_device + ) + commit_gdn_replayssm_spec( self.spec_write_pos, self.spec_cache_base, self.spec_is_flush, - first_decode_d, + num_accepted_tokens.to(torch.int32), sbi, max_cache_len=self.spec_flush_threshold, max_spec_len=self.max_spec_len, + cache_buf_len=self.spec_cache_buf_len, ) - spec_write_pos_d = self.spec_write_pos - spec_cache_base_d = self.spec_cache_base - spec_is_flush_d = self.spec_is_flush + if first_decode_d is not None: + reset_gdn_replayssm_spec_cursors( + self.spec_write_pos, + self.spec_cache_base, + self.spec_is_flush, + first_decode_d, + sbi, + max_cache_len=self.spec_flush_threshold, + max_spec_len=self.max_spec_len, + ) + spec_write_pos_d = self.spec_write_pos + spec_cache_base_d = self.spec_cache_base + spec_is_flush_d = self.spec_is_flush # Prepare per-request tensors for cudagraph. m.num_actual_tokens is # token-padded for FULL graph replay, but the GDN state/query/accepted @@ -640,6 +719,20 @@ def build( # type: ignore[override] num_accepted_tokens = self.num_accepted_tokens[:batch_size] num_accepted_tokens[num_spec_decodes:].fill_(1) + if self.gdn_spec_backend == "flashinfer_ucache": + # Padded rows: negative sentinel -> the kernel retires the + # whole CTA at entry (pad-skip), so under-bucket batches cost + # ~nothing instead of a full T-step verify per padded row. + # Requires kernel commit with _exit_cta_if_neg (455b0f6+); + # on older kernels use 0 (P=0 verify against null page 0). + # Fill to the END of the fixed buffers: rows beyond the live + # batch must read hist=0 / idx=-1 on graph replay. One eager + # fill over <=max_bs int32 per step — free. + self.spec_hist_len_gathered[num_spec_decodes:].fill_(0) + self.spec_state_indices_col0[num_spec_decodes:].fill_( + self.ucache_pad_row_id + ) + if ( self.use_full_cuda_graph and num_prefills == 0 @@ -697,6 +790,8 @@ def build( # type: ignore[override] spec_write_pos_d=spec_write_pos_d, spec_cache_base_d=spec_cache_base_d, spec_is_flush_d=spec_is_flush_d, + spec_hist_len_d=spec_hist_len_d, + spec_state_indices_col0_d=spec_state_indices_col0_d, nums_dict=nums_dict, batch_ptr=batch_ptr, token_chunk_offset_ptr=token_chunk_offset_ptr, From 271f0bfc70a7a36ad863091b2ba4645a01a95175 Mon Sep 17 00:00:00 2001 From: Amey Naik <212485788+ameynaik-hub@users.noreply.github.com> Date: Fri, 24 Jul 2026 15:34:55 -0700 Subject: [PATCH 21/23] GDN ucache spec backend: launch-path perf optimizations + kernel test 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). --- .../kernels/test_replayssm_ucache_spec_gdn.py | 463 ++++++++++++++++++ .../layers/fla/ops/gdn_ucache_spec.py | 132 +++-- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 35 +- vllm/v1/attention/backends/gdn_attn.py | 14 +- 4 files changed, 608 insertions(+), 36 deletions(-) create mode 100644 tests/kernels/test_replayssm_ucache_spec_gdn.py diff --git a/tests/kernels/test_replayssm_ucache_spec_gdn.py b/tests/kernels/test_replayssm_ucache_spec_gdn.py new file mode 100644 index 000000000000..c050122cd833 --- /dev/null +++ b/tests/kernels/test_replayssm_ucache_spec_gdn.py @@ -0,0 +1,463 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Integration tests for the flashinfer_ucache GDN spec-decode backend. + +Covers what is NEW in the vLLM integration (the kernel's own numerics are +anchored by the kernel repo's fp32-reference gates: _ucache_check.py, +_ucache_flush_check.py, _ucache_g3_ring.py): + +1. commit_gdn_ucache_hist semantics (flush restart folded into commit, + first-decode reset, null-block masking). +2. Intra-step shared hist_len: with restart_hist_on_flush=False, N "layers" + sharing one hist tensor within a step all see the same P (the bug the + flag exists to prevent). +3. Protocol equivalence: the vLLM bookkeeping (builder-owned commit, + restart_hist_on_flush=False) must produce bit-identical outputs and + checkpoints to the kernel-repo protocol (hist += accepted each step, + wrapper masked_fill_ restart) over many steps crossing several flushes. +4. Strided packed-qkv slices (production layout) match dense inputs. +5. Null-page rows only scribble the reserved page 0. + +Run inside the vLLM container with: + VLLM_GDN_UCACHE_MODULE= +""" + +import os + +import pytest +import torch + +from vllm.model_executor.layers.fla.ops.gdn_ucache_spec import ( + UCACHE_W_RING, + commit_gdn_ucache_hist, + load_ucache_kernel_module, + ucache_flush_min, +) + +DEV = "cuda" +HK, HV, K, V = 16, 64, 128, 128 # qwen122b geometry + + +def _kmod(): + if not torch.cuda.is_available(): + return None + try: + return load_ucache_kernel_module(strided_qkv=True) + except Exception: + return None + + +pytestmark = pytest.mark.skipif( + _kmod() is None, + reason="CUDA + ucache kernel module required " + "(set VLLM_GDN_UCACHE_MODULE=)", +) + + +def _gating_params(): + torch.manual_seed(7) + A_log = torch.randn(HV, device=DEV, dtype=torch.float32) * 0.1 + dt_bias = torch.randn(HV, device=DEV, dtype=torch.float32) * 0.1 + return A_log, dt_bias + + +def _rand_inputs(B, T, seed): + g = torch.Generator(device=DEV).manual_seed(seed) + mk = lambda *s: torch.randn(*s, generator=g, device=DEV, dtype=torch.bfloat16) + q = mk(B, T, HK, K) + k = mk(B, T, HK, K) + v = mk(B, T, HV, V) * 0.5 + a = mk(B, T, HV) * 0.5 + b = mk(B, T, HV) + return q, k, v, a, b + + +def _pools(num_blocks, seed=11): + g = torch.Generator(device=DEV).manual_seed(seed) + ckpt = ( + torch.randn( + num_blocks, HV, V, K, generator=g, device=DEV, dtype=torch.float32 + ) + * 0.05 + ).to(torch.bfloat16) + k_cache = torch.zeros( + num_blocks, HK, UCACHE_W_RING, K, device=DEV, dtype=torch.bfloat16 + ) + u_cache = torch.zeros( + num_blocks, HV, UCACHE_W_RING, V, device=DEV, dtype=torch.bfloat16 + ) + g_cache = torch.zeros( + num_blocks, HV, UCACHE_W_RING, device=DEV, dtype=torch.float32 + ) + return ckpt, k_cache, u_cache, g_cache + + +def _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt, sbi, kc, uc, gc, hist, + T, restart): + return mod.gated_delta_rule_mtp_ucache_flush( + A_log=A_log, + a=a, + dt_bias=dt_bias, + q=q, + k=k, + v=v, + b=b, + initial_state_source=ckpt, + initial_state_indices=sbi, + k_cache=kc, + u_cache=uc, + g_cache=gc, + hist_len=hist, + scale=K**-0.5, + use_qk_l2norm_in_kernel=True, + output=None, + flush_min=ucache_flush_min(T), + restart_hist_on_flush=restart, + ) + + +def test_commit_kernel_semantics(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + n_blocks = 32 + flush_min = ucache_flush_min(4) # 13 + hist = torch.zeros(n_blocks, dtype=torch.int32, device=DEV) + # rows -> blocks 3, 5, 7, 9 (block 0 = NULL among indices) + sbi = torch.tensor([3, 5, 7, 0, 9], dtype=torch.int32, device=DEV) + hist[3] = 6 # below threshold: 6 + acc + hist[5] = 13 # at threshold: flushed last step -> acc + hist[7] = 16 # max legal: flushed -> acc + hist[9] = 2 # first-decode row: reset to 0 regardless + hist[0] = 99 # null block must never be touched + acc = torch.tensor([3, 2, 4, 1, 4], dtype=torch.int32, device=DEV) + first_decode = torch.tensor([0, 0, 0, 0, 1], dtype=torch.int8, device=DEV) + commit_gdn_ucache_hist(hist, acc, sbi, first_decode, flush_min=flush_min) + torch.cuda.synchronize() + assert hist[3].item() == 9 # 6 + 3 + assert hist[5].item() == 2 # flush restart + 2 + assert hist[7].item() == 4 # flush restart + 4 + assert hist[9].item() == 0 # first-decode reset + assert hist[0].item() == 99 # null untouched + assert (hist[1:] <= UCACHE_W_RING).all() # non-null blocks only + + +@pytest.mark.parametrize("T", [4, 8]) +def test_intra_step_shared_hist_and_flag(T): + """Two 'layers' share one hist tensor in a step; with restart=False the + second layer must see the same P and produce the same fold as the first + (independent pools). With restart=True the second layer would see P=0.""" + mod = _kmod() + A_log, dt_bias = _gating_params() + B, n_blocks = 3, 8 + flush_min = ucache_flush_min(T) + q, k, v, a, b = _rand_inputs(B, T, seed=23) + sbi = torch.tensor([1, 4, 6], dtype=torch.int32, device=DEV) + + hist_master = torch.tensor( + [flush_min, 5, flush_min + 1], dtype=torch.int32, device=DEV + ) + # Two independent "layers" with identical pools and inputs. + outs, ckpts, hists = [], [], [] + hist = hist_master.clone() + # Pre-fill rings identically for both layers so P>0 rows have history. + for layer in range(2): + ckpt, kc, uc, gc = _pools(n_blocks, seed=31) + # Prime the ring: run one step from P=0 (appends T entries), then + # set hist to the master values for the step under test. + hist0 = torch.zeros(B, dtype=torch.int32, device=DEV) + _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt, sbi, kc, uc, gc, + hist0, T, restart=False) + hist_layer = hist if layer == 0 else hist # SHARED tensor + out = _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt, sbi, kc, uc, + gc, hist_layer, T, restart=False) + torch.cuda.synchronize() + outs.append(out.clone()) + ckpts.append(ckpt.clone()) + hists.append(hist_layer.clone()) + # hist untouched by both calls + assert torch.equal(hists[0], hist_master) + assert torch.equal(hists[1], hist_master) + # both layers saw identical P -> identical outputs and identical folds + assert torch.equal(outs[0], outs[1]) + assert torch.equal(ckpts[0], ckpts[1]) + + +@pytest.mark.parametrize("T", [4]) +@pytest.mark.parametrize("nreq", [1, 3, 8]) +def test_protocol_equivalence_multi_step(T, nreq): + """vLLM bookkeeping (commit kernel + restart=False) vs kernel-repo + bookkeeping (hist += accepted + wrapper restart) over 24 steps crossing + several flush cycles: outputs and checkpoints must match bit-for-bit.""" + mod = _kmod() + A_log, dt_bias = _gating_params() + n_blocks = 16 + flush_min = ucache_flush_min(T) + # permuted, non-trivial block assignment (block 0 reserved) + perm = torch.randperm(n_blocks - 1)[:nreq] + 1 + sbi = perm.to(torch.int32).to(DEV) + + ckpt_a, kc_a, uc_a, gc_a = _pools(n_blocks, seed=41) + ckpt_b, kc_b, uc_b, gc_b = _pools(n_blocks, seed=41) + + # Protocol A (vLLM): block-keyed master + commit kernel. + hist_blocks = torch.zeros(n_blocks, dtype=torch.int32, device=DEV) + # Protocol B (kernel repo): request-keyed hist, wrapper restart. + hist_req = torch.zeros(nreq, dtype=torch.int32, device=DEV) + + gen = torch.Generator().manual_seed(97) + prev_acc = torch.zeros(nreq, dtype=torch.int32, device=DEV) + first = torch.zeros(nreq, dtype=torch.int8, device=DEV) + for step in range(24): + q, k, v, a, b = _rand_inputs(nreq, T, seed=1000 + step) + # A: commit (prev step's acceptance), gather, call with restart=False + commit_gdn_ucache_hist( + hist_blocks, prev_acc, sbi, first if step == 0 else None, + flush_min=flush_min, + ) + gathered = hist_blocks.index_select(0, sbi.to(torch.int64)).contiguous() + out_a = _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt_a, sbi, + kc_a, uc_a, gc_a, gathered, T, restart=False) + # B: kernel-repo protocol on request-keyed hist + out_b = _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt_b, sbi, + kc_b, uc_b, gc_b, hist_req, T, restart=True) + torch.cuda.synchronize() + assert torch.equal(out_a, out_b), f"outputs diverged at step {step}" + assert torch.equal(ckpt_a, ckpt_b), f"checkpoints diverged at step {step}" + assert (gathered <= UCACHE_W_RING).all() + acc = torch.randint(1, T + 1, (nreq,), generator=gen).to( + torch.int32 + ).to(DEV) + prev_acc = acc + hist_req += acc # protocol B commit (wrapper already restarted) + + +def test_strided_packed_qkv_matches_dense(): + """Production layout: q/k/v as last-dim slices of one packed row must + match dense contiguous copies of the same data.""" + mod = _kmod() + A_log, dt_bias = _gating_params() + B, T, n_blocks = 4, 4, 8 + q, k, v, a, b = _rand_inputs(B, T, seed=53) + packed = torch.cat( + [q.flatten(2), k.flatten(2), v.flatten(2)], dim=-1 + ).contiguous() # [B, T, 2*HK*K + HV*V], token stride shared by slices + qs = packed[..., : HK * K].unflatten(-1, (HK, K)) + ks = packed[..., HK * K : 2 * HK * K].unflatten(-1, (HK, K)) + vs = packed[..., 2 * HK * K :].unflatten(-1, (HV, V)) + sbi = torch.arange(1, B + 1, dtype=torch.int32, device=DEV) + + outs = [] + for (qq, kk, vv) in [(q, k, v), (qs, ks, vs)]: + ckpt, kc, uc, gc = _pools(n_blocks, seed=61) + hist = torch.zeros(B, dtype=torch.int32, device=DEV) + out = _call(mod, A_log, dt_bias, qq, kk, vv, a, b, ckpt, sbi, + kc, uc, gc, hist, T, restart=False) + torch.cuda.synchronize() + outs.append(out.clone()) + assert torch.equal(outs[0], outs[1]) + + +def test_null_page_rows_only_touch_page_zero(): + mod = _kmod() + A_log, dt_bias = _gating_params() + B, T, n_blocks = 3, 4, 8 + q, k, v, a, b = _rand_inputs(B, T, seed=71) + # row 1 -> null page 0; rows 0/2 -> real pages + sbi = torch.tensor([2, 0, 5], dtype=torch.int32, device=DEV) + ckpt, kc, uc, gc = _pools(n_blocks, seed=79) + snap = ckpt.clone() + hist = torch.tensor([13, 0, 13], dtype=torch.int32, device=DEV) + # prime rows 0/2 rings so their flush folds something + hist0 = torch.zeros(B, dtype=torch.int32, device=DEV) + _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt, sbi, kc, uc, gc, hist0, + T, restart=False) + snap_after_prime = ckpt.clone() + _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt, sbi, kc, uc, gc, hist, + T, restart=False) + torch.cuda.synchronize() + # pages not referenced by any row are byte-stable + untouched = [i for i in range(n_blocks) if i not in (0, 2, 5)] + for i in untouched: + assert torch.equal(ckpt[i], snap[i]), f"page {i} was modified" + # flushed real rows actually changed their pages + assert not torch.equal(ckpt[2], snap_after_prime[2]) + assert not torch.equal(ckpt[5], snap_after_prime[5]) + + +def _vllm_style_strided_pools(num_blocks, seed=11): + """Carve (ckpt, k, u, g) as block-strided views of one page-major backing, + exactly like vLLM's _reshape_kv_cache_tensors (inner dims dense, dim-0 + stride = whole page).""" + shapes = [(HV, V, K), (HK, UCACHE_W_RING, K), (HV, UCACHE_W_RING, V), + (HV, UCACHE_W_RING)] + dtypes = [torch.bfloat16, torch.bfloat16, torch.bfloat16, torch.float32] + page_bytes = sum( + int(torch.empty(s, device="meta").numel()) * t.itemsize + for s, t in zip(shapes, dtypes) + ) + raw = torch.zeros(num_blocks * page_bytes, dtype=torch.int8, device=DEV) + out, off = [], 0 + for s, t in zip(shapes, dtypes): + n = int(torch.empty(s, device="meta").numel()) + view = torch.as_strided( + raw.view(t), + size=(num_blocks, *s), + stride=(page_bytes // t.itemsize, + *torch.empty(s, device="meta").stride()), + storage_offset=off // t.itemsize, + ) + out.append(view) + off += n * t.itemsize + ckpt, kc, uc, gc = out + g = torch.Generator(device=DEV).manual_seed(seed) + ckpt.copy_((torch.randn(ckpt.shape, generator=g, device=DEV, + dtype=torch.float32) * 0.05).to(torch.bfloat16)) + return ckpt, kc, uc, gc + + +def test_block_strided_pools_match_dense(): + """vLLM paged layout: block-strided pool views must produce bit-identical + outputs, ring contents, and folds to dense contiguous pools.""" + mod = _kmod() + A_log, dt_bias = _gating_params() + B, T, n_blocks = 4, 4, 6 + sbi = torch.tensor([1, 3, 4, 5], dtype=torch.int32, device=DEV) + + # dense reference pools with identical initial checkpoint values + ckpt_d, kc_d, uc_d, gc_d = _pools(n_blocks, seed=91) + ckpt_s, kc_s, uc_s, gc_s = _vllm_style_strided_pools(n_blocks, seed=91) + ckpt_s.copy_(ckpt_d) + assert not ckpt_s.is_contiguous() and not kc_s.is_contiguous() + + hist_d = torch.zeros(B, dtype=torch.int32, device=DEV) + hist_s = torch.zeros(B, dtype=torch.int32, device=DEV) + gen = torch.Generator().manual_seed(3) + prev_acc_d = torch.zeros(B, dtype=torch.int32, device=DEV) + hist_blocks_d = torch.zeros(n_blocks, dtype=torch.int32, device=DEV) + hist_blocks_s = torch.zeros(n_blocks, dtype=torch.int32, device=DEV) + fm = ucache_flush_min(T) + for step in range(12): # crosses >= 2 flush cycles + qd, kd, vd, a, b = _rand_inputs(B, T, seed=500 + step) + # packed strided q/k/v (production layout) so the wrapper's + # static-descriptor mode engages — required for strided pools + packed = torch.cat( + [qd.flatten(2), kd.flatten(2), vd.flatten(2)], dim=-1 + ).contiguous() + q = packed[..., : HK * K].unflatten(-1, (HK, K)) + k = packed[..., HK * K : 2 * HK * K].unflatten(-1, (HK, K)) + v = packed[..., 2 * HK * K :].unflatten(-1, (HV, V)) + commit_gdn_ucache_hist(hist_blocks_d, prev_acc_d, sbi, None, flush_min=fm) + commit_gdn_ucache_hist(hist_blocks_s, prev_acc_d, sbi, None, flush_min=fm) + hd = hist_blocks_d.index_select(0, sbi.to(torch.int64)).contiguous() + hs = hist_blocks_s.index_select(0, sbi.to(torch.int64)).contiguous() + out_d = _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt_d, sbi, + kc_d, uc_d, gc_d, hd, T, restart=False) + out_s = _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt_s, sbi, + kc_s, uc_s, gc_s, hs, T, restart=False) + torch.cuda.synchronize() + assert torch.equal(out_d, out_s), f"outputs diverged at step {step}" + assert torch.equal(ckpt_d[1:], ckpt_s[1:]), f"ckpt diverged at step {step}" + assert torch.equal(kc_d[1:], kc_s[1:]) and torch.equal(uc_d[1:], uc_s[1:]) + prev_acc_d = torch.randint(1, T + 1, (B,), generator=gen).to( + torch.int32 + ).to(DEV) + + +def test_block_strided_pools_past_2gb(): + """64-bit pool addressing: rows on high blocks whose element offsets + exceed 2^31 must match dense-pool results (page ~4.7MB -> block ~950 + sits past the old 32-bit ceiling).""" + if torch.cuda.get_device_properties(0).total_memory < 30 * 2**30: + pytest.skip("needs ~10GB free GPU memory") + mod = _kmod() + A_log, dt_bias = _gating_params() + B, T, n_blocks = 3, 4, 2000 + sbi = torch.tensor([1951, 1975, 1999], dtype=torch.int32, device=DEV) + + ckpt_d, kc_d, uc_d, gc_d = _pools(n_blocks, seed=101) + ckpt_s, kc_s, uc_s, gc_s = _vllm_style_strided_pools(n_blocks, seed=101) + ckpt_s.copy_(ckpt_d) + # sanity: the high blocks really are past 2^31 elements in the backing + assert 1999 * ckpt_s.stride(0) > 2**31, ( + f"test page too small to cross 2^31: stride0={ckpt_s.stride(0)}") + + hist_blocks = torch.zeros(n_blocks, dtype=torch.int32, device=DEV) + prev_acc = torch.zeros(B, dtype=torch.int32, device=DEV) + fm = ucache_flush_min(T) + gen = torch.Generator().manual_seed(5) + hist_blocks_d = hist_blocks.clone() + for step in range(8): # crosses a flush cycle + qd, kd, vd, a, b = _rand_inputs(B, T, seed=800 + step) + packed = torch.cat( + [qd.flatten(2), kd.flatten(2), vd.flatten(2)], dim=-1 + ).contiguous() + q = packed[..., : HK * K].unflatten(-1, (HK, K)) + k = packed[..., HK * K : 2 * HK * K].unflatten(-1, (HK, K)) + v = packed[..., 2 * HK * K :].unflatten(-1, (HV, V)) + commit_gdn_ucache_hist(hist_blocks, prev_acc, sbi, None, flush_min=fm) + commit_gdn_ucache_hist(hist_blocks_d, prev_acc, sbi, None, flush_min=fm) + hs = hist_blocks.index_select(0, sbi.to(torch.int64)).contiguous() + hd = hist_blocks_d.index_select(0, sbi.to(torch.int64)).contiguous() + out_s = _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt_s, sbi, + kc_s, uc_s, gc_s, hs, T, restart=False) + out_d = _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt_d, sbi, + kc_d, uc_d, gc_d, hd, T, restart=False) + torch.cuda.synchronize() + assert torch.equal(out_s, out_d), f"outputs diverged at step {step}" + assert torch.equal(ckpt_s[1900:], ckpt_d[1900:]), f"high-block ckpt diverged at {step}" + prev_acc = torch.randint(1, T + 1, (B,), generator=gen).to( + torch.int32 + ).to(DEV) + + +def test_pad_skip_negative_rows_exit_early(): + """Padded rows (sentinel idx < 0) retire their CTAs at kernel entry. + + (a) real-row outputs and pages are bit-identical whether pad rows carry + -1 (pad-skip) or the legacy null-page 0 (P=0 verify) -- pad rows can + never influence real rows; + (b) with -1, pad rows write NOTHING: page 0 (ckpt and all three rings) + stays byte-stable, unlike legacy 0-padding which scribbles it. + """ + mod = _kmod() + A_log, dt_bias = _gating_params() + B, T, n_blocks = 6, 4, 8 + q, k, v, a, b = _rand_inputs(B, T, seed=101) + # rows 0..2 real -> pages 2/5/7 (row 0 at flush threshold), rows 3..5 pad + sbi_neg = torch.tensor([2, 5, 7, -1, -1, -1], dtype=torch.int32, device=DEV) + sbi_nul = torch.tensor([2, 5, 7, 0, 0, 0], dtype=torch.int32, device=DEV) + hist = torch.tensor([13, 7, 0, 0, 0, 0], dtype=torch.int32, device=DEV) + + ckpt1, kc1, uc1, gc1 = _pools(n_blocks, seed=103) + ckpt2, kc2, uc2, gc2 = _pools(n_blocks, seed=103) + page0_snap = (ckpt1[0].clone(), kc1[0].clone(), uc1[0].clone(), + gc1[0].clone()) + + out1 = _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt1, sbi_neg, + kc1, uc1, gc1, hist.clone(), T, restart=False) + out2 = _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt2, sbi_nul, + kc2, uc2, gc2, hist.clone(), T, restart=False) + torch.cuda.synchronize() + + # (a) real rows bit-identical across the two pad conventions. The wrapper + # returns [B, T, HV, V]; slice BATCH rows (pad rows are uninitialized + # torch.empty garbage in the pad-skip arm by design — never compare them). + assert out1.shape[0] == B + assert torch.equal(out1[:3], out2[:3]) + for p in (2, 5, 7): + assert torch.equal(ckpt1[p], ckpt2[p]), f"real page {p} ckpt differs" + assert torch.equal(kc1[p], kc2[p]), f"real page {p} k-ring differs" + assert torch.equal(uc1[p], uc2[p]), f"real page {p} u-ring differs" + assert torch.equal(gc1[p], gc2[p]), f"real page {p} g-ring differs" + + # (b) pad-skip rows leave page 0 byte-stable + assert torch.equal(ckpt1[0], page0_snap[0]), "pad rows wrote ckpt page 0" + assert torch.equal(kc1[0], page0_snap[1]), "pad rows wrote k-ring page 0" + assert torch.equal(uc1[0], page0_snap[2]), "pad rows wrote u-ring page 0" + assert torch.equal(gc1[0], page0_snap[3]), "pad rows wrote g-ring page 0" + # ...and legacy 0-padding does scribble the null page (sanity that the + # comparison above is meaningful) + scribbled = (not torch.equal(kc2[0], page0_snap[1]) + or not torch.equal(uc2[0], page0_snap[2]) + or not torch.equal(gc2[0], page0_snap[3])) + assert scribbled, "expected legacy 0-padding to write the null page" diff --git a/vllm/model_executor/layers/fla/ops/gdn_ucache_spec.py b/vllm/model_executor/layers/fla/ops/gdn_ucache_spec.py index 60542c16a270..39b6e992deb8 100644 --- a/vllm/model_executor/layers/fla/ops/gdn_ucache_spec.py +++ b/vllm/model_executor/layers/fla/ops/gdn_ucache_spec.py @@ -111,6 +111,23 @@ def load_ucache_kernel_module(strided_qkv: bool = True): return mod +# A_log/dt_bias must reach the kernel wrappers as STABLE tensor objects: +# their fp32->bf16 cast cache is keyed by object identity, and calling +# .detach() per call creates a fresh object every time -> cache miss -> +# one cast kernel per tensor per layer per step (measured ~2 x 1.9us x 36 +# layers/step). Detach ONCE per parameter here; the parameter object owns +# the entry lifetime. +_DETACHED: dict = {} + + +def _detached(t: torch.Tensor) -> torch.Tensor: + d = _DETACHED.get(id(t)) + if d is None: + d = t.detach() + _DETACHED[id(t)] = d + return d + + _PAD_BUFS: dict = {} @@ -143,6 +160,7 @@ def gdn_ucache_spec_verify( scale: float, strided_qkv: bool = True, pad_to: int | None = None, + output: torch.Tensor | None = None, ) -> torch.Tensor: """Run the fused verify+flush kernel; returns [total_spec, HV, V] view. @@ -175,32 +193,77 @@ def gdn_ucache_spec_verify( if _io is not mixed_qkv_spec.dtype: mixed_qkv_spec = mixed_qkv_spec.to(_io) + # Caller-provided destination ([real_total, HV, V], the layer's + # core_attn_out slice): kernels STG directly into it (wrapper + # alias) and the layer skips its slice-assign (the 2.7us + # DtoD per layer). Only usable when the padded batch exactly covers it + # and dtypes match (fp16-IO mode casts, so it falls back there). + _dest = output + if ( + _dest is not None + and (_io is not _out_dtype or _dest.dtype is not _out_dtype + or not _dest.is_contiguous()) + ): + _dest = None + real_total = total_spec if pad_to is not None and B < pad_to: dev = mixed_qkv_spec.device qkv_dim = mixed_qkv_spec.shape[1] - packed = _pad_scratch( - ("qkv", dev), (pad_to * T, qkv_dim), mixed_qkv_spec.dtype, dev - ) - packed[:total_spec].copy_(mixed_qkv_spec) - a_buf = _pad_scratch(("a", dev), (pad_to * T, HV), a.dtype, dev) - b_buf = _pad_scratch(("b", dev), (pad_to * T, HV), b.dtype, dev) - a_buf[:total_spec].copy_(a[:total_spec]) - b_buf[:total_spec].copy_(b[:total_spec]) - mixed_qkv_spec, a, b = packed, a_buf, b_buf - # Stage hist/idx into fixed pad_to-length scratch (pad rows: - # hist=0, state index=-1 pad-skip sentinel). - hist_buf = _pad_scratch(("hist", dev), (pad_to,), torch.int32, dev) - hist_buf[:B].copy_(hist_len) - hist_buf[B:].fill_(0) - hist_len = hist_buf - idx_buf = _pad_scratch( - ("idx", dev), (pad_to,), state_indices.dtype, dev, - fill=UCACHE_PAD_ROW_ID - ) - idx_buf[:B].copy_(state_indices) - idx_buf[B:].fill_(UCACHE_PAD_ROW_ID) - state_indices = idx_buf + def _claimable(t, rows, cols): + need = (t.storage_offset() + (rows - 1) * t.stride(0) + + (cols - 1) * t.stride(1) + 1) + return (t.untyped_storage().nbytes() // t.element_size()) >= need + + if ( + os.environ.get("VLLM_GDN_UCACHE_ZEROCOPY_PAD") == "1" + and _claimable(mixed_qkv_spec, pad_to * T, qkv_dim) + and _claimable(a, pad_to * T, HV) + and _claimable(b, pad_to * T, HV) + ): + # ZERO-COPY padding: claim bucket-shaped views over the REAL + # (possibly strided) tensors via as_strided. Rows beyond + # total_spec are phantom — their addresses land in descriptors + # but are NEVER dereferenced: the pad sentinel (-1) in the + # bucket-length index buffers retires those CTAs at kernel + # entry before any load/TMA/cp.async. Kills the per-layer + # direct_copy (packed qkv) + 2 DtoD (a/b) in drain-bucket + # steps. Env-gated: out-of-bounds-by-contract addressing. + mixed_qkv_spec = mixed_qkv_spec.as_strided( + (pad_to * T, qkv_dim), mixed_qkv_spec.stride() + ) + a = a.as_strided((pad_to * T, HV), a.stride()) + b = b.as_strided((pad_to * T, HV), b.stride()) + else: + packed = _pad_scratch( + ("qkv", dev), (pad_to * T, qkv_dim), mixed_qkv_spec.dtype, dev + ) + packed[:total_spec].copy_(mixed_qkv_spec) + a_buf = _pad_scratch(("a", dev), (pad_to * T, HV), a.dtype, dev) + b_buf = _pad_scratch(("b", dev), (pad_to * T, HV), b.dtype, dev) + a_buf[:total_spec].copy_(a[:total_spec]) + b_buf[:total_spec].copy_(b[:total_spec]) + mixed_qkv_spec, a, b = packed, a_buf, b_buf + # hist/idx tensors arriving PRE-PADDED at bucket length (the + # builder fills pad rows each step) are used as-is — no per-layer + # copies/fills. Shorter tensors take the legacy staging path. + if hist_len.shape[0] >= pad_to: + hist_len = hist_len[:pad_to] + else: + hist_buf = _pad_scratch(("hist", dev), (pad_to,), torch.int32, dev) + hist_buf[:B].copy_(hist_len) + hist_buf[B:].fill_(0) + hist_len = hist_buf + if state_indices.shape[0] >= pad_to: + state_indices = state_indices[:pad_to] + else: + idx_buf = _pad_scratch( + ("idx", dev), (pad_to,), state_indices.dtype, dev, + fill=UCACHE_PAD_ROW_ID + ) + idx_buf[:B].copy_(state_indices) + idx_buf[B:].fill_(UCACHE_PAD_ROW_ID) + state_indices = idx_buf B = pad_to total_spec = B * T @@ -210,22 +273,29 @@ def gdn_ucache_spec_verify( q = qkv[..., : HK * K].unflatten(-1, (HK, K)) k = qkv[..., HK * K : 2 * HK * K].unflatten(-1, (HK, K)) v = qkv[..., 2 * HK * K :].unflatten(-1, (HV, V)) - # a/b are chunk() views (token stride 2*HV); rows [0:total_spec] is the - # exact window the Triton spec kernel reads (parity). Materialize - # contiguous copies for the kernel wrappers. - a_spec = a[:total_spec].contiguous().reshape(B, T, HV) - b_spec = b[:total_spec].contiguous().reshape(B, T, HV) + # a/b are chunk() views with token stride 2*HV; rows [0:total_spec] is the + # exact window the Triton spec kernel reads (parity). reshape is a pure + # VIEW here (regular token stride), and the wrappers' strided-a/b mode + # (ab_t_stride, kernel repo) reads it directly -- the two per-layer + # .contiguous() copies (~3us x 36 layers/step) are gone. Wrappers fall + # back to staging automatically if the stride pattern is irregular. + a_spec = a[:total_spec].reshape(B, T, HV) + b_spec = b[:total_spec].reshape(B, T, HV) if _io is not a_spec.dtype: a_spec = a_spec.to(_io) b_spec = b_spec.to(_io) + _fused_out = None + if _dest is not None and _dest.shape[0] == B * T: + _fused_out = _dest.view(B, T, HV, V) out = mod.gated_delta_rule_mtp_ucache_flush( # nn.Parameters carry requires_grad=True, which DLPack refuses to - # export; detach() shares storage. - A_log=A_log.detach(), + # export; detach() shares storage so the wrapper's identity-keyed + # bf16 cast cache still hits. + A_log=_detached(A_log), a=a_spec, - dt_bias=dt_bias.detach(), + dt_bias=_detached(dt_bias), q=q, k=k, v=v, @@ -238,6 +308,8 @@ def gdn_ucache_spec_verify( hist_len=hist_len, scale=scale, use_qk_l2norm_in_kernel=True, + # caller-destination alias when available, else zero-copy view + output=_fused_out, flush_min=ucache_flush_min(T), restart_hist_on_flush=False, # builder-owned restart (see module doc) ) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 2abf3f23397e..7dc9c09d7cdd 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -1556,6 +1556,23 @@ def _forward_core( assert attn_metadata.spec_hist_len_d is not None assert attn_metadata.spec_state_indices_col0_d is not None n = attn_metadata.num_spec_decodes + _pb = attn_metadata.spec_padded_rows + _slice_rows = max( + n, _pb if _pb is not None + else (self.gdn_spec_ucache_pad_b or 0) + ) + _spec_dest = None + if ( + attn_metadata.num_prefills == 0 + and attn_metadata.num_decodes == 0 + ): + # Graph steps: hand the PADDED extent so the kernels write + # in place at the bucket size (padded rows are discarded + # downstream, same as every other padded activation). + _spec_dest = core_attn_out[ + : (_pb * self.max_spec_len) + if _pb is not None else num_actual_tokens + ] core_attn_out_spec = gdn_ucache_spec_verify( mixed_qkv_spec=mixed_qkv_spec, a=a, @@ -1566,8 +1583,14 @@ def _forward_core( u_cache=self_kv_cache[2], k_cache=self_kv_cache[3], g_cache=self_kv_cache[4], - hist_len=attn_metadata.spec_hist_len_d[:n], - state_indices=attn_metadata.spec_state_indices_col0_d[:n], + # Bucket-length PRE-PADDED slices (builder fills pad rows: + # hist=0, idx=-1): the adapter skips its per-layer hist/idx + # staging copies+fills entirely. + hist_len=attn_metadata.spec_hist_len_d[:_slice_rows], + state_indices=attn_metadata.spec_state_indices_col0_d[ + :_slice_rows + ], + output=_spec_dest, num_spec_decodes=n, max_spec_len=self.max_spec_len, num_k_heads=self.num_k_heads // self.tp_size, @@ -1575,7 +1598,9 @@ def _forward_core( head_v_dim=self.head_v_dim, scale=self.head_k_dim**-0.5, strided_qkv=self.gdn_spec_ucache_strided, - pad_to=self.gdn_spec_ucache_pad_b, + pad_to=( + _pb if _pb is not None else self.gdn_spec_ucache_pad_b + ), ).unsqueeze(0) last_recurrent_state = None elif spec_sequence_masks is not None and self.use_cache_spec_kernel: @@ -1776,7 +1801,9 @@ def _forward_core( merged_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) core_attn_out[:num_actual_tokens] = merged_out.squeeze(0) elif spec_sequence_masks is not None: - core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0) + _spec_flat = core_attn_out_spec.squeeze(0) + if _spec_flat.data_ptr() != core_attn_out.data_ptr(): + core_attn_out[:num_actual_tokens] = _spec_flat else: core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0) diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index ad017b8d9d38..cf8ab828d7fb 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -82,6 +82,11 @@ class GDNAttentionMetadata: # for the CuTeDSL verify+flush kernel. None on the triton backend. spec_hist_len_d: torch.Tensor | None = None # shape: [batch,] int32 spec_state_indices_col0_d: torch.Tensor | None = None # shape: [batch,] int32 + # CUDA-graph padded row count of this step's bucket (None when eager). + # The layer pads the ucache spec call to THIS (bucket-aware) instead of + # the static max: zero-copy pad claims then fit the bucket's buffers and + # the kernel grid matches the bucket. + spec_padded_rows: int | None = None # Pre-computed FLA chunk metadata (avoids GPU->CPU sync in prepare_chunk_indices) chunk_indices: torch.Tensor | None = None @@ -525,6 +530,7 @@ def build( # type: ignore[override] spec_is_flush_d = None spec_hist_len_d = None spec_state_indices_col0_d = None + spec_padded_rows = None if self.use_cached_kernel and spec_sequence_masks is None and num_decodes > 0: num_prompt_tokens_cpu = m.num_prompt_tokens_cpu num_computed_tokens_cpu = m._num_computed_tokens_cpu @@ -720,13 +726,16 @@ def build( # type: ignore[override] num_accepted_tokens[num_spec_decodes:].fill_(1) if self.gdn_spec_backend == "flashinfer_ucache": + spec_padded_rows = batch_size # Padded rows: negative sentinel -> the kernel retires the # whole CTA at entry (pad-skip), so under-bucket batches cost # ~nothing instead of a full T-step verify per padded row. # Requires kernel commit with _exit_cta_if_neg (455b0f6+); # on older kernels use 0 (P=0 verify against null page 0). - # Fill to the END of the fixed buffers: rows beyond the live - # batch must read hist=0 / idx=-1 on graph replay. One eager + # Fill to the END of the fixed buffers (not just the graph + # bucket): the layer passes bucket-length PRE-PADDED slices + # to the adapter (skipping its per-layer re-staging), and the + # adapter's pad_to may exceed the current bucket. One eager # fill over <=max_bs int32 per step — free. self.spec_hist_len_gathered[num_spec_decodes:].fill_(0) self.spec_state_indices_col0[num_spec_decodes:].fill_( @@ -792,6 +801,7 @@ def build( # type: ignore[override] spec_is_flush_d=spec_is_flush_d, spec_hist_len_d=spec_hist_len_d, spec_state_indices_col0_d=spec_state_indices_col0_d, + spec_padded_rows=spec_padded_rows, nums_dict=nums_dict, batch_ptr=batch_ptr, token_chunk_offset_ptr=token_chunk_offset_ptr, From a44058a052bae3b664b52e2845e00f4730d2268f Mon Sep 17 00:00:00 2001 From: Amey Naik <212485788+ameynaik-hub@users.noreply.github.com> Date: Fri, 24 Jul 2026 20:07:52 -0700 Subject: [PATCH 22/23] gdn ucache backend: adopt the Triton ring cursor model (RING_SLOTS=32) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../layers/fla/ops/gdn_ucache_spec.py | 47 ++++++++---- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 9 ++- vllm/v1/attention/backends/gdn_attn.py | 72 ++++++++++++++----- 3 files changed, 89 insertions(+), 39 deletions(-) diff --git a/vllm/model_executor/layers/fla/ops/gdn_ucache_spec.py b/vllm/model_executor/layers/fla/ops/gdn_ucache_spec.py index 39b6e992deb8..3330006fa52c 100644 --- a/vllm/model_executor/layers/fla/ops/gdn_ucache_spec.py +++ b/vllm/model_executor/layers/fla/ops/gdn_ucache_spec.py @@ -9,21 +9,21 @@ checkpoint (and restarts the ring at slots ``[0, T)``) for rows at or past the threshold. One kernel per layer per step; CUDA-graph capturable. -Ring page layout on this backend (allocated via ``ring_slots=16``): +Ring page layout — IDENTICAL to the Triton backend (pow2 ring, 32 slots): page[1] = checkpoint [blocks, HV, V, K] fp16 default (GDN_UCACHE_STATE_DTYPE) - page[2] = u_cache [blocks, HV, 16, V] fp16 default (GDN_UCACHE_RING_DTYPE) - page[3] = k_cache [blocks, HK, 16, K] fp16 default (GDN_UCACHE_RING_DTYPE) - page[4] = g_cache [blocks, HV, 16] f32 (abs cumulative log-decay) - -Cursor model: ONE block-keyed persistent ``hist_len`` buffer, committed -eagerly in the metadata builder (outside any captured region) by -``commit_gdn_ucache_hist`` with the previous step's acceptance: -``new = (old >= flush_min ? 0 : old) + accepted`` — the ring restart of a -flush step is folded into the commit, so ``hist_len`` is strictly read-only -inside the captured forward (the kernel wrapper is called with -``restart_hist_on_flush=False``; see the kernel-repo flag). The builder -gathers the block-keyed values into fixed-address request-keyed buffers that -the captured kernel reads. + page[2] = u_cache [blocks, HV, 32, V] fp16 default (GDN_UCACHE_RING_DTYPE) + page[3] = k_cache [blocks, HK, 32, K] fp16 default (GDN_UCACHE_RING_DTYPE) + page[4] = g_cache [blocks, HV, 32] f32 (abs cumulative log-decay) + +Cursor model — IDENTICAL to the Triton backend: block-keyed ``write_pos`` / +``cache_base`` / ``is_flush``, committed eagerly in the metadata builder +(outside any captured region) by ``commit_gdn_replayssm_spec`` with the +previous step's acceptance. The live window is +``[cache_base, cache_base + write_pos) mod 32``; the kernel appends at +``(cache_base + write_pos + s) & 31`` — past the window, so a flush never +overwrites rows a sibling CTA still reads. The builder gathers the +block-keyed cursor values into fixed-address request-keyed buffers that the +captured kernel reads (the CuTeDSL kernel wants request-keyed inputs). """ import importlib.util @@ -39,6 +39,9 @@ # The kernel's hardcoded ring depth; replayssm_buffer_len must equal this. UCACHE_W_RING = 16 +# Physical ring depth — must match the kernel's RING_SLOTS AND the Triton +# backend's pow2 ring (next_pow2(replayssm_buffer_len + num_spec) = 32). +UCACHE_RING_SLOTS = 32 _KMOD: Any = None @@ -101,6 +104,11 @@ def load_ucache_kernel_module(strided_qkv: bool = True): assert mod.W_RING == UCACHE_W_RING, ( f"ucache kernel W_RING={mod.W_RING} != expected {UCACHE_W_RING}" ) + assert getattr(mod, "RING_SLOTS", None) == UCACHE_RING_SLOTS, ( + f"ucache kernel RING_SLOTS={getattr(mod, 'RING_SLOTS', None)} != " + f"expected {UCACHE_RING_SLOTS} (pre-ring kernel builds are " + "incompatible with this backend's Triton-ring cursor model)" + ) logger.info_once( "GDN spec backend flashinfer_ucache: loaded kernel module from %s " "(strided_qkv=%s)", @@ -151,6 +159,7 @@ def gdn_ucache_spec_verify( k_cache: torch.Tensor, # [blocks, HK, 16, K] bf16 g_cache: torch.Tensor, # [blocks, HV, 16] f32 hist_len: torch.Tensor, # [B] int32, request-keyed (gathered by builder) + cache_base: torch.Tensor, # [B] int32 ring window origin (gathered) state_indices: torch.Tensor, # [B] int32 physical block per request num_spec_decodes: int, max_spec_len: int, @@ -254,6 +263,13 @@ def _claimable(t, rows, cols): hist_buf[:B].copy_(hist_len) hist_buf[B:].fill_(0) hist_len = hist_buf + if cache_base.shape[0] >= pad_to: + cache_base = cache_base[:pad_to] + else: + base_buf = _pad_scratch(("base", dev), (pad_to,), torch.int32, dev) + base_buf[:B].copy_(cache_base) + base_buf[B:].fill_(0) + cache_base = base_buf if state_indices.shape[0] >= pad_to: state_indices = state_indices[:pad_to] else: @@ -306,12 +322,13 @@ def _claimable(t, rows, cols): u_cache=u_cache, g_cache=g_cache, hist_len=hist_len, + cache_base=cache_base, scale=scale, use_qk_l2norm_in_kernel=True, # caller-destination alias when available, else zero-copy view output=_fused_out, flush_min=ucache_flush_min(T), - restart_hist_on_flush=False, # builder-owned restart (see module doc) + restart_hist_on_flush=False, # builder-owned cursor commit (module doc) ) if out.dtype is not _out_dtype: out = out.to(_out_dtype) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 7dc9c09d7cdd..9044c81ae196 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -434,10 +434,8 @@ def get_state_shape( self.conv_kernel_size, self.cache_config.replayssm_buffer_len, self.num_spec, - # ucache kernel: fixed 16-slot linear ring (u/k/g), no pow2. - ring_slots=16 - if self.gdn_spec_backend == "flashinfer_ucache" - else None, + # Both backends share the Triton ring layout: + # next_pow2(B + num_spec) = 32 slots (kernel RING_SLOTS). ) elif self.cache_config.use_replayssm: return MambaStateShapeCalculator.gated_delta_net_replayssm_state_shape( @@ -1584,9 +1582,10 @@ def _forward_core( k_cache=self_kv_cache[3], g_cache=self_kv_cache[4], # Bucket-length PRE-PADDED slices (builder fills pad rows: - # hist=0, idx=-1): the adapter skips its per-layer hist/idx + # hist=0, base=0, idx=-1): the adapter skips its per-layer # staging copies+fills entirely. hist_len=attn_metadata.spec_hist_len_d[:_slice_rows], + cache_base=attn_metadata.spec_ring_base_d[:_slice_rows], state_indices=attn_metadata.spec_state_indices_col0_d[ :_slice_rows ], diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index cf8ab828d7fb..b3c70cdb1c88 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -81,6 +81,7 @@ class GDNAttentionMetadata: # each build from the block-keyed master) + contiguous col-0 state indices # for the CuTeDSL verify+flush kernel. None on the triton backend. spec_hist_len_d: torch.Tensor | None = None # shape: [batch,] int32 + spec_ring_base_d: torch.Tensor | None = None # shape: [batch,] int32 spec_state_indices_col0_d: torch.Tensor | None = None # shape: [batch,] int32 # CUDA-graph padded row count of this step's bucket (None when eager). # The layer pads the ucache spec call to THIS (bucket-aware) instead of @@ -225,14 +226,11 @@ def __init__( if self.use_cache_spec_kernel else "triton" ) - self.spec_hist_len: torch.Tensor | None = None if self.gdn_spec_backend == "flashinfer_ucache": from vllm.model_executor.layers.fla.ops.gdn_ucache_spec import ( UCACHE_PAD_ROW_ID, - ucache_flush_min, ) - self.ucache_flush_min = ucache_flush_min(self.max_spec_len) self.ucache_pad_row_id = UCACHE_PAD_ROW_ID max_reqs = max( vllm_config.scheduler_config.max_num_seqs, @@ -241,6 +239,9 @@ def __init__( self.spec_hist_len_gathered: torch.Tensor = torch.zeros( (max_reqs,), dtype=torch.int32, device=device ) + self.spec_ring_base_gathered: torch.Tensor = torch.zeros( + (max_reqs,), dtype=torch.int32, device=device + ) self.spec_state_indices_col0: torch.Tensor = torch.full( (max_reqs,), UCACHE_PAD_ROW_ID, dtype=torch.int32, device=device ) @@ -529,6 +530,7 @@ def build( # type: ignore[override] spec_cache_base_d = None spec_is_flush_d = None spec_hist_len_d = None + spec_ring_base_d = None spec_state_indices_col0_d = None spec_padded_rows = None if self.use_cached_kernel and spec_sequence_masks is None and num_decodes > 0: @@ -595,39 +597,69 @@ def build( # type: ignore[override] first_decode_d = first_decode_full.index_select(0, spec_row_idx) if self.gdn_spec_backend == "flashinfer_ucache": - from vllm.model_executor.layers.fla.ops.gdn_ucache_spec import ( - commit_gdn_ucache_hist, + # RING cursors — identical block-keyed write_pos/cache_base + # model and commit kernel as the Triton backend (the CuTeDSL + # kernel consumes the same ring; its flush predicate + # P >= flush_min coincides with the commit's arming rule + # wp + 2T > L). The kernel wants REQUEST-keyed cursor values, + # so we additionally gather into fixed-address buffers the + # captured graph reads (padded rows filled below). + from vllm.model_executor.layers.fla.ops.gdn_replayssm_spec_decode import ( + commit_gdn_replayssm_spec, + reset_gdn_replayssm_spec_cursors, ) - if self.spec_hist_len is None: + if self.spec_write_pos is None: n_blocks = self.vllm_config.cache_config.num_gpu_blocks assert n_blocks is not None and n_blocks > 0, ( "--use-replayssm-spec needs num_gpu_blocks at " - "build time to size the block-keyed hist buffer" + "build time to size the block-keyed cursor buffers" + ) + self.spec_write_pos = torch.zeros( + n_blocks, dtype=torch.int32, device=self.cursor_device ) - self.spec_hist_len = torch.zeros( + self.spec_cache_base = torch.zeros( n_blocks, dtype=torch.int32, device=self.cursor_device ) - # Commit-at-start with the previous step's acceptance; the - # flush-step ring restart is folded in (hist >= flush_min -> 0 - # before adding accepted). Eager, outside any captured region. - commit_gdn_ucache_hist( - self.spec_hist_len, + self.spec_is_flush = torch.zeros( + n_blocks, dtype=torch.int8, device=self.cursor_device + ) + commit_gdn_replayssm_spec( + self.spec_write_pos, + self.spec_cache_base, + self.spec_is_flush, num_accepted_tokens.to(torch.int32), sbi, - first_decode_d, - flush_min=self.ucache_flush_min, + max_cache_len=self.spec_flush_threshold, + max_spec_len=self.max_spec_len, + cache_buf_len=self.spec_cache_buf_len, ) - # Gather block-keyed -> request-keyed fixed-address buffers - # (the captured kernel reads these; padded rows filled below). + if first_decode_d is not None: + reset_gdn_replayssm_spec_cursors( + self.spec_write_pos, + self.spec_cache_base, + self.spec_is_flush, + first_decode_d, + sbi, + max_cache_len=self.spec_flush_threshold, + max_spec_len=self.max_spec_len, + ) + sbi64 = sbi.to(torch.int64) torch.index_select( - self.spec_hist_len, + self.spec_write_pos, 0, - sbi.to(torch.int64), + sbi64, out=self.spec_hist_len_gathered[:num_spec_decodes], ) + torch.index_select( + self.spec_cache_base, + 0, + sbi64, + out=self.spec_ring_base_gathered[:num_spec_decodes], + ) self.spec_state_indices_col0[:num_spec_decodes].copy_(sbi) spec_hist_len_d = self.spec_hist_len_gathered + spec_ring_base_d = self.spec_ring_base_gathered spec_state_indices_col0_d = self.spec_state_indices_col0 else: from vllm.model_executor.layers.fla.ops.gdn_replayssm_spec_decode import ( @@ -738,6 +770,7 @@ def build( # type: ignore[override] # adapter's pad_to may exceed the current bucket. One eager # fill over <=max_bs int32 per step — free. self.spec_hist_len_gathered[num_spec_decodes:].fill_(0) + self.spec_ring_base_gathered[num_spec_decodes:].fill_(0) self.spec_state_indices_col0[num_spec_decodes:].fill_( self.ucache_pad_row_id ) @@ -800,6 +833,7 @@ def build( # type: ignore[override] spec_cache_base_d=spec_cache_base_d, spec_is_flush_d=spec_is_flush_d, spec_hist_len_d=spec_hist_len_d, + spec_ring_base_d=spec_ring_base_d, spec_state_indices_col0_d=spec_state_indices_col0_d, spec_padded_rows=spec_padded_rows, nums_dict=nums_dict, From b597ed8a514eb757ae807410809980722e679e09 Mon Sep 17 00:00:00 2001 From: Amey Naik <212485788+ameynaik-hub@users.noreply.github.com> Date: Sat, 25 Jul 2026 22:07:00 -0700 Subject: [PATCH 23/23] =?UTF-8?q?gdn=20ucache=20ring:=20review=20fixes=20?= =?UTF-8?q?=E2=80=94=20pad=20fill=20on=20every=20spec=20step,=20init-time?= =?UTF-8?q?=20ring=20check,=20ring-era=20kernel=20test?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .../kernels/test_replayssm_ucache_spec_gdn.py | 283 +++++++++++------- .../layers/fla/ops/gdn_ucache_spec.py | 91 ++---- .../layers/mamba/mamba_utils.py | 77 +++-- vllm/v1/attention/backends/gdn_attn.py | 33 +- 4 files changed, 260 insertions(+), 224 deletions(-) diff --git a/tests/kernels/test_replayssm_ucache_spec_gdn.py b/tests/kernels/test_replayssm_ucache_spec_gdn.py index c050122cd833..928007caeb59 100644 --- a/tests/kernels/test_replayssm_ucache_spec_gdn.py +++ b/tests/kernels/test_replayssm_ucache_spec_gdn.py @@ -3,20 +3,26 @@ """Integration tests for the flashinfer_ucache GDN spec-decode backend. Covers what is NEW in the vLLM integration (the kernel's own numerics are -anchored by the kernel repo's fp32-reference gates: _ucache_check.py, -_ucache_flush_check.py, _ucache_g3_ring.py): - -1. commit_gdn_ucache_hist semantics (flush restart folded into commit, - first-decode reset, null-block masking). -2. Intra-step shared hist_len: with restart_hist_on_flush=False, N "layers" - sharing one hist tensor within a step all see the same P (the bug the - flag exists to prevent). -3. Protocol equivalence: the vLLM bookkeeping (builder-owned commit, - restart_hist_on_flush=False) must produce bit-identical outputs and - checkpoints to the kernel-repo protocol (hist += accepted each step, - wrapper masked_fill_ restart) over many steps crossing several flushes. +anchored by the kernel repo's fp32-reference suite, tests/gdn/): + +1. Shared Triton cursor semantics (commit_gdn_replayssm_spec / + reset_gdn_replayssm_spec_cursors driving the ucache kernel): flush rows + slide cache_base past the folded window mod RING_SLOTS, first-decode + reset clears BOTH cursors, null-block rows are never touched. +2. Intra-step shared cursors: N "layers" reading one gathered + (hist_len, cache_base) pair within a step all see the same window (the + kernel treats cursors as read-only). +3. Protocol equivalence: the vLLM bookkeeping (block-keyed cursors + + commit_gdn_replayssm_spec + gather, restart_hist_on_flush=False) must + produce bit-identical outputs, checkpoints, and window origins to the + kernel-repo standalone protocol (request-keyed cursors, wrapper commit + via restart_hist_on_flush=True) over many steps crossing several + flushes and ring wrap-arounds. 4. Strided packed-qkv slices (production layout) match dense inputs. 5. Null-page rows only scribble the reserved page 0. +6. Block-strided (vLLM paged) pools match dense pools, incl. past 2^31 + elements (64-bit addressing). +7. Padded rows (negative sentinel) retire their CTAs at kernel entry. Run inside the vLLM container with: VLLM_GDN_UCACHE_MODULE= @@ -27,15 +33,26 @@ import pytest import torch +# This suite allocates bf16 checkpoint/ring pools; pin the kernel module's +# env-selected dtypes to match BEFORE it loads (the adapter defaults both to +# fp16 for serving). Must run before the first load_ucache_kernel_module(). +os.environ["GDN_UCACHE_STATE_DTYPE"] = "bf16" +os.environ["GDN_UCACHE_RING_DTYPE"] = "bf16" + +from vllm.model_executor.layers.fla.ops.gdn_replayssm_spec_decode import ( + commit_gdn_replayssm_spec, + reset_gdn_replayssm_spec_cursors, +) from vllm.model_executor.layers.fla.ops.gdn_ucache_spec import ( + UCACHE_RING_SLOTS, UCACHE_W_RING, - commit_gdn_ucache_hist, load_ucache_kernel_module, ucache_flush_min, ) DEV = "cuda" HK, HV, K, V = 16, 64, 128, 128 # qwen122b geometry +RING = UCACHE_RING_SLOTS # 32 def _kmod(): @@ -80,20 +97,37 @@ def _pools(num_blocks, seed=11): ) * 0.05 ).to(torch.bfloat16) - k_cache = torch.zeros( - num_blocks, HK, UCACHE_W_RING, K, device=DEV, dtype=torch.bfloat16 - ) - u_cache = torch.zeros( - num_blocks, HV, UCACHE_W_RING, V, device=DEV, dtype=torch.bfloat16 - ) - g_cache = torch.zeros( - num_blocks, HV, UCACHE_W_RING, device=DEV, dtype=torch.float32 - ) + k_cache = torch.zeros(num_blocks, HK, RING, K, device=DEV, dtype=torch.bfloat16) + u_cache = torch.zeros(num_blocks, HV, RING, V, device=DEV, dtype=torch.bfloat16) + g_cache = torch.zeros(num_blocks, HV, RING, device=DEV, dtype=torch.float32) return ckpt, k_cache, u_cache, g_cache +def _cursors(n_blocks): + """Block-keyed cursor triple, exactly as the metadata builder allocates.""" + wp = torch.zeros(n_blocks, dtype=torch.int32, device=DEV) + cb = torch.zeros(n_blocks, dtype=torch.int32, device=DEV) + fl = torch.zeros(n_blocks, dtype=torch.int8, device=DEV) + return wp, cb, fl + + +def _commit(wp, cb, fl, acc, sbi, T): + """The builder's per-step commit with its exact parameters + (L = buffer_len + T = W_RING + T, physical ring = next_pow2(L) = 32).""" + commit_gdn_replayssm_spec( + wp, cb, fl, acc, sbi, + max_cache_len=UCACHE_W_RING + T, + max_spec_len=T, + cache_buf_len=RING, + ) + + +def _gather(t, sbi): + return t.index_select(0, sbi.to(torch.int64)).contiguous() + + def _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt, sbi, kc, uc, gc, hist, - T, restart): + base, T, restart): return mod.gated_delta_rule_mtp_ucache_flush( A_log=A_log, a=a, @@ -108,6 +142,7 @@ def _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt, sbi, kc, uc, gc, hist, u_cache=uc, g_cache=gc, hist_len=hist, + cache_base=base, scale=K**-0.5, use_qk_l2norm_in_kernel=True, output=None, @@ -116,68 +151,82 @@ def _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt, sbi, kc, uc, gc, hist, ) -def test_commit_kernel_semantics(): +def test_shared_cursor_commit_semantics(): + """The shared Triton commit drives the ucache backend: flushed rows + (is_flush armed == the kernel just folded, wp >= flush_min) slide + cache_base past the folded window mod 32 and restart wp at the accepted + count; verify rows just grow wp; first-decode reset clears BOTH cursors; + the null block is never touched.""" if not torch.cuda.is_available(): pytest.skip("CUDA required") + T = 4 n_blocks = 32 - flush_min = ucache_flush_min(4) # 13 - hist = torch.zeros(n_blocks, dtype=torch.int32, device=DEV) - # rows -> blocks 3, 5, 7, 9 (block 0 = NULL among indices) - sbi = torch.tensor([3, 5, 7, 0, 9], dtype=torch.int32, device=DEV) - hist[3] = 6 # below threshold: 6 + acc - hist[5] = 13 # at threshold: flushed last step -> acc - hist[7] = 16 # max legal: flushed -> acc - hist[9] = 2 # first-decode row: reset to 0 regardless - hist[0] = 99 # null block must never be touched + wp, cb, fl = _cursors(n_blocks) + # rows -> blocks 3, 5, 7, 9 (+ a null row on block 0) + sbi = torch.tensor([3, 5, 7, 9, 0], dtype=torch.int32, device=DEV) + wp[3], fl[3] = 6, 0 # verify row: 6 + acc + wp[5], fl[5] = 13, 1 # flushed at threshold: base += 13, wp = acc + wp[7], cb[7], fl[7] = 16, 20, 1 # flushed at max, wrapping base: (20+16)&31 + wp[9], fl[9] = 12, 0 # verify row landing ON the arm point + wp[0] = 99 # null block must never be touched acc = torch.tensor([3, 2, 4, 1, 4], dtype=torch.int32, device=DEV) - first_decode = torch.tensor([0, 0, 0, 0, 1], dtype=torch.int8, device=DEV) - commit_gdn_ucache_hist(hist, acc, sbi, first_decode, flush_min=flush_min) + _commit(wp, cb, fl, acc, sbi, T) + torch.cuda.synchronize() + assert wp[3].item() == 9 and cb[3].item() == 0 and fl[3].item() == 0 + assert wp[5].item() == 2 and cb[5].item() == 13 and fl[5].item() == 0 + assert wp[7].item() == 4 and cb[7].item() == (20 + 16) % RING # == 4 + assert fl[7].item() == 0 + # wp 12+1=13 == flush_min(4): is_flush arms for the NEXT step + assert wp[9].item() == 13 and fl[9].item() == 1 + assert wp[0].item() == 99 and cb[0].item() == 0 # null untouched + assert (wp[1:] <= UCACHE_W_RING).all() + assert (cb >= 0).all() and (cb < RING).all() + + # first-decode reset (prefill->decode handoff) clears BOTH cursors + do_reset = torch.tensor([0, 0, 0, 1, 0], dtype=torch.int8, device=DEV) + reset_gdn_replayssm_spec_cursors( + wp, cb, fl, do_reset, sbi, + max_cache_len=UCACHE_W_RING + T, max_spec_len=T, + ) torch.cuda.synchronize() - assert hist[3].item() == 9 # 6 + 3 - assert hist[5].item() == 2 # flush restart + 2 - assert hist[7].item() == 4 # flush restart + 4 - assert hist[9].item() == 0 # first-decode reset - assert hist[0].item() == 99 # null untouched - assert (hist[1:] <= UCACHE_W_RING).all() # non-null blocks only + assert wp[9].item() == 0 and cb[9].item() == 0 and fl[9].item() == 0 + assert wp[5].item() == 2 and cb[5].item() == 13 # others untouched @pytest.mark.parametrize("T", [4, 8]) -def test_intra_step_shared_hist_and_flag(T): - """Two 'layers' share one hist tensor in a step; with restart=False the - second layer must see the same P and produce the same fold as the first - (independent pools). With restart=True the second layer would see P=0.""" +def test_intra_step_shared_cursors(T): + """Two 'layers' share one gathered (hist, base) pair in a step; the + kernel treats cursors as read-only, so both layers must see the same + window and produce identical outputs and folds (independent pools).""" mod = _kmod() A_log, dt_bias = _gating_params() B, n_blocks = 3, 8 - flush_min = ucache_flush_min(T) + fm = ucache_flush_min(T) q, k, v, a, b = _rand_inputs(B, T, seed=23) sbi = torch.tensor([1, 4, 6], dtype=torch.int32, device=DEV) - hist_master = torch.tensor( - [flush_min, 5, flush_min + 1], dtype=torch.int32, device=DEV - ) - # Two independent "layers" with identical pools and inputs. - outs, ckpts, hists = [], [], [] + hist_master = torch.tensor([fm, 5, fm + 1], dtype=torch.int32, device=DEV) + base_master = torch.tensor([0, 28, 5], dtype=torch.int32, device=DEV) hist = hist_master.clone() - # Pre-fill rings identically for both layers so P>0 rows have history. - for layer in range(2): + base = base_master.clone() + outs, ckpts = [], [] + for _layer in range(2): ckpt, kc, uc, gc = _pools(n_blocks, seed=31) - # Prime the ring: run one step from P=0 (appends T entries), then - # set hist to the master values for the step under test. + # Prime the ring (appends T entries from P=0), then run the step + # under test with the SHARED master cursors. hist0 = torch.zeros(B, dtype=torch.int32, device=DEV) + base0 = torch.zeros(B, dtype=torch.int32, device=DEV) _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt, sbi, kc, uc, gc, - hist0, T, restart=False) - hist_layer = hist if layer == 0 else hist # SHARED tensor + hist0, base0, T, restart=False) out = _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt, sbi, kc, uc, - gc, hist_layer, T, restart=False) + gc, hist, base, T, restart=False) torch.cuda.synchronize() outs.append(out.clone()) ckpts.append(ckpt.clone()) - hists.append(hist_layer.clone()) - # hist untouched by both calls - assert torch.equal(hists[0], hist_master) - assert torch.equal(hists[1], hist_master) - # both layers saw identical P -> identical outputs and identical folds + # cursors untouched by both calls (restart=False -> kernel read-only) + assert torch.equal(hist, hist_master) + assert torch.equal(base, base_master) + # both layers saw identical windows -> identical outputs and folds assert torch.equal(outs[0], outs[1]) assert torch.equal(ckpts[0], ckpts[1]) @@ -185,50 +234,56 @@ def test_intra_step_shared_hist_and_flag(T): @pytest.mark.parametrize("T", [4]) @pytest.mark.parametrize("nreq", [1, 3, 8]) def test_protocol_equivalence_multi_step(T, nreq): - """vLLM bookkeeping (commit kernel + restart=False) vs kernel-repo - bookkeeping (hist += accepted + wrapper restart) over 24 steps crossing - several flush cycles: outputs and checkpoints must match bit-for-bit.""" + """vLLM bookkeeping (block-keyed cursors + shared Triton commit + + gather, restart=False) vs kernel-repo standalone bookkeeping + (request-keyed cursors, wrapper commit via restart=True) over 24 steps + crossing several flush cycles and ring wraps: outputs, checkpoints, and + window origins must match bit-for-bit.""" mod = _kmod() A_log, dt_bias = _gating_params() n_blocks = 16 - flush_min = ucache_flush_min(T) # permuted, non-trivial block assignment (block 0 reserved) - perm = torch.randperm(n_blocks - 1)[:nreq] + 1 - sbi = perm.to(torch.int32).to(DEV) + perm = torch.randperm(n_blocks - 1, generator=torch.Generator().manual_seed(13)) + sbi = (perm[:nreq] + 1).to(torch.int32).to(DEV) ckpt_a, kc_a, uc_a, gc_a = _pools(n_blocks, seed=41) ckpt_b, kc_b, uc_b, gc_b = _pools(n_blocks, seed=41) - # Protocol A (vLLM): block-keyed master + commit kernel. - hist_blocks = torch.zeros(n_blocks, dtype=torch.int32, device=DEV) - # Protocol B (kernel repo): request-keyed hist, wrapper restart. + # Protocol A (vLLM): block-keyed cursors + shared Triton commit. + wp, cb, fl = _cursors(n_blocks) + # Protocol B (kernel repo): request-keyed cursors, wrapper commit. hist_req = torch.zeros(nreq, dtype=torch.int32, device=DEV) + base_req = torch.zeros(nreq, dtype=torch.int32, device=DEV) gen = torch.Generator().manual_seed(97) prev_acc = torch.zeros(nreq, dtype=torch.int32, device=DEV) - first = torch.zeros(nreq, dtype=torch.int8, device=DEV) for step in range(24): q, k, v, a, b = _rand_inputs(nreq, T, seed=1000 + step) - # A: commit (prev step's acceptance), gather, call with restart=False - commit_gdn_ucache_hist( - hist_blocks, prev_acc, sbi, first if step == 0 else None, - flush_min=flush_min, - ) - gathered = hist_blocks.index_select(0, sbi.to(torch.int64)).contiguous() + # A: commit the previous step's acceptance, gather, run restart=False + _commit(wp, cb, fl, prev_acc, sbi, T) + hd = _gather(wp, sbi) + bd = _gather(cb, sbi) out_a = _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt_a, sbi, - kc_a, uc_a, gc_a, gathered, T, restart=False) - # B: kernel-repo protocol on request-keyed hist + kc_a, uc_a, gc_a, hd, bd, T, restart=False) + # B: kernel-repo protocol; the wrapper commits flushed rows itself + # DURING the call (A's builder slides at the NEXT step's commit), so + # the phase-aligned comparison point is B's PRE-call cursors. + hist_b_pre = hist_req.clone() + base_b_pre = base_req.clone() out_b = _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt_b, sbi, - kc_b, uc_b, gc_b, hist_req, T, restart=True) + kc_b, uc_b, gc_b, hist_req, base_req, T, restart=True) torch.cuda.synchronize() assert torch.equal(out_a, out_b), f"outputs diverged at step {step}" assert torch.equal(ckpt_a, ckpt_b), f"checkpoints diverged at step {step}" - assert (gathered <= UCACHE_W_RING).all() + assert torch.equal(hd, hist_b_pre), f"windows diverged at step {step}" + assert torch.equal(bd, base_b_pre), f"window origins diverged at step {step}" + assert (hd <= UCACHE_W_RING).all() + assert (bd >= 0).all() and (bd < RING).all() acc = torch.randint(1, T + 1, (nreq,), generator=gen).to( torch.int32 ).to(DEV) prev_acc = acc - hist_req += acc # protocol B commit (wrapper already restarted) + hist_req += acc # protocol B commit (wrapper already slid the base) def test_strided_packed_qkv_matches_dense(): @@ -250,8 +305,9 @@ def test_strided_packed_qkv_matches_dense(): for (qq, kk, vv) in [(q, k, v), (qs, ks, vs)]: ckpt, kc, uc, gc = _pools(n_blocks, seed=61) hist = torch.zeros(B, dtype=torch.int32, device=DEV) + base = torch.zeros(B, dtype=torch.int32, device=DEV) out = _call(mod, A_log, dt_bias, qq, kk, vv, a, b, ckpt, sbi, - kc, uc, gc, hist, T, restart=False) + kc, uc, gc, hist, base, T, restart=False) torch.cuda.synchronize() outs.append(out.clone()) assert torch.equal(outs[0], outs[1]) @@ -267,13 +323,15 @@ def test_null_page_rows_only_touch_page_zero(): ckpt, kc, uc, gc = _pools(n_blocks, seed=79) snap = ckpt.clone() hist = torch.tensor([13, 0, 13], dtype=torch.int32, device=DEV) + base = torch.zeros(B, dtype=torch.int32, device=DEV) # prime rows 0/2 rings so their flush folds something hist0 = torch.zeros(B, dtype=torch.int32, device=DEV) + base0 = torch.zeros(B, dtype=torch.int32, device=DEV) _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt, sbi, kc, uc, gc, hist0, - T, restart=False) + base0, T, restart=False) snap_after_prime = ckpt.clone() _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt, sbi, kc, uc, gc, hist, - T, restart=False) + base, T, restart=False) torch.cuda.synchronize() # pages not referenced by any row are byte-stable untouched = [i for i in range(n_blocks) if i not in (0, 2, 5)] @@ -288,8 +346,7 @@ def _vllm_style_strided_pools(num_blocks, seed=11): """Carve (ckpt, k, u, g) as block-strided views of one page-major backing, exactly like vLLM's _reshape_kv_cache_tensors (inner dims dense, dim-0 stride = whole page).""" - shapes = [(HV, V, K), (HK, UCACHE_W_RING, K), (HV, UCACHE_W_RING, V), - (HV, UCACHE_W_RING)] + shapes = [(HV, V, K), (HK, RING, K), (HV, RING, V), (HV, RING)] dtypes = [torch.bfloat16, torch.bfloat16, torch.bfloat16, torch.float32] page_bytes = sum( int(torch.empty(s, device="meta").numel()) * t.itemsize @@ -329,13 +386,10 @@ def test_block_strided_pools_match_dense(): ckpt_s.copy_(ckpt_d) assert not ckpt_s.is_contiguous() and not kc_s.is_contiguous() - hist_d = torch.zeros(B, dtype=torch.int32, device=DEV) - hist_s = torch.zeros(B, dtype=torch.int32, device=DEV) + wp_d, cb_d, fl_d = _cursors(n_blocks) + wp_s, cb_s, fl_s = _cursors(n_blocks) gen = torch.Generator().manual_seed(3) - prev_acc_d = torch.zeros(B, dtype=torch.int32, device=DEV) - hist_blocks_d = torch.zeros(n_blocks, dtype=torch.int32, device=DEV) - hist_blocks_s = torch.zeros(n_blocks, dtype=torch.int32, device=DEV) - fm = ucache_flush_min(T) + prev_acc = torch.zeros(B, dtype=torch.int32, device=DEV) for step in range(12): # crosses >= 2 flush cycles qd, kd, vd, a, b = _rand_inputs(B, T, seed=500 + step) # packed strided q/k/v (production layout) so the wrapper's @@ -346,19 +400,19 @@ def test_block_strided_pools_match_dense(): q = packed[..., : HK * K].unflatten(-1, (HK, K)) k = packed[..., HK * K : 2 * HK * K].unflatten(-1, (HK, K)) v = packed[..., 2 * HK * K :].unflatten(-1, (HV, V)) - commit_gdn_ucache_hist(hist_blocks_d, prev_acc_d, sbi, None, flush_min=fm) - commit_gdn_ucache_hist(hist_blocks_s, prev_acc_d, sbi, None, flush_min=fm) - hd = hist_blocks_d.index_select(0, sbi.to(torch.int64)).contiguous() - hs = hist_blocks_s.index_select(0, sbi.to(torch.int64)).contiguous() + _commit(wp_d, cb_d, fl_d, prev_acc, sbi, T) + _commit(wp_s, cb_s, fl_s, prev_acc, sbi, T) + hd, bd = _gather(wp_d, sbi), _gather(cb_d, sbi) + hs, bs = _gather(wp_s, sbi), _gather(cb_s, sbi) out_d = _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt_d, sbi, - kc_d, uc_d, gc_d, hd, T, restart=False) + kc_d, uc_d, gc_d, hd, bd, T, restart=False) out_s = _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt_s, sbi, - kc_s, uc_s, gc_s, hs, T, restart=False) + kc_s, uc_s, gc_s, hs, bs, T, restart=False) torch.cuda.synchronize() assert torch.equal(out_d, out_s), f"outputs diverged at step {step}" assert torch.equal(ckpt_d[1:], ckpt_s[1:]), f"ckpt diverged at step {step}" assert torch.equal(kc_d[1:], kc_s[1:]) and torch.equal(uc_d[1:], uc_s[1:]) - prev_acc_d = torch.randint(1, T + 1, (B,), generator=gen).to( + prev_acc = torch.randint(1, T + 1, (B,), generator=gen).to( torch.int32 ).to(DEV) @@ -381,11 +435,10 @@ def test_block_strided_pools_past_2gb(): assert 1999 * ckpt_s.stride(0) > 2**31, ( f"test page too small to cross 2^31: stride0={ckpt_s.stride(0)}") - hist_blocks = torch.zeros(n_blocks, dtype=torch.int32, device=DEV) + wp_d, cb_d, fl_d = _cursors(n_blocks) + wp_s, cb_s, fl_s = _cursors(n_blocks) prev_acc = torch.zeros(B, dtype=torch.int32, device=DEV) - fm = ucache_flush_min(T) gen = torch.Generator().manual_seed(5) - hist_blocks_d = hist_blocks.clone() for step in range(8): # crosses a flush cycle qd, kd, vd, a, b = _rand_inputs(B, T, seed=800 + step) packed = torch.cat( @@ -394,17 +447,18 @@ def test_block_strided_pools_past_2gb(): q = packed[..., : HK * K].unflatten(-1, (HK, K)) k = packed[..., HK * K : 2 * HK * K].unflatten(-1, (HK, K)) v = packed[..., 2 * HK * K :].unflatten(-1, (HV, V)) - commit_gdn_ucache_hist(hist_blocks, prev_acc, sbi, None, flush_min=fm) - commit_gdn_ucache_hist(hist_blocks_d, prev_acc, sbi, None, flush_min=fm) - hs = hist_blocks.index_select(0, sbi.to(torch.int64)).contiguous() - hd = hist_blocks_d.index_select(0, sbi.to(torch.int64)).contiguous() + _commit(wp_d, cb_d, fl_d, prev_acc, sbi, T) + _commit(wp_s, cb_s, fl_s, prev_acc, sbi, T) + hd, bd = _gather(wp_d, sbi), _gather(cb_d, sbi) + hs, bs = _gather(wp_s, sbi), _gather(cb_s, sbi) out_s = _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt_s, sbi, - kc_s, uc_s, gc_s, hs, T, restart=False) + kc_s, uc_s, gc_s, hs, bs, T, restart=False) out_d = _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt_d, sbi, - kc_d, uc_d, gc_d, hd, T, restart=False) + kc_d, uc_d, gc_d, hd, bd, T, restart=False) torch.cuda.synchronize() assert torch.equal(out_s, out_d), f"outputs diverged at step {step}" - assert torch.equal(ckpt_s[1900:], ckpt_d[1900:]), f"high-block ckpt diverged at {step}" + assert torch.equal(ckpt_s[1900:], ckpt_d[1900:]), ( + f"high-block ckpt diverged at {step}") prev_acc = torch.randint(1, T + 1, (B,), generator=gen).to( torch.int32 ).to(DEV) @@ -427,6 +481,7 @@ def test_pad_skip_negative_rows_exit_early(): sbi_neg = torch.tensor([2, 5, 7, -1, -1, -1], dtype=torch.int32, device=DEV) sbi_nul = torch.tensor([2, 5, 7, 0, 0, 0], dtype=torch.int32, device=DEV) hist = torch.tensor([13, 7, 0, 0, 0, 0], dtype=torch.int32, device=DEV) + base = torch.zeros(B, dtype=torch.int32, device=DEV) ckpt1, kc1, uc1, gc1 = _pools(n_blocks, seed=103) ckpt2, kc2, uc2, gc2 = _pools(n_blocks, seed=103) @@ -434,9 +489,9 @@ def test_pad_skip_negative_rows_exit_early(): gc1[0].clone()) out1 = _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt1, sbi_neg, - kc1, uc1, gc1, hist.clone(), T, restart=False) + kc1, uc1, gc1, hist.clone(), base.clone(), T, restart=False) out2 = _call(mod, A_log, dt_bias, q, k, v, a, b, ckpt2, sbi_nul, - kc2, uc2, gc2, hist.clone(), T, restart=False) + kc2, uc2, gc2, hist.clone(), base.clone(), T, restart=False) torch.cuda.synchronize() # (a) real rows bit-identical across the two pad conventions. The wrapper diff --git a/vllm/model_executor/layers/fla/ops/gdn_ucache_spec.py b/vllm/model_executor/layers/fla/ops/gdn_ucache_spec.py index 3330006fa52c..69aab33aef79 100644 --- a/vllm/model_executor/layers/fla/ops/gdn_ucache_spec.py +++ b/vllm/model_executor/layers/fla/ops/gdn_ucache_spec.py @@ -33,7 +33,6 @@ import torch from vllm.logger import init_logger -from vllm.triton_utils import tl, triton logger = init_logger(__name__) @@ -45,9 +44,6 @@ _KMOD: Any = None -# Reserved null page id (mirrors vllm/v1/attention/backends/utils.py). -NULL_BLOCK_ID = 0 - # Padded-row sentinel: the kernel retires the whole CTA at entry for rows # with state index < 0 (pad-skip), so padded rows cost ~nothing instead of a # full T-step verify. Requires kernel commit 455b0f6+ (_exit_cta_if_neg); @@ -101,14 +97,21 @@ def load_ucache_kernel_module(strided_qkv: bool = True): from flashinfer.gdn_kernels import ( # type: ignore[import-not-found] gdn_decode_bf16_wy_ucache_flush as mod, ) - assert mod.W_RING == UCACHE_W_RING, ( - f"ucache kernel W_RING={mod.W_RING} != expected {UCACHE_W_RING}" - ) - assert getattr(mod, "RING_SLOTS", None) == UCACHE_RING_SLOTS, ( - f"ucache kernel RING_SLOTS={getattr(mod, 'RING_SLOTS', None)} != " - f"expected {UCACHE_RING_SLOTS} (pre-ring kernel builds are " - "incompatible with this backend's Triton-ring cursor model)" - ) + # RuntimeError (not assert): stripped asserts under `python -O` would + # let a pre-ring flat-layout kernel run against ring cursors and + # silently corrupt state. resolve_gdn_spec_backend pre-screens this at + # engine init by scanning the module source; this is the authoritative + # post-import check. + if mod.W_RING != UCACHE_W_RING or ( + getattr(mod, "RING_SLOTS", None) != UCACHE_RING_SLOTS + ): + raise RuntimeError( + f"ucache kernel module incompatible: W_RING={mod.W_RING} " + f"(expected {UCACHE_W_RING}), " + f"RING_SLOTS={getattr(mod, 'RING_SLOTS', None)} " + f"(expected {UCACHE_RING_SLOTS}). Pre-ring kernel builds are " + "incompatible with this backend's Triton-ring cursor model." + ) logger.info_once( "GDN spec backend flashinfer_ucache: loaded kernel module from %s " "(strided_qkv=%s)", @@ -335,63 +338,7 @@ def _claimable(t, rows, cols): return out.reshape(B * T, HV, V)[:real_total] -@triton.jit -def _commit_gdn_ucache_hist_kernel( - hist_ptr, # [num_gpu_blocks] int32, block-keyed - num_accepted_ptr, # [n_rows] int32 (previous step's acceptance) - sbi_ptr, # base ptr of spec_state_indices_tensor[:, 0] - first_decode_ptr, # [n_rows] int8 (dummy when HAS_RESET == False) - n_rows, - sbi_stride, - FLUSH_MIN: tl.constexpr, - HAS_RESET: tl.constexpr, - NULL_BLOCK: tl.constexpr, - BLOCK: tl.constexpr, -): - offs = tl.arange(0, BLOCK) - m = offs < n_rows - blk = tl.load(sbi_ptr + offs * sbi_stride, mask=m, other=NULL_BLOCK).to( - tl.int64 - ) - valid = m & (blk > NULL_BLOCK) - old = tl.load(hist_ptr + blk, mask=valid, other=0).to(tl.int32) - acc = tl.load(num_accepted_ptr + offs, mask=valid, other=0).to(tl.int32) - # A row whose previous verify launched with hist >= FLUSH_MIN flushed - # in-kernel (every layer): its ring restarted at [0, T), so committed - # history restarts at 0 before adding the accepted count. - new = tl.where(old >= FLUSH_MIN, 0, old) + acc - if HAS_RESET: - fd = tl.load(first_decode_ptr + offs, mask=valid, other=0).to(tl.int32) - new = tl.where(fd != 0, 0, new) # prefill->decode / block-recycle - tl.store(hist_ptr + blk, new, mask=valid) - - -def commit_gdn_ucache_hist( - hist_len: torch.Tensor, # [num_gpu_blocks] int32, block-keyed - num_accepted_tokens: torch.Tensor, # [n_rows] int32 - state_indices: torch.Tensor, # [n_rows] view of block ids (may be strided) - first_decode: torch.Tensor | None, # [n_rows] int8 or None - *, - flush_min: int, -) -> None: - """Eager (outside-capture) hist commit; mirrors commit_gdn_replayssm_spec. - - Invariant afterwards: hist <= (flush_min - 1) + T == W_RING, the kernel's - legal [0, W_RING] range. - """ - n_rows = state_indices.shape[0] - if n_rows == 0: - return - BLOCK = max(triton.next_power_of_2(n_rows), 16) - _commit_gdn_ucache_hist_kernel[(1,)]( - hist_len, - num_accepted_tokens, - state_indices, - first_decode if first_decode is not None else hist_len, - n_rows, - state_indices.stride(0), - FLUSH_MIN=flush_min, - HAS_RESET=first_decode is not None, - NULL_BLOCK=NULL_BLOCK_ID, - BLOCK=BLOCK, - ) +# (The legacy block-keyed hist-only commit, commit_gdn_ucache_hist, was +# removed: the backend shares the Triton backend's cursor machinery — +# commit_gdn_replayssm_spec / reset_gdn_replayssm_spec_cursors — which +# commits write_pos AND cache_base together.) diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index a597cf850e12..034d288b238f 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -692,36 +692,67 @@ def kda_state_copy_func(cls): def _ucache_kernel_available() -> tuple[bool, str]: - """Init-time check that the ucache CuTeDSL kernel module is loadable. + """Init-time check that the ucache CuTeDSL kernel module is loadable + AND is a ring build (RING_SLOTS == 32). Mirrors load_ucache_kernel_module's resolution order without importing (and JIT-compiling) the module: an explicit VLLM_GDN_UCACHE_MODULE path - must exist, else flashinfer.gdn_kernels must provide the module. - """ - path = os.environ.get("VLLM_GDN_UCACHE_MODULE") - if path: - return ( - os.path.isfile(path), - f"VLLM_GDN_UCACHE_MODULE points to a missing file: {path!r}", - ) + must exist, else flashinfer.gdn_kernels must provide the module. The + ring check scans the module SOURCE for the RING_SLOTS constant so that + a pre-ring kernel fails here, at engine init, rather than at the first + speculative-decode step mid-serving (load_ucache_kernel_module re-checks + the imported module authoritatively on first use).""" import importlib.util - try: - found = ( - importlib.util.find_spec( + path = os.environ.get("VLLM_GDN_UCACHE_MODULE") + if path: + if not os.path.isfile(path): + return ( + False, + f"VLLM_GDN_UCACHE_MODULE points to a missing file: {path!r}", + ) + src_path = path + else: + try: + spec = importlib.util.find_spec( "flashinfer.gdn_kernels.gdn_decode_bf16_wy_ucache_flush" ) - is not None - ) - except ModuleNotFoundError: - found = False - return ( - found, - "ucache CuTeDSL kernel module not found: set VLLM_GDN_UCACHE_MODULE=" - "/abs/path/to/gdn_decode_bf16_wy_ucache_flush.py or install a " - "FlashInfer build that provides " - "flashinfer.gdn_kernels.gdn_decode_bf16_wy_ucache_flush", - ) + except ModuleNotFoundError: + spec = None + if spec is None: + return ( + False, + "ucache CuTeDSL kernel module not found: set " + "VLLM_GDN_UCACHE_MODULE=/abs/path/to/" + "gdn_decode_bf16_wy_ucache_flush.py or install a FlashInfer " + "build that provides " + "flashinfer.gdn_kernels.gdn_decode_bf16_wy_ucache_flush", + ) + src_path = spec.origin or "" + + if src_path and os.path.isfile(src_path): + import re + + try: + with open(src_path, encoding="utf-8", errors="replace") as f: + src = f.read() + except OSError: + src = "" + if src: + m = re.search( + r"^RING_SLOTS(?:\s*:\s*\w+)?\s*=\s*(\d+)", src, re.MULTILINE + ) + ring_slots = int(m.group(1)) if m else None + if ring_slots != 32: + return ( + False, + f"kernel module at {src_path!r} is not a ring build " + f"(RING_SLOTS={ring_slots}, need 32): pre-ring ucache " + "kernels are incompatible with this backend's " + "Triton-ring cursor model — update the FlashInfer " + "ucache kernel", + ) + return (True, "") def resolve_gdn_spec_backend(vllm_config) -> str: diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index b3c70cdb1c88..4489fb066f49 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -658,6 +658,21 @@ def build( # type: ignore[override] out=self.spec_ring_base_gathered[:num_spec_decodes], ) self.spec_state_indices_col0[:num_spec_decodes].copy_(sbi) + # Fill the pad tail of all three fixed buffers EVERY spec + # step, not only in the full-CG pure-spec branch below: the + # layer hands the adapter bucket-length PRE-PADDED slices on + # eager mixed (prefill+spec) steps too, so a stale tail row + # from an earlier step would reach the kernel carrying a + # possibly-reallocated block id plus old ring cursors (ghost + # CTAs appending to a live request's block). idx uses the + # negative sentinel -> the kernel retires the CTA at entry + # (pad-skip; kernels with _exit_cta_if_neg, 455b0f6+). One + # eager fill over <=max_bs int32 per step - free. + self.spec_hist_len_gathered[num_spec_decodes:].fill_(0) + self.spec_ring_base_gathered[num_spec_decodes:].fill_(0) + self.spec_state_indices_col0[num_spec_decodes:].fill_( + self.ucache_pad_row_id + ) spec_hist_len_d = self.spec_hist_len_gathered spec_ring_base_d = self.spec_ring_base_gathered spec_state_indices_col0_d = self.spec_state_indices_col0 @@ -759,21 +774,9 @@ def build( # type: ignore[override] if self.gdn_spec_backend == "flashinfer_ucache": spec_padded_rows = batch_size - # Padded rows: negative sentinel -> the kernel retires the - # whole CTA at entry (pad-skip), so under-bucket batches cost - # ~nothing instead of a full T-step verify per padded row. - # Requires kernel commit with _exit_cta_if_neg (455b0f6+); - # on older kernels use 0 (P=0 verify against null page 0). - # Fill to the END of the fixed buffers (not just the graph - # bucket): the layer passes bucket-length PRE-PADDED slices - # to the adapter (skipping its per-layer re-staging), and the - # adapter's pad_to may exceed the current bucket. One eager - # fill over <=max_bs int32 per step — free. - self.spec_hist_len_gathered[num_spec_decodes:].fill_(0) - self.spec_ring_base_gathered[num_spec_decodes:].fill_(0) - self.spec_state_indices_col0[num_spec_decodes:].fill_( - self.ucache_pad_row_id - ) + # Pad rows (hist=0, base=0, idx=negative sentinel) were + # already filled at gather time above — done every spec + # step, including eager mixed ones, so nothing to do here. if ( self.use_full_cuda_graph