From 866fea2b9900bf49d552c205d2eaac4716fb63ac Mon Sep 17 00:00:00 2001 From: Johnny-Liou <75152852+Johnny-Liou@users.noreply.github.com> Date: Fri, 24 Jul 2026 12:39:49 -0400 Subject: [PATCH] [Kernel] ReplaySSM: cache SSM inputs for faster Mamba2 standard decode (#48018) Signed-off-by: Johnny-Liou Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: tomeras91 <57313761+tomeras91@users.noreply.github.com> Co-authored-by: Cyrus Leung --- benchmarks/replayssm/e2e_decode_speedup.py | 267 +++++++ ...ayssm_prefill_decode_equivalence_mamba2.py | 244 ++++++ .../test_replayssm_standard_decode_mamba2.py | 673 +++++++++++++++++ tests/kernels/mamba/utils.py | 160 ++++ tests/models/test_registry.py | 16 + .../test_mamba_update_block_table.py | 2 + .../test_replayssm_metadata_builder.py | 256 +++++++ tests/v1/e2e/test_replayssm_decode.py | 138 ++++ vllm/config/cache.py | 11 + vllm/config/model.py | 4 + vllm/config/vllm.py | 33 +- vllm/engine/arg_utils.py | 8 + .../layers/mamba/mamba_mixer2.py | 120 ++- .../layers/mamba/mamba_utils.py | 34 + .../layers/mamba/ops/replayssm_config.py | 66 ++ ...tive_state_update_replayssm_output_only.py | 709 ++++++++++++++++++ vllm/model_executor/models/interfaces.py | 25 + vllm/model_executor/models/nemotron_h.py | 31 +- vllm/model_executor/models/registry.py | 3 + vllm/v1/attention/backend.py | 6 + vllm/v1/attention/backends/mamba_attn.py | 144 ++++ vllm/v1/worker/gpu_input_batch.py | 26 + vllm/v1/worker/gpu_model_runner.py | 9 + 23 files changed, 2954 insertions(+), 31 deletions(-) create mode 100644 benchmarks/replayssm/e2e_decode_speedup.py create mode 100644 tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py create mode 100644 tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py create mode 100644 tests/v1/attention/test_replayssm_metadata_builder.py create mode 100644 tests/v1/e2e/test_replayssm_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 diff --git a/benchmarks/replayssm/e2e_decode_speedup.py b/benchmarks/replayssm/e2e_decode_speedup.py new file mode 100644 index 0000000000..b7084a2131 --- /dev/null +++ b/benchmarks/replayssm/e2e_decode_speedup.py @@ -0,0 +1,267 @@ +# 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 Mamba2 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. The two modes run in +separate subprocesses so each gets 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: + python e2e_decode_speedup.py --model-id nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 + python e2e_decode_speedup.py --dtype auto --buffer-len 16 \ + --model-id nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 # B300 NVFP4 +""" + +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." + ) + 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( + "--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( + "--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, + ) + 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) + + _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.accelerator.is_available(): + torch.accelerator.synchronize() + t0 = time.perf_counter() + outs = llm.generate(prompts, sp, use_tqdm=False) + if torch.accelerator.is_available(): + torch.accelerator.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), + "--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/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 0000000000..5319d8964d --- /dev/null +++ b/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py @@ -0,0 +1,244 @@ +# 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, + * the ReplaySSM output_only decode kernel, + +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.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( + *, + 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, + ) + 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, + ) + 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( + "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( + 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( + 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_standard_decode_mamba2.py b/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py new file mode 100644 index 0000000000..588ce52dbe --- /dev/null +++ b/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py @@ -0,0 +1,673 @@ +# 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, the ReplaySSM output_only kernel reproduces 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, +) +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.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( + *, + 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, + ) + 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, + ) + + # 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("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( + 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( + 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("precision", _PRECISIONS) +@pytest.mark.parametrize("geometry", _REAL_GEOMETRIES) +def test_replayssm_standard_decode_real_geometry( + 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( + 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( + "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( + 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( + 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("precision", _PRECISIONS) +@pytest.mark.parametrize("with_padding", [False, True]) +def test_replayssm_standard_decode_with_batch_indices( + 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, + ) + selective_state_update_replayssm_output_only( + state_cached, x, dt, A, B, C, bc_pre=bc_pre, **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]) + + +# 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/kernels/mamba/utils.py b/tests/kernels/mamba/utils.py index fb8a4b0a28..b82c4d82c3 100644 --- a/tests/kernels/mamba/utils.py +++ b/tests/kernels/mamba/utils.py @@ -76,3 +76,163 @@ def selective_state_update_ref( 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: + 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/models/test_registry.py b/tests/models/test_registry.py index 1418704036..45db16e879 100644 --- a/tests/models/test_registry.py +++ b/tests/models/test_registry.py @@ -134,6 +134,22 @@ def test_registry_is_pp(model_arch, is_pp, init_cuda): ) +@create_new_process_for_each_test() +@pytest.mark.parametrize( + "model_arch,supported", + [ + # ReplaySSM is opt-in per model; only Nemotron-H sets the flag today. + ("NemotronHForCausalLM", True), + ("Mamba2ForCausalLM", False), + ("Zamba2ForCausalLM", False), + ], +) +def test_registry_supports_replayssm(model_arch, supported): + model_info = ModelRegistry._try_inspect_model_cls(model_arch) + assert model_info is not None + assert model_info.supports_replayssm is supported + + def test_lazy_modelinfo_package_hash_includes_submodules(tmp_path): package_dir = tmp_path / "model_package" package_dir.mkdir() diff --git a/tests/v1/attention/test_mamba_update_block_table.py b/tests/v1/attention/test_mamba_update_block_table.py index 99dcb09ab1..4ec1382702 100644 --- a/tests/v1/attention/test_mamba_update_block_table.py +++ b/tests/v1/attention/test_mamba_update_block_table.py @@ -42,6 +42,8 @@ def _make_vllm_config( cache_config=SimpleNamespace( block_size=block_size, mamba_cache_mode="all", + use_replayssm=False, + replayssm_buffer_len=16, ), compilation_config=SimpleNamespace( cudagraph_mode=CUDAGraphMode.FULL, diff --git a/tests/v1/attention/test_replayssm_metadata_builder.py b/tests/v1/attention/test_replayssm_metadata_builder.py new file mode 100644 index 0000000000..cbbfedee96 --- /dev/null +++ b/tests/v1/attention/test_replayssm_metadata_builder.py @@ -0,0 +1,256 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Mamba2 ReplaySSM decode write-position derivation in +BaseMambaAttentionMetadataBuilder: write_pos and is_flush computed from the +per-request ring origin (replayssm_decode_base) and num_computed. +""" + +from dataclasses import dataclass + +import pytest +import torch + +from tests.v1.attention.utils import ( + BatchSpec, + MockMambaBuilder, + create_common_attn_metadata, + create_vllm_config, +) +from vllm.v1.kv_cache_interface import MambaSpec + +BLOCK_SIZE = 16 +DEVICE = torch.device("cpu") + + +@dataclass +class ReplaySSMBuildCase: + """A decode batch and its expected per-row write_pos / is_flush. + + num_computed = seq_len - query_len; write_pos = + (num_computed - decode_base) % buffer_len; is_flush = write_pos == + buffer_len - 1 (or a forced one-token flush when num_computed < decode_base). + """ + + seq_lens: list[int] + query_lens: list[int] + is_prefilling: list[bool] + decode_base: list[int] + buffer_len: int + expected_write_pos: list[int] + expected_is_flush: list[int] + mamba_cache_mode: str = "none" + + +REPLAYSSM_BUILD_CASES = { + # decode_base == num_prompt (fresh request). + "fresh_decode": ReplaySSMBuildCase( + seq_lens=[106], + query_lens=[1], + is_prefilling=[False], + decode_base=[100], + buffer_len=16, + expected_write_pos=[5], + expected_is_flush=[0], + ), + # decode_base > num_prompt anchors write_pos at the resume point. + "resumed_reanchors_to_zero": ReplaySSMBuildCase( + seq_lens=[106], + query_lens=[1], + is_prefilling=[False], + decode_base=[105], + buffer_len=16, + expected_write_pos=[0], + expected_is_flush=[0], + ), + # write_pos == buffer_len - 1 flushes. + "flush_boundary": ReplaySSMBuildCase( + seq_lens=[116], + query_lens=[1], + is_prefilling=[False], + decode_base=[100], + buffer_len=16, + expected_write_pos=[15], + expected_is_flush=[1], + ), + # Resumed request landing on a flush boundary. + "resumed_flush_boundary": ReplaySSMBuildCase( + seq_lens=[121], + query_lens=[1], + is_prefilling=[False], + decode_base=[105], + buffer_len=16, + expected_write_pos=[15], + expected_is_flush=[1], + ), + # Per-row write_pos / is_flush are independent. + "mixed_rows": ReplaySSMBuildCase( + seq_lens=[104, 106, 216], + query_lens=[1, 1, 1], + is_prefilling=[False, False, False], + decode_base=[100, 105, 200], + buffer_len=16, + expected_write_pos=[3, 0, 15], + expected_is_flush=[0, 0, 1], + ), + # write_pos wraps within the buffer (6 % 4 == 2). + "small_buffer_wrap": ReplaySSMBuildCase( + seq_lens=[112], + query_lens=[1], + is_prefilling=[False], + decode_base=[105], + buffer_len=4, + expected_write_pos=[2], + expected_is_flush=[0], + ), + # Single-token prefill-as-decode still in the prompt (num_computed < + # decode_base): forced one-token flush. + "leftover_prompt_one_token_flush": ReplaySSMBuildCase( + seq_lens=[100], + query_lens=[1], + is_prefilling=[True], + decode_base=[100], + buffer_len=16, + expected_write_pos=[0], + expected_is_flush=[1], + ), + # Align mode (block_size 16). Past the first boundary the ring re-anchors at + # the block start: num_computed 117 -> block_start 112, write_pos 5 (vs 1 in + # none mode). + "align_reanchor_past_boundary": ReplaySSMBuildCase( + seq_lens=[118], + query_lens=[1], + is_prefilling=[False], + decode_base=[100], + buffer_len=16, + expected_write_pos=[5], + expected_is_flush=[0], + mamba_cache_mode="align", + ), + # First-block boundary (num_computed+1 == 112) forces a flush even though + # write_pos (11) != buffer_len - 1. + "align_first_block_boundary_flush": ReplaySSMBuildCase( + seq_lens=[112], + query_lens=[1], + is_prefilling=[False], + decode_base=[100], + buffer_len=16, + expected_write_pos=[11], + expected_is_flush=[1], + mamba_cache_mode="align", + ), + # First step of a new block re-anchors write_pos to 0. + "align_new_block_start_zero": ReplaySSMBuildCase( + seq_lens=[113], + query_lens=[1], + is_prefilling=[False], + decode_base=[100], + buffer_len=16, + expected_write_pos=[0], + expected_is_flush=[0], + mamba_cache_mode="align", + ), + # block_size % buffer_len == 0: a later boundary lands on write_pos == + # buffer_len - 1, so the boundary flush coincides with the natural flush. + "align_boundary_coincides_natural_flush": ReplaySSMBuildCase( + seq_lens=[128], + query_lens=[1], + is_prefilling=[False], + decode_base=[100], + buffer_len=16, + expected_write_pos=[15], + expected_is_flush=[1], + mamba_cache_mode="align", + ), + # block_size % buffer_len != 0 (buffer_len 6): the boundary step still flushes + # although write_pos (3) != buffer_len - 1. + "align_unaligned_buffer_forces_flush": ReplaySSMBuildCase( + seq_lens=[128], + query_lens=[1], + is_prefilling=[False], + decode_base=[100], + buffer_len=6, + expected_write_pos=[3], + expected_is_flush=[1], + mamba_cache_mode="align", + ), + # Per-row independence in align mode: partial-block / new-block / boundary. + "align_mixed_rows": ReplaySSMBuildCase( + seq_lens=[105, 113, 112], + query_lens=[1, 1, 1], + is_prefilling=[False, False, False], + decode_base=[100, 100, 100], + buffer_len=16, + expected_write_pos=[4, 0, 11], + expected_is_flush=[0, 0, 1], + mamba_cache_mode="align", + ), +} + + +def _make_mamba_spec(buffer_len: int) -> MambaSpec: + # Five-tensor ReplaySSM page; the builder only reads shapes[4][0] (bc groups). + return MambaSpec( + block_size=BLOCK_SIZE, + shapes=( + (1, 1), + (1, 1, 1), + (1, buffer_len, 1), + (1, buffer_len), + (1, buffer_len, 1), + ), + dtypes=(torch.float32,), + ) + + +def _create_replayssm_builder( + buffer_len: int, mamba_cache_mode: str = "none" +) -> MockMambaBuilder: + vllm_config = create_vllm_config( + model_name="Qwen/Qwen3.5-0.8B", block_size=BLOCK_SIZE + ) + # Set the flags after construction to skip validate_mamba_cached_kernel + # (it requires a Triton backend) on the mock model. + vllm_config.cache_config.use_replayssm = True + vllm_config.cache_config.replayssm_buffer_len = buffer_len + vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode + return MockMambaBuilder( + _make_mamba_spec(buffer_len), ["layer0"], vllm_config, DEVICE + ) + + +def _build(builder: MockMambaBuilder, case: ReplaySSMBuildCase): + batch = BatchSpec(seq_lens=case.seq_lens, query_lens=case.query_lens) + common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE).replace( + is_prefilling=torch.tensor(case.is_prefilling, dtype=torch.bool), + replayssm_decode_base_cpu=torch.tensor(case.decode_base, dtype=torch.int32), + ) + return builder.build(0, common) + + +@pytest.mark.parametrize( + "case", REPLAYSSM_BUILD_CASES.values(), ids=REPLAYSSM_BUILD_CASES.keys() +) +def test_replayssm_write_pos(case: ReplaySSMBuildCase): + builder = _create_replayssm_builder(case.buffer_len, case.mamba_cache_mode) + meta = _build(builder, case) + + assert meta.write_pos_d is not None + assert meta.is_flush_d is not None + n = len(case.expected_write_pos) + assert meta.write_pos_d[:n].tolist() == case.expected_write_pos + assert meta.is_flush_d[:n].tolist() == case.expected_is_flush + + +def test_resumed_request_differs_from_fresh(): + """Same token count, different decode_base: fresh (base 100) -> write_pos 5, + resumed (base 105) -> write_pos 0.""" + builder = _create_replayssm_builder(16) + batch = BatchSpec(seq_lens=[106, 106], query_lens=[1, 1]) + common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE).replace( + is_prefilling=torch.tensor([False, False]), + replayssm_decode_base_cpu=torch.tensor([100, 105], dtype=torch.int32), + ) + meta = builder.build(0, common) + + assert meta.write_pos_d.tolist()[:2] == [5, 0] + assert meta.is_flush_d.tolist()[:2] == [0, 0] diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py new file mode 100644 index 0000000000..4fc6768e74 --- /dev/null +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -0,0 +1,138 @@ +# 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 vllm.v1.metrics.reader import Counter + +from ...models.utils import check_logprobs_close +from ...utils import large_gpu_mark, multi_gpu_test + +# Mamba2 (Nemotron-3) hybrid. +MAMBA2_MODEL = "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16" +MODELS = [ + pytest.param(MAMBA2_MODEL, marks=large_gpu_mark(min_gb=40)), +] + +PROMPTS = [ + "The capital of France is", + "Once upon a time, in a small village,", +] + + +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) + 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", + ) + + +@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]) +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) + + +# Prefix spans several mamba blocks; prefix caching only reuses full blocks. +_PC_SENTENCE = ( + "In a detailed survey of state space models, the authors compared many " + "architectures across a wide range of long-context language tasks and " + "measured their throughput, memory use, and accuracy in careful detail. " +) +_PC_PREFIX = _PC_SENTENCE * 120 +PREFIX_CACHING_PROMPTS = [ + _PC_PREFIX + "The most important conclusion was that", + _PC_PREFIX + "Surprisingly, the experiments showed that", + _PC_PREFIX + "The most important conclusion was that", +] + + +def _prefix_cache_hits(llm) -> int: + return sum( + m.value + for m in llm.llm.get_metrics() + if isinstance(m, Counter) and m.name == "vllm:prefix_cache_hits" + ) + + +def _check_replayssm_prefix_caching_parity( + vllm_runner, model_name, *, tensor_parallel_size=1 +): + # align mode materializes the exact SSM state at each block boundary, so + # ReplaySSM's cached prefixes must match the always-materialized baseline. + common = dict( + max_model_len=8192, + trust_remote_code=True, + enable_prefix_caching=True, + enable_chunked_prefill=True, + mamba_cache_mode="align", + disable_log_stats=False, # required for llm.get_metrics() + tensor_parallel_size=tensor_parallel_size, + ) + with vllm_runner(model_name, **common) as llm: + baseline = llm.generate_greedy_logprobs( + PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 + ) + with vllm_runner( + model_name, use_replayssm=True, replayssm_buffer_len=16, **common + ) as llm: + # Prime the cache, then measure, so cache hits are deterministic. + llm.generate_greedy_logprobs( + PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 + ) + replay = llm.generate_greedy_logprobs( + PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 + ) + replay_hits = _prefix_cache_hits(llm) + + # Without real cache hits the cached path is never exercised. + assert replay_hits > 0, ( + "ReplaySSM align-mode run produced no prefix-cache hits; the shared " + "prefix may be shorter than one mamba block, so prefix caching is inert" + ) + check_logprobs_close( + outputs_0_lst=baseline, + outputs_1_lst=replay, + name_0="baseline_align_pc", + name_1="replayssm_align_pc", + ) + + +@pytest.mark.parametrize("model_name", MODELS) +def test_replayssm_prefix_caching_matches_baseline(vllm_runner, model_name): + _check_replayssm_prefix_caching_parity(vllm_runner, model_name) + + +@multi_gpu_test(num_gpus=2) +@pytest.mark.parametrize("model_name", [MAMBA2_MODEL]) +def test_replayssm_prefix_caching_matches_baseline_tp2(vllm_runner, model_name): + _check_replayssm_prefix_caching_parity( + vllm_runner, model_name, tensor_parallel_size=2 + ) diff --git a/vllm/config/cache.py b/vllm/config/cache.py index a628e7d7cd..45c9362418 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -145,6 +145,17 @@ 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 buffer length B: with use_replayssm, standard decode + caches recent SSM inputs in a size-B ring buffer and flushes the checkpoint + state to HBM every B steps. Default 16.""" + use_replayssm: bool = False + """Use the ReplaySSM Mamba2 decode kernel: cache recent SSM inputs and skip + the per-step full-state store, writing the checkpoint back only on flush. + Requires mamba_cache_mode 'none' or 'align' (prefix caching) and the Triton + mamba backend; standard (non-speculative) decode only. In align mode flushes + are most efficient when mamba_block_size is a multiple of replayssm_buffer_len, + but this is not required.""" # Will be set after profiling. num_gpu_blocks: int | None = field(default=None, init=False) diff --git a/vllm/config/model.py b/vllm/config/model.py index d9d2f57dc4..0e7621c2b1 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1696,6 +1696,10 @@ class ModelConfig: def supports_mamba_prefix_caching(self) -> bool: return self._model_info.supports_mamba_prefix_caching + @property + def supports_replayssm(self) -> bool: + return self._model_info.supports_replayssm + @property def use_mla(self) -> bool: return self.is_deepseek_mla and not envs.VLLM_MLA_DISABLE diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index c05e35cbbc..a7f185cad4 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -39,7 +39,7 @@ from .kv_events import KVEventsConfig 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 @@ -2308,6 +2308,37 @@ class VllmConfig: ) return self + @model_validator(mode="after") + def validate_mamba_cached_kernel(self) -> "VllmConfig": + if not self.cache_config.use_replayssm: + return self + # ReplaySSM adds a 3-tensor ring to the mamba state; only models that + # opt in (supports_replayssm) build a consistent shape on both the layer + # and config paths. Reject others so the mamba page size cannot desync. + if self.model_config is not None and not self.model_config.supports_replayssm: + raise ValueError( + "--use-replayssm is only supported for Nemotron-H models " + f"(got architecture {self.model_config.architecture!r})" + ) + if self.cache_config.mamba_cache_mode == "all": + raise ValueError( + "--use-replayssm supports prefix caching only in align mode; " + "pass --mamba-cache-mode align" + ) + 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.kv_transfer_config is not None + and self.kv_transfer_config.is_kv_transfer_instance + ): + raise ValueError( + "--use-replayssm is incompatible with KV connectors " + "(P/D disaggregation, KV cache offload)" + ) + 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 29ebb988e9..a14de27190 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -695,6 +695,8 @@ class EngineArgs: mamba_block_size: int | None = get_field(CacheConfig, "mamba_block_size") prefix_match_unit: int | None = get_field(CacheConfig, "prefix_match_unit") mamba_cache_mode: MambaCacheMode = CacheConfig.mamba_cache_mode + replayssm_buffer_len: int = CacheConfig.replayssm_buffer_len + use_replayssm: bool = CacheConfig.use_replayssm mamba_backend: MambaBackendEnum = MambaBackendEnum.TRITON enable_mamba_cache_stochastic_rounding: bool = ( @@ -1201,6 +1203,10 @@ class EngineArgs: 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( "--kv-offloading-size", **cache_kwargs["kv_offloading_size"] ) @@ -1910,6 +1916,8 @@ class EngineArgs: mamba_block_size=self.mamba_block_size, prefix_match_unit=self.prefix_match_unit, mamba_cache_mode=self.mamba_cache_mode, + replayssm_buffer_len=self.replayssm_buffer_len, + use_replayssm=self.use_replayssm, kv_offloading_size=self.kv_offloading_size, kv_offloading_backend=self.kv_offloading_backend, ) diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index a6524961ea..7538a2b6b4 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -32,6 +32,9 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( causal_conv1d_update, ) from vllm.model_executor.layers.mamba.ops.layernorm_gated import rms_norm_gated +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.ssd_combined import ( mamba_chunk_scan_combined_varlen, ) @@ -494,12 +497,27 @@ class MambaMixer2(MambaBase, PluggableLayer): 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_replayssm = ( + cache_config.use_replayssm if cache_config is not None else False + ) + self.replayssm_buffer_len = ( + cache_config.replayssm_buffer_len + if cache_config is not None and cache_config.use_replayssm + else None + ) + self.mamba_config = vllm_config.mamba_config + if self.use_replayssm and self.num_heads % self.tp_size != 0: + raise ValueError( + "--use-replayssm requires tensor-parallel heads to divide evenly" + ) + # The tuple is (conv_state, ssm_state); with the cached (ReplaySSM) decode + # kernel enabled it is (conv_state, ssm_state, x_cache, dt_cache, B_cache). + _n_state = 5 if self.use_replayssm else 2 + self.kv_cache = tuple(torch.tensor([]) for _ in range(_n_state)) self.num_spec = vllm_config.num_speculative_tokens if self.num_spec > 0: @@ -591,7 +609,7 @@ class MambaMixer2(MambaBase, PluggableLayer): # 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 +720,10 @@ class MambaMixer2(MambaBase, PluggableLayer): else self.kv_cache[0].transpose(-1, -2) ) ssm_state = self.kv_cache[1] + if self.use_replayssm: + 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 @@ -1027,37 +1049,78 @@ class MambaMixer2(MambaBase, PluggableLayer): # - 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_replayssm: + assert self.replayssm_buffer_len is not None + 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.replayssm_buffer_len, + state_batch_indices=state_indices_tensor_d_input, + out=preallocated_ssm_out_d, + # Stochastic Rounding for the vanilla decode path is read + # from mamba_config inside ssu_dispatch; the replay kernel + # isn't a dispatch backend, so pass it here. + enable_stochastic_rounding=( + self.mamba_config.enable_stochastic_rounding + ), + cache_philox_rounds=( + self.mamba_config.stochastic_rounding_philox_rounds + ), + ) + 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( + base_dtype = MambaStateDtypeCalculator.mamba2_state_dtype( self.model_config.dtype, self.cache_config.mamba_cache_dtype, self.cache_config.mamba_ssm_cache_dtype, ) + if self.use_replayssm: + return MambaStateDtypeCalculator.append_replayssm_ring( + base_dtype, self.model_config.dtype + ) + return base_dtype - def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]: - return MambaStateShapeCalculator.mamba2_state_shape( + def get_state_shape(self) -> tuple[tuple[int, ...], ...]: + tp_world_size = get_tensor_model_parallel_world_size() + base_shape = MambaStateShapeCalculator.mamba2_state_shape( intermediate_size=self.intermediate_size, - tp_world_size=get_tensor_model_parallel_world_size(), + tp_world_size=tp_world_size, n_groups=self.n_groups, num_heads=self.num_heads, head_dim=self.head_dim, @@ -1065,6 +1128,15 @@ class MambaMixer2(MambaBase, PluggableLayer): conv_kernel=self.conv_kernel_size, num_spec=self.num_spec, ) + if self.use_replayssm: + assert self.replayssm_buffer_len is not None + return MambaStateShapeCalculator.append_replayssm_ring( + base_shape, + self.n_groups, + tp_world_size, + self.replayssm_buffer_len, + ) + return base_shape @property def mamba_type(self) -> MambaAttentionBackendEnum: diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 9e78b82228..d974da5c27 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -80,6 +80,18 @@ class MambaStateDtypeCalculator: model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype ) + @classmethod + def append_replayssm_ring( + cls, + base_dtypes: tuple[torch.dtype, ...], + model_dtype: ModelDType | torch.dtype, + ) -> tuple[torch.dtype, ...]: + """Append the ReplaySSM ring dtypes to a base ``(conv, ssm)`` tuple: + ``(x_cache, dt_cache, B_cache)`` = ``(activation, fp32, activation)``. + """ + activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype) + return (*base_dtypes, activation_dtype, torch.float32, activation_dtype) + @classmethod def _mamba_state_dtype( cls, @@ -186,6 +198,28 @@ class MambaStateShapeCalculator: temporal_state_shape = (divide(num_heads, tp_world_size), head_dim, state_size) return conv_state_shape, temporal_state_shape + @classmethod + def append_replayssm_ring( + cls, + base_shapes: tuple[tuple[int, ...], ...], + n_groups: int, + tp_world_size: int, + replayssm_buffer_len: int, + ) -> tuple[tuple[int, ...], ...]: + """Append the ReplaySSM ring shapes (x_cache, dt_cache, B_cache) to a + base ``(conv, ssm)`` tuple. ``base_shapes[1]`` is the ssm shape + ``(nheads // tp, head_dim, state_size)``; B_cache uses the un-extended + ``n_groups``. + """ + local_nheads, head_dim, state_size = base_shapes[1] + local_ngroups = divide(n_groups, tp_world_size) + return ( + *base_shapes, + (local_nheads, replayssm_buffer_len, head_dim), + (local_nheads, replayssm_buffer_len), + (local_ngroups, replayssm_buffer_len, state_size), + ) + @classmethod def short_conv_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 0000000000..aff4fc1c3f --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/replayssm_config.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Launch-config selection for the ReplaySSM Mamba2 output_only decode kernel. + +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. +""" + +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_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 get_replayssm_config(kernel: str, **shape) -> tuple: + """Return the launch config for ``kernel`` (override > tuned default). + + kernel: "mamba2_output_only". ``shape`` carries the keying dims (dstate; + ``L`` for the buffer length, default 16); hardware is auto-detected. + """ + if kernel in _overrides: + return _overrides[kernel] + bw = _is_blackwell() + if kernel == "mamba2_output_only": + return _mamba2_output_only(shape["dstate"], shape.get("L", 16), 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 0000000000..c6046b06a5 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py @@ -0,0 +1,709 @@ +# 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 convert_rs_fp16x2, 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, + rand_seed_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, + NF_DSTATE_TILE: tl.constexpr, + NF_NDS: tl.constexpr, + FL_DSTATE_TILE: tl.constexpr, + FL_NDS: tl.constexpr, + USE_RS_ROUNDING: tl.constexpr, + PHILOX_ROUNDS: 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_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) + 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, 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) + 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) + + 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 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 + + # 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_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 -- 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) + 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 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 + ) + 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 + ) + if USE_RS_ROUNDING: + # Stochastic-round fp32->fp16 (Blackwell cvt.rs), mirroring the + # baseline. Only the flush step stores state, so this runs at 1/L + # the baseline's per-step rate. Absolute per-element offsets seed + # the RNG so each state element draws independently. + rand_seed = tl.load(rand_seed_ptr) + rand_offsets = ( + state_batch_idx * stride_state_batch + + pid_h * stride_state_head + + offs_m[:, None] * stride_state_dim + + offs_n_f[None, :] * stride_state_dstate + ) + if PHILOX_ROUNDS > 0: + rand = tl.randint(rand_seed, rand_offsets, PHILOX_ROUNDS) + else: + rand = tl.randint(rand_seed, rand_offsets) + tl.static_assert( + state_ptrs.dtype.element_ty == tl.float16, + "stochastic rounding requires an fp16 SSM state cache", + ) + state_store = convert_rs_fp16x2(state_new, rand) + else: + state_store = state_new.to(st_f.dtype) + tl.store(state_ptrs, state_store, 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: + 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, + enable_stochastic_rounding: bool = False, + cache_philox_rounds: int = 0, +) -> 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, 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) + state_indices_strides = ( + (state_batch_indices.stride(0), state_batch_indices.stride(1)) + if state_batch_indices is not None + else (0, 0) + ) + rand_seed = ( + torch.randint(0, 2**32, (1,), device=state.device) + if enable_stochastic_rounding + else None + ) + + with torch.accelerator.device_index(x.device.index): + # Both kernels always launch: the precompute kernel self-skips flush + # rows per row (the branch can't be hoisted out under CUDA graphs), so + # it is a no-op when every row is flushing. + _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, + rand_seed, + 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, + nf_dstate_tile, + nf_nds, + fl_dstate_tile, + fl_nds, + enable_stochastic_rounding, + cache_philox_rounds, + num_warps=num_warps, + num_stages=num_stages, + ) + + if not has_heads: + out = out.squeeze(1) + return out diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 32534f8fbb..df9c18ce53 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -984,6 +984,31 @@ def supports_mamba_prefix_caching( return getattr(model, "supports_mamba_prefix_caching", False) +@runtime_checkable +class SupportsReplaySSM(Protocol): + """The interface for models whose Mamba2 layers support ReplaySSM cached + standard decode. + + This is currently experimental. + """ + + supports_replayssm: ClassVar[Literal[True]] = True + + +@overload +def supports_replayssm(model: object) -> TypeIs[SupportsReplaySSM]: ... + + +@overload +def supports_replayssm(model: type[object]) -> TypeIs[type[SupportsReplaySSM]]: ... + + +def supports_replayssm( + model: type[object] | object, +) -> TypeIs[type[SupportsReplaySSM]] | TypeIs[SupportsReplaySSM]: + return getattr(model, "supports_replayssm", False) + + @runtime_checkable class SupportsCrossEncoding(Protocol): """The interface required for all models that support cross encoding.""" diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index 84f24094a3..50dc5c0222 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -69,6 +69,7 @@ from vllm.model_executor.models.interfaces import ( SupportsMambaPrefixCaching, SupportsPP, SupportsQuant, + SupportsReplaySSM, ) from vllm.model_executor.models.utils import ( AutoWeightsLoader, @@ -707,6 +708,7 @@ class NemotronHForCausalLM( SupportsQuant, MixtureOfExperts, SupportsMambaPrefixCaching, + SupportsReplaySSM, ): # Relevant only if self.has_moe is True is_non_gated_moe: bool = True @@ -742,18 +744,24 @@ 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 + base_dtype = MambaStateDtypeCalculator.mamba2_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, ) + if cache_config.use_replayssm: + return MambaStateDtypeCalculator.append_replayssm_ring( + base_dtype, vllm_config.model_config.dtype + ) + return base_dtype @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: @@ -763,12 +771,15 @@ class NemotronHForCausalLM( Tuple containing: - conv_state_shape: Shape for convolutional state cache - temporal_state_shape: Shape for state space model cache + - (when use_replayssm 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( + base_shape = MambaStateShapeCalculator.mamba2_state_shape( intermediate_size=intermediate_size, tp_world_size=parallel_config.tensor_parallel_size, n_groups=hf_config.n_groups, @@ -778,6 +789,14 @@ class NemotronHForCausalLM( conv_kernel=hf_config.conv_kernel, num_spec=vllm_config.num_speculative_tokens, ) + if cache_config.use_replayssm: + return MambaStateShapeCalculator.append_replayssm_ring( + base_shape, + hf_config.n_groups, + parallel_config.tensor_parallel_size, + cache_config.replayssm_buffer_len, + ) + return base_shape @classmethod def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index b17a41834d..a4fd44fc2c 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -55,6 +55,7 @@ from .interfaces import ( supports_multimodal_encoder_tp_data, supports_multimodal_raw_input_only, supports_pp, + supports_replayssm, supports_transcription, ) from .interfaces_base import ( @@ -789,6 +790,7 @@ class _ModelInfo: is_hybrid: bool has_noops: bool supports_mamba_prefix_caching: bool + supports_replayssm: bool supports_transcription: bool supports_transcription_only: bool @@ -815,6 +817,7 @@ class _ModelInfo: is_attention_free=is_attention_free(model), is_hybrid=is_hybrid(model), supports_mamba_prefix_caching=supports_mamba_prefix_caching(model), + supports_replayssm=supports_replayssm(model), supports_transcription=supports_transcription(model), supports_transcription_only=( supports_transcription(model) and model.supports_transcription_only diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 6f23cff0b1..b479aab581 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -480,6 +480,11 @@ class CommonAttentionMetadata: fixed sliding window. None disables R-SWA. The attention backend copies this into its own persistent buffer and reads ``rswa_window`` from model config.""" + replayssm_decode_base_cpu: torch.Tensor | None = None + """(batch_size,) CPU ring origin for Mamba2 ReplaySSM decode: num_computed + at the current decode run's last full-state write. write_pos counts from + here, so a preemption-resumed request re-anchors past the prompt boundary.""" + # WARNING: Deprecated fields. Will be removed in a future release (v0.15.0) _seq_lens_cpu: torch.Tensor | None = None _num_computed_tokens_cpu: torch.Tensor | None = None @@ -591,6 +596,7 @@ class CommonAttentionMetadata: 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), + replayssm_decode_base_cpu=maybe_slice_reqs(self.replayssm_decode_base_cpu), ) diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index 6fce7d0dc7..28928f2894 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -74,6 +74,12 @@ class BaseMambaAttentionMetadata: nums_dict: dict | None = None batch_ptr: torch.Tensor | None = None token_chunk_offset_ptr: torch.Tensor | None = None + # ReplaySSM standard decode: per-row ring cursor and flush flag, plus the + # per-step (decode_rows, ngroups, replayssm_buffer_len) fp32 scratch for the + # precomputed k^T q products. All None when use_replayssm is disabled. + write_pos_d: torch.Tensor | None = None + is_flush_d: torch.Tensor | None = None + bc_pre_scratch: torch.Tensor | None = None class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): @@ -98,6 +104,8 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): 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_replayssm = vllm_config.cache_config.use_replayssm + self.replayssm_buffer_len = vllm_config.cache_config.replayssm_buffer_len assert isinstance(kv_cache_spec, MambaSpec) scheduler_config = vllm_config.scheduler_config @@ -158,6 +166,36 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): dtype=torch.int32, device=device, ) + # ReplaySSM standard-decode CUDA-graph buffers: per-row ring cursor, + # flush flag, and the k^T q precompute scratch. + if self.use_replayssm: + 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, + ) + # B_cache shape = (ngroups, replayssm_buffer_len, dstate); the page + # layout is (conv_state, ssm_state, 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( + ( + bc_scratch_bs, + bc_ngroups, + self.replayssm_buffer_len, + ), + dtype=torch.float32, + device=device, + ) + else: + self.decode_bc_pre_scratch = None self._init_reorder_batch_threshold(1, self.use_spec_decode) if self.use_spec_decode: @@ -414,6 +452,8 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): has_prior_state = seq_lens_cpu > 1 prefill_to_decode = single_token_prefill_rows & has_prior_state if torch.any(prefill_to_decode).item(): + # 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( @@ -444,6 +484,8 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): # 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() @@ -530,6 +572,79 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): num_reqs - num_prefills : num_reqs ] + if self.use_replayssm and num_decodes > 0: + decode_base_cpu = common_attn_metadata.replayssm_decode_base_cpu + num_computed_tokens_cpu = common_attn_metadata._num_computed_tokens_cpu + if decode_base_cpu is None or num_computed_tokens_cpu is None: + raise ValueError( + "--use-replayssm requires CPU decode-base and " + "computed-token counts to derive decode write positions" + ) + num_computed_d = num_computed_tokens_cpu[:num_decodes] + decode_base_d = decode_base_cpu[:num_decodes] + align_mode = self.vllm_config.cache_config.mamba_cache_mode == "align" + block_size = self.kv_cache_spec.block_size + if align_mode: + # After a boundary the align copy leaves an exact checkpoint at + # the block start and the new block's ring restarts empty, so + # re-anchor there; max() keeps the prompt-end anchor for the + # first (partial) block. + effective_base = torch.maximum( + decode_base_d, (num_computed_d // block_size) * block_size + ) + else: + effective_base = decode_base_d + # write_pos counts decode steps since the ring's last full-state + # write (the anchor), so a resumed request re-anchors correctly. + decode_steps_cpu = num_computed_d - effective_base + 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 + # 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 & ~leftover_prompt, + decode_steps_cpu, + torch.zeros_like(decode_steps_cpu), + ) + write_pos_cpu = torch.remainder(decode_steps_cpu, self.replayssm_buffer_len) + is_flush_cpu = ( + write_pos_cpu == self.replayssm_buffer_len - 1 + ) | leftover_prompt + if align_mode: + # Force a flush on the step completing a mamba block so the exact + # boundary state is materialized for prefix caching. + is_flush_cpu = is_flush_cpu | ( + valid_decode_rows + & ((num_computed_d + query_lens_cpu) % block_size == 0) + ) + is_flush_cpu = is_flush_cpu.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_replayssm + and self.decode_bc_pre_scratch is not None + and num_decodes > 0 + ): + bc_pre_scratch = self.decode_bc_pre_scratch[:num_decodes] + metadata = self.metadata_cls( num_prefills=num_prefills, num_prefill_tokens=num_prefill_tokens, @@ -539,6 +654,9 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): 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, num_accepted_tokens=num_accepted_tokens, query_start_loc_d=query_start_loc_d, block_idx_last_scheduled_token=block_idx_last_scheduled_token, @@ -573,6 +691,9 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): 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 if ( metadata.num_prefills == 0 and metadata.num_decodes <= self.decode_cudagraph_max_bs @@ -634,11 +755,34 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): ) block_idx_last_scheduled_token_prev_step[metadata.num_decodes :] = 0 + if self.use_replayssm: + 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.decode_bc_pre_scratch is not None: + bc_pre_scratch = self.decode_bc_pre_scratch[: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, 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_input_batch.py b/vllm/v1/worker/gpu_input_batch.py index baa502ee81..0b805fca76 100644 --- a/vllm/v1/worker/gpu_input_batch.py +++ b/vllm/v1/worker/gpu_input_batch.py @@ -106,6 +106,7 @@ class InputBatch: is_pooling_model: bool = False, cp_kv_cache_interleave_size: int = 1, reasoning_config: ReasoningConfig | None = None, + use_replayssm: bool = False, slot_mapping_modes: list[SlotMappingMode] | None = None, ): self.thinking_budget_state_holder = maybe_create_thinking_budget_state_holder( @@ -170,6 +171,17 @@ class InputBatch: ) self.num_computed_tokens_cpu = self.num_computed_tokens_cpu_tensor.numpy() + # Mamba2 ReplaySSM decode ring origin (num_computed at each request's + # last full-state write); populated only when the feature is on. + self.use_replayssm = use_replayssm + self.replayssm_decode_base_cpu_tensor = torch.zeros( + (max_num_reqs,), + device="cpu", + dtype=torch.int32, + pin_memory=PIN_MEMORY, + ) + self.replayssm_decode_base = self.replayssm_decode_base_cpu_tensor.numpy() + # Block table. self.block_table = MultiGroupBlockTable( max_num_reqs=max_num_reqs, @@ -377,6 +389,11 @@ class InputBatch: # Number of tokens without spec decode tokens. self.num_tokens_no_spec[req_index] = request.num_tokens + if self.use_replayssm: + # Ring origin = full context at (re)admission (prompt + any resumed + # output), so a resumed request re-anchors past the prompt. + self.replayssm_decode_base[req_index] = request.num_tokens + self.num_computed_tokens_cpu[req_index] = request.num_computed_tokens self.block_table.add_row(request.block_ids, req_index) @@ -597,6 +614,11 @@ class InputBatch: self.num_prompt_tokens[i2], self.num_prompt_tokens[i1], ) + if self.use_replayssm: + self.replayssm_decode_base[i1], self.replayssm_decode_base[i2] = ( + self.replayssm_decode_base[i2], + self.replayssm_decode_base[i1], + ) self.num_computed_tokens_cpu[i1], self.num_computed_tokens_cpu[i2] = ( self.num_computed_tokens_cpu[i2], self.num_computed_tokens_cpu[i1], @@ -754,6 +776,10 @@ class InputBatch: last_req_index ] self.num_prompt_tokens[empty_index] = self.num_prompt_tokens[last_req_index] + if self.use_replayssm: + self.replayssm_decode_base[empty_index] = self.replayssm_decode_base[ + last_req_index + ] self.num_computed_tokens_cpu[empty_index] = self.num_computed_tokens_cpu[ last_req_index ] diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 95221c8693..dd31ddab67 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -727,6 +727,7 @@ class GPUModelRunner( is_pooling_model=self.is_pooling_model, cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size, reasoning_config=self.vllm_config.reasoning_config, + use_replayssm=self.cache_config.use_replayssm, ) # Separate cuda stream for overlapping transfer of sampled token ids from @@ -2413,6 +2414,12 @@ class GPUModelRunner( if self.model_config.rswa_window is not None: rswa_prefix_lens = num_prompt_tokens_cpu + replayssm_decode_base_cpu = None + if self.cache_config.use_replayssm: + replayssm_decode_base_cpu = ( + self.input_batch.replayssm_decode_base_cpu_tensor[:num_reqs_padded] + ) + cm_base = CommonAttentionMetadata( query_start_loc=self.query_start_loc.gpu[: num_reqs_padded + 1], query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs_padded + 1], @@ -2420,6 +2427,7 @@ class GPUModelRunner( _seq_lens_cpu=seq_lens_cpu, _num_computed_tokens_cpu=num_computed_tokens_cpu, seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, + replayssm_decode_base_cpu=replayssm_decode_base_cpu, num_reqs=num_reqs_padded, num_actual_tokens=num_tokens_padded, max_query_len=max_query_len, @@ -7262,6 +7270,7 @@ class GPUModelRunner( is_pooling_model=self.is_pooling_model, cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size, reasoning_config=self.vllm_config.reasoning_config, + use_replayssm=self.cache_config.use_replayssm, slot_mapping_modes=slot_mapping_modes, )