| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """End-to-end autoregressive decode benchmark: ReplaySSM vs the standard SSM kernel. |
| 4 | |
| 5 | Loads a hybrid Mamba2 model, replicates one prompt across the batch, and times a |
| 6 | long greedy decode (CUDA graphs on) once with the standard kernel and once with |
| 7 | ReplaySSM, then reports the per-step / throughput speedup. The two modes run in |
| 8 | separate subprocesses so each gets a clean CUDA context. |
| 9 | |
| 10 | The FlashInfer FP4-MoE autotuner is disabled by default (it is unstable under |
| 11 | CUDA-graph capture on the pre-release Blackwell FP4 path); pass |
| 12 | --no-disable-flashinfer-autotune for non-FP4 models. |
| 13 | |
| 14 | Examples: |
| 15 | python e2e_decode_speedup.py --model-id nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 |
| 16 | python e2e_decode_speedup.py --dtype auto --buffer-len 16 \ |
| 17 | --model-id nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 # B300 NVFP4 |
| 18 | |
| 19 | """ |
| 20 | |
| 21 | import argparse |
| 22 | import json |
| 23 | import os |
| 24 | import subprocess |
| 25 | import sys |
| 26 | import time |
| 27 | |
| 28 | DEFAULT_PROMPT = "My cat wrote all this CUDA code for a new language model and" |
| 29 | |
| 30 | MODE_LABEL = {"standard": "standard", "replayssm": "ReplaySSM"} |
| 31 | |
| 32 | |
| 33 | def parse_args(): |
| 34 | p = argparse.ArgumentParser( |
| 35 | description="E2E decode speedup: ReplaySSM vs the standard SSM kernel." |
| 36 | ) |
| 37 | p.add_argument("--model-id", default="nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16") |
| 38 | p.add_argument("--prompt", default=DEFAULT_PROMPT) |
| 39 | p.add_argument("--batch-size", type=int, default=256) |
| 40 | p.add_argument("--num-steps", type=int, default=1000) |
| 41 | p.add_argument("--warmup-steps", type=int, default=128) |
| 42 | p.add_argument("--repeats", type=int, default=1) |
| 43 | p.add_argument( |
| 44 | "--buffer-len", type=int, default=16, help="ReplaySSM input-buffer length." |
| 45 | ) |
| 46 | p.add_argument( |
| 47 | "--dtype", |
| 48 | default="bfloat16", |
| 49 | choices=["bfloat16", "float16", "float32", "auto"], |
| 50 | ) |
| 51 | p.add_argument("--gpu-memory-utilization", type=float, default=0.9) |
| 52 | p.add_argument("--max-model-len", type=int, default=None) |
| 53 | p.add_argument( |
| 54 | "--disable-flashinfer-autotune", |
| 55 | action=argparse.BooleanOptionalAction, |
| 56 | default=True, |
| 57 | help="Disable the FlashInfer FP4-MoE autotuner (default: on). " |
| 58 | "It is unstable under CUDA-graph capture on the " |
| 59 | "pre-release Blackwell FP4 path; pass " |
| 60 | "--no-disable-flashinfer-autotune for non-FP4 models.", |
| 61 | ) |
| 62 | p.add_argument( |
| 63 | "--mamba-ssm-cache-dtype", |
| 64 | default="auto", |
| 65 | choices=["auto", "float32", "float16", "bfloat16"], |
| 66 | help="SSM state dtype (both modes). 'auto' = config-driven; " |
| 67 | "'float32' = fp32 state, 'bfloat16' = s16 state.", |
| 68 | ) |
| 69 | p.add_argument( |
| 70 | "--baseline-ssm-config", |
| 71 | default="", |
| 72 | help="Pin the STANDARD baseline's SSM launch config as " |
| 73 | "'bsm,nw' via override_ssm_config (forces the in-process " |
| 74 | "engine so the override reaches the kernel). Empty = off.", |
| 75 | ) |
| 76 | p.add_argument( |
| 77 | "--worker", |
| 78 | choices=["standard", "replayssm"], |
| 79 | default=None, |
| 80 | help=argparse.SUPPRESS, |
| 81 | ) |
| 82 | return p.parse_args() |
| 83 | |
| 84 | |
| 85 | def resolve_max_model_len(args) -> int: |
| 86 | if args.max_model_len is not None: |
| 87 | return args.max_model_len |
| 88 | return args.num_steps + 256 |
| 89 | |
| 90 | |
| 91 | def run_worker(args): |
| 92 | # override_ssm_config is a module global; it only reaches the model if the |
| 93 | # engine runs in-process (default V1 spawns a separate EngineCore). Force it. |
| 94 | if args.worker == "standard" and args.baseline_ssm_config: |
| 95 | os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" |
| 96 | |
| 97 | import torch |
| 98 | |
| 99 | from vllm import LLM, SamplingParams |
| 100 | |
| 101 | mode = args.worker |
| 102 | max_model_len = resolve_max_model_len(args) |
| 103 | |
| 104 | llm_kwargs = dict( |
| 105 | model=args.model_id, |
| 106 | tensor_parallel_size=1, |
| 107 | dtype=args.dtype, |
| 108 | max_model_len=max_model_len, |
| 109 | trust_remote_code=True, |
| 110 | enable_prefix_caching=False, |
| 111 | enable_chunked_prefill=False, |
| 112 | max_num_seqs=args.batch_size, |
| 113 | max_num_batched_tokens=max(max_model_len, args.batch_size * 64), |
| 114 | enforce_eager=False, |
| 115 | disable_log_stats=True, |
| 116 | gpu_memory_utilization=args.gpu_memory_utilization, |
| 117 | # SSM state dtype (applies to both standard and ReplaySSM). |
| 118 | mamba_ssm_cache_dtype=args.mamba_ssm_cache_dtype, |
| 119 | ) |
| 120 | if args.disable_flashinfer_autotune: |
| 121 | # FP4-MoE autotuner is unstable under CUDA-graph capture on Blackwell; |
| 122 | # re-enable (--no-disable-flashinfer-autotune) only for non-FP4 models. |
| 123 | llm_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False} |
| 124 | if mode == "replayssm": |
| 125 | llm_kwargs.update(use_replayssm=True, replayssm_buffer_len=args.buffer_len) |
| 126 | |
| 127 | _ssm_cm = None |
| 128 | if mode == "standard" and args.baseline_ssm_config: |
| 129 | from vllm.model_executor.layers.mamba.ops.mamba_ssm import override_ssm_config |
| 130 | |
| 131 | _bsm, _nw = (int(x) for x in args.baseline_ssm_config.split(",")) |
| 132 | _ssm_cm = override_ssm_config((_bsm, _nw)) |
| 133 | _ssm_cm.__enter__() # active through LLM() graph capture + decode |
| 134 | print( |
| 135 | f"[{mode}] override_ssm_config -> (BLOCK_SIZE_M={_bsm}, num_warps={_nw})", |
| 136 | flush=True, |
| 137 | ) |
| 138 | |
| 139 | llm = LLM(**llm_kwargs) |
| 140 | prompts = [args.prompt] * args.batch_size |
| 141 | |
| 142 | def timed_generate(n_tokens): |
| 143 | sp = SamplingParams( |
| 144 | n=1, |
| 145 | temperature=0.0, |
| 146 | ignore_eos=True, |
| 147 | min_tokens=n_tokens, |
| 148 | max_tokens=n_tokens, |
| 149 | ) |
| 150 | if torch.accelerator.is_available(): |
| 151 | torch.accelerator.synchronize() |
| 152 | t0 = time.perf_counter() |
| 153 | outs = llm.generate(prompts, sp, use_tqdm=False) |
| 154 | if torch.accelerator.is_available(): |
| 155 | torch.accelerator.synchronize() |
| 156 | elapsed = time.perf_counter() - t0 |
| 157 | produced = min(len(o.outputs[0].token_ids) for o in outs) |
| 158 | assert produced == n_tokens, f"expected {n_tokens} tokens, got {produced}" |
| 159 | return elapsed |
| 160 | |
| 161 | timed_generate(args.warmup_steps) |
| 162 | |
| 163 | best = None |
| 164 | for _ in range(args.repeats): |
| 165 | elapsed = timed_generate(args.num_steps) |
| 166 | tok_s = args.batch_size * args.num_steps / elapsed |
| 167 | per_step_ms = elapsed / args.num_steps * 1e3 |
| 168 | print( |
| 169 | f"[{mode}] {elapsed:.3f}s {tok_s:,.0f} tok/s {per_step_ms:.3f} ms/step", |
| 170 | flush=True, |
| 171 | ) |
| 172 | if best is None or elapsed < best["elapsed_s"]: |
| 173 | best = { |
| 174 | "mode": mode, |
| 175 | "elapsed_s": elapsed, |
| 176 | "tok_s": tok_s, |
| 177 | "per_step_ms": per_step_ms, |
| 178 | } |
| 179 | |
| 180 | print("RESULT_JSON " + json.dumps(best), flush=True) |
| 181 | if _ssm_cm is not None: |
| 182 | _ssm_cm.__exit__(None, None, None) |
| 183 | |
| 184 | |
| 185 | def run_one_mode(args, mode) -> dict: |
| 186 | cmd = [ |
| 187 | sys.executable, |
| 188 | __file__, |
| 189 | "--worker", |
| 190 | mode, |
| 191 | "--model-id", |
| 192 | args.model_id, |
| 193 | "--prompt", |
| 194 | args.prompt, |
| 195 | "--batch-size", |
| 196 | str(args.batch_size), |
| 197 | "--num-steps", |
| 198 | str(args.num_steps), |
| 199 | "--warmup-steps", |
| 200 | str(args.warmup_steps), |
| 201 | "--repeats", |
| 202 | str(args.repeats), |
| 203 | "--buffer-len", |
| 204 | str(args.buffer_len), |
| 205 | "--dtype", |
| 206 | args.dtype, |
| 207 | "--gpu-memory-utilization", |
| 208 | str(args.gpu_memory_utilization), |
| 209 | "--mamba-ssm-cache-dtype", |
| 210 | args.mamba_ssm_cache_dtype, |
| 211 | "--baseline-ssm-config", |
| 212 | args.baseline_ssm_config, |
| 213 | ] |
| 214 | cmd.append( |
| 215 | "--disable-flashinfer-autotune" |
| 216 | if args.disable_flashinfer_autotune |
| 217 | else "--no-disable-flashinfer-autotune" |
| 218 | ) |
| 219 | if args.max_model_len is not None: |
| 220 | cmd += ["--max-model-len", str(args.max_model_len)] |
| 221 | |
| 222 | result = None |
| 223 | proc = subprocess.Popen( |
| 224 | cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1 |
| 225 | ) |
| 226 | for line in proc.stdout: |
| 227 | sys.stdout.write(line) |
| 228 | sys.stdout.flush() |
| 229 | if line.startswith("RESULT_JSON "): |
| 230 | result = json.loads(line[len("RESULT_JSON ") :]) |
| 231 | proc.wait() |
| 232 | if proc.returncode != 0: |
| 233 | raise RuntimeError(f"mode '{mode}' worker exited with {proc.returncode}") |
| 234 | if result is None: |
| 235 | raise RuntimeError(f"mode '{mode}' produced no RESULT_JSON line") |
| 236 | return result |
| 237 | |
| 238 | |
| 239 | def main(): |
| 240 | args = parse_args() |
| 241 | if args.worker is not None: |
| 242 | run_worker(args) |
| 243 | return |
| 244 | |
| 245 | print( |
| 246 | f"model={args.model_id} batch_size={args.batch_size} " |
| 247 | f"steps={args.num_steps} buffer_len={args.buffer_len} dtype={args.dtype}" |
| 248 | ) |
| 249 | |
| 250 | std = run_one_mode(args, "standard") |
| 251 | fla = run_one_mode(args, "replayssm") |
| 252 | speedup = std["per_step_ms"] / fla["per_step_ms"] |
| 253 | |
| 254 | print() |
| 255 | header = f"{'mode':<10}{'ms/step':>12}{'tok/s':>16}{'wall (s)':>12}" |
| 256 | print(header) |
| 257 | print("-" * len(header)) |
| 258 | for r in (std, fla): |
| 259 | print( |
| 260 | f"{MODE_LABEL[r['mode']]:<10}{r['per_step_ms']:>12.3f}" |
| 261 | f"{r['tok_s']:>16,.0f}{r['elapsed_s']:>12.3f}" |
| 262 | ) |
| 263 | print("-" * len(header)) |
| 264 | print(f"speedup (standard / ReplaySSM, per step): {speedup:.3f}x") |
| 265 | |
| 266 | |
| 267 | if __name__ == "__main__": |
| 268 | main() |