code/models/experimental/gr00t/benchmarks/mb1_sdpa.py

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