wehub-resource-sync
775 lines
25 KiB
Python
775 lines
25 KiB
Python
#!/usr/bin/env python3
|
||
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
"""
|
||
Benchmark and tuning script for the Mamba selective_state_update kernel.
|
||
|
||
Mirrors the fused MoE tuning workflow: sweeps (BLOCK_SIZE_M, num_warps) across
|
||
an effective_batch grid for a given (headdim, dstate, ngroups, cache_dtype) and
|
||
saves the best config per effective_batch to JSON. Generated configs are picked
|
||
up by selective_state_update at runtime.
|
||
|
||
Usage:
|
||
python -m benchmarks.kernels.benchmark_selective_state_update \
|
||
--all-dstates --save-configs --compare
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
import sys
|
||
from io import StringIO
|
||
from itertools import product
|
||
from typing import Any
|
||
|
||
import torch
|
||
|
||
from tests.kernels.mamba.utils import selective_state_update_ref
|
||
from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
|
||
_CONFIGS_DIR,
|
||
_canonical_cache_dtype,
|
||
_get_default_ssm_launch_config,
|
||
get_ssm_config_file_name,
|
||
get_ssm_device_name,
|
||
override_ssm_config,
|
||
selective_state_update,
|
||
)
|
||
from vllm.triton_utils import triton
|
||
|
||
# bf16 shares configs with fp16 - same bit width.
|
||
_SSM_CACHE_DTYPE_MAP: dict[str, torch.dtype] = {
|
||
"float32": torch.float32,
|
||
"float16": torch.float16,
|
||
"bfloat16": torch.float16,
|
||
}
|
||
|
||
_RESULTS_DIR = os.path.dirname(os.path.realpath(__file__))
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Tuning search space
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_BSM_CHOICES_ALL = [4, 8, 16, 32, 64, 128, 256]
|
||
|
||
NUM_WARPS_CHOICES = [1, 2, 4, 8]
|
||
|
||
|
||
def _block_size_m_choices(headdim: int) -> list[int]:
|
||
"""BLOCK_SIZE_M candidates worth sweeping for a given headdim.
|
||
|
||
BLOCK_SIZE_M > next_pow2(headdim) wastes >=50% of each tile via masking
|
||
(offs_m >= dim rows are zeroed out), so we cap the sweep there.
|
||
"""
|
||
ceiling = 1
|
||
while ceiling < headdim:
|
||
ceiling <<= 1
|
||
return [b for b in _BSM_CHOICES_ALL if b <= ceiling]
|
||
|
||
|
||
# Default deployment shapes. effective_batch = batch * nheads scales the
|
||
# kernel grid, so configs transfer across (model, TP) combos sharing
|
||
# (headdim, dstate, cache_dtype).
|
||
DEFAULT_BATCH_SIZES = [1, 8, 16, 32, 64, 128, 256, 512, 1024, 1536, 2048]
|
||
DEFAULT_NHEADS = [128, 256]
|
||
|
||
ALL_DSTATES = [16, 32, 64, 128, 256]
|
||
|
||
# Default tuning shape — matches Nemotron-3-Super and Nemotron-3-Nano Mamba layers.
|
||
# Override with CLI flags for other architectures.
|
||
DEFAULT_HEADDIM = 64
|
||
DEFAULT_NGROUPS = 8
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Benchmark helper
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _make_inputs(
|
||
batch: int,
|
||
nheads: int,
|
||
dim: int,
|
||
dstate: int,
|
||
ngroups: int,
|
||
dtype: torch.dtype,
|
||
state_dtype: torch.dtype | None = None,
|
||
device: str = "cuda",
|
||
):
|
||
if state_dtype is None:
|
||
state_dtype = dtype
|
||
state = torch.randn(batch, nheads, dim, dstate, dtype=state_dtype, device=device)
|
||
x = torch.randn(batch, nheads, dim, dtype=dtype, device=device)
|
||
dt = torch.randn(batch, nheads, dim, dtype=dtype, device=device)
|
||
A = -torch.rand(nheads, dim, dstate, dtype=torch.float32, device=device)
|
||
B = torch.randn(batch, ngroups, dstate, dtype=dtype, device=device)
|
||
C = torch.randn(batch, ngroups, dstate, dtype=dtype, device=device)
|
||
D = torch.randn(nheads, dim, dtype=dtype, device=device)
|
||
dt_bias = torch.randn(nheads, dim, dtype=dtype, device=device)
|
||
out = torch.zeros(batch, nheads, dim, dtype=dtype, device=device)
|
||
return state, x, dt, A, B, C, D, dt_bias, out
|
||
|
||
|
||
def benchmark_config(
|
||
batch: int,
|
||
nheads: int,
|
||
dim: int,
|
||
dstate: int,
|
||
ngroups: int,
|
||
block_size_m: int,
|
||
num_warps_val: int,
|
||
dtype: torch.dtype,
|
||
state_dtype: torch.dtype | None = None,
|
||
num_iters: int = 100,
|
||
num_warmup: int = 20,
|
||
graph_batch_size: int = 10,
|
||
) -> float | None:
|
||
"""
|
||
Time one (BLOCK_SIZE_M, num_warps) config for selective_state_update.
|
||
Returns elapsed time in microseconds, or None on error.
|
||
|
||
Uses CUDA graph capture-and-replay to isolate kernel time from Python
|
||
eager-mode dispatch / kwarg-resolution overhead, mirroring the timing
|
||
methodology in benchmarks/kernels/benchmark_moe.py.
|
||
"""
|
||
state, x, dt, A, B, C, D, dt_bias, out = _make_inputs(
|
||
batch, nheads, dim, dstate, ngroups, dtype, state_dtype=state_dtype
|
||
)
|
||
|
||
def _call_kernel() -> None:
|
||
selective_state_update(
|
||
state,
|
||
x,
|
||
dt,
|
||
A,
|
||
B,
|
||
C,
|
||
D=D,
|
||
z=None,
|
||
dt_bias=dt_bias,
|
||
dt_softplus=True,
|
||
out=out,
|
||
)
|
||
|
||
try:
|
||
with override_ssm_config((block_size_m, num_warps_val)):
|
||
# Eager-mode warmup: triggers Triton autotune / JIT, primes caches.
|
||
for _ in range(num_warmup):
|
||
_call_kernel()
|
||
torch.accelerator.synchronize()
|
||
|
||
# Capture graph_batch_size invocations into a CUDA graph so the
|
||
# timed region runs without Python dispatch overhead per call.
|
||
graph = torch.cuda.CUDAGraph()
|
||
with torch.cuda.graph(graph):
|
||
for _ in range(graph_batch_size):
|
||
_call_kernel()
|
||
torch.accelerator.synchronize()
|
||
|
||
# Warmup graph replays (let the runtime stabilize).
|
||
for _ in range(5):
|
||
graph.replay()
|
||
torch.accelerator.synchronize()
|
||
|
||
start = torch.cuda.Event(enable_timing=True)
|
||
end = torch.cuda.Event(enable_timing=True)
|
||
latencies: list[float] = []
|
||
for _ in range(num_iters):
|
||
start.record()
|
||
graph.replay()
|
||
end.record()
|
||
end.synchronize()
|
||
latencies.append(start.elapsed_time(end))
|
||
graph.reset()
|
||
# elapsed_time returns ms; each replay runs graph_batch_size kernels,
|
||
# so divide by (num_iters * graph_batch_size) and convert ms -> us.
|
||
return sum(latencies) / (num_iters * graph_batch_size) * 1000
|
||
except Exception as e:
|
||
if "OutOfResources" not in str(e):
|
||
print(
|
||
f" Warning: config M={block_size_m},w={num_warps_val} "
|
||
f"raised {type(e).__name__}: {e}"
|
||
)
|
||
return None
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Tuning loop
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
# CUDA grid Y/Z dim limit — both `batch` and `nheads` must fit individually.
|
||
_CUDA_MAX_GRID_DIM = 65535
|
||
|
||
# Above this, kernel state-offset arithmetic (batch * nheads * headdim * dstate)
|
||
# overflows int32 and the launch raises cudaErrorIllegalAddress.
|
||
# 262144 covers Nemotron Super TP1 BS=2048.
|
||
_MAX_EFFECTIVE_BATCH = 262144
|
||
|
||
|
||
def expand_batch_x_nheads(
|
||
batch_sizes: list[int],
|
||
nheads_list: list[int],
|
||
ngroups: int,
|
||
) -> list[tuple[int, int, int]]:
|
||
"""Cross-product batch_sizes × nheads_list → sorted [(effective_batch,
|
||
batch, nheads)], deduped by effective_batch. Filters pairs that exceed
|
||
the CUDA grid dim limit, the effective_batch ceiling, or where nheads is
|
||
not a positive multiple of ngroups.
|
||
"""
|
||
seen: dict[int, tuple[int, int]] = {}
|
||
skipped_grid: list[tuple[int, int]] = []
|
||
skipped_ngroups: list[tuple[int, int]] = []
|
||
skipped_eb: list[tuple[int, int]] = []
|
||
for b, n in product(batch_sizes, nheads_list):
|
||
if b <= 0 or n <= 0:
|
||
continue
|
||
if b > _CUDA_MAX_GRID_DIM or n > _CUDA_MAX_GRID_DIM:
|
||
skipped_grid.append((b, n))
|
||
continue
|
||
if n % ngroups != 0:
|
||
skipped_ngroups.append((b, n))
|
||
continue
|
||
if b * n > _MAX_EFFECTIVE_BATCH:
|
||
skipped_eb.append((b, n))
|
||
continue
|
||
seen.setdefault(b * n, (b, n))
|
||
if skipped_grid:
|
||
print(
|
||
f" Note: skipping (batch, nheads) pairs exceeding CUDA grid dim "
|
||
f"{_CUDA_MAX_GRID_DIM}: {skipped_grid}"
|
||
)
|
||
if skipped_ngroups:
|
||
print(
|
||
f" Note: skipping (batch, nheads) pairs where nheads % ngroups != 0 "
|
||
f"for ngroups={ngroups}: {skipped_ngroups}"
|
||
)
|
||
if skipped_eb:
|
||
print(
|
||
f" Note: skipping (batch, nheads) pairs whose effective_batch "
|
||
f"exceeds {_MAX_EFFECTIVE_BATCH}: {skipped_eb}"
|
||
)
|
||
return sorted((eb, b, n) for eb, (b, n) in seen.items())
|
||
|
||
|
||
def tune_dstate(
|
||
dstate: int,
|
||
headdim: int,
|
||
ngroups: int,
|
||
dtype: torch.dtype,
|
||
num_iters: int,
|
||
verbose: bool,
|
||
active: list[tuple[int, int, int]],
|
||
state_dtype: torch.dtype | None = None,
|
||
) -> tuple[dict[int, dict], dict[int, dict[tuple[int, int], float]]]:
|
||
"""For each (effective_batch, batch, nheads) in *active*, sweep
|
||
(BLOCK_SIZE_M, num_warps) and return
|
||
({effective_batch: best_config}, {effective_batch: {(bsm, nw): us}}).
|
||
The second map is the full timing grid, used downstream so we don't
|
||
re-measure the same config in the comparison phase.
|
||
"""
|
||
best_per_eb: dict[int, dict] = {}
|
||
timings: dict[int, dict[tuple[int, int], float]] = {}
|
||
|
||
print(f"\n{'=' * 74}")
|
||
effective_state_dtype = state_dtype if state_dtype is not None else dtype
|
||
print(
|
||
f"Tuning headdim={headdim} dstate={dstate} ngroups={ngroups} "
|
||
f"dtype={dtype} ssm_cache_dtype={effective_state_dtype}"
|
||
)
|
||
print(f"{'=' * 74}")
|
||
|
||
bsm_choices = _block_size_m_choices(headdim)
|
||
print(f"BSM candidates (capped at next_pow2(headdim={headdim})): {bsm_choices}")
|
||
|
||
hdr = f"{'EffBatch':>8} | {'BLOCK_M':>7} | {'warps':>5} | {'us':>10} | note"
|
||
print(hdr)
|
||
print("-" * 52)
|
||
|
||
for eb, batch, nheads in active:
|
||
best_time = float("inf")
|
||
best_cfg: dict = {}
|
||
eb_timings: dict[tuple[int, int], float] = {}
|
||
|
||
for bsm, nw in product(bsm_choices, NUM_WARPS_CHOICES):
|
||
t = benchmark_config(
|
||
batch=batch,
|
||
nheads=nheads,
|
||
dim=headdim,
|
||
dstate=dstate,
|
||
ngroups=ngroups,
|
||
block_size_m=bsm,
|
||
num_warps_val=nw,
|
||
dtype=dtype,
|
||
state_dtype=state_dtype,
|
||
num_iters=num_iters,
|
||
)
|
||
if t is None:
|
||
continue
|
||
eb_timings[(bsm, nw)] = t
|
||
is_best = t < best_time
|
||
if is_best:
|
||
best_time = t
|
||
best_cfg = {"BLOCK_SIZE_M": bsm, "num_warps": nw}
|
||
if verbose:
|
||
marker = " <-- best" if is_best else ""
|
||
print(f"{eb:>8} | {bsm:>7} | {nw:>5} | {t:>10.2f} |{marker}")
|
||
|
||
timings[eb] = eb_timings
|
||
|
||
if not best_cfg:
|
||
print(
|
||
f"{eb:>8} | {'-':>7} | {'-':>5} | {'-':>10} | "
|
||
f"no working config (skipped)"
|
||
)
|
||
continue
|
||
|
||
if not verbose:
|
||
print(
|
||
f"{eb:>8} | {best_cfg['BLOCK_SIZE_M']:>7} | "
|
||
f"{best_cfg['num_warps']:>5} | {best_time:>10.2f} | best"
|
||
)
|
||
|
||
best_per_eb[eb] = best_cfg
|
||
|
||
return best_per_eb, timings
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Correctness validation
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def validate_configs(
|
||
dstate: int,
|
||
headdim: int,
|
||
ngroups: int,
|
||
tuned: dict[int, dict],
|
||
active: list[tuple[int, int, int]],
|
||
dtype: torch.dtype,
|
||
atol: float = 1e-2,
|
||
rtol: float = 1e-2,
|
||
state_dtype: torch.dtype | None = None,
|
||
) -> dict[int, bool]:
|
||
"""
|
||
For every (effective_batch, batch, nheads) in *active* that has a tuned
|
||
config, run the kernel with that config and compare against the reference.
|
||
Returns {effective_batch: passed}.
|
||
"""
|
||
# Disable TF32 so the reference's matmul matches the Triton kernel's
|
||
# fp32 accumulation; otherwise large ebs show bf16 rounding mismatches.
|
||
torch.set_float32_matmul_precision("highest")
|
||
|
||
print(f"\n{'=' * 74}")
|
||
effective_state_dtype = state_dtype if state_dtype is not None else dtype
|
||
print(
|
||
f"Validation headdim={headdim} dstate={dstate} ngroups={ngroups} "
|
||
f"dtype={dtype} ssm_cache_dtype={effective_state_dtype} atol={atol}"
|
||
)
|
||
print(f"{'=' * 74}")
|
||
print(f"{'EffBatch':>8} | {'MaxAbsErr':>12} | {'Status':>8}")
|
||
print("-" * 36)
|
||
|
||
results: dict[int, bool] = {}
|
||
|
||
for eb, batch, nheads in active:
|
||
cfg = tuned.get(eb)
|
||
if cfg is None:
|
||
continue
|
||
state, x, dt, A, B, C, D, dt_bias, out = _make_inputs(
|
||
batch=batch,
|
||
nheads=nheads,
|
||
dim=headdim,
|
||
dstate=dstate,
|
||
ngroups=ngroups,
|
||
dtype=dtype,
|
||
state_dtype=state_dtype,
|
||
)
|
||
# Clone state before GPU kernel modifies it in-place
|
||
state_ref = state.clone()
|
||
|
||
with override_ssm_config((cfg["BLOCK_SIZE_M"], cfg["num_warps"])):
|
||
selective_state_update(
|
||
state,
|
||
x,
|
||
dt,
|
||
A,
|
||
B,
|
||
C,
|
||
D=D,
|
||
z=None,
|
||
dt_bias=dt_bias,
|
||
dt_softplus=True,
|
||
out=out,
|
||
)
|
||
torch.accelerator.synchronize()
|
||
gpu_out = out.detach().cpu()
|
||
|
||
# Reference uses the original (unmodified) state
|
||
# Upcast to fp32 so the reference sums in fp32 (matches the Triton
|
||
# kernel); summing in bf16 over `dstate` blows up the error.
|
||
ref_out = (
|
||
selective_state_update_ref(
|
||
state_ref.float(),
|
||
x.float(),
|
||
dt.float(),
|
||
A.float(),
|
||
B.float(),
|
||
C.float(),
|
||
D=D.float(),
|
||
dt_bias=dt_bias.float(),
|
||
dt_softplus=True,
|
||
)
|
||
.to(out.dtype)
|
||
.cpu()
|
||
)
|
||
|
||
passed = torch.allclose(gpu_out.float(), ref_out.float(), atol=atol, rtol=rtol)
|
||
max_err = (gpu_out.float() - ref_out.float()).abs().max().item()
|
||
status = "PASS" if passed else "FAIL"
|
||
results[eb] = passed
|
||
print(f"{eb:>8} | {max_err:>12.6f} | {status:>8}")
|
||
|
||
n_pass = sum(results.values())
|
||
n_total = len(results)
|
||
print(f"\n {n_pass}/{n_total} configs passed validation for dstate={dstate}")
|
||
return results
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Save configs
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def save_configs(
|
||
headdim: int,
|
||
dstate: int,
|
||
cache_dtype: str,
|
||
configs: dict[int, dict],
|
||
save_dir: str | None = None,
|
||
) -> str:
|
||
# bf16 shares configs with fp16, use common filename for both
|
||
cache_dtype = _canonical_cache_dtype(cache_dtype)
|
||
|
||
base_dir = save_dir if save_dir else _CONFIGS_DIR
|
||
os.makedirs(base_dir, exist_ok=True)
|
||
file_path = os.path.join(
|
||
base_dir,
|
||
get_ssm_config_file_name(headdim, dstate, cache_dtype, get_ssm_device_name()),
|
||
)
|
||
# triton_version is informational only, the loader ignores it
|
||
payload: dict[str, Any] = {
|
||
"triton_version": triton.__version__,
|
||
**{str(k): v for k, v in sorted(configs.items())},
|
||
}
|
||
with open(file_path, "w") as f:
|
||
json.dump(payload, f, indent=4)
|
||
return file_path
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Comparison table
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def current_heuristic(dstate: int, is_blackwell: bool = False) -> dict:
|
||
"""Return the current hard-coded BLOCK_SIZE_M / num_warps for dstate."""
|
||
bsm, nw = _get_default_ssm_launch_config(dstate, is_blackwell)
|
||
return {"BLOCK_SIZE_M": bsm, "num_warps": nw}
|
||
|
||
|
||
def compare_heuristic_vs_tuned(
|
||
dstate: int,
|
||
headdim: int,
|
||
ngroups: int,
|
||
tuned: dict[int, dict],
|
||
timings: dict[int, dict[tuple[int, int], float]],
|
||
active: list[tuple[int, int, int]],
|
||
dtype: torch.dtype,
|
||
num_iters: int,
|
||
is_blackwell: bool,
|
||
state_dtype: torch.dtype | None = None,
|
||
):
|
||
heur_cfg = current_heuristic(dstate, is_blackwell)
|
||
heur_key = (heur_cfg["BLOCK_SIZE_M"], heur_cfg["num_warps"])
|
||
|
||
print(f"\n{'=' * 74}")
|
||
print(
|
||
f"Comparison headdim={headdim} dstate={dstate} "
|
||
f"ngroups={ngroups} — heuristic vs tuned"
|
||
)
|
||
print(
|
||
f"Heuristic: BLOCK_SIZE_M={heur_cfg['BLOCK_SIZE_M']}, "
|
||
f"num_warps={heur_cfg['num_warps']}"
|
||
)
|
||
print(f"{'=' * 74}")
|
||
hdr = (
|
||
f"{'EffBatch':>8} | {'Heur(us)':>10} | {'Tuned(us)':>10} | "
|
||
f"{'Speedup':>8} | Best config"
|
||
)
|
||
print(hdr)
|
||
print("-" * len(hdr))
|
||
|
||
for eb, batch, nheads in active:
|
||
eb_timings = timings.get(eb, {})
|
||
|
||
# Heuristic timing: reuse the tuning measurement if the heuristic
|
||
# config was in the swept grid; otherwise measure it once.
|
||
t_h = eb_timings.get(heur_key)
|
||
if t_h is None:
|
||
t_h = benchmark_config(
|
||
batch=batch,
|
||
nheads=nheads,
|
||
dim=headdim,
|
||
dstate=dstate,
|
||
ngroups=ngroups,
|
||
block_size_m=heur_cfg["BLOCK_SIZE_M"],
|
||
num_warps_val=heur_cfg["num_warps"],
|
||
dtype=dtype,
|
||
state_dtype=state_dtype,
|
||
num_iters=num_iters,
|
||
)
|
||
|
||
# `tuned[eb]` may be missing if all configs failed in tune_dstate;
|
||
# in that case fall back to the heuristic so the table still prints.
|
||
best = tuned.get(eb) or heur_cfg
|
||
t_t = eb_timings.get((best["BLOCK_SIZE_M"], best["num_warps"]))
|
||
|
||
if t_h is None or t_t is None:
|
||
print(f"{eb:>8} | {'N/A':>10} | {'N/A':>10} | {'N/A':>8} |")
|
||
continue
|
||
speedup = t_h / t_t
|
||
marker = " <--" if speedup > 1.05 else ""
|
||
print(
|
||
f"{eb:>8} | {t_h:>10.2f} | {t_t:>10.2f} | "
|
||
f"{speedup:>7.2f}x | "
|
||
f"M={best['BLOCK_SIZE_M']},w={best['num_warps']}{marker}"
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# CLI
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def save_results(device_name: str, output: str, results_file: str | None = None) -> str:
|
||
"""Save the full benchmark output to a results text file."""
|
||
if results_file is None:
|
||
results_file = os.path.join(
|
||
_RESULTS_DIR, f"ssm_benchmark_results_{device_name}.txt"
|
||
)
|
||
with open(results_file, "w") as f:
|
||
f.write(output)
|
||
return results_file
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(
|
||
description="Tune selective_state_update kernel for Mamba SSM"
|
||
)
|
||
parser.add_argument(
|
||
"--dstate",
|
||
type=int,
|
||
default=128,
|
||
help="SSM state size to tune for (default: 128)",
|
||
)
|
||
parser.add_argument(
|
||
"--all-dstates",
|
||
action="store_true",
|
||
help="Tune all common dstate values: " + str(ALL_DSTATES),
|
||
)
|
||
parser.add_argument(
|
||
"--dtype",
|
||
type=str,
|
||
default="bfloat16",
|
||
choices=["float16", "bfloat16"],
|
||
help="Activation / input data type (default: bfloat16)",
|
||
)
|
||
parser.add_argument(
|
||
"--mamba-ssm-cache-dtype",
|
||
type=str,
|
||
default="float32",
|
||
choices=list(_SSM_CACHE_DTYPE_MAP.keys()),
|
||
help="SSM state cache dtype (default: float32)",
|
||
)
|
||
parser.add_argument(
|
||
"--num-iters",
|
||
type=int,
|
||
default=100,
|
||
help="Number of timing iterations (default: 100)",
|
||
)
|
||
parser.add_argument(
|
||
"--save-configs",
|
||
action="store_true",
|
||
help=f"Save best configs to JSON in {_CONFIGS_DIR}",
|
||
)
|
||
parser.add_argument(
|
||
"--compare",
|
||
action="store_true",
|
||
help="Show comparison table: heuristic vs tuned",
|
||
)
|
||
parser.add_argument(
|
||
"--verbose",
|
||
action="store_true",
|
||
help="Print every (BLOCK_SIZE_M, num_warps) result, not just best",
|
||
)
|
||
parser.add_argument(
|
||
"--results-file",
|
||
type=str,
|
||
default=None,
|
||
help="Path to save the benchmark results text file "
|
||
"(default: ssm_benchmark_results_<device>.txt alongside this script)",
|
||
)
|
||
parser.add_argument(
|
||
"--save-dir",
|
||
type=str,
|
||
default=None,
|
||
help=f"Directory to save JSON configs (default: {_CONFIGS_DIR})",
|
||
)
|
||
parser.add_argument(
|
||
"--headdim",
|
||
type=int,
|
||
default=DEFAULT_HEADDIM,
|
||
help=f"Per-head feature dim (default: {DEFAULT_HEADDIM})",
|
||
)
|
||
parser.add_argument(
|
||
"--ngroups",
|
||
type=int,
|
||
default=DEFAULT_NGROUPS,
|
||
help=f"Number of B/C groups (default: {DEFAULT_NGROUPS})",
|
||
)
|
||
parser.add_argument(
|
||
"--batch-sizes",
|
||
type=int,
|
||
nargs="+",
|
||
default=DEFAULT_BATCH_SIZES,
|
||
metavar="B",
|
||
help=f"Decoder batch sizes to sweep (default: {DEFAULT_BATCH_SIZES})",
|
||
)
|
||
parser.add_argument(
|
||
"--nheads",
|
||
type=int,
|
||
nargs="+",
|
||
default=DEFAULT_NHEADS,
|
||
metavar="N",
|
||
help=f"Number of heads per rank to sweep (default: {DEFAULT_NHEADS}). "
|
||
"effective_batch = batch * nheads; cross-product is deduped by eb.",
|
||
)
|
||
parser.add_argument(
|
||
"--validate",
|
||
action="store_true",
|
||
help="After tuning, verify each best config against a CPU reference "
|
||
"implementation. Configs that fail are flagged in the output.",
|
||
)
|
||
parser.add_argument(
|
||
"--atol",
|
||
type=float,
|
||
default=1e-2,
|
||
help="Absolute tolerance for --validate (default: 1e-2)",
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
|
||
state_dtype = _SSM_CACHE_DTYPE_MAP[args.mamba_ssm_cache_dtype]
|
||
device_name = get_ssm_device_name()
|
||
cap = torch.cuda.get_device_capability()
|
||
is_blackwell = cap[0] >= 10
|
||
|
||
# Mirror all output to a results file (like Unix tee).
|
||
buf = StringIO()
|
||
|
||
class _Tee:
|
||
"""Writes to both the original stdout and an in-memory buffer."""
|
||
|
||
def write(self, s):
|
||
buf.write(s)
|
||
sys.__stdout__.write(s)
|
||
|
||
def flush(self):
|
||
sys.__stdout__.flush()
|
||
|
||
sys.stdout = _Tee() # type: ignore[assignment]
|
||
|
||
try:
|
||
print(f"Device : {device_name} (sm_{cap[0]}{cap[1]})")
|
||
print(f"Blackwell: {is_blackwell}")
|
||
print(f"dtype : {args.dtype}")
|
||
print(f"ssm_cache_dtype: {args.mamba_ssm_cache_dtype}")
|
||
print(f"headdim: {args.headdim}")
|
||
print(f"ngroups: {args.ngroups}")
|
||
print(f"triton : {triton.__version__}")
|
||
|
||
dstates = ALL_DSTATES if args.all_dstates else [args.dstate]
|
||
active = expand_batch_x_nheads(args.batch_sizes, args.nheads, args.ngroups)
|
||
|
||
for dstate in dstates:
|
||
tuned, timings = tune_dstate(
|
||
dstate=dstate,
|
||
headdim=args.headdim,
|
||
ngroups=args.ngroups,
|
||
dtype=dtype,
|
||
num_iters=args.num_iters,
|
||
verbose=args.verbose,
|
||
active=active,
|
||
state_dtype=state_dtype,
|
||
)
|
||
|
||
if args.compare:
|
||
compare_heuristic_vs_tuned(
|
||
dstate=dstate,
|
||
headdim=args.headdim,
|
||
ngroups=args.ngroups,
|
||
tuned=tuned,
|
||
timings=timings,
|
||
active=active,
|
||
dtype=dtype,
|
||
num_iters=args.num_iters,
|
||
is_blackwell=is_blackwell,
|
||
state_dtype=state_dtype,
|
||
)
|
||
|
||
if args.validate:
|
||
validity = validate_configs(
|
||
dstate=dstate,
|
||
headdim=args.headdim,
|
||
ngroups=args.ngroups,
|
||
tuned=tuned,
|
||
active=active,
|
||
dtype=dtype,
|
||
atol=args.atol,
|
||
state_dtype=state_dtype,
|
||
)
|
||
# Filter out any configs that failed correctness check
|
||
failed = [eb for eb, ok in validity.items() if not ok]
|
||
if failed:
|
||
print(
|
||
f"\n WARNING: {len(failed)} config(s) failed validation "
|
||
f"for dstate={dstate}: effective_batches {failed}"
|
||
)
|
||
print(" These will NOT be saved even with --save-configs.")
|
||
tuned = {
|
||
eb: cfg for eb, cfg in tuned.items() if validity.get(eb, True)
|
||
}
|
||
|
||
if args.save_configs:
|
||
path = save_configs(
|
||
headdim=args.headdim,
|
||
dstate=dstate,
|
||
cache_dtype=args.mamba_ssm_cache_dtype,
|
||
configs=tuned,
|
||
save_dir=args.save_dir,
|
||
)
|
||
print(f"\nSaved: {path}")
|
||
else:
|
||
print(f"\nBest configs for dstate={dstate}:")
|
||
for eb, cfg in sorted(tuned.items()):
|
||
print(f" effective_batch={eb:>6}: {cfg}")
|
||
print("\n(Re-run with --save-configs to persist to JSON)")
|
||
finally:
|
||
sys.stdout = sys.__stdout__
|
||
results_path = save_results(device_name, buf.getvalue(), args.results_file)
|
||
print(f"\nResults saved to: {results_path}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|