iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Kernel] ReplaySSM: cache SSM inputs for faster Mamba2 standard decode (#48018)

Browse Source 
Signed-off-by: Johnny-Liou <a897111@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
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 <tlleungac@connect.ust.hk>
This commit is contained in:
Johnny-Liou
2026-07-24 09:39:49 -07:00
committed by GitHub
co-authored by Claude Opus 4.8 mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> tomeras91 Cyrus Leung
parent d02df748bf
commit 866fea2b99
23 changed files with 2954 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmarks/replayssm/e2e_decode_speedup.py
+267
View File
@@ -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()
tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py
+244
View File
@@ -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,
)
tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py
+673
View File
@@ -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)
tests/kernels/mamba/utils.py
+160
View File
@@ -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
tests/models/test_registry.py
+16
View File
@@ -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()
tests/v1/attention/test_mamba_update_block_table.py
+2
View File
@@ -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,
tests/v1/attention/test_replayssm_metadata_builder.py
+256
View File
@@ -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]
tests/v1/e2e/test_replayssm_decode.py
+138
View File
@@ -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
)
vllm/config/cache.py
+11
View File
@@ -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)
vllm/config/model.py
+4
View File
@@ -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
vllm/config/vllm.py
+32 -1
View File
@@ -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
vllm/engine/arg_utils.py
+8
View File
@@ -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,
)
vllm/model_executor/layers/mamba/mamba_mixer2.py
+96 -24
View File
@@ -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:
vllm/model_executor/layers/mamba/mamba_utils.py
+34
View File
@@ -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,
vllm/model_executor/layers/mamba/ops/replayssm_config.py
+66
View File
@@ -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}")
vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py
+709
View File
@@ -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
vllm/model_executor/models/interfaces.py
+25
View File
@@ -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."""
vllm/model_executor/models/nemotron_h.py
+25 -6
View File
@@ -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]:
vllm/model_executor/models/registry.py
+3
View File
@@ -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
vllm/v1/attention/backend.py
+6
View File
@@ -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),
)
vllm/v1/attention/backends/mamba_attn.py
+144
View File
@@ -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=(
vllm/v1/worker/gpu_input_batch.py
+26
View File
@@ -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
]
vllm/v1/worker/gpu_model_runner.py
+9
View File
@@ -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,
)