@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d56285c747
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 2, 22:00 UTC

Tuning script and configs for Triton Mamba SSU kernel (#43083)

Signed-off-by: Banani Ghosh <bg2502@nyu.edu> Signed-off-by: Daniel Serebrenik <daserebrenik@nvidia.com> Co-authored-by: Banani Ghosh <bg2502@nyu.edu>

danisereb committed May 24, 2026 at 20:12 UTC d56285c747ec3133163c740adc3c68f3f5feb4e7
13 files changed +1770 -91
benchmarks/kernels/benchmark_selective_state_update.py new
+774
@@ -0,0 +1,774 @@
1 +#!/usr/bin/env python3
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +"""
5 +Benchmark and tuning script for the Mamba selective_state_update kernel.
6 +
7 +Mirrors the fused MoE tuning workflow: sweeps (BLOCK_SIZE_M, num_warps) across
8 +an effective_batch grid for a given (headdim, dstate, ngroups, cache_dtype) and
9 +saves the best config per effective_batch to JSON. Generated configs are picked
10 +up by selective_state_update at runtime.
11 +
12 +Usage:
13 + python -m benchmarks.kernels.benchmark_selective_state_update \
14 + --all-dstates --save-configs --compare
15 +"""
16 +
17 +import argparse
18 +import json
19 +import os
20 +import sys
21 +from io import StringIO
22 +from itertools import product
23 +from typing import Any
24 +
25 +import torch
26 +
27 +from tests.kernels.mamba.utils import selective_state_update_ref
28 +from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
29 + _CONFIGS_DIR,
30 + _canonical_cache_dtype,
31 + _get_default_ssm_launch_config,
32 + get_ssm_config_file_name,
33 + get_ssm_device_name,
34 + override_ssm_config,
35 + selective_state_update,
36 +)
37 +from vllm.triton_utils import triton
38 +
39 +# bf16 shares configs with fp16 - same bit width.
40 +_SSM_CACHE_DTYPE_MAP: dict[str, torch.dtype] = {
41 + "float32": torch.float32,
42 + "float16": torch.float16,
43 + "bfloat16": torch.float16,
44 +}
45 +
46 +_RESULTS_DIR = os.path.dirname(os.path.realpath(__file__))
47 +
48 +# ---------------------------------------------------------------------------
49 +# Tuning search space
50 +# ---------------------------------------------------------------------------
51 +
52 +_BSM_CHOICES_ALL = [4, 8, 16, 32, 64, 128, 256]
53 +
54 +NUM_WARPS_CHOICES = [1, 2, 4, 8]
55 +
56 +
57 +def _block_size_m_choices(headdim: int) -> list[int]:
58 + """BLOCK_SIZE_M candidates worth sweeping for a given headdim.
59 +
60 + BLOCK_SIZE_M > next_pow2(headdim) wastes >=50% of each tile via masking
61 + (offs_m >= dim rows are zeroed out), so we cap the sweep there.
62 + """
63 + ceiling = 1
64 + while ceiling < headdim:
65 + ceiling <<= 1
66 + return [b for b in _BSM_CHOICES_ALL if b <= ceiling]
67 +
68 +
69 +# Default deployment shapes. effective_batch = batch * nheads scales the
70 +# kernel grid, so configs transfer across (model, TP) combos sharing
71 +# (headdim, dstate, cache_dtype).
72 +DEFAULT_BATCH_SIZES = [1, 8, 16, 32, 64, 128, 256, 512, 1024, 1536, 2048]
73 +DEFAULT_NHEADS = [128, 256]
74 +
75 +ALL_DSTATES = [16, 32, 64, 128, 256]
76 +
77 +# Default tuning shape — matches Nemotron-3-Super and Nemotron-3-Nano Mamba layers.
78 +# Override with CLI flags for other architectures.
79 +DEFAULT_HEADDIM = 64
80 +DEFAULT_NGROUPS = 8
81 +
82 +
83 +# ---------------------------------------------------------------------------
84 +# Benchmark helper
85 +# ---------------------------------------------------------------------------
86 +
87 +
88 +def _make_inputs(
89 + batch: int,
90 + nheads: int,
91 + dim: int,
92 + dstate: int,
93 + ngroups: int,
94 + dtype: torch.dtype,
95 + state_dtype: torch.dtype | None = None,
96 + device: str = "cuda",
97 +):
98 + if state_dtype is None:
99 + state_dtype = dtype
100 + state = torch.randn(batch, nheads, dim, dstate, dtype=state_dtype, device=device)
101 + x = torch.randn(batch, nheads, dim, dtype=dtype, device=device)
102 + dt = torch.randn(batch, nheads, dim, dtype=dtype, device=device)
103 + A = -torch.rand(nheads, dim, dstate, dtype=torch.float32, device=device)
104 + B = torch.randn(batch, ngroups, dstate, dtype=dtype, device=device)
105 + C = torch.randn(batch, ngroups, dstate, dtype=dtype, device=device)
106 + D = torch.randn(nheads, dim, dtype=dtype, device=device)
107 + dt_bias = torch.randn(nheads, dim, dtype=dtype, device=device)
108 + out = torch.zeros(batch, nheads, dim, dtype=dtype, device=device)
109 + return state, x, dt, A, B, C, D, dt_bias, out
110 +
111 +
112 +def benchmark_config(
113 + batch: int,
114 + nheads: int,
115 + dim: int,
116 + dstate: int,
117 + ngroups: int,
118 + block_size_m: int,
119 + num_warps_val: int,
120 + dtype: torch.dtype,
121 + state_dtype: torch.dtype | None = None,
122 + num_iters: int = 100,
123 + num_warmup: int = 20,
124 + graph_batch_size: int = 10,
125 +) -> float | None:
126 + """
127 + Time one (BLOCK_SIZE_M, num_warps) config for selective_state_update.
128 + Returns elapsed time in microseconds, or None on error.
129 +
130 + Uses CUDA graph capture-and-replay to isolate kernel time from Python
131 + eager-mode dispatch / kwarg-resolution overhead, mirroring the timing
132 + methodology in benchmarks/kernels/benchmark_moe.py.
133 + """
134 + state, x, dt, A, B, C, D, dt_bias, out = _make_inputs(
135 + batch, nheads, dim, dstate, ngroups, dtype, state_dtype=state_dtype
136 + )
137 +
138 + def _call_kernel() -> None:
139 + selective_state_update(
140 + state,
141 + x,
142 + dt,
143 + A,
144 + B,
145 + C,
146 + D=D,
147 + z=None,
148 + dt_bias=dt_bias,
149 + dt_softplus=True,
150 + out=out,
151 + )
152 +
153 + try:
154 + with override_ssm_config((block_size_m, num_warps_val)):
155 + # Eager-mode warmup: triggers Triton autotune / JIT, primes caches.
156 + for _ in range(num_warmup):
157 + _call_kernel()
158 + torch.accelerator.synchronize()
159 +
160 + # Capture graph_batch_size invocations into a CUDA graph so the
161 + # timed region runs without Python dispatch overhead per call.
162 + graph = torch.cuda.CUDAGraph()
163 + with torch.cuda.graph(graph):
164 + for _ in range(graph_batch_size):
165 + _call_kernel()
166 + torch.accelerator.synchronize()
167 +
168 + # Warmup graph replays (let the runtime stabilize).
169 + for _ in range(5):
170 + graph.replay()
171 + torch.accelerator.synchronize()
172 +
173 + start = torch.cuda.Event(enable_timing=True)
174 + end = torch.cuda.Event(enable_timing=True)
175 + latencies: list[float] = []
176 + for _ in range(num_iters):
177 + start.record()
178 + graph.replay()
179 + end.record()
180 + end.synchronize()
181 + latencies.append(start.elapsed_time(end))
182 + graph.reset()
183 + # elapsed_time returns ms; each replay runs graph_batch_size kernels,
184 + # so divide by (num_iters * graph_batch_size) and convert ms -> us.
185 + return sum(latencies) / (num_iters * graph_batch_size) * 1000
186 + except Exception as e:
187 + if "OutOfResources" not in str(e):
188 + print(
189 + f" Warning: config M={block_size_m},w={num_warps_val} "
190 + f"raised {type(e).__name__}: {e}"
191 + )
192 + return None
193 +
194 +
195 +# ---------------------------------------------------------------------------
196 +# Tuning loop
197 +# ---------------------------------------------------------------------------
198 +
199 +
200 +# CUDA grid Y/Z dim limit — both `batch` and `nheads` must fit individually.
201 +_CUDA_MAX_GRID_DIM = 65535
202 +
203 +# Above this, kernel state-offset arithmetic (batch * nheads * headdim * dstate)
204 +# overflows int32 and the launch raises cudaErrorIllegalAddress.
205 +# 262144 covers Nemotron Super TP1 BS=2048.
206 +_MAX_EFFECTIVE_BATCH = 262144
207 +
208 +
209 +def expand_batch_x_nheads(
210 + batch_sizes: list[int],
211 + nheads_list: list[int],
212 + ngroups: int,
213 +) -> list[tuple[int, int, int]]:
214 + """Cross-product batch_sizes × nheads_list → sorted [(effective_batch,
215 + batch, nheads)], deduped by effective_batch. Filters pairs that exceed
216 + the CUDA grid dim limit, the effective_batch ceiling, or where nheads is
217 + not a positive multiple of ngroups.
218 + """
219 + seen: dict[int, tuple[int, int]] = {}
220 + skipped_grid: list[tuple[int, int]] = []
221 + skipped_ngroups: list[tuple[int, int]] = []
222 + skipped_eb: list[tuple[int, int]] = []
223 + for b, n in product(batch_sizes, nheads_list):
224 + if b <= 0 or n <= 0:
225 + continue
226 + if b > _CUDA_MAX_GRID_DIM or n > _CUDA_MAX_GRID_DIM:
227 + skipped_grid.append((b, n))
228 + continue
229 + if n % ngroups != 0:
230 + skipped_ngroups.append((b, n))
231 + continue
232 + if b * n > _MAX_EFFECTIVE_BATCH:
233 + skipped_eb.append((b, n))
234 + continue
235 + seen.setdefault(b * n, (b, n))
236 + if skipped_grid:
237 + print(
238 + f" Note: skipping (batch, nheads) pairs exceeding CUDA grid dim "
239 + f"{_CUDA_MAX_GRID_DIM}: {skipped_grid}"
240 + )
241 + if skipped_ngroups:
242 + print(
243 + f" Note: skipping (batch, nheads) pairs where nheads % ngroups != 0 "
244 + f"for ngroups={ngroups}: {skipped_ngroups}"
245 + )
246 + if skipped_eb:
247 + print(
248 + f" Note: skipping (batch, nheads) pairs whose effective_batch "
249 + f"exceeds {_MAX_EFFECTIVE_BATCH}: {skipped_eb}"
250 + )
251 + return sorted((eb, b, n) for eb, (b, n) in seen.items())
252 +
253 +
254 +def tune_dstate(
255 + dstate: int,
256 + headdim: int,
257 + ngroups: int,
258 + dtype: torch.dtype,
259 + num_iters: int,
260 + verbose: bool,
261 + active: list[tuple[int, int, int]],
262 + state_dtype: torch.dtype | None = None,
263 +) -> tuple[dict[int, dict], dict[int, dict[tuple[int, int], float]]]:
264 + """For each (effective_batch, batch, nheads) in *active*, sweep
265 + (BLOCK_SIZE_M, num_warps) and return
266 + ({effective_batch: best_config}, {effective_batch: {(bsm, nw): us}}).
267 + The second map is the full timing grid, used downstream so we don't
268 + re-measure the same config in the comparison phase.
269 + """
270 + best_per_eb: dict[int, dict] = {}
271 + timings: dict[int, dict[tuple[int, int], float]] = {}
272 +
273 + print(f"\n{'=' * 74}")
274 + effective_state_dtype = state_dtype if state_dtype is not None else dtype
275 + print(
276 + f"Tuning headdim={headdim} dstate={dstate} ngroups={ngroups} "
277 + f"dtype={dtype} ssm_cache_dtype={effective_state_dtype}"
278 + )
279 + print(f"{'=' * 74}")
280 +
281 + bsm_choices = _block_size_m_choices(headdim)
282 + print(f"BSM candidates (capped at next_pow2(headdim={headdim})): {bsm_choices}")
283 +
284 + hdr = f"{'EffBatch':>8} | {'BLOCK_M':>7} | {'warps':>5} | {'us':>10} | note"
285 + print(hdr)
286 + print("-" * 52)
287 +
288 + for eb, batch, nheads in active:
289 + best_time = float("inf")
290 + best_cfg: dict = {}
291 + eb_timings: dict[tuple[int, int], float] = {}
292 +
293 + for bsm, nw in product(bsm_choices, NUM_WARPS_CHOICES):
294 + t = benchmark_config(
295 + batch=batch,
296 + nheads=nheads,
297 + dim=headdim,
298 + dstate=dstate,
299 + ngroups=ngroups,
300 + block_size_m=bsm,
301 + num_warps_val=nw,
302 + dtype=dtype,
303 + state_dtype=state_dtype,
304 + num_iters=num_iters,
305 + )
306 + if t is None:
307 + continue
308 + eb_timings[(bsm, nw)] = t
309 + is_best = t < best_time
310 + if is_best:
311 + best_time = t
312 + best_cfg = {"BLOCK_SIZE_M": bsm, "num_warps": nw}
313 + if verbose:
314 + marker = " <-- best" if is_best else ""
315 + print(f"{eb:>8} | {bsm:>7} | {nw:>5} | {t:>10.2f} |{marker}")
316 +
317 + timings[eb] = eb_timings
318 +
319 + if not best_cfg:
320 + print(
321 + f"{eb:>8} | {'-':>7} | {'-':>5} | {'-':>10} | "
322 + f"no working config (skipped)"
323 + )
324 + continue
325 +
326 + if not verbose:
327 + print(
328 + f"{eb:>8} | {best_cfg['BLOCK_SIZE_M']:>7} | "
329 + f"{best_cfg['num_warps']:>5} | {best_time:>10.2f} | best"
330 + )
331 +
332 + best_per_eb[eb] = best_cfg
333 +
334 + return best_per_eb, timings
335 +
336 +
337 +# ---------------------------------------------------------------------------
338 +# Correctness validation
339 +# ---------------------------------------------------------------------------
340 +
341 +
342 +def validate_configs(
343 + dstate: int,
344 + headdim: int,
345 + ngroups: int,
346 + tuned: dict[int, dict],
347 + active: list[tuple[int, int, int]],
348 + dtype: torch.dtype,
349 + atol: float = 1e-2,
350 + rtol: float = 1e-2,
351 + state_dtype: torch.dtype | None = None,
352 +) -> dict[int, bool]:
353 + """
354 + For every (effective_batch, batch, nheads) in *active* that has a tuned
355 + config, run the kernel with that config and compare against the reference.
356 + Returns {effective_batch: passed}.
357 + """
358 + # Disable TF32 so the reference's matmul matches the Triton kernel's
359 + # fp32 accumulation; otherwise large ebs show bf16 rounding mismatches.
360 + torch.set_float32_matmul_precision("highest")
361 +
362 + print(f"\n{'=' * 74}")
363 + effective_state_dtype = state_dtype if state_dtype is not None else dtype
364 + print(
365 + f"Validation headdim={headdim} dstate={dstate} ngroups={ngroups} "
366 + f"dtype={dtype} ssm_cache_dtype={effective_state_dtype} atol={atol}"
367 + )
368 + print(f"{'=' * 74}")
369 + print(f"{'EffBatch':>8} | {'MaxAbsErr':>12} | {'Status':>8}")
370 + print("-" * 36)
371 +
372 + results: dict[int, bool] = {}
373 +
374 + for eb, batch, nheads in active:
375 + cfg = tuned.get(eb)
376 + if cfg is None:
377 + continue
378 + state, x, dt, A, B, C, D, dt_bias, out = _make_inputs(
379 + batch=batch,
380 + nheads=nheads,
381 + dim=headdim,
382 + dstate=dstate,
383 + ngroups=ngroups,
384 + dtype=dtype,
385 + state_dtype=state_dtype,
386 + )
387 + # Clone state before GPU kernel modifies it in-place
388 + state_ref = state.clone()
389 +
390 + with override_ssm_config((cfg["BLOCK_SIZE_M"], cfg["num_warps"])):
391 + selective_state_update(
392 + state,
393 + x,
394 + dt,
395 + A,
396 + B,
397 + C,
398 + D=D,
399 + z=None,
400 + dt_bias=dt_bias,
401 + dt_softplus=True,
402 + out=out,
403 + )
404 + torch.accelerator.synchronize()
405 + gpu_out = out.detach().cpu()
406 +
407 + # Reference uses the original (unmodified) state
408 + # Upcast to fp32 so the reference sums in fp32 (matches the Triton
409 + # kernel); summing in bf16 over `dstate` blows up the error.
410 + ref_out = (
411 + selective_state_update_ref(
412 + state_ref.float(),
413 + x.float(),
414 + dt.float(),
415 + A.float(),
416 + B.float(),
417 + C.float(),
418 + D=D.float(),
419 + dt_bias=dt_bias.float(),
420 + dt_softplus=True,
421 + )
422 + .to(out.dtype)
423 + .cpu()
424 + )
425 +
426 + passed = torch.allclose(gpu_out.float(), ref_out.float(), atol=atol, rtol=rtol)
427 + max_err = (gpu_out.float() - ref_out.float()).abs().max().item()
428 + status = "PASS" if passed else "FAIL"
429 + results[eb] = passed
430 + print(f"{eb:>8} | {max_err:>12.6f} | {status:>8}")
431 +
432 + n_pass = sum(results.values())
433 + n_total = len(results)
434 + print(f"\n {n_pass}/{n_total} configs passed validation for dstate={dstate}")
435 + return results
436 +
437 +
438 +# ---------------------------------------------------------------------------
439 +# Save configs
440 +# ---------------------------------------------------------------------------
441 +
442 +
443 +def save_configs(
444 + headdim: int,
445 + dstate: int,
446 + cache_dtype: str,
447 + configs: dict[int, dict],
448 + save_dir: str | None = None,
449 +) -> str:
450 + # bf16 shares configs with fp16, use common filename for both
451 + cache_dtype = _canonical_cache_dtype(cache_dtype)
452 +
453 + base_dir = save_dir if save_dir else _CONFIGS_DIR
454 + os.makedirs(base_dir, exist_ok=True)
455 + file_path = os.path.join(
456 + base_dir,
457 + get_ssm_config_file_name(headdim, dstate, cache_dtype, get_ssm_device_name()),
458 + )
459 + # triton_version is informational only, the loader ignores it
460 + payload: dict[str, Any] = {
461 + "triton_version": triton.__version__,
462 + **{str(k): v for k, v in sorted(configs.items())},
463 + }
464 + with open(file_path, "w") as f:
465 + json.dump(payload, f, indent=4)
466 + return file_path
467 +
468 +
469 +# ---------------------------------------------------------------------------
470 +# Comparison table
471 +# ---------------------------------------------------------------------------
472 +
473 +
474 +def current_heuristic(dstate: int, is_blackwell: bool = False) -> dict:
475 + """Return the current hard-coded BLOCK_SIZE_M / num_warps for dstate."""
476 + bsm, nw = _get_default_ssm_launch_config(dstate, is_blackwell)
477 + return {"BLOCK_SIZE_M": bsm, "num_warps": nw}
478 +
479 +
480 +def compare_heuristic_vs_tuned(
481 + dstate: int,
482 + headdim: int,
483 + ngroups: int,
484 + tuned: dict[int, dict],
485 + timings: dict[int, dict[tuple[int, int], float]],
486 + active: list[tuple[int, int, int]],
487 + dtype: torch.dtype,
488 + num_iters: int,
489 + is_blackwell: bool,
490 + state_dtype: torch.dtype | None = None,
491 +):
492 + heur_cfg = current_heuristic(dstate, is_blackwell)
493 + heur_key = (heur_cfg["BLOCK_SIZE_M"], heur_cfg["num_warps"])
494 +
495 + print(f"\n{'=' * 74}")
496 + print(
497 + f"Comparison headdim={headdim} dstate={dstate} "
498 + f"ngroups={ngroups} — heuristic vs tuned"
499 + )
500 + print(
501 + f"Heuristic: BLOCK_SIZE_M={heur_cfg['BLOCK_SIZE_M']}, "
502 + f"num_warps={heur_cfg['num_warps']}"
503 + )
504 + print(f"{'=' * 74}")
505 + hdr = (
506 + f"{'EffBatch':>8} | {'Heur(us)':>10} | {'Tuned(us)':>10} | "
507 + f"{'Speedup':>8} | Best config"
508 + )
509 + print(hdr)
510 + print("-" * len(hdr))
511 +
512 + for eb, batch, nheads in active:
513 + eb_timings = timings.get(eb, {})
514 +
515 + # Heuristic timing: reuse the tuning measurement if the heuristic
516 + # config was in the swept grid; otherwise measure it once.
517 + t_h = eb_timings.get(heur_key)
518 + if t_h is None:
519 + t_h = benchmark_config(
520 + batch=batch,
521 + nheads=nheads,
522 + dim=headdim,
523 + dstate=dstate,
524 + ngroups=ngroups,
525 + block_size_m=heur_cfg["BLOCK_SIZE_M"],
526 + num_warps_val=heur_cfg["num_warps"],
527 + dtype=dtype,
528 + state_dtype=state_dtype,
529 + num_iters=num_iters,
530 + )
531 +
532 + # `tuned[eb]` may be missing if all configs failed in tune_dstate;
533 + # in that case fall back to the heuristic so the table still prints.
534 + best = tuned.get(eb) or heur_cfg
535 + t_t = eb_timings.get((best["BLOCK_SIZE_M"], best["num_warps"]))
536 +
537 + if t_h is None or t_t is None:
538 + print(f"{eb:>8} | {'N/A':>10} | {'N/A':>10} | {'N/A':>8} |")
539 + continue
540 + speedup = t_h / t_t
541 + marker = " <--" if speedup > 1.05 else ""
542 + print(
543 + f"{eb:>8} | {t_h:>10.2f} | {t_t:>10.2f} | "
544 + f"{speedup:>7.2f}x | "
545 + f"M={best['BLOCK_SIZE_M']},w={best['num_warps']}{marker}"
546 + )
547 +
548 +
549 +# ---------------------------------------------------------------------------
550 +# CLI
551 +# ---------------------------------------------------------------------------
552 +
553 +
554 +def save_results(device_name: str, output: str, results_file: str | None = None) -> str:
555 + """Save the full benchmark output to a results text file."""
556 + if results_file is None:
557 + results_file = os.path.join(
558 + _RESULTS_DIR, f"ssm_benchmark_results_{device_name}.txt"
559 + )
560 + with open(results_file, "w") as f:
561 + f.write(output)
562 + return results_file
563 +
564 +
565 +def main():
566 + parser = argparse.ArgumentParser(
567 + description="Tune selective_state_update kernel for Mamba SSM"
568 + )
569 + parser.add_argument(
570 + "--dstate",
571 + type=int,
572 + default=128,
573 + help="SSM state size to tune for (default: 128)",
574 + )
575 + parser.add_argument(
576 + "--all-dstates",
577 + action="store_true",
578 + help="Tune all common dstate values: " + str(ALL_DSTATES),
579 + )
580 + parser.add_argument(
581 + "--dtype",
582 + type=str,
583 + default="bfloat16",
584 + choices=["float16", "bfloat16"],
585 + help="Activation / input data type (default: bfloat16)",
586 + )
587 + parser.add_argument(
588 + "--mamba-ssm-cache-dtype",
589 + type=str,
590 + default="float32",
591 + choices=list(_SSM_CACHE_DTYPE_MAP.keys()),
592 + help="SSM state cache dtype (default: float32)",
593 + )
594 + parser.add_argument(
595 + "--num-iters",
596 + type=int,
597 + default=100,
598 + help="Number of timing iterations (default: 100)",
599 + )
600 + parser.add_argument(
601 + "--save-configs",
602 + action="store_true",
603 + help=f"Save best configs to JSON in {_CONFIGS_DIR}",
604 + )
605 + parser.add_argument(
606 + "--compare",
607 + action="store_true",
608 + help="Show comparison table: heuristic vs tuned",
609 + )
610 + parser.add_argument(
611 + "--verbose",
612 + action="store_true",
613 + help="Print every (BLOCK_SIZE_M, num_warps) result, not just best",
614 + )
615 + parser.add_argument(
616 + "--results-file",
617 + type=str,
618 + default=None,
619 + help="Path to save the benchmark results text file "
620 + "(default: ssm_benchmark_results_<device>.txt alongside this script)",
621 + )
622 + parser.add_argument(
623 + "--save-dir",
624 + type=str,
625 + default=None,
626 + help=f"Directory to save JSON configs (default: {_CONFIGS_DIR})",
627 + )
628 + parser.add_argument(
629 + "--headdim",
630 + type=int,
631 + default=DEFAULT_HEADDIM,
632 + help=f"Per-head feature dim (default: {DEFAULT_HEADDIM})",
633 + )
634 + parser.add_argument(
635 + "--ngroups",
636 + type=int,
637 + default=DEFAULT_NGROUPS,
638 + help=f"Number of B/C groups (default: {DEFAULT_NGROUPS})",
639 + )
640 + parser.add_argument(
641 + "--batch-sizes",
642 + type=int,
643 + nargs="+",
644 + default=DEFAULT_BATCH_SIZES,
645 + metavar="B",
646 + help=f"Decoder batch sizes to sweep (default: {DEFAULT_BATCH_SIZES})",
647 + )
648 + parser.add_argument(
649 + "--nheads",
650 + type=int,
651 + nargs="+",
652 + default=DEFAULT_NHEADS,
653 + metavar="N",
654 + help=f"Number of heads per rank to sweep (default: {DEFAULT_NHEADS}). "
655 + "effective_batch = batch * nheads; cross-product is deduped by eb.",
656 + )
657 + parser.add_argument(
658 + "--validate",
659 + action="store_true",
660 + help="After tuning, verify each best config against a CPU reference "
661 + "implementation. Configs that fail are flagged in the output.",
662 + )
663 + parser.add_argument(
664 + "--atol",
665 + type=float,
666 + default=1e-2,
667 + help="Absolute tolerance for --validate (default: 1e-2)",
668 + )
669 + args = parser.parse_args()
670 +
671 + dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
672 + state_dtype = _SSM_CACHE_DTYPE_MAP[args.mamba_ssm_cache_dtype]
673 + device_name = get_ssm_device_name()
674 + cap = torch.cuda.get_device_capability()
675 + is_blackwell = cap[0] >= 10
676 +
677 + # Mirror all output to a results file (like Unix tee).
678 + buf = StringIO()
679 +
680 + class _Tee:
681 + """Writes to both the original stdout and an in-memory buffer."""
682 +
683 + def write(self, s):
684 + buf.write(s)
685 + sys.__stdout__.write(s)
686 +
687 + def flush(self):
688 + sys.__stdout__.flush()
689 +
690 + sys.stdout = _Tee() # type: ignore[assignment]
691 +
692 + try:
693 + print(f"Device : {device_name} (sm_{cap[0]}{cap[1]})")
694 + print(f"Blackwell: {is_blackwell}")
695 + print(f"dtype : {args.dtype}")
696 + print(f"ssm_cache_dtype: {args.mamba_ssm_cache_dtype}")
697 + print(f"headdim: {args.headdim}")
698 + print(f"ngroups: {args.ngroups}")
699 + print(f"triton : {triton.__version__}")
700 +
701 + dstates = ALL_DSTATES if args.all_dstates else [args.dstate]
702 + active = expand_batch_x_nheads(args.batch_sizes, args.nheads, args.ngroups)
703 +
704 + for dstate in dstates:
705 + tuned, timings = tune_dstate(
706 + dstate=dstate,
707 + headdim=args.headdim,
708 + ngroups=args.ngroups,
709 + dtype=dtype,
710 + num_iters=args.num_iters,
711 + verbose=args.verbose,
712 + active=active,
713 + state_dtype=state_dtype,
714 + )
715 +
716 + if args.compare:
717 + compare_heuristic_vs_tuned(
718 + dstate=dstate,
719 + headdim=args.headdim,
720 + ngroups=args.ngroups,
721 + tuned=tuned,
722 + timings=timings,
723 + active=active,
724 + dtype=dtype,
725 + num_iters=args.num_iters,
726 + is_blackwell=is_blackwell,
727 + state_dtype=state_dtype,
728 + )
729 +
730 + if args.validate:
731 + validity = validate_configs(
732 + dstate=dstate,
733 + headdim=args.headdim,
734 + ngroups=args.ngroups,
735 + tuned=tuned,
736 + active=active,
737 + dtype=dtype,
738 + atol=args.atol,
739 + state_dtype=state_dtype,
740 + )
741 + # Filter out any configs that failed correctness check
742 + failed = [eb for eb, ok in validity.items() if not ok]
743 + if failed:
744 + print(
745 + f"\n WARNING: {len(failed)} config(s) failed validation "
746 + f"for dstate={dstate}: effective_batches {failed}"
747 + )
748 + print(" These will NOT be saved even with --save-configs.")
749 + tuned = {
750 + eb: cfg for eb, cfg in tuned.items() if validity.get(eb, True)
751 + }
752 +
753 + if args.save_configs:
754 + path = save_configs(
755 + headdim=args.headdim,
756 + dstate=dstate,
757 + cache_dtype=args.mamba_ssm_cache_dtype,
758 + configs=tuned,
759 + save_dir=args.save_dir,
760 + )
761 + print(f"\nSaved: {path}")
762 + else:
763 + print(f"\nBest configs for dstate={dstate}:")
764 + for eb, cfg in sorted(tuned.items()):
765 + print(f" effective_batch={eb:>6}: {cfg}")
766 + print("\n(Re-run with --save-configs to persist to JSON)")
767 + finally:
768 + sys.stdout = sys.__stdout__
769 + results_path = save_results(device_name, buf.getvalue(), args.results_file)
770 + print(f"\nResults saved to: {results_path}")
771 +
772 +
773 +if __name__ == "__main__":
774 + main()
tests/kernels/mamba/__init__.py
tests/kernels/mamba/test_mamba_ssm.py
+1 -72
@@ -6,6 +6,7 @@ import torch
6 import torch.nn.functional as F
7 from einops import rearrange, repeat
8
9 +from tests.kernels.mamba.utils import selective_state_update_ref
10 from tests.kernels.utils import opcheck
11 from vllm import _custom_ops as ops # noqa: F401
12 from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
@@ -17,78 +18,6 @@ from vllm.utils.torch_utils import set_random_seed
18 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
19
20
20 -def selective_state_update_ref(
21 - state, x, dt, A, B, C, D=None, z=None, dt_bias=None, dt_softplus=False
22 -):
23 - """
24 - Argument:
25 - state: (batch, dim, dstate) or (batch, nheads, dim, dstate)
26 - x: (batch, dim) or (batch, nheads, dim)
27 - dt: (batch, dim) or (batch, nheads, dim)
28 - A: (dim, dstate) or (nheads, dim, dstate)
29 - B: (batch, dstate) or (batch, ngroups, dstate)
30 - C: (batch, dstate) or (batch, ngroups, dstate)
31 - D: (dim,) or (nheads, dim)
32 - z: (batch, dim) or (batch, nheads, dim)
33 - dt_bias: (dim,) or (nheads, dim)
34 - Return:
35 - out: (batch, dim) or (batch, nheads, dim)
36 - """
37 - has_heads = state.dim() > 3
38 - if state.dim() == 3:
39 - state = state.unsqueeze(1)
40 - if x.dim() == 2:
41 - x = x.unsqueeze(1)
42 - if dt.dim() == 2:
43 - dt = dt.unsqueeze(1)
44 - if A.dim() == 2:
45 - A = A.unsqueeze(0)
46 - if B.dim() == 2:
47 - B = B.unsqueeze(1)
48 - if C.dim() == 2:
49 - C = C.unsqueeze(1)
50 - if D is not None and D.dim() == 1:
51 - D = D.unsqueeze(0)
52 - if z is not None and z.dim() == 2:
53 - z = z.unsqueeze(1)
54 - if dt_bias is not None and dt_bias.dim() == 1:
55 - dt_bias = dt_bias.unsqueeze(0)
56 - batch, nheads, dim, dstate = state.shape
57 - assert x.shape == (batch, nheads, dim)
58 - assert dt.shape == x.shape
59 - assert A.shape == (nheads, dim, dstate)
60 - ngroups = B.shape[1]
61 - assert nheads % ngroups == 0, "nheads must be divisible by ngroups"
62 - assert B.shape == (batch, ngroups, dstate)
63 - assert C.shape == B.shape
64 - if D is not None:
65 - assert D.shape == (nheads, dim)
66 - if z is not None:
67 - assert z.shape == x.shape
68 - if dt_bias is not None:
69 - assert dt_bias.shape == (nheads, dim)
70 - dt = dt + dt_bias
71 - dt = F.softplus(dt) if dt_softplus else dt
72 - dA = torch.exp(
73 - rearrange(dt, "b h d -> b h d 1") * A
74 - ) # (batch, nheads, dim, dstate)
75 - B = repeat(B, "b g n -> b (g h) n", h=nheads // ngroups) # (batch, nheads, dstate)
76 - C = repeat(C, "b g n -> b (g h) n", h=nheads // ngroups) # (batch, nheads, dstate)
77 - dB = rearrange(dt, "b h d -> b h d 1") * rearrange(
78 - B, "b h n -> b h 1 n"
79 - ) # (batch, nheads, dim, dstate)
80 - state.copy_(
81 - state * dA + dB * rearrange(x, "b h d -> b h d 1")
82 - ) # (batch, dim, dstate
83 - out = torch.einsum("bhdn,bhn->bhd", state.to(C.dtype), C)
84 - if D is not None:
85 - out += (x * D).to(out.dtype)
86 - out = (out if z is None else out * F.silu(z)).to(x.dtype)
87 - if not has_heads:
88 - out = out.squeeze(1)
89 - return out
90 -
91 -
21 def selective_scan_ref(
22 u,
23 delta,
tests/kernels/mamba/test_mamba_ssm_configs.py new
+212
@@ -0,0 +1,212 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Unit tests for the JSON-based config loader added to selective_state_update.
5 +
6 +Tests cover:
7 + - Flat MoE-style filename generation
8 + - VLLM_TUNED_CONFIG_FOLDER env-var override
9 + - Fallback to heuristic when no config file exists
10 + - Nearest effective_batch interpolation
11 + - Edge cases: non-dict JSON, empty config
12 +"""
13 +
14 +import json
15 +
16 +from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
17 + _get_default_ssm_launch_config,
18 + _try_get_optimal_ssm_config_cached,
19 + get_ssm_config_file_name,
20 + get_ssm_configs,
21 + get_ssm_device_name,
22 + try_get_optimal_ssm_config,
23 +)
24 +
25 +# Common kwargs for try_get_optimal_ssm_config. Tests pick (batch, nheads) so
26 +# their product (effective_batch) matches the value being probed.
27 +_HEADDIM = 64
28 +_CACHE_DTYPE = "float32"
29 +
30 +
31 +def _clear_caches() -> None:
32 + get_ssm_configs.cache_clear()
33 + _try_get_optimal_ssm_config_cached.cache_clear()
34 +
35 +
36 +def _write_config(tmp_path, dstate: int, payload: dict) -> None:
37 + """Write payload as the bundled config for (headdim, dstate, cache_dtype)."""
38 + device_name = get_ssm_device_name()
39 + config_path = tmp_path / get_ssm_config_file_name(
40 + _HEADDIM, dstate, _CACHE_DTYPE, device_name
41 + )
42 + with open(config_path, "w") as f:
43 + json.dump(payload, f)
44 +
45 +
46 +# ---------------------------------------------------------------------------
47 +# Config filename generation
48 +# ---------------------------------------------------------------------------
49 +
50 +
51 +def test_config_file_name_format():
52 + name = get_ssm_config_file_name(
53 + headdim=64, dstate=128, cache_dtype="float32", device_name="NVIDIA_B200"
54 + )
55 + assert name == (
56 + "headdim=64,dstate=128,device_name=NVIDIA_B200,cache_dtype=float32.json"
57 + )
58 +
59 +
60 +# ---------------------------------------------------------------------------
61 +# VLLM_TUNED_CONFIG_FOLDER override
62 +# ---------------------------------------------------------------------------
63 +
64 +
65 +def test_env_override_loads_custom_config(monkeypatch, tmp_path):
66 + """VLLM_TUNED_CONFIG_FOLDER should take precedence over the bundled dir."""
67 + _write_config(
68 + tmp_path,
69 + dstate=16,
70 + payload={
71 + "1": {"BLOCK_SIZE_M": 4, "num_warps": 1},
72 + },
73 + )
74 +
75 + monkeypatch.setenv("VLLM_TUNED_CONFIG_FOLDER", str(tmp_path))
76 + _clear_caches()
77 +
78 + cfg = get_ssm_configs(_HEADDIM, 16, _CACHE_DTYPE)
79 + assert cfg is not None
80 + assert cfg[1] == {"BLOCK_SIZE_M": 4, "num_warps": 1}
81 +
82 + _clear_caches()
83 +
84 +
85 +# ---------------------------------------------------------------------------
86 +# Fallback to heuristic when no config file exists
87 +# ---------------------------------------------------------------------------
88 +
89 +
90 +def test_fallback_when_no_config(monkeypatch, tmp_path):
91 + """try_get_optimal_ssm_config must fall back to _get_default_ssm_launch_config
92 + when no JSON file is found for the current
93 + (device, headdim, dstate, cache_dtype) combination.
94 + """
95 + monkeypatch.setenv("VLLM_TUNED_CONFIG_FOLDER", str(tmp_path))
96 + monkeypatch.setattr(
97 + "vllm.model_executor.layers.mamba.ops.mamba_ssm._CONFIGS_DIR",
98 + str(tmp_path),
99 + )
100 +
101 + for dstate in (8, 16, 32, 64, 128, 256):
102 + for is_blackwell in (False, True):
103 + _clear_caches()
104 + block_m, warps = try_get_optimal_ssm_config(
105 + headdim=_HEADDIM,
106 + dstate=dstate,
107 + batch=1,
108 + nheads=1,
109 + cache_dtype=_CACHE_DTYPE,
110 + is_blackwell=is_blackwell,
111 + )
112 + assert (block_m, warps) == _get_default_ssm_launch_config(
113 + dstate, is_blackwell=is_blackwell
114 + )
115 +
116 + _clear_caches()
117 +
118 +
119 +# ---------------------------------------------------------------------------
120 +# Nearest effective_batch interpolation
121 +# ---------------------------------------------------------------------------
122 +
123 +
124 +def test_nearest_effective_batch_interpolation(monkeypatch, tmp_path):
125 + """When effective_batch = batch*nheads is not an exact key, the closest
126 + key should be selected."""
127 + _write_config(
128 + tmp_path,
129 + dstate=32,
130 + payload={
131 + "64": {"BLOCK_SIZE_M": 8, "num_warps": 1},
132 + "4096": {"BLOCK_SIZE_M": 32, "num_warps": 4},
133 + },
134 + )
135 +
136 + monkeypatch.setenv("VLLM_TUNED_CONFIG_FOLDER", str(tmp_path))
137 + _clear_caches()
138 +
139 + # effective_batch = 1*128 = 128 -> closer to 64 than to 4096
140 + block_m, warps = try_get_optimal_ssm_config(
141 + headdim=_HEADDIM,
142 + dstate=32,
143 + batch=1,
144 + nheads=128,
145 + cache_dtype=_CACHE_DTYPE,
146 + is_blackwell=False,
147 + )
148 + assert block_m == 8 and warps == 1
149 +
150 + # effective_batch = 4*1024 = 4096 -> exact match on 4096
151 + block_m, warps = try_get_optimal_ssm_config(
152 + headdim=_HEADDIM,
153 + dstate=32,
154 + batch=4,
155 + nheads=1024,
156 + cache_dtype=_CACHE_DTYPE,
157 + is_blackwell=False,
158 + )
159 + assert block_m == 32 and warps == 4
160 +
161 + _clear_caches()
162 +
163 +
164 +# ---------------------------------------------------------------------------
165 +# Edge cases: malformed / empty config files
166 +# ---------------------------------------------------------------------------
167 +
168 +
169 +def test_non_dict_json_returns_none(monkeypatch, tmp_path):
170 + """A valid JSON file that is not a dict (e.g. a list) must be ignored
171 + and return None rather than raising AttributeError."""
172 + device_name = get_ssm_device_name()
173 + config_path = tmp_path / get_ssm_config_file_name(
174 + _HEADDIM, 16, _CACHE_DTYPE, device_name
175 + )
176 + with open(config_path, "w") as f:
177 + json.dump([1, 2, 3], f)
178 +
179 + monkeypatch.setenv("VLLM_TUNED_CONFIG_FOLDER", str(tmp_path))
180 + monkeypatch.setattr(
181 + "vllm.model_executor.layers.mamba.ops.mamba_ssm._CONFIGS_DIR",
182 + str(tmp_path),
183 + )
184 + _clear_caches()
185 +
186 + assert get_ssm_configs(_HEADDIM, 16, _CACHE_DTYPE) is None
187 +
188 + _clear_caches()
189 +
190 +
191 +def test_empty_config_falls_back_to_heuristic(monkeypatch, tmp_path):
192 + """An empty JSON object {} must not crash min() — should fall back
193 + to the hard-coded heuristic."""
194 + _write_config(tmp_path, dstate=64, payload={})
195 +
196 + monkeypatch.setenv("VLLM_TUNED_CONFIG_FOLDER", str(tmp_path))
197 + _clear_caches()
198 +
199 + dstate = 64
200 + block_m, warps = try_get_optimal_ssm_config(
201 + headdim=_HEADDIM,
202 + dstate=dstate,
203 + batch=1,
204 + nheads=64,
205 + cache_dtype=_CACHE_DTYPE,
206 + is_blackwell=False,
207 + )
208 + assert (block_m, warps) == _get_default_ssm_launch_config(
209 + dstate=dstate, is_blackwell=False
210 + )
211 +
212 + _clear_caches()
tests/kernels/mamba/utils.py new
+78
@@ -0,0 +1,78 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +import torch.nn.functional as F
6 +from einops import rearrange, repeat
7 +
8 +
9 +def selective_state_update_ref(
10 + state, x, dt, A, B, C, D=None, z=None, dt_bias=None, dt_softplus=False
11 +):
12 + """
13 + Argument:
14 + state: (batch, dim, dstate) or (batch, nheads, dim, dstate)
15 + x: (batch, dim) or (batch, nheads, dim)
16 + dt: (batch, dim) or (batch, nheads, dim)
17 + A: (dim, dstate) or (nheads, dim, dstate)
18 + B: (batch, dstate) or (batch, ngroups, dstate)
19 + C: (batch, dstate) or (batch, ngroups, dstate)
20 + D: (dim,) or (nheads, dim)
21 + z: (batch, dim) or (batch, nheads, dim)
22 + dt_bias: (dim,) or (nheads, dim)
23 + Return:
24 + out: (batch, dim) or (batch, nheads, dim)
25 + """
26 + has_heads = state.dim() > 3
27 + if state.dim() == 3:
28 + state = state.unsqueeze(1)
29 + if x.dim() == 2:
30 + x = x.unsqueeze(1)
31 + if dt.dim() == 2:
32 + dt = dt.unsqueeze(1)
33 + if A.dim() == 2:
34 + A = A.unsqueeze(0)
35 + if B.dim() == 2:
36 + B = B.unsqueeze(1)
37 + if C.dim() == 2:
38 + C = C.unsqueeze(1)
39 + if D is not None and D.dim() == 1:
40 + D = D.unsqueeze(0)
41 + if z is not None and z.dim() == 2:
42 + z = z.unsqueeze(1)
43 + if dt_bias is not None and dt_bias.dim() == 1:
44 + dt_bias = dt_bias.unsqueeze(0)
45 + batch, nheads, dim, dstate = state.shape
46 + assert x.shape == (batch, nheads, dim)
47 + assert dt.shape == x.shape
48 + assert A.shape == (nheads, dim, dstate)
49 + ngroups = B.shape[1]
50 + assert nheads % ngroups == 0, "nheads must be divisible by ngroups"
51 + assert B.shape == (batch, ngroups, dstate)
52 + assert C.shape == B.shape
53 + if D is not None:
54 + assert D.shape == (nheads, dim)
55 + if z is not None:
56 + assert z.shape == x.shape
57 + if dt_bias is not None:
58 + assert dt_bias.shape == (nheads, dim)
59 + dt = dt + dt_bias
60 + dt = F.softplus(dt) if dt_softplus else dt
61 + dA = torch.exp(
62 + rearrange(dt, "b h d -> b h d 1") * A
63 + ) # (batch, nheads, dim, dstate)
64 + B = repeat(B, "b g n -> b (g h) n", h=nheads // ngroups) # (batch, nheads, dstate)
65 + C = repeat(C, "b g n -> b (g h) n", h=nheads // ngroups) # (batch, nheads, dstate)
66 + dB = rearrange(dt, "b h d -> b h d 1") * rearrange(
67 + B, "b h n -> b h 1 n"
68 + ) # (batch, nheads, dim, dstate)
69 + state.copy_(
70 + state * dA + dB * rearrange(x, "b h d -> b h d 1")
71 + ) # (batch, dim, dstate
72 + out = torch.einsum("bhdn,bhn->bhd", state.to(C.dtype), C)
73 + if D is not None:
74 + out += (x * D).to(out.dtype)
75 + out = (out if z is None else out * F.silu(z)).to(x.dtype)
76 + if not has_heads:
77 + out = out.squeeze(1)
78 + return out
vllm/envs.py
+4 -1
@@ -1744,7 +1744,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
1744 "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT": lambda: bool(
1745 int(os.getenv("VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", "0"))
1746 ),
1747 - # Allows vllm to find tuned config under customized folder
1747 + # User override folder for tuned Triton-kernel configs. Shared by MoE,
1748 + # Mamba SSU, and LoRA. Filenames are distinct so one folder can hold all.
1749 + # Each component first checks this folder, then the configs shipped with
1750 + # vLLM (if any). If no JSON matches, it uses a hard-coded heuristic.
1751 "VLLM_TUNED_CONFIG_FOLDER": lambda: os.getenv("VLLM_TUNED_CONFIG_FOLDER", None),
1752 # Valid values are container,code_interpreter,web_search_preview
1753 # ex VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS=container,code_interpreter
vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=NVIDIA_B200,cache_dtype=float16.json new
+87
@@ -0,0 +1,87 @@
1 +{
2 + "triton_version": "3.6.0",
3 + "8": {
4 + "BLOCK_SIZE_M": 4,
5 + "num_warps": 4
6 + },
7 + "16": {
8 + "BLOCK_SIZE_M": 8,
9 + "num_warps": 8
10 + },
11 + "32": {
12 + "BLOCK_SIZE_M": 4,
13 + "num_warps": 1
14 + },
15 + "64": {
16 + "BLOCK_SIZE_M": 16,
17 + "num_warps": 4
18 + },
19 + "128": {
20 + "BLOCK_SIZE_M": 32,
21 + "num_warps": 4
22 + },
23 + "256": {
24 + "BLOCK_SIZE_M": 16,
25 + "num_warps": 2
26 + },
27 + "512": {
28 + "BLOCK_SIZE_M": 32,
29 + "num_warps": 4
30 + },
31 + "1024": {
32 + "BLOCK_SIZE_M": 16,
33 + "num_warps": 2
34 + },
35 + "2048": {
36 + "BLOCK_SIZE_M": 32,
37 + "num_warps": 2
38 + },
39 + "4096": {
40 + "BLOCK_SIZE_M": 16,
41 + "num_warps": 2
42 + },
43 + "8192": {
44 + "BLOCK_SIZE_M": 16,
45 + "num_warps": 2
46 + },
47 + "12288": {
48 + "BLOCK_SIZE_M": 16,
49 + "num_warps": 2
50 + },
51 + "16384": {
52 + "BLOCK_SIZE_M": 16,
53 + "num_warps": 2
54 + },
55 + "24576": {
56 + "BLOCK_SIZE_M": 16,
57 + "num_warps": 2
58 + },
59 + "32768": {
60 + "BLOCK_SIZE_M": 16,
61 + "num_warps": 2
62 + },
63 + "49152": {
64 + "BLOCK_SIZE_M": 16,
65 + "num_warps": 2
66 + },
67 + "65536": {
68 + "BLOCK_SIZE_M": 16,
69 + "num_warps": 2
70 + },
71 + "98304": {
72 + "BLOCK_SIZE_M": 16,
73 + "num_warps": 2
74 + },
75 + "131072": {
76 + "BLOCK_SIZE_M": 16,
77 + "num_warps": 2
78 + },
79 + "196608": {
80 + "BLOCK_SIZE_M": 16,
81 + "num_warps": 2
82 + },
83 + "262144": {
84 + "BLOCK_SIZE_M": 16,
85 + "num_warps": 2
86 + }
87 +}
vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=NVIDIA_B200,cache_dtype=float32.json new
+87
@@ -0,0 +1,87 @@
1 +{
2 + "triton_version": "3.6.0",
3 + "8": {
4 + "BLOCK_SIZE_M": 4,
5 + "num_warps": 4
6 + },
7 + "16": {
8 + "BLOCK_SIZE_M": 4,
9 + "num_warps": 1
10 + },
11 + "32": {
12 + "BLOCK_SIZE_M": 4,
13 + "num_warps": 1
14 + },
15 + "64": {
16 + "BLOCK_SIZE_M": 4,
17 + "num_warps": 1
18 + },
19 + "128": {
20 + "BLOCK_SIZE_M": 4,
21 + "num_warps": 1
22 + },
23 + "256": {
24 + "BLOCK_SIZE_M": 4,
25 + "num_warps": 1
26 + },
27 + "512": {
28 + "BLOCK_SIZE_M": 4,
29 + "num_warps": 1
30 + },
31 + "1024": {
32 + "BLOCK_SIZE_M": 4,
33 + "num_warps": 1
34 + },
35 + "2048": {
36 + "BLOCK_SIZE_M": 4,
37 + "num_warps": 1
38 + },
39 + "4096": {
40 + "BLOCK_SIZE_M": 4,
41 + "num_warps": 1
42 + },
43 + "8192": {
44 + "BLOCK_SIZE_M": 4,
45 + "num_warps": 1
46 + },
47 + "12288": {
48 + "BLOCK_SIZE_M": 4,
49 + "num_warps": 1
50 + },
51 + "16384": {
52 + "BLOCK_SIZE_M": 4,
53 + "num_warps": 1
54 + },
55 + "24576": {
56 + "BLOCK_SIZE_M": 4,
57 + "num_warps": 1
58 + },
59 + "32768": {
60 + "BLOCK_SIZE_M": 4,
61 + "num_warps": 1
62 + },
63 + "49152": {
64 + "BLOCK_SIZE_M": 4,
65 + "num_warps": 1
66 + },
67 + "65536": {
68 + "BLOCK_SIZE_M": 4,
69 + "num_warps": 1
70 + },
71 + "98304": {
72 + "BLOCK_SIZE_M": 4,
73 + "num_warps": 1
74 + },
75 + "131072": {
76 + "BLOCK_SIZE_M": 4,
77 + "num_warps": 1
78 + },
79 + "196608": {
80 + "BLOCK_SIZE_M": 4,
81 + "num_warps": 1
82 + },
83 + "262144": {
84 + "BLOCK_SIZE_M": 4,
85 + "num_warps": 1
86 + }
87 +}
vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=NVIDIA_GB200,cache_dtype=float16.json new
+87
@@ -0,0 +1,87 @@
1 +{
2 + "triton_version": "3.6.0",
3 + "8": {
4 + "BLOCK_SIZE_M": 4,
5 + "num_warps": 4
6 + },
7 + "16": {
8 + "BLOCK_SIZE_M": 8,
9 + "num_warps": 8
10 + },
11 + "32": {
12 + "BLOCK_SIZE_M": 4,
13 + "num_warps": 1
14 + },
15 + "64": {
16 + "BLOCK_SIZE_M": 16,
17 + "num_warps": 4
18 + },
19 + "128": {
20 + "BLOCK_SIZE_M": 32,
21 + "num_warps": 4
22 + },
23 + "256": {
24 + "BLOCK_SIZE_M": 16,
25 + "num_warps": 2
26 + },
27 + "512": {
28 + "BLOCK_SIZE_M": 32,
29 + "num_warps": 1
30 + },
31 + "1024": {
32 + "BLOCK_SIZE_M": 16,
33 + "num_warps": 2
34 + },
35 + "2048": {
36 + "BLOCK_SIZE_M": 16,
37 + "num_warps": 2
38 + },
39 + "4096": {
40 + "BLOCK_SIZE_M": 16,
41 + "num_warps": 2
42 + },
43 + "8192": {
44 + "BLOCK_SIZE_M": 16,
45 + "num_warps": 2
46 + },
47 + "12288": {
48 + "BLOCK_SIZE_M": 16,
49 + "num_warps": 2
50 + },
51 + "16384": {
52 + "BLOCK_SIZE_M": 16,
53 + "num_warps": 2
54 + },
55 + "24576": {
56 + "BLOCK_SIZE_M": 16,
57 + "num_warps": 2
58 + },
59 + "32768": {
60 + "BLOCK_SIZE_M": 16,
61 + "num_warps": 2
62 + },
63 + "49152": {
64 + "BLOCK_SIZE_M": 16,
65 + "num_warps": 2
66 + },
67 + "65536": {
68 + "BLOCK_SIZE_M": 16,
69 + "num_warps": 2
70 + },
71 + "98304": {
72 + "BLOCK_SIZE_M": 16,
73 + "num_warps": 2
74 + },
75 + "131072": {
76 + "BLOCK_SIZE_M": 16,
77 + "num_warps": 2
78 + },
79 + "196608": {
80 + "BLOCK_SIZE_M": 16,
81 + "num_warps": 2
82 + },
83 + "262144": {
84 + "BLOCK_SIZE_M": 16,
85 + "num_warps": 2
86 + }
87 +}
vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=NVIDIA_GB200,cache_dtype=float32.json new
+87
@@ -0,0 +1,87 @@
1 +{
2 + "triton_version": "3.6.0",
3 + "8": {
4 + "BLOCK_SIZE_M": 4,
5 + "num_warps": 1
6 + },
7 + "16": {
8 + "BLOCK_SIZE_M": 4,
9 + "num_warps": 1
10 + },
11 + "32": {
12 + "BLOCK_SIZE_M": 4,
13 + "num_warps": 1
14 + },
15 + "64": {
16 + "BLOCK_SIZE_M": 4,
17 + "num_warps": 1
18 + },
19 + "128": {
20 + "BLOCK_SIZE_M": 4,
21 + "num_warps": 1
22 + },
23 + "256": {
24 + "BLOCK_SIZE_M": 4,
25 + "num_warps": 1
26 + },
27 + "512": {
28 + "BLOCK_SIZE_M": 4,
29 + "num_warps": 1
30 + },
31 + "1024": {
32 + "BLOCK_SIZE_M": 4,
33 + "num_warps": 1
34 + },
35 + "2048": {
36 + "BLOCK_SIZE_M": 4,
37 + "num_warps": 1
38 + },
39 + "4096": {
40 + "BLOCK_SIZE_M": 4,
41 + "num_warps": 1
42 + },
43 + "8192": {
44 + "BLOCK_SIZE_M": 4,
45 + "num_warps": 1
46 + },
47 + "12288": {
48 + "BLOCK_SIZE_M": 4,
49 + "num_warps": 1
50 + },
51 + "16384": {
52 + "BLOCK_SIZE_M": 4,
53 + "num_warps": 1
54 + },
55 + "24576": {
56 + "BLOCK_SIZE_M": 4,
57 + "num_warps": 1
58 + },
59 + "32768": {
60 + "BLOCK_SIZE_M": 4,
61 + "num_warps": 1
62 + },
63 + "49152": {
64 + "BLOCK_SIZE_M": 4,
65 + "num_warps": 1
66 + },
67 + "65536": {
68 + "BLOCK_SIZE_M": 4,
69 + "num_warps": 1
70 + },
71 + "98304": {
72 + "BLOCK_SIZE_M": 4,
73 + "num_warps": 1
74 + },
75 + "131072": {
76 + "BLOCK_SIZE_M": 4,
77 + "num_warps": 1
78 + },
79 + "196608": {
80 + "BLOCK_SIZE_M": 4,
81 + "num_warps": 1
82 + },
83 + "262144": {
84 + "BLOCK_SIZE_M": 4,
85 + "num_warps": 1
86 + }
87 +}
vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=NVIDIA_H100_80GB_HBM3,cache_dtype=float16.json new
+87
@@ -0,0 +1,87 @@
1 +{
2 + "triton_version": "3.6.0",
3 + "8": {
4 + "BLOCK_SIZE_M": 4,
5 + "num_warps": 1
6 + },
7 + "16": {
8 + "BLOCK_SIZE_M": 16,
9 + "num_warps": 4
10 + },
11 + "32": {
12 + "BLOCK_SIZE_M": 4,
13 + "num_warps": 1
14 + },
15 + "64": {
16 + "BLOCK_SIZE_M": 16,
17 + "num_warps": 4
18 + },
19 + "128": {
20 + "BLOCK_SIZE_M": 8,
21 + "num_warps": 2
22 + },
23 + "256": {
24 + "BLOCK_SIZE_M": 8,
25 + "num_warps": 2
26 + },
27 + "512": {
28 + "BLOCK_SIZE_M": 8,
29 + "num_warps": 2
30 + },
31 + "1024": {
32 + "BLOCK_SIZE_M": 8,
33 + "num_warps": 1
34 + },
35 + "2048": {
36 + "BLOCK_SIZE_M": 16,
37 + "num_warps": 2
38 + },
39 + "4096": {
40 + "BLOCK_SIZE_M": 16,
41 + "num_warps": 2
42 + },
43 + "8192": {
44 + "BLOCK_SIZE_M": 16,
45 + "num_warps": 2
46 + },
47 + "12288": {
48 + "BLOCK_SIZE_M": 16,
49 + "num_warps": 4
50 + },
51 + "16384": {
52 + "BLOCK_SIZE_M": 16,
53 + "num_warps": 2
54 + },
55 + "24576": {
56 + "BLOCK_SIZE_M": 16,
57 + "num_warps": 2
58 + },
59 + "32768": {
60 + "BLOCK_SIZE_M": 16,
61 + "num_warps": 2
62 + },
63 + "49152": {
64 + "BLOCK_SIZE_M": 16,
65 + "num_warps": 2
66 + },
67 + "65536": {
68 + "BLOCK_SIZE_M": 16,
69 + "num_warps": 2
70 + },
71 + "98304": {
72 + "BLOCK_SIZE_M": 16,
73 + "num_warps": 2
74 + },
75 + "131072": {
76 + "BLOCK_SIZE_M": 16,
77 + "num_warps": 2
78 + },
79 + "196608": {
80 + "BLOCK_SIZE_M": 16,
81 + "num_warps": 2
82 + },
83 + "262144": {
84 + "BLOCK_SIZE_M": 16,
85 + "num_warps": 2
86 + }
87 +}
vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=NVIDIA_H100_80GB_HBM3,cache_dtype=float32.json new
+87
@@ -0,0 +1,87 @@
1 +{
2 + "triton_version": "3.6.0",
3 + "8": {
4 + "BLOCK_SIZE_M": 4,
5 + "num_warps": 1
6 + },
7 + "16": {
8 + "BLOCK_SIZE_M": 4,
9 + "num_warps": 1
10 + },
11 + "32": {
12 + "BLOCK_SIZE_M": 4,
13 + "num_warps": 1
14 + },
15 + "64": {
16 + "BLOCK_SIZE_M": 4,
17 + "num_warps": 1
18 + },
19 + "128": {
20 + "BLOCK_SIZE_M": 4,
21 + "num_warps": 1
22 + },
23 + "256": {
24 + "BLOCK_SIZE_M": 4,
25 + "num_warps": 1
26 + },
27 + "512": {
28 + "BLOCK_SIZE_M": 4,
29 + "num_warps": 1
30 + },
31 + "1024": {
32 + "BLOCK_SIZE_M": 4,
33 + "num_warps": 1
34 + },
35 + "2048": {
36 + "BLOCK_SIZE_M": 4,
37 + "num_warps": 1
38 + },
39 + "4096": {
40 + "BLOCK_SIZE_M": 4,
41 + "num_warps": 1
42 + },
43 + "8192": {
44 + "BLOCK_SIZE_M": 16,
45 + "num_warps": 2
46 + },
47 + "12288": {
48 + "BLOCK_SIZE_M": 64,
49 + "num_warps": 8
50 + },
51 + "16384": {
52 + "BLOCK_SIZE_M": 16,
53 + "num_warps": 1
54 + },
55 + "24576": {
56 + "BLOCK_SIZE_M": 64,
57 + "num_warps": 8
58 + },
59 + "32768": {
60 + "BLOCK_SIZE_M": 16,
61 + "num_warps": 1
62 + },
63 + "49152": {
64 + "BLOCK_SIZE_M": 64,
65 + "num_warps": 8
66 + },
67 + "65536": {
68 + "BLOCK_SIZE_M": 16,
69 + "num_warps": 2
70 + },
71 + "98304": {
72 + "BLOCK_SIZE_M": 16,
73 + "num_warps": 1
74 + },
75 + "131072": {
76 + "BLOCK_SIZE_M": 16,
77 + "num_warps": 2
78 + },
79 + "196608": {
80 + "BLOCK_SIZE_M": 16,
81 + "num_warps": 1
82 + },
83 + "262144": {
84 + "BLOCK_SIZE_M": 16,
85 + "num_warps": 2
86 + }
87 +}
vllm/model_executor/layers/mamba/ops/mamba_ssm.py
+179 -18
@@ -4,16 +4,190 @@
4 # Copyright (c) 2024, Tri Dao, Albert Gu.
5 # Adapted from https://github.com/state-spaces/mamba/blob/v2.2.4/mamba_ssm/ops/triton/selective_state_update.py
6
7 +import functools
8 +import json
9 +import os
10 +from contextlib import contextmanager
11 +from typing import Any
12 +
13 import torch
14 from packaging import version
15
16 +import vllm.envs as envs
17 from vllm import _custom_ops as ops
18 +from vllm.logger import init_logger
19 from vllm.model_executor.layers.mamba.ops.triton_helpers import fast_exp
20 +from vllm.platforms import current_platform
21 from vllm.triton_utils import HAS_TRITON, tl, triton
22 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
23
24 +logger = init_logger(__name__)
25 +
26 TRITON3 = HAS_TRITON and (version.parse(triton.__version__) >= version.parse("3.0.0"))
27
28 +
29 +# ---------------------------------------------------------------------------
30 +# JSON config loading
31 +# ---------------------------------------------------------------------------
32 +
33 +_CONFIGS_DIR = os.path.join(
34 + os.path.dirname(os.path.realpath(__file__)), "configs", "selective_state_update"
35 +)
36 +
37 +
38 +def get_ssm_config_file_name(
39 + headdim: int, dstate: int, cache_dtype: str, device_name: str
40 +) -> str:
41 + """Return the JSON filename for the given kernel shape.
42 +
43 + Layout: ``configs/selective_state_update/
44 + headdim=<H>,dstate=<D>,device_name=<dev>,cache_dtype=<dt>.json``.
45 + """
46 + return (
47 + f"headdim={headdim},dstate={dstate},"
48 + f"device_name={device_name},cache_dtype={cache_dtype}.json"
49 + )
50 +
51 +
52 +def get_ssm_device_name() -> str:
53 + return current_platform.get_device_name().replace(" ", "_")
54 +
55 +
56 +def _canonical_cache_dtype(cache_dtype: str) -> str:
57 + """Canonical key for config lookup. bf16 and fp16 share the same tuned
58 + configs because the kernel only sees bit width when accessing state."""
59 + return "float16" if cache_dtype == "bfloat16" else cache_dtype
60 +
61 +
62 +@functools.cache
63 +def get_ssm_configs(
64 + headdim: int, dstate: int, cache_dtype: str
65 +) -> dict[int, Any] | None:
66 + """
67 + Return tuned (BLOCK_SIZE_M, num_warps) configs for *selective_state_update*
68 + keyed by ``effective_batch = batch * nheads``, or ``None`` if no config
69 + file is found for the (headdim, dstate, cache_dtype, device) combination.
70 +
71 + They can be generated with:
72 + benchmarks/kernels/benchmark_selective_state_update.py --save-configs
73 + """
74 + cache_dtype = _canonical_cache_dtype(cache_dtype)
75 + device_name = get_ssm_device_name()
76 + json_file_name = get_ssm_config_file_name(headdim, dstate, cache_dtype, device_name)
77 +
78 + config_file_paths: list[str] = []
79 +
80 + # User-supplied override
81 + user_defined_config_folder = envs.VLLM_TUNED_CONFIG_FOLDER
82 + if user_defined_config_folder is not None:
83 + config_file_paths.append(
84 + os.path.join(user_defined_config_folder, json_file_name)
85 + )
86 +
87 + # Bundled default
88 + config_file_paths.append(os.path.join(_CONFIGS_DIR, json_file_name))
89 +
90 + for path in config_file_paths:
91 + if os.path.exists(path):
92 + with open(path) as f:
93 + logger.info_once(
94 + "Using SSM config from %s for selective_state_update.",
95 + path,
96 + scope="global",
97 + )
98 + raw = json.load(f)
99 + if isinstance(raw, dict):
100 + # triton_version included in the config file only for reference
101 + raw.pop("triton_version", None)
102 + return {int(k): v for k, v in raw.items() if k.isdigit()}
103 +
104 + logger.warning_once(
105 + "Using default Mamba SSU config. Performance might be sub-optimal! "
106 + "Config file not found at %s",
107 + ", ".join(config_file_paths),
108 + )
109 + return None
110 +
111 +
112 +def _get_default_ssm_launch_config(
113 + dstate: int,
114 + is_blackwell: bool,
115 +) -> tuple[int, int]:
116 + """Hard-coded fallback heuristic used when no tuned config is available."""
117 + BLOCK_SIZE_M, num_warps = 4, 8
118 + if dstate <= 16:
119 + BLOCK_SIZE_M, num_warps = 32, 4
120 + elif dstate <= 32:
121 + BLOCK_SIZE_M, num_warps = 16, 4
122 + elif dstate <= 64:
123 + BLOCK_SIZE_M, num_warps = 8, 4
124 + else:
125 + if is_blackwell:
126 + BLOCK_SIZE_M, num_warps = 32, 8
127 + elif dstate <= 128:
128 + BLOCK_SIZE_M, num_warps = 4, 4
129 + return BLOCK_SIZE_M, num_warps
130 +
131 +
132 +@functools.cache
133 +def _try_get_optimal_ssm_config_cached(
134 + headdim: int,
135 + dstate: int,
136 + batch: int,
137 + nheads: int,
138 + cache_dtype: str,
139 + is_blackwell: bool,
140 +) -> tuple[int, int]:
141 + """Cached resolution. See :func:`try_get_optimal_ssm_config`."""
142 + effective_batch = batch * nheads
143 + configs = get_ssm_configs(headdim, dstate, cache_dtype)
144 + if configs:
145 + # Pick the closest effective_batch in the tuned grid (MoE strategy).
146 + closest = min(configs.keys(), key=lambda x: abs(x - effective_batch))
147 + cfg = configs[closest]
148 + return cfg["BLOCK_SIZE_M"], cfg["num_warps"]
149 +
150 + return _get_default_ssm_launch_config(dstate, is_blackwell)
151 +
152 +
153 +# Override hook for benchmarks/tests, see `override_ssm_config`.
154 +_ssm_config_override: tuple[int, int] | None = None
155 +
156 +
157 +@contextmanager
158 +def override_ssm_config(config: tuple[int, int]):
159 + """Pin ``try_get_optimal_ssm_config`` to ``config`` for the duration of
160 + the context. Used by the tuning benchmark to time specific configs."""
161 + global _ssm_config_override
162 + prev = _ssm_config_override
163 + _ssm_config_override = config
164 + try:
165 + yield
166 + finally:
167 + _ssm_config_override = prev
168 +
169 +
170 +def try_get_optimal_ssm_config(
171 + headdim: int,
172 + dstate: int,
173 + batch: int,
174 + nheads: int,
175 + cache_dtype: str,
176 + is_blackwell: bool,
177 +) -> tuple[int, int]:
178 + """Return (BLOCK_SIZE_M, num_warps) for the given kernel shape.
179 +
180 + Tuning is keyed on ``effective_batch = batch * nheads`` (the kernel grid
181 + scales with the product), so configs transfer across (model, TP) combos
182 + sharing ``(headdim, dstate, cache_dtype)``.
183 + """
184 + if _ssm_config_override is not None:
185 + return _ssm_config_override
186 + return _try_get_optimal_ssm_config_cached(
187 + headdim, dstate, batch, nheads, cache_dtype, is_blackwell
188 + )
189 +
190 +
191 if TRITON3:
192
193 @triton.jit
@@ -440,24 +614,11 @@ def selective_state_update(
614 else (0, 0)
615 )
616 # We don't want autotune since it will overwrite the state.
443 - # We instead tune by hand based on dstate.
444 -
445 - # Default
446 - BLOCK_SIZE_M, num_warps = 4, 8
447 -
448 - if dstate <= 16:
449 - BLOCK_SIZE_M, num_warps = 32, 4
450 - elif dstate <= 32:
451 - BLOCK_SIZE_M, num_warps = 16, 4
452 - elif dstate <= 64:
453 - BLOCK_SIZE_M, num_warps = 8, 4
454 - else:
455 - # dstate > 64
456 - if is_blackwell:
457 - # Optimized for B200 with dstate>64
458 - BLOCK_SIZE_M, num_warps = 32, 8
459 - elif dstate <= 128:
460 - BLOCK_SIZE_M, num_warps = 4, 4
617 + # Load from JSON config if available, otherwise fall back to heuristic.
618 + cache_dtype = str(state.dtype).removeprefix("torch.")
619 + BLOCK_SIZE_M, num_warps = try_get_optimal_ssm_config(
620 + dim, dstate, N, nheads, cache_dtype, is_blackwell
621 + )
622
623 tie_hdim = (
624 A.stride(-1) == 0