bench: add autoregressive video DiT SDPA config + GB200/GB300 results…

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 <path>`.
+
+### 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/<config>/gb200/`.

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()

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
 

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",
+)

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

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

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