diff --git a/docs/inference/optimizations.md b/docs/inference/optimizations.md index 3df7af111..a8707856b 100644 --- a/docs/inference/optimizations.md +++ b/docs/inference/optimizations.md @@ -286,6 +286,11 @@ Counts accumulate per shape; the file is flushed periodically and at exit. Hand the JSONL files to the kernel folks. The `SAGE_ATTN_THREE` backend records to the same log, so either implementation can be used for collection. +A collected log can directly drive the NVFP4/FP8 attention kernel benchmark +(per-shape accuracy, TFLOPS, and latency for the local SageAttention3 port, +FlashInfer NVFP4, SageAttention FP8, and torch SDPA): +`python fastvideo-kernel/benchmarks/bench_nvfp4_fp8_attention.py --shapes-jsonl /tmp/attn_shapes.jsonl.pid`. + ### V-MoBA / SLA / SageSLA These backends are model-specific and require the corresponding kernels and diff --git a/fastvideo-kernel/benchmarks/bench_nvfp4_fp8_attention.py b/fastvideo-kernel/benchmarks/bench_nvfp4_fp8_attention.py new file mode 100644 index 000000000..5a67782cd --- /dev/null +++ b/fastvideo-kernel/benchmarks/bench_nvfp4_fp8_attention.py @@ -0,0 +1,420 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Benchmark NVFP4 vs FP8 attention kernels: accuracy, TFLOPS, latency. + +Lanes (each skips cleanly, with a printed reason, where it cannot run): + + nvfp4_local fastvideo-kernel/attn_qat_infer (local SageAttention3 + port, ``sageattn_blackwell``). Built for sm_120a only. + nvfp4_flashinfer flashinfer.nvfp4_attention_sm120_quantize_qkv + + nvfp4_attention_sm120_fwd. Compute capability 12.0 only, + head_dim 64/128. + fp8_sageattn ``sageattention.sageattn`` — the attention op activated by + examples/inference/optimizations/fp8_wan2_1_1_3b.py + (PR #1496, FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN → + fastvideo/attention/backends/sage_attn.py). sageattn + auto-dispatches per-arch; on sm89+ it selects the + qk-int8 / pv-fp8 CUDA kernels. Note: that script's *FP8 + weight* quantization (fastvideo.layers.quantization FP8 → + torch._scaled_mm) applies to DiT linears, not attention, + so it is out of scope here. + sdpa_ torch SDPA in bf16/fp16 — latency/TFLOPS baseline. + +Accuracy for every lane is measured against an fp32 torch-SDPA ground truth +computed on the same fixed-seed inputs (fp32 SDPA is accuracy-only and gets +no speed row). Reported: cosine similarity, max abs err, mean abs err. + +All lanes run at the [B, H, L, D] interface. + +Usage: + python bench_nvfp4_fp8_attention.py # default shapes + python bench_nvfp4_fp8_attention.py --causal --json out.json + python bench_nvfp4_fp8_attention.py --shapes-jsonl /tmp/attn_shapes.jsonl.pid1234 + # shapes collected via FASTVIDEO_ATTN_SHAPE_LOG (fastvideo/attention/ + # shape_logger.py); recorded causal/dtype/sm_scale are honored. +""" + +from __future__ import annotations + +import argparse +import dataclasses +import inspect +import json +import math +import statistics +import sys +from collections.abc import Callable +from pathlib import Path + +import torch + +_REPO_ROOT = Path(__file__).resolve().parents[2] + +# --------------------------------------------------------------------------- +# Pure helpers (CPU-unit-tested in fastvideo/tests/attention/ +# test_nvfp4_fp8_attention_bench.py) — keep them torch-free where possible. +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class Shape: + batch: int + num_q_heads: int + num_kv_heads: int + seq_len_q: int + seq_len_kv: int + head_dim: int + causal: bool + dtype: str = "torch.bfloat16" + sm_scale: float | None = None + + def label(self) -> str: + base = f"{self.batch}x{self.num_q_heads}x{self.seq_len_q}x{self.head_dim}" + if self.seq_len_kv != self.seq_len_q or self.num_kv_heads != self.num_q_heads: + base += f" (kv {self.num_kv_heads}x{self.seq_len_kv})" + if self.causal: + base += " causal" + return base + + +def attn_tflops(shape: Shape, ms: float) -> float: + """TFLOPS of one dense attention forward. + + FLOPs = 4 * B * H * Lq * Lk * D: + QK^T is B*H*Lq*Lk*D multiply-adds = 2*B*H*Lq*Lk*D FLOPs, and P@V is + another 2*B*H*Lq*Lk*D. Causal masking computes ~half the score matrix, + so FLOPs are halved. TFLOPS = FLOPs / (ms * 1e9). + """ + flops = 4.0 * shape.batch * shape.num_q_heads * shape.seq_len_q * shape.seq_len_kv * shape.head_dim + if shape.causal: + flops *= 0.5 + return flops / (ms * 1e9) + + +def parse_shapes_jsonl(paths: list[str], default_dtype: str) -> list[Shape]: + """Parse fastvideo/attention/shape_logger.py output (one JSON object per + line; keys: batch, num_q_heads, num_kv_heads, seq_len_q, seq_len_kv, + head_dim, dtype, causal, sm_scale, count, ...). Dedupes across files, + preserving first-seen order.""" + shapes: list[Shape] = [] + seen: set[Shape] = set() + for path in paths: + with open(path) as f: + for line in f: + line = line.strip() + if not line: + continue + rec = json.loads(line) + shape = Shape( + batch=int(rec["batch"]), + num_q_heads=int(rec["num_q_heads"]), + num_kv_heads=int(rec.get("num_kv_heads", rec["num_q_heads"])), + seq_len_q=int(rec["seq_len_q"]), + seq_len_kv=int(rec.get("seq_len_kv", rec["seq_len_q"])), + head_dim=int(rec["head_dim"]), + causal=bool(rec.get("causal", False)), + dtype=str(rec.get("dtype") or default_dtype), + sm_scale=rec.get("sm_scale"), + ) + if shape not in seen: + seen.add(shape) + shapes.append(shape) + return shapes + + +@dataclasses.dataclass(frozen=True) +class LaneSpec: + name: str + exact_cap: tuple[int, int] | None = None # kernel runs on exactly this SM + min_cap: tuple[int, int] | None = None + head_dims: tuple[int, ...] | None = None # None = any + same_shape: bool = False # requires q/k/v equal shapes (self-attention) + same_heads: bool = False # requires num_q_heads == num_kv_heads + + +def skip_reason(spec: LaneSpec, shape: Shape, cap: tuple[int, int] | None, + unavailable: str | None) -> str | None: + """Why this lane cannot run this shape (None = it can).""" + if unavailable: + return unavailable + if cap is None: + return "no CUDA device" + if spec.exact_cap is not None and cap != spec.exact_cap: + return (f"requires SM {spec.exact_cap[0]}.{spec.exact_cap[1]}, " + f"device is SM {cap[0]}.{cap[1]}") + if spec.min_cap is not None and cap < spec.min_cap: + return (f"requires SM >= {spec.min_cap[0]}.{spec.min_cap[1]}, " + f"device is SM {cap[0]}.{cap[1]}") + if spec.head_dims is not None and shape.head_dim not in spec.head_dims: + return f"head_dim {shape.head_dim} not in {list(spec.head_dims)}" + if spec.same_shape and (shape.seq_len_q != shape.seq_len_kv or shape.num_q_heads != shape.num_kv_heads): + return "requires equal q/k/v shapes (self-attention)" + if spec.same_heads and shape.num_q_heads != shape.num_kv_heads: + return "requires num_q_heads == num_kv_heads" + return None + + +def format_table(headers: list[str], rows: list[list[str]]) -> str: + """Aligned plain-text table (left-justified columns).""" + cols = list(zip(headers, *rows)) if rows else [(h,) for h in headers] + widths = [max(len(str(c)) for c in col) for col in cols] + lines = [ + " ".join(str(h).ljust(w) for h, w in zip(headers, widths)), + " ".join("-" * w for w in widths), + ] + for row in rows: + lines.append(" ".join(str(c).ljust(w) for c, w in zip(row, widths))) + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# Lanes (GPU) +# --------------------------------------------------------------------------- + +AttnFn = Callable[[torch.Tensor, torch.Tensor, torch.Tensor, bool, float], torch.Tensor] + + +@dataclasses.dataclass +class Lane: + spec: LaneSpec + fn: AttnFn | None # None when unavailable + unavailable: str | None = None + + +def _run_sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool, scale: float) -> torch.Tensor: + kwargs: dict = {"is_causal": causal, "scale": scale} + if q.shape[1] != k.shape[1]: + kwargs["enable_gqa"] = True + return torch.nn.functional.scaled_dot_product_attention(q, k, v, **kwargs) + + +def _load_nvfp4_local(single_level_p_quant: bool) -> Lane: + spec = LaneSpec("nvfp4_local", exact_cap=(12, 0), head_dims=(64, 128), same_shape=True) + for p in (_REPO_ROOT, _REPO_ROOT / "fastvideo-kernel", _REPO_ROOT / "fastvideo-kernel" / "python"): + if str(p) not in sys.path: + sys.path.insert(0, str(p)) + try: + from attn_qat_infer import sageattn_blackwell # noqa: PLC0415 + except ImportError as e: + return Lane(spec, None, f"attn_qat_infer import failed ({e}); build fastvideo-kernel with sm_120a") + + def fn(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool, scale: float) -> torch.Tensor: + # sageattn_blackwell pads L to a multiple of 128 internally and trims. + return sageattn_blackwell(q, k, v, is_causal=causal, + single_level_p_quant=single_level_p_quant, sm_scale=scale) + + return Lane(spec, fn) + + +def _load_nvfp4_flashinfer() -> Lane: + spec = LaneSpec("nvfp4_flashinfer", exact_cap=(12, 0), head_dims=(64, 128), same_shape=True) + try: + from flashinfer import ( # noqa: PLC0415 + nvfp4_attention_sm120_fwd, nvfp4_attention_sm120_quantize_qkv) + except ImportError as e: + return Lane(spec, None, f"flashinfer nvfp4_attention_sm120 not installed ({e})") + + quant_kwargs = {} + if "per_block_mean" in inspect.signature(nvfp4_attention_sm120_quantize_qkv).parameters: + quant_kwargs["per_block_mean"] = True + + def fn(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool, scale: float) -> torch.Tensor: + seq_len = q.shape[2] + packed = nvfp4_attention_sm120_quantize_qkv(q, k, v, **quant_kwargs) + out, _lse = nvfp4_attention_sm120_fwd(*packed, sm_scale=scale, causal=causal, out_dtype=q.dtype) + # flashinfer pads seq_len to a multiple of 128; trim it back off. + return out[:, :, :seq_len, :] + + return Lane(spec, fn) + + +def _load_fp8_sageattn() -> Lane: + # sageattn's fp8 (qk-int8 / pv-fp8) kernels need sm89+; the package + # dispatches per-arch internally. + spec = LaneSpec("fp8_sageattn", min_cap=(8, 9), same_heads=True) + try: + from sageattention import sageattn # noqa: PLC0415 + except ImportError as e: + return Lane(spec, None, f"sageattention not installed ({e})") + + def fn(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool, scale: float) -> torch.Tensor: + return sageattn(q, k, v, tensor_layout="HND", is_causal=causal, sm_scale=scale) + + return Lane(spec, fn) + + +def build_lanes(args: argparse.Namespace) -> list[Lane]: + return [ + _load_nvfp4_local(single_level_p_quant=args.p_quant == "single"), + _load_nvfp4_flashinfer(), + _load_fp8_sageattn(), + Lane(LaneSpec(f"sdpa_{args.dtype}"), _run_sdpa), + ] + + +# --------------------------------------------------------------------------- +# Measurement +# --------------------------------------------------------------------------- + +_DTYPES = {"torch.bfloat16": torch.bfloat16, "torch.float16": torch.float16, + "bf16": torch.bfloat16, "fp16": torch.float16} + + +def make_qkv(shape: Shape, seed: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + gen = torch.Generator(device="cuda").manual_seed(seed) # fixed seed per shape + dtype = _DTYPES.get(shape.dtype, torch.bfloat16) + + def rand(heads: int, seq: int) -> torch.Tensor: + return torch.randn(shape.batch, heads, seq, shape.head_dim, + generator=gen, device="cuda", dtype=dtype) + + return rand(shape.num_q_heads, shape.seq_len_q), rand(shape.num_kv_heads, shape.seq_len_kv), \ + rand(shape.num_kv_heads, shape.seq_len_kv) + + +def accuracy_stats(out: torch.Tensor, ref: torch.Tensor) -> dict[str, float]: + out = out.float() + diff = (out - ref).abs() + cos = torch.nn.functional.cosine_similarity(out.flatten(), ref.flatten(), dim=0) + return {"cos_sim": cos.item(), "max_abs_err": diff.max().item(), "mean_abs_err": diff.mean().item()} + + +def time_fn(fn: Callable[[], torch.Tensor], warmup: int, iters: int) -> float: + """Median CUDA-event latency in ms.""" + for _ in range(warmup): + fn() + torch.cuda.synchronize() + times = [] + for _ in range(iters): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + torch.cuda.synchronize() + times.append(start.elapsed_time(end)) + return statistics.median(times) + + +def bench_shape(shape: Shape, lanes: list[Lane], cap: tuple[int, int] | None, + args: argparse.Namespace) -> list[dict]: + results = [] + runnable = [(lane, skip_reason(lane.spec, shape, cap, lane.unavailable)) for lane in lanes] + q = k = v = ref = None + if any(reason is None for _, reason in runnable): + q, k, v = make_qkv(shape, args.seed) + scale = shape.sm_scale or shape.head_dim**-0.5 + # fp32 SDPA ground truth (accuracy baseline only; never timed). + ref = _run_sdpa(q.float(), k.float(), v.float(), shape.causal, scale) + for lane, reason in runnable: + rec: dict = {"shape": shape.label(), "lane": lane.spec.name, **dataclasses.asdict(shape)} + if reason is not None: + rec.update(status="skip", reason=reason) + results.append(rec) + continue + scale = shape.sm_scale or shape.head_dim**-0.5 + try: + out = lane.fn(q, k, v, shape.causal, scale) + rec.update(accuracy_stats(out, ref)) + del out + ms = time_fn(lambda: lane.fn(q, k, v, shape.causal, scale), args.warmup, args.iters) + rec.update(status="ok", median_ms=ms, tflops=attn_tflops(shape, ms)) + except Exception as e: # lane failed at runtime: report, keep going + rec.update(status="fail", reason=f"{type(e).__name__}: {e}") + results.append(rec) + del q, k, v, ref + if torch.cuda.is_available(): + torch.cuda.empty_cache() + return results + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +# Wan2.1-1.3B DiT self-attention is 12 heads x head_dim 128; 81-frame video +# gives latent grids 21x30x52 (480p) = 32760 tokens, 21x45x80 (720p) = 75600. +DEFAULT_SHAPE_DIMS = [ + (1, 8, 512, 64), # smoke + (1, 12, 1024, 128), # smoke + (1, 12, 32760, 128), # Wan2.1-1.3B 480p + (1, 12, 75600, 128), # Wan2.1-1.3B 720p +] + + +def default_shapes(causal: bool, dtype: str) -> list[Shape]: + return [ + Shape(batch=b, num_q_heads=h, num_kv_heads=h, seq_len_q=length, seq_len_kv=length, + head_dim=d, causal=causal, dtype=dtype) for b, h, length, d in DEFAULT_SHAPE_DIMS + ] + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0], + formatter_class=argparse.ArgumentDefaultsHelpFormatter) + p.add_argument("--shapes-jsonl", nargs="+", metavar="PATH", + help="attention shape log(s) written via FASTVIDEO_ATTN_SHAPE_LOG " + "(fastvideo/attention/shape_logger.py); replaces the default shapes") + p.add_argument("--causal", action="store_true", + help="run default shapes causally (jsonl shapes keep their recorded flag)") + p.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16", + help="input dtype for default shapes / jsonl records without one") + p.add_argument("--p-quant", choices=["single", "two"], default="single", + help="nvfp4_local P-matrix quantization: single-level or two-level " + "(sageattn_blackwell single_level_p_quant knob)") + p.add_argument("--warmup", type=int, default=10, help="warmup iterations per lane") + p.add_argument("--iters", type=int, default=50, help="timed iterations per lane (median reported)") + p.add_argument("--seed", type=int, default=0, help="input seed (fixed per shape)") + p.add_argument("--json", metavar="PATH", help="also dump results as JSON") + return p.parse_args(argv) + + +def _fmt(rec: dict, key: str, spec: str) -> str: + return format(rec[key], spec) if key in rec else "-" + + +def main(argv: list[str] | None = None) -> None: + args = parse_args(argv) + dtype_name = f"torch.{'bfloat16' if args.dtype == 'bf16' else 'float16'}" + shapes = (parse_shapes_jsonl(args.shapes_jsonl, dtype_name) + if args.shapes_jsonl else default_shapes(args.causal, dtype_name)) + + cap = None + meta: dict = {"torch": torch.__version__, "warmup": args.warmup, "iters": args.iters, "seed": args.seed} + if torch.cuda.is_available(): + cap = torch.cuda.get_device_capability(0) + meta.update(device=torch.cuda.get_device_name(0), sm=f"{cap[0]}.{cap[1]}") + print(f"Device: {meta['device']} (SM {meta['sm']}), torch {torch.__version__}") + else: + print("No CUDA device: all lanes will be skipped (shape/skip dry run).") + + lanes = build_lanes(args) + results = [] + for shape in shapes: + print(f"\n== {shape.label()} dtype={shape.dtype} ==") + results.extend(bench_shape(shape, lanes, cap, args)) + + headers = ["shape", "lane", "status", "median_ms", "TFLOPS", "cos_sim", "max_abs_err", "mean_abs_err"] + rows = [] + for rec in results: + status = rec["status"] if rec["status"] == "ok" else f"{rec['status']}: {rec['reason']}" + rows.append([ + rec["shape"], rec["lane"], status, + _fmt(rec, "median_ms", ".3f"), + _fmt(rec, "tflops", ".1f"), + _fmt(rec, "cos_sim", ".6f"), + _fmt(rec, "max_abs_err", ".4e"), + _fmt(rec, "mean_abs_err", ".4e"), + ]) + print("\n" + format_table(headers, rows)) + print("\nAccuracy is vs fp32 torch-SDPA ground truth on identical inputs " + "(fp32 SDPA itself is excluded from the speed columns).") + + if args.json: + with open(args.json, "w") as f: + json.dump({"meta": meta, "results": results}, f, indent=2) + print(f"Wrote {args.json}") + + +if __name__ == "__main__": + main() diff --git a/fastvideo/tests/attention/test_nvfp4_fp8_attention_bench.py b/fastvideo/tests/attention/test_nvfp4_fp8_attention_bench.py new file mode 100644 index 000000000..b5dbec034 --- /dev/null +++ b/fastvideo/tests/attention/test_nvfp4_fp8_attention_bench.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only tests for the pure parts of +fastvideo-kernel/benchmarks/bench_nvfp4_fp8_attention.py (flops formula, +shapes-jsonl parsing, lane-skip logic, table formatting). GPU lanes are +exercised manually on real hardware.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from pathlib import Path + +import pytest + +_SCRIPT = Path(__file__).resolve().parents[3] / "fastvideo-kernel" / "benchmarks" / "bench_nvfp4_fp8_attention.py" + +spec = importlib.util.spec_from_file_location("bench_nvfp4_fp8_attention", _SCRIPT) +bench = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = bench # dataclasses resolve annotations via sys.modules +spec.loader.exec_module(bench) + + +def _shape(**overrides): + base = dict(batch=1, num_q_heads=12, num_kv_heads=12, seq_len_q=1024, + seq_len_kv=1024, head_dim=128, causal=False) + base.update(overrides) + return bench.Shape(**base) + + +def test_attn_tflops_formula(): + # 4*B*H*Lq*Lk*D FLOPs at 1 ms. + shape = _shape(batch=2, num_q_heads=3, seq_len_q=100, seq_len_kv=200, head_dim=64) + expected = 4 * 2 * 3 * 100 * 200 * 64 / 1e9 # TFLOPS at 1 ms + assert bench.attn_tflops(shape, 1.0) == pytest.approx(expected) + # Causal halves the FLOPs; 2 ms halves the rate again. + causal = _shape(batch=2, num_q_heads=3, seq_len_q=100, seq_len_kv=200, head_dim=64, causal=True) + assert bench.attn_tflops(causal, 2.0) == pytest.approx(expected / 4) + + +def test_parse_shapes_jsonl_dedupe_and_defaults(tmp_path): + rec = {"backend": "FLASHINFER_SAGE_ATTN3", "batch": 1, "num_q_heads": 12, + "num_kv_heads": 12, "seq_len_q": 75600, "seq_len_kv": 75600, + "head_dim": 128, "dtype": "torch.bfloat16", "causal": False, + "sm_scale": 0.0884, "count": 1200, "torch": "2.8.0", + "device": "NVIDIA GB10", "sm": "12.1"} + minimal = {"batch": 2, "num_q_heads": 8, "seq_len_q": 512, "head_dim": 64} + path = tmp_path / "shapes.jsonl" + path.write_text(json.dumps(rec) + "\n\n" + json.dumps(rec) + "\n" + json.dumps(minimal) + "\n") + + shapes = bench.parse_shapes_jsonl([str(path)], default_dtype="torch.float16") + assert len(shapes) == 2 # duplicate line deduped + full, mini = shapes + assert (full.batch, full.num_q_heads, full.seq_len_q, full.head_dim) == (1, 12, 75600, 128) + assert full.sm_scale == pytest.approx(0.0884) + assert full.dtype == "torch.bfloat16" + # Missing keys fall back: kv defaults to q, dtype to the CLI default. + assert (mini.num_kv_heads, mini.seq_len_kv) == (8, 512) + assert mini.dtype == "torch.float16" + assert mini.causal is False and mini.sm_scale is None + + +def test_skip_reason(): + sm120 = bench.LaneSpec("nvfp4", exact_cap=(12, 0), head_dims=(64, 128), same_shape=True) + ok = _shape() + assert bench.skip_reason(sm120, ok, (12, 0), None) is None + assert "not installed" in bench.skip_reason(sm120, ok, (12, 0), "not installed") + assert "no CUDA" in bench.skip_reason(sm120, ok, None, None) + assert "SM 12.0" in bench.skip_reason(sm120, ok, (10, 0), None) # GB200 + assert "SM 12.0" in bench.skip_reason(sm120, ok, (12, 1), None) # Spark + assert "head_dim 96" in bench.skip_reason(sm120, _shape(head_dim=96), (12, 0), None) + cross = _shape(seq_len_kv=256) + assert "equal q/k/v" in bench.skip_reason(sm120, cross, (12, 0), None) + + fp8 = bench.LaneSpec("fp8", min_cap=(8, 9), same_heads=True) + assert bench.skip_reason(fp8, ok, (9, 0), None) is None + assert bench.skip_reason(fp8, cross, (12, 0), None) is None # cross-attn ok + assert "SM >= 8.9" in bench.skip_reason(fp8, ok, (8, 6), None) + assert "num_q_heads" in bench.skip_reason(fp8, _shape(num_kv_heads=4), (9, 0), None) + + anywhere = bench.LaneSpec("sdpa_bf16") + assert bench.skip_reason(anywhere, cross, (7, 5), None) is None + + +def test_format_table_alignment(): + headers = ["shape", "lane", "TFLOPS"] + rows = [["1x12x75600x128", "nvfp4_local", "812.3"], + ["1x8x512x64", "sdpa_bf16", "9.1"]] + table = bench.format_table(headers, rows) + lines = table.splitlines() + assert len(lines) == 4 # header, rule, 2 rows + assert len({len(line) for line in lines}) == 1 # columns aligned + assert "nvfp4_local" in lines[2] and "sdpa_bf16" in lines[3] + assert bench.format_table(headers, []).splitlines()[0].startswith("shape") + + +def test_cli_defaults_and_shapes(): + args = bench.parse_args([]) + assert (args.warmup, args.iters, args.p_quant) == (10, 50, "single") + shapes = bench.default_shapes(causal=False, dtype="torch.bfloat16") + assert bench.Shape(1, 12, 12, 75600, 75600, 128, False) in shapes # Wan2.1 720p + with pytest.raises(SystemExit) as exc: + bench.parse_args(["--help"]) + assert exc.value.code == 0