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

# SPDX-License-Identifier: Apache-2.0
"""mb1_prefill_matmul — prefill-shaped matmuls (LLM rows S=128..1024, ViT rows M=256/352/1792) on the p150a.

Closes research §6 unknown 4 (matmul efficiency at LLM rows S=116-556 and ViT rows M=704-2688 on the 11x10 grid).
Shapes follow blackhole-perf-model-v2 §2.2/§3:

* Qwen3 LLM layer (2048-d, 16 q / 8 kv heads x 128, ff 6144): fused qkv [2048, 4096], o_proj [2048, 2048],
  fused gate_up [2048, 12288], down [6144, 2048]; S in {128, 256, 512, 1024}.
* SigLIP ViT (1152-d, ff 4304): fused qkv [1152, 3456], out [1152, 1152], fc1 [1152, 4304], fc2 [4304, 1152];
  Qwen3-VL ViT (1024-d, ff 4096): [1024, 3072], [1024, 1024], [1024, 4096], [4096, 1024]; M in {256, 352, 1792}.

For every shape the script times ``ttnn.matmul`` (auto config), a 2D-mcast ``MatmulMultiCoreReuseMultiCastProgramConfig``
on the full 11x10 grid (tt_transformers ``matmul_config`` recipe) and ``ttnn.experimental.minimal_matmul`` with the
tt_dit heuristic blocking, for bf16 and bfloat8_b weights (bf16 activations), HiFi2 and LoFi fidelity, and fp32
accumulation on/off.  TFLOPS = 2*M*K*N / t; peak = 110 cores x 4096 FLOP/cycle x 1.35 GHz (/2 for HiFi2)
(perf-model-v2 §2.1).

Usage:
    with-device.sh python models/experimental/gr00t/benchmarks/mb1_prefill_matmul.py --out results/mb1_prefill_matmul.json
"""

from __future__ import annotations

import argparse
import json
import math
import os
import sys
import time
from typing import Any, Dict, List, Optional, Tuple

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import mb1_common as C  # noqa: E402

GRID = (11, 10)

LLM_SHAPES = [("qkv", 2048, 4096), ("o_proj", 2048, 2048), ("gate_up", 2048, 12288), ("down", 6144, 2048)]
SIGLIP_SHAPES = [("qkv", 1152, 3456), ("out", 1152, 1152), ("fc1", 1152, 4304), ("fc2", 4304, 1152)]
QWENVIT_SHAPES = [("qkv", 1024, 3072), ("out", 1024, 1024), ("fc1", 1024, 4096), ("fc2", 4096, 1024)]


def largest_divisor(n: int, cap: int = 8) -> int:
    for d in range(cap, 0, -1):
        if n % d == 0:
            return d
    return 1


def mcast2d_candidates(ttnn, M: int, K: int, N: int, fp32_acc: bool):
    """2D-mcast program configs to try for one shape (tt_transformers ``matmul_config`` recipe + row-fitted variants).

    * ``full``      : 11x10 grid, M split over the 10 rows, N over the 11 columns (per_core_M = ceil(Mt/10)).
    * ``full_T``    : the same grid with ``transpose_mcast=True`` (M over the 11 columns, N over the 10 rows).
    * ``rows=Mt``   : grid (11, Mt) with per_core_M = 1 when Mt <= 10 (no padded rows, fewer cores).
    * ``cols=Mt_T`` : grid (Mt, 10) transposed with per_core_M = 1 when Mt <= 11.
    Each with in0_block_w = largest divisor of Kt <= 8 (and 4 for the full grid).
    """
    Mt, Kt, Nt = math.ceil(M / 32), math.ceil(K / 32), math.ceil(N / 32)
    max_dst = 4 if fp32_acc else 8

    def subblock(per_core_M: int, per_core_N: int):
        h, w = 1, max(d for d in C.divisors(per_core_N) if d <= max_dst)
        for hh in (4, 2):
            if per_core_M % hh == 0:
                ww = max(d for d in C.divisors(per_core_N) if d * hh <= max_dst)
                if ww * hh > w * h:
                    h, w = hh, ww
                break
        return h, w

    ibw_main = largest_divisor(Kt, 8)
    cands = []
    layouts = [("full", (11, 10), False), ("full_T", (11, 10), True)]
    if Mt <= 10:
        layouts.append((f"rows={Mt}", (11, Mt), False))
    if Mt <= 11:
        layouts.append((f"cols={Mt}_T", (Mt, 10), True))
    for name, grid, transpose in layouts:
        gm, gn = (grid[0], grid[1]) if transpose else (grid[1], grid[0])  # cores along M / along N
        per_core_M = math.ceil(Mt / gm)
        per_core_N = math.ceil(Nt / gn)
        ibws = [ibw_main] + ([4] if name == "full" and ibw_main != 4 and Kt % 4 == 0 else [])
        for ibw in ibws:
            h, w = subblock(per_core_M, per_core_N)
            pc = ttnn.MatmulMultiCoreReuseMultiCastProgramConfig(
                compute_with_storage_grid_size=ttnn.CoreCoord(grid[0], grid[1]),
                in0_block_w=ibw,
                out_subblock_h=h,
                out_subblock_w=w,
                per_core_M=per_core_M,
                per_core_N=per_core_N,
                transpose_mcast=transpose,
                fused_activation=None,
                fuse_batch=True,
            )
            desc = dict(
                layout=name,
                grid=list(grid),
                transpose_mcast=transpose,
                in0_block_w=ibw,
                per_core_M=per_core_M,
                per_core_N=per_core_N,
                out_subblock=[h, w],
                cores_used=math.ceil(Mt / per_core_M) * math.ceil(Nt / per_core_N),
            )
            cands.append((f"mcast2d:{name}:ibw{ibw}", pc, desc))
    return cands


def run_shape(ttnn, device, args, group: str, name: str, M: int, K: int, N: int, w_dtype: str) -> List[Dict[str, Any]]:
    import torch

    rows: List[Dict[str, Any]] = []
    g = torch.Generator().manual_seed(0)
    a_t = torch.randn((1, 1, M, K), generator=g) * 0.5
    w_t = torch.randn((1, 1, K, N), generator=g) * (1.0 / math.sqrt(K))
    ref = (a_t.float() @ w_t.float())[0, 0]
    a = ttnn.from_torch(
        a_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device, memory_config=ttnn.DRAM_MEMORY_CONFIG
    )
    w = ttnn.from_torch(
        w_t,
        dtype=C.ttnn_dtype(ttnn, w_dtype),
        layout=ttnn.TILE_LAYOUT,
        device=device,
        memory_config=ttnn.DRAM_MEMORY_CONFIG,
    )
    flop = 2.0 * M * K * N
    wb = C.weight_bytes(K, N, w_dtype)
    families = args.families.split(",")
    fid_opts = [("HiFi2", False), ("LoFi", False), ("HiFi2", True)]
    if args.quick:
        fid_opts = fid_opts[:2]
    variants = []  # (family_label, builder) ; builder() -> (fn, desc) for a given compute config
    if "auto" in families:
        variants.append(("auto", None))
    if "mcast2d" in families:
        for label, pc, desc in mcast2d_candidates(ttnn, M, K, N, False):
            variants.append((label, (pc, desc)))
    if "minimal" in families:
        variants.append(("minimal", None))
    try:
        for label, payload in variants:
            for fidelity, fp32_acc in fid_opts:
                ck = C.compute_config(ttnn, fidelity, fp32_acc)
                family = label.split(":")[0]
                row: Dict[str, Any] = dict(
                    group=group,
                    name=name,
                    M=M,
                    K=K,
                    N=N,
                    w_dtype=w_dtype,
                    family=family,
                    variant=label,
                    fidelity=fidelity,
                    fp32_acc=fp32_acc,
                    gflop=round(flop / 1e9, 2),
                    weight_MB=round(wb / 1e6, 2),
                )
                try:
                    if family == "auto":

                        def fn():
                            return ttnn.matmul(a, w, memory_config=ttnn.DRAM_MEMORY_CONFIG, compute_kernel_config=ck)

                        row["config"] = {}
                    elif family == "mcast2d":
                        pc, desc = payload
                        if fp32_acc and desc["out_subblock"][0] * desc["out_subblock"][1] > 4:
                            # dst holds 4 fp32 tiles: rebuild the config with a smaller subblock
                            pc2 = [c for c in mcast2d_candidates(ttnn, M, K, N, True) if c[0] == label]
                            if not pc2:
                                raise RuntimeError("no fp32-acc variant of this config")
                            pc, desc = pc2[0][1], pc2[0][2]

                        def fn(pc=pc):
                            return ttnn.matmul(
                                a, w, program_config=pc, memory_config=ttnn.DRAM_MEMORY_CONFIG, compute_kernel_config=ck
                            )

                        row["config"] = desc
                    elif family == "minimal":
                        from models.tt_dit.utils.matmul import get_matmul_config

                        mc = get_matmul_config(M, K, N, ttnn.CoreCoord(GRID[0], GRID[1]), use_heuristic=True)

                        def fn():
                            return ttnn.experimental.minimal_matmul(
                                a, w, config=mc, memory_config=ttnn.DRAM_MEMORY_CONFIG, compute_kernel_config=ck
                            )

                        row["config"] = {
                            "M_block": mc.M_block_size,
                            "K_block": mc.K_block_size,
                            "N_block": mc.N_block_size,
                            "subblock": [mc.subblock_h, mc.subblock_w],
                        }
                    else:
                        raise ValueError(family)
                    if not args.no_check:
                        out = fn()
                        got = ttnn.to_torch(out)[0, 0].float()
                        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)),
                        pct_peak_trace=C.r(
                            C.pct(C.tflops(flop, t.trace_us), C.peak_tflops(fidelity, args.aiclk_mhz)), 1
                        ),
                        tflops_eager=C.r(C.tflops(flop, t.eager_us)),
                        weight_gbs_trace=C.r(C.gbs(wb, t.trace_us)),
                        error=t.error,
                    )
                except Exception as e:  # noqa: BLE001
                    row["error"] = C.short_error(e)
                rows.append(row)
                print(
                    f"  {group}/{name} M={M} K={K} N={N} {w_dtype} {label} {fidelity} fp32acc={fp32_acc}: trace {row.get('trace_us')} us = {row.get('tflops_trace')} TFLOPS ({row.get('pct_peak_trace')} %) pcc={row.get('pcc')} {row.get('error') or ''}"
                )
                sys.stdout.flush()
    finally:
        ttnn.deallocate(a)
        ttnn.deallocate(w)
    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_prefill_matmul.json"))
    ap.add_argument("--llm-s", default="128,256,512,1024")
    ap.add_argument("--vit-m", default="256,352,1792")
    ap.add_argument("--dtypes", default="bf16,bfp8")
    ap.add_argument("--families", default="auto,mcast2d,minimal")
    ap.add_argument("--groups", default="llm,siglip,qwenvit")
    ap.add_argument("--iters", type=int, default=16)
    ap.add_argument("--reps", type=int, default=5)
    ap.add_argument("--aiclk-mhz", type=float, default=C.DEFAULT_AICLK_MHZ)
    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:
        groups = args.groups.split(",")
        for w_dtype in args.dtypes.split(","):
            if "llm" in groups:
                for S in [int(x) for x in args.llm_s.split(",")]:
                    for name, K, N in LLM_SHAPES:
                        result["rows"].extend(run_shape(ttnn, device, args, "llm", name, S, K, N, w_dtype))
                    C.write_json(args.out, result)
            for grp, shapes in (("siglip", SIGLIP_SHAPES), ("qwenvit", QWENVIT_SHAPES)):
                if grp not in groups:
                    continue
                for M in [int(x) for x in args.vit_m.split(",")]:
                    for name, K, N in shapes:
                        result["rows"].extend(run_shape(ttnn, device, args, grp, name, M, K, N, w_dtype))
                    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)

    # best per shape summary
    bests = []
    seen = set()
    for rw in result["rows"]:
        key = (rw["group"], rw["name"], rw["M"], rw["w_dtype"], rw["fidelity"], rw["fp32_acc"])
        if key in seen:
            continue
        seen.add(key)
        cands = [
            x
            for x in result["rows"]
            if (x["group"], x["name"], x["M"], x["w_dtype"], x["fidelity"], x["fp32_acc"]) == key and x.get("trace_us")
        ]
        if cands:
            bests.append(min(cands, key=lambda x: x["trace_us"]))
    C.print_table(
        bests,
        [
            "group",
            "name",
            "M",
            "K",
            "N",
            "w_dtype",
            "fidelity",
            "fp32_acc",
            "variant",
            "trace_us",
            "tflops_trace",
            "pct_peak_trace",
            "pcc",
        ],
        "best variant per shape",
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())