diff --git a/benchmark/sdpa_benchmark_training/README.md b/benchmark/sdpa_benchmark_training/README.md index 5aa64643c..a212bc4ec 100755 --- a/benchmark/sdpa_benchmark_training/README.md +++ b/benchmark/sdpa_benchmark_training/README.md @@ -16,6 +16,7 @@ This directory contains benchmarking tools for Scaled Dot Product Attention (SDP - `ltx2.py` - LTX-2 video DiT self-attention benchmarks (bidirectional, no mask) - `gpt_oss.py` - GPT-OSS sliding-window-attention GQA benchmarks (causal, SWA=128) - `qwen35.py` - Qwen 3.5 GQA benchmarks (head_dim=256, causal, bf16 bidirectional — Blackwell fp8/fa4 limits) + - `auto_regressive_dit.py` - Autoregressive video DiT (short Q, long cached KV, bf16/mxfp8, no_mask) - `runner.py` - Configuration-based benchmark runner - `config_types.py` - Data types for benchmark configuration - `charts.py` - Chart generation utilities @@ -55,6 +56,9 @@ python -m benchmark.sdpa_benchmark_training.runner --config ltx2 # Run Qwen 3.5 benchmark suite (cuDNN bf16 at head_dim=256) python -m benchmark.sdpa_benchmark_training.runner --config qwen35 +# Run Autoregressive video DiT benchmark suite (short Q, long cached KV) +python -m benchmark.sdpa_benchmark_training.runner --config auto_regressive_dit + # Dry run (show what would be executed) python -m benchmark.sdpa_benchmark_training.runner --config llama --dry-run @@ -337,4 +341,14 @@ Runs were captured on GB200 and GB300 with cuDNN 9.23.0 and FAv4 4.0.0b15. ![Qwen 3.5 Causal on GB300](results/qwen35/gb300/qwen35_top_left.png) - `batch=2; num_q_heads=32; num_kv_heads=2; head_dim=256` — cuDNN BF16 at head_dim=256 on Blackwell +### GB300 - Autoregressive video DiT (short Q, long cached KV) +![Autoregressive DiT on GB300](results/auto_regressive_dit/gb300/auto_regressive_dit_no_mask.png) +- `batch=1; num_q_heads=9; num_kv_heads=9; head_dim=128; s_q ∈ {985..8192}; s_kv=62208` +- Forward-only (autoregressive inference). cuDNN 9.30.0 with prefill split-K on bf16/fp8/mxfp8; FAv4 BF16 swept over `num_splits ∈ {1, 2, 4, 8, 16, 32}` with the best annotated on each bar (`ks=`). FAv4 FP8/MXFP8 are absent — the CuTe-DSL FAv4 build rejects those input types. +- Reproduce with `python -m benchmark.sdpa_benchmark_training.bench_ar_dit_peak --out `. + +### GB200 - Autoregressive video DiT +![Autoregressive DiT on GB200](results/auto_regressive_dit/gb200/auto_regressive_dit_no_mask.png) +- Same configuration as the GB300 chart above, captured on GB200. + GB200 results are available under the same layout at `results//gb200/`. diff --git a/benchmark/sdpa_benchmark_training/bench_ar_dit_peak.py b/benchmark/sdpa_benchmark_training/bench_ar_dit_peak.py new file mode 100644 index 000000000..f1f2f5b8c --- /dev/null +++ b/benchmark/sdpa_benchmark_training/bench_ar_dit_peak.py @@ -0,0 +1,132 @@ +"""Peak-vs-peak SDPA bench for the autoregressive DiT shape. + +Sweeps FAv4 ``num_splits`` and reports the best, paired against cuDNN +(whatever the linked libcudnn provides; this script is intended to run +against a cuDNN build that has split-K so the comparison is split-KV +on both sides). + +CSV schema matches the rest of ``benchmark.sdpa_benchmark_training`` +plus an extra ``num_splits`` column so the per-seqlen winners are +visible. +""" + +from __future__ import annotations + +import argparse +import csv +import logging +import os +from pathlib import Path + +logging.basicConfig(level=logging.INFO, format="%(message)s") +log = logging.getLogger(__name__) + +import torch +import cudnn + +from benchmark.sdpa_benchmark_training.benchmark_single_sdpa import run_benchmark + +H = 9 +D = 128 +S_KV = 62208 +S_Q_LIST = [985, 1024, 2048, 4096, 8192] +DTYPES_CUDNN = ["bfloat16", "fp8", "mxfp8"] +FA4_SPLITS_SWEEP = [1, 2, 4, 8, 16, 32] +WARMUP = 5 +ITERS = 30 + + +def run(backend, dtype, s_q, fa4_num_splits=None): + return run_benchmark( + batch_size=1, + q_seqlen=s_q, + kv_seqlen=S_KV, + num_q_heads=H, + num_kv_heads=H, + head_dim=D, + data_type=dtype, + backend=backend, + attn_mask="no_mask", + profile_pass="fwd", + num_iterations=ITERS, + num_warmup_iterations=WARMUP, + skip_ref=True, + deterministic_bwd=False, + sliding_window_size=None, + fa4_num_splits=fa4_num_splits, + ) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--out", required=True, help="output CSV path") + args = p.parse_args() + + rows = [] + gpu_name = torch.cuda.get_device_name(0) + cudnn_be = cudnn.backend_version() + cudnn_fe = cudnn.__version__ + log.info(f"gpu={gpu_name} cudnn-frontend={cudnn_fe} backend={cudnn_be}") + + for s_q in S_Q_LIST: + # cuDNN on all dtypes + for dt in DTYPES_CUDNN: + try: + r = run("cudnn", dt, s_q) + log.info(f"cudnn {dt:>8} s_q={s_q:>5}: {r['time_ms']:.3f} ms {r['tflops']:.1f} TF") + rows.append( + dict( + backend="cudnn", + data_type=dt, + q_seqlen=s_q, + kv_seqlen=S_KV, + num_splits=0, + time_ms=r["time_ms"], + tflops=r["tflops"], + gpu_name=gpu_name, + cudnn_backend_version=cudnn_be, + ) + ) + except Exception as e: + log.info(f"cudnn {dt} s_q={s_q} FAILED: {e}") + + # FAv4 BF16, sweep num_splits + best = None + for ks in FA4_SPLITS_SWEEP: + try: + r = run("flash_attention_4", "bfloat16", s_q, fa4_num_splits=ks) + log.info(f" fa4 bfloat16 s_q={s_q:>5} num_splits={ks:>2}: {r['time_ms']:.3f} ms {r['tflops']:.1f} TF") + if best is None or r["time_ms"] < best["time_ms"]: + best = dict(r) + best["num_splits"] = ks + except Exception as e: + log.info(f" fa4 bfloat16 s_q={s_q} num_splits={ks} FAILED: {e}") + if best is not None: + log.info(f"fa4 bfloat16 s_q={s_q:>5} BEST num_splits={best['num_splits']:>2}: {best['time_ms']:.3f} ms {best['tflops']:.1f} TF") + rows.append( + dict( + backend="flash_attention_4", + data_type="bfloat16", + q_seqlen=s_q, + kv_seqlen=S_KV, + num_splits=best["num_splits"], + time_ms=best["time_ms"], + tflops=best["tflops"], + gpu_name=gpu_name, + cudnn_backend_version=cudnn_be, + ) + ) + + fields = ["backend", "data_type", "q_seqlen", "kv_seqlen", "num_splits", "time_ms", "tflops", "gpu_name", "cudnn_backend_version"] + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + with out.open("w", newline="") as f: + w = csv.DictWriter(f, fieldnames=fields) + w.writeheader() + for r in rows: + w.writerow(r) + log.info(f"wrote {len(rows)} rows -> {out}") + + +if __name__ == "__main__": + main() diff --git a/benchmark/sdpa_benchmark_training/benchmark_single_sdpa.py b/benchmark/sdpa_benchmark_training/benchmark_single_sdpa.py index e2c8fc991..c3dcbd163 100755 --- a/benchmark/sdpa_benchmark_training/benchmark_single_sdpa.py +++ b/benchmark/sdpa_benchmark_training/benchmark_single_sdpa.py @@ -117,6 +117,12 @@ def parse_args(): help="Number of warmup iterations to run before measuring performance", ) parser.add_argument("--verbose", action="store_true", help="Verbose output") + parser.add_argument( + "--fa4_num_splits", + default=None, + type=int, + help="FlashAttention-4 only: force num_splits (KV split count). " "Default is None (FA4 picks automatically).", + ) parser.add_argument( "--fwd_bwd", action="store_true", @@ -197,6 +203,7 @@ def run_benchmark( deterministic_bwd: bool = False, sliding_window_size: Optional[int] = None, verbose: bool = False, + fa4_num_splits: Optional[int] = None, ) -> Dict[str, Any]: """ Run a single SDPA benchmark. @@ -288,6 +295,8 @@ def run_benchmark( cmd.extend(["--sliding_window_size", str(sliding_window_size)]) if verbose: cmd.append("--verbose") + if fa4_num_splits is not None: + cmd.extend(["--fa4_num_splits", str(fa4_num_splits)]) # Run benchmark result = subprocess.run( @@ -1102,13 +1111,18 @@ def flash_attention_3_sdpa(query, key, value): def flash_attention_4_sdpa(query, key, value): window_size = (args.sliding_window_size, 0) if args.sliding_window_size else (None, None) + kwargs = dict( + causal=args.attn_mask != "no_mask", + window_size=window_size, + deterministic=args.deterministic_bwd, + ) + if args.fa4_num_splits is not None: + kwargs["num_splits"] = args.fa4_num_splits output, _ = flash_attn_interface.flash_attn_func( query, key, value, - causal=args.attn_mask != "no_mask", - window_size=window_size, - deterministic=args.deterministic_bwd, + **kwargs, ) return output diff --git a/benchmark/sdpa_benchmark_training/configs/auto_regressive_dit.py b/benchmark/sdpa_benchmark_training/configs/auto_regressive_dit.py new file mode 100644 index 000000000..13ec21ac2 --- /dev/null +++ b/benchmark/sdpa_benchmark_training/configs/auto_regressive_dit.py @@ -0,0 +1,56 @@ +""" +Autoregressive video DiT SDPA Benchmark Configuration + +Benchmarks the self-attention in an autoregressive (world-model / next-frame) +video DiT. Unlike the bidirectional video DiTs (LTX-2, Wan 2.2), an +autoregressive DiT generates one frame at a time conditioned on a growing +frame-history KV cache, so each attention call has a short query (the new +frame's patch tokens) against a long key/value tensor (concatenated past +frames): + + s_q in {985, 1024, 2048, 4096, 8192} # one new frame, varied resolution + s_kv = 62208 # cached history of past frames + num_heads = 9 + head_dim = 128 + +Tokens for the new frame = (H/p) * (W/p) for typical 720p..1440p patchified +inputs at patch sizes 1..4. The KV cache size of 62208 (= 486 * 128) is +representative of ~30s of generation history at the resolutions above. + +Within a single attention call, the new-frame query tokens attend to the +entire past-frame KV cache (no causality across frames is enforced inside +the operator; frame-level autoregression is the outer loop), so the +operator-level mask is ``no_mask``. + +Usage: + python -m benchmark.sdpa_benchmark_training.runner --config auto_regressive_dit + python -m benchmark.sdpa_benchmark_training.runner --config auto_regressive_dit --dry-run +""" + +from ..config_types import ModelPreset, BenchmarkConfig + +AR_DIT = ModelPreset( + name="auto_regressive_dit", + num_q_heads=9, + num_kv_heads=9, + head_dim=128, +) + +CONFIG = BenchmarkConfig( + name="auto_regressive_dit", + models=[AR_DIT], + seqlens=[ + (985, 62208), + (1024, 62208), + (2048, 62208), + (4096, 62208), + (8192, 62208), + ], + backends=["cudnn", "flash_attention_4"], + data_types=["bfloat16", "fp8", "mxfp8"], + attn_masks=["no_mask"], + profile_pass="fwd", + batch_size=1, + num_iterations=10, + output_dir="results", +) diff --git a/benchmark/sdpa_benchmark_training/results/auto_regressive_dit/gb200/auto_regressive_dit_no_mask.csv b/benchmark/sdpa_benchmark_training/results/auto_regressive_dit/gb200/auto_regressive_dit_no_mask.csv new file mode 100644 index 000000000..1c81aba8a --- /dev/null +++ b/benchmark/sdpa_benchmark_training/results/auto_regressive_dit/gb200/auto_regressive_dit_no_mask.csv @@ -0,0 +1,21 @@ +backend,data_type,q_seqlen,kv_seqlen,num_splits,time_ms,tflops,gpu_name,cudnn_backend_version +cudnn,bfloat16,985,62208,0,0.205,1380.0,NVIDIA GB200,92400 +cudnn,fp8,985,62208,0,0.157,1796.0,NVIDIA GB200,92400 +cudnn,mxfp8,985,62208,0,0.164,1717.0,NVIDIA GB200,92400 +flash_attention_4,bfloat16,985,62208,4,0.212,1332.0,NVIDIA GB200,92400 +cudnn,bfloat16,1024,62208,0,0.205,1429.0,NVIDIA GB200,92400 +cudnn,fp8,1024,62208,0,0.157,1870.0,NVIDIA GB200,92400 +cudnn,mxfp8,1024,62208,0,0.164,1785.0,NVIDIA GB200,92400 +flash_attention_4,bfloat16,1024,62208,4,0.211,1389.0,NVIDIA GB200,92400 +cudnn,bfloat16,2048,62208,0,0.373,1573.0,NVIDIA GB200,92400 +cudnn,fp8,2048,62208,0,0.294,1996.0,NVIDIA GB200,92400 +cudnn,mxfp8,2048,62208,0,0.307,1915.0,NVIDIA GB200,92400 +flash_attention_4,bfloat16,2048,62208,2,0.388,1513.0,NVIDIA GB200,92400 +cudnn,bfloat16,4096,62208,0,0.692,1697.0,NVIDIA GB200,92400 +cudnn,fp8,4096,62208,0,0.568,2066.0,NVIDIA GB200,92400 +cudnn,mxfp8,4096,62208,0,0.596,1971.0,NVIDIA GB200,92400 +flash_attention_4,bfloat16,4096,62208,1,0.673,1746.0,NVIDIA GB200,92400 +cudnn,bfloat16,8192,62208,0,1.333,1762.0,NVIDIA GB200,92400 +cudnn,fp8,8192,62208,0,1.129,2080.0,NVIDIA GB200,92400 +cudnn,mxfp8,8192,62208,0,1.181,1988.0,NVIDIA GB200,92400 +flash_attention_4,bfloat16,8192,62208,1,1.303,1802.0,NVIDIA GB200,92400 diff --git a/benchmark/sdpa_benchmark_training/results/auto_regressive_dit/gb200/auto_regressive_dit_no_mask.png b/benchmark/sdpa_benchmark_training/results/auto_regressive_dit/gb200/auto_regressive_dit_no_mask.png new file mode 100644 index 000000000..0ac474400 Binary files /dev/null and b/benchmark/sdpa_benchmark_training/results/auto_regressive_dit/gb200/auto_regressive_dit_no_mask.png differ diff --git a/benchmark/sdpa_benchmark_training/results/auto_regressive_dit/gb300/auto_regressive_dit_no_mask.csv b/benchmark/sdpa_benchmark_training/results/auto_regressive_dit/gb300/auto_regressive_dit_no_mask.csv new file mode 100644 index 000000000..cdc80f8cc --- /dev/null +++ b/benchmark/sdpa_benchmark_training/results/auto_regressive_dit/gb300/auto_regressive_dit_no_mask.csv @@ -0,0 +1,21 @@ +backend,data_type,q_seqlen,kv_seqlen,num_splits,time_ms,tflops,gpu_name,cudnn_backend_version +cudnn,bfloat16,985,62208,0,0.161,1752.0,NVIDIA GB300,92400 +cudnn,fp8,985,62208,0,0.112,2519.0,NVIDIA GB300,92400 +cudnn,mxfp8,985,62208,0,0.12,2359.0,NVIDIA GB300,92400 +flash_attention_4,bfloat16,985,62208,4,0.195,1451.0,NVIDIA GB300,92400 +cudnn,bfloat16,1024,62208,0,0.162,1813.0,NVIDIA GB300,92400 +cudnn,fp8,1024,62208,0,0.112,2619.0,NVIDIA GB300,92400 +cudnn,mxfp8,1024,62208,0,0.12,2447.0,NVIDIA GB300,92400 +flash_attention_4,bfloat16,1024,62208,4,0.194,1515.0,NVIDIA GB300,92400 +cudnn,bfloat16,2048,62208,0,0.305,1923.0,NVIDIA GB300,92400 +cudnn,fp8,2048,62208,0,0.212,2768.0,NVIDIA GB300,92400 +cudnn,mxfp8,2048,62208,0,0.226,2598.0,NVIDIA GB300,92400 +flash_attention_4,bfloat16,2048,62208,2,0.364,1613.0,NVIDIA GB300,92400 +cudnn,bfloat16,4096,62208,0,0.573,2050.0,NVIDIA GB300,92400 +cudnn,fp8,4096,62208,0,0.394,2978.0,NVIDIA GB300,92400 +cudnn,mxfp8,4096,62208,0,0.437,2687.0,NVIDIA GB300,92400 +flash_attention_4,bfloat16,4096,62208,1,0.571,2055.0,NVIDIA GB300,92400 +cudnn,bfloat16,8192,62208,0,1.126,2085.0,NVIDIA GB300,92400 +cudnn,fp8,8192,62208,0,0.782,3002.0,NVIDIA GB300,92400 +cudnn,mxfp8,8192,62208,0,0.867,2707.0,NVIDIA GB300,92400 +flash_attention_4,bfloat16,8192,62208,1,1.134,2071.0,NVIDIA GB300,92400 diff --git a/benchmark/sdpa_benchmark_training/results/auto_regressive_dit/gb300/auto_regressive_dit_no_mask.png b/benchmark/sdpa_benchmark_training/results/auto_regressive_dit/gb300/auto_regressive_dit_no_mask.png new file mode 100644 index 000000000..4caca3e73 Binary files /dev/null and b/benchmark/sdpa_benchmark_training/results/auto_regressive_dit/gb300/auto_regressive_dit_no_mask.png differ diff --git a/benchmark/sdpa_benchmark_training/runner.py b/benchmark/sdpa_benchmark_training/runner.py index 20fe905fb..7c934da15 100644 --- a/benchmark/sdpa_benchmark_training/runner.py +++ b/benchmark/sdpa_benchmark_training/runner.py @@ -166,8 +166,10 @@ def expand_config(self, config: BenchmarkConfig) -> Iterator[Dict[str, Any]]: det_values = [False] if profile_pass == "fwd" else list(config.deterministic_bwd) for det_bwd in det_values: - # MXFP8 constraints: cudnn-only - if data_type == "mxfp8": + # FP8/MXFP8 constraints: cudnn-only. FA4's cute interface + # rejects non-fp16/bf16 inputs outright, so emitting those + # cases just produces tracebacks in the CSV. + if data_type in ("fp8", "mxfp8"): if backend != "cudnn": continue