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