# SPDX-License-Identifier: Apache-2.0 """mb1_sdpa — ``ttnn.transformer.scaled_dot_product_attention`` timings for the GR00T attention shapes. Shape classes (arch synthesis §1 / perf-model-v2 §2.3): * ViT (SigLIP / Qwen3-VL vision): non-causal, B = images as batch (1, 2, 4), 16 heads, S in {256, 352, 448}, head_dim 96 (SigLIP 72 padded to 96, ttnn-blocks-vlm.md line 114) and 64 (Qwen3-VL ViT). * LLM (Qwen3 decoder prefill): causal GQA, 16 q / 8 kv heads, S in {128, 256, 512, 1024}, head_dim 128. * DiT (action head): 32 heads, Sq = 64 (49/51/41 padded), Sk in {64, 128, 320, 512}, head_dim 64 (48 padded), non-causal (self-attention Sk = 64, cross-attention Sk = VL tokens). For each shape the default program config and a sweep of ``SDPAProgramConfig(q_chunk_size, k_chunk_size)`` over {32, 64, 128, 256, 512} are timed (trace replay of ``--iters`` calls); FLOP = 4 * B * H * Sq * Sk * d (halved for causal). Usage: with-device.sh python models/experimental/gr00t/benchmarks/mb1_sdpa.py --out results/mb1_sdpa.json """ from __future__ import annotations import argparse import itertools import math import os import sys import time from typing import Any, Dict, List, Optional sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) import mb1_common as C # noqa: E402 CHUNKS = [32, 64, 128, 256, 512] def cases(args) -> List[Dict[str, Any]]: out = [] if "vit" in args.groups: for B, S, d in itertools.product([1, 2, 4], [256, 352, 448], [96, 64]): out.append(dict(group="vit", B=B, H=16, Hkv=16, Sq=S, Sk=S, d=d, causal=False)) if "llm" in args.groups: for S in [128, 256, 512, 1024]: out.append(dict(group="llm", B=1, H=16, Hkv=8, Sq=S, Sk=S, d=128, causal=True)) if "dit" in args.groups: for Sk in [64, 128, 320, 512]: out.append(dict(group="dit", B=1, H=32, Hkv=32, Sq=64, Sk=Sk, d=64, causal=False)) if args.quick: out = [c for c in out if c["B"] == 1][:6] return out def run_case(ttnn, device, args, case: Dict[str, Any]) -> List[Dict[str, Any]]: import torch B, H, Hkv, Sq, Sk, d = case["B"], case["H"], case["Hkv"], case["Sq"], case["Sk"], case["d"] g = torch.Generator().manual_seed(0) q_t = torch.randn((B, H, Sq, d), generator=g) k_t = torch.randn((B, Hkv, Sk, d), generator=g) v_t = torch.randn((B, Hkv, Sk, d), generator=g) q = ttnn.from_torch(q_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device, pad_value=0.0) k = ttnn.from_torch(k_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device, pad_value=0.0) v = ttnn.from_torch(v_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device, pad_value=0.0) ck = C.compute_config(ttnn, args.fidelity, args.fp32_acc) flop = 4.0 * B * H * Sq * Sk * d * (0.5 if case["causal"] else 1.0) # reference if Hkv != H: rep = H // Hkv k_ref = k_t.reshape(B, Hkv, 1, Sk, d).repeat(1, 1, rep, 1, 1).reshape(B, H, Sk, d) v_ref = v_t.reshape(B, Hkv, 1, Sk, d).repeat(1, 1, rep, 1, 1).reshape(B, H, Sk, d) else: k_ref, v_ref = k_t, v_t ref = torch.nn.functional.scaled_dot_product_attention(q_t, k_ref, v_ref, is_causal=case["causal"]) rows = [] grid = device.compute_with_storage_grid_size() combos: List[Optional[Dict[str, int]]] = [None] q_opts = [c for c in CHUNKS if c <= max(32, Sq)] k_opts = [c for c in CHUNKS if c <= max(32, Sk)] if args.quick: q_opts, k_opts = q_opts[-2:], k_opts[-2:] combos += [dict(q=qc, k=kc) for qc in q_opts for kc in k_opts] try: for combo in combos: row = dict(case) row.update(fidelity=args.fidelity, fp32_acc=args.fp32_acc, gflop=round(flop / 1e9, 3)) try: pc = None if combo is not None: pc = ttnn.SDPAProgramConfig( compute_with_storage_grid_size=grid, q_chunk_size=combo["q"], k_chunk_size=combo["k"], exp_approx_mode=args.exp_approx, ) row["q_chunk"], row["k_chunk"] = combo["q"], combo["k"] else: row["q_chunk"], row["k_chunk"] = "default", "default" def fn(): return ttnn.transformer.scaled_dot_product_attention( q, k, v, is_causal=case["causal"], program_config=pc, compute_kernel_config=ck ) if not args.no_check: out = fn() got = ttnn.to_torch(out).float()[:, :, :Sq, :] row["pcc"] = round(pcc(got, ref), 5) ttnn.deallocate(out) t = C.measure(ttnn, device, fn, iters=args.iters, warmup=2, reps=args.reps) row.update( eager_us=C.r(t.eager_us), trace_us=C.r(t.trace_us), tflops_trace=C.r(C.tflops(flop, t.trace_us)), error=t.error, ) except Exception as e: # noqa: BLE001 row["error"] = C.short_error(e) rows.append(row) print( f" {case['group']} B={B} H={H}/{Hkv} Sq={Sq} Sk={Sk} d={d} q={row['q_chunk']} k={row['k_chunk']}: trace {row.get('trace_us')} us eager {row.get('eager_us')} pcc={row.get('pcc')} {row.get('error') or ''}" ) finally: for t_ in (q, k, v): ttnn.deallocate(t_) return rows def pcc(a, b) -> float: import torch a = a.flatten().double() b = b.flatten().double() return float(torch.corrcoef(torch.stack([a, b]))[0, 1]) def main() -> int: ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--out", default=os.path.join(C.default_out_dir(), "mb1_sdpa.json")) ap.add_argument("--groups", default="vit,llm,dit") ap.add_argument("--fidelity", default="HiFi2") ap.add_argument("--fp32-acc", action="store_true") ap.add_argument("--exp-approx", action="store_true", default=False) ap.add_argument("--iters", type=int, default=16) ap.add_argument("--reps", type=int, default=5) ap.add_argument("--quick", action="store_true") ap.add_argument("--no-check", action="store_true") args = ap.parse_args() import ttnn device = C.open_device(ttnn) result: Dict[str, Any] = { "env": C.env_report(), "device": C.device_report(ttnn, device), "args": vars(args), "rows": [], } t0 = time.perf_counter() try: for case in cases(args): result["rows"].extend(run_case(ttnn, device, args, case)) C.write_json(args.out, result) finally: result["wall_s"] = round(time.perf_counter() - t0, 1) C.write_json(args.out, result) ttnn.close_device(device) bests = [] seen = set() for rw in result["rows"]: key = (rw["group"], rw["B"], rw["Sq"], rw["Sk"], rw["d"]) if key in seen: continue seen.add(key) cands = [ x for x in result["rows"] if (x["group"], x["B"], x["Sq"], x["Sk"], x["d"]) == key and x.get("trace_us") ] if cands: bests.append(min(cands, key=lambda x: x["trace_us"])) C.print_table( bests, [ "group", "B", "H", "Hkv", "Sq", "Sk", "d", "q_chunk", "k_chunk", "trace_us", "eager_us", "tflops_trace", "pcc", ], "best chunk config per SDPA shape", ) return 0 if __name__ == "__main__": sys.exit(main())