Skip to content

41. Compiling Mountain Goat to GA-1: Fusion, Tiling and Why Decoding Is Memory-Bound

Mountain goats in space helmets on Pluto

What you will understand: what it takes to send real Mountain Goat programs, including a transformer block, to Chapter 40's model accelerator, and what the compiler's decisions are worth there. The back end reads the same mg dialect text that mgc mlir writes for the CPU path, and produces a GA-1 program: a matmul becomes Chapter 40's blocked schedule (block shape chosen by comparing the analytic model), a run of elementwise operations becomes one fused kernel that keeps its intermediates in the scratchpad, and reductions and transposes become tile loops. Every program it accepts is run on the simulator and its printed matrices are compared with mgc run's. Then the chapter measures three things on GA-1: what common-subexpression elimination, fusion and double buffering are each worth; where a transformer block spends its time; and why running one token at a time (decode) leaves the matrix unit nearly idle while many tokens at once (prefill) do not.

What you need to know first: Chapter 40 (GA-1, tiles, double buffering, the analytic model) and Chapters 29 and 30 (softmax, attention, the transformer block). The Mountain Goat compiler is used only for its front end (mgc mlir); nothing in the compiler changes.

This is a model, not a chip

GA-1's parameters are invented (Chapter 40 says so), so every cycle count here describes GA-1 only. What transfers to real hardware is the shape of the argument: fusion removes round trips to memory, a one-token matmul is limited by streaming the weights, and the same weights serve many tokens for free. The ratios are not calibrated against any device.

Compile and run

cd docs/part41/code                                   # plain Python 3; needs the Chapter 38 build for `mgc mlir` (../../part38/code/build.sh)
./show_program.py > show_program_out.txt             # kernels of the stable softmax, with and without fusion
./compare_with_cpu.py > compare_out.txt              # all book examples: GA-1 result against `mgc run` (about 2 minutes)
./studies.py > studies_out.txt                       # options, decode versus prefill, a transformer block (about 50 seconds)
./check_backend.py                                   # the checks of this page (about 45 seconds)
./backend_mutation.py > backend_mutation_out.txt     # 14 WRONG back ends the checker must catch (about 12 minutes)
cd ../../part15/code && ./run_lit.sh                 # the whole test suite
Every listing and output on this page comes from these commands.

Why this page shows a mutation run, and why 'caught' is the wanted result

backend_mutation.py writes broken versions of the back end and the simulator and expects the checker to fail on each. "Caught" is good; "NOT CAUGHT" would be a gap. Nothing on this page is a failing test.

What the back end does

It is a Python program over the text of the IR, not an MLIR pass (a limit, below). The steps:

  1. Parse and flatten. Read the func.func bodies of mgc mlir, and inline every call, so the program is one list of operations with unique names.
  2. CSE and dead code. Merge identical pure operations (the softmax definition mentions exp(x - row_max(x)) twice, so every call of it carries two identical copies) and drop what no print needs.
  3. Group. With fusion on, consecutive elementwise operations of the same shape become one group; a broadcast of a value made inside the group starts a new one.
  4. Emit, kernel by kernel. Each tensor lives in DRAM. A kernel loads the tiles it needs, computes, and stores what later kernels need.
mg operation GA-1 code
matmul Chapter 40's output-stationary blocked schedule; the block shape and double buffering are the ones the analytic model of Chapter 40 predicts fastest among all that fit the scratchpad
add sub mul div ge neg relu exp sqrt log, scalar operations, broadcast one vector-unit instruction per tile (exp, sqrt, log go to the special-function unit); a size-1 axis is broadcast by loading its only tile and expanding it
reduce (sum or max, either axis) accumulate tile by tile; the padding of a partial tile is hidden behind a mask (fill 0 for sum, minus infinity for max)
transpose load the tile at swapped coordinates, transpose it on the vector unit, store
constant, print a DRAM array; the printed value is read back after the run

Everything else (dynamic shapes, rank other than 2, reshape, permute, contract) is refused with a message rather than guessed at.

The whole back end:

#!/usr/bin/env python3
"""A back end that compiles Mountain Goat to GA-1 (Chapter 40's model accelerator). It reads the `mg` dialect text that the real front end writes (`mgc mlir prog.mg`: the same
IR that mg-opt lowers to LLVM), flattens the calls, and emits a GA-1 program with the data layout and schedule decisions made here:

  * every tensor lives in DRAM (a name -> rows of numbers); every kernel loads the tiles it needs into the scratchpad and stores its results;
  * matmul: Chapter 40's output-stationary blocked schedule, with the block shape and double buffering CHOSEN by comparing the analytic model of Chapter 40 over every shape that fits;
  * elementwise operations (add sub mul div ge neg relu exp sqrt log, scalar operations and broadcasts) are emitted tile by tile; with fusion on, a run of them becomes ONE kernel
    that keeps its intermediate tiles in the scratchpad and stores only what later code needs;
  * reductions (row or column, sum or max) accumulate tile by tile, hiding the padding of a partial tile behind a mask;  transpose is a vector-unit transpose of each tile;
  * options: cse (merge identical pure operations: the front end emits some twice), fuse, double_buffer (software-pipeline the loads of elementwise, reduction and transpose kernels).

Not a MLIR pass: a Python program over the text of the IR. Unsupported (it says so): dynamic shapes, rank other than 2, reshape, permute, contract."""
import math, os, re, sys
from dataclasses import dataclass, field
here = os.path.dirname(os.path.abspath(__file__)); sys.path.insert(0, os.path.join(here, "..", "..", "part40", "code")); sys.path.insert(0, here)
from ga import Machine
from ga_ext import Sim2
import schedules as S

class Unsupported(Exception): pass

# ------------------------------------------------------------------ parsing the mg dialect text
def shape_of(text):
    m = re.findall(r"tensor<([\dx?]+)xf64>", text)
    if not m: raise Unsupported("no tensor type in: " + text)
    dims = m[-1].split("x")
    if "?" in dims: raise Unsupported("dynamic shapes are not supported by the GA-1 back end")
    if len(dims) != 2: raise Unsupported(f"rank {len(dims)} tensors are not supported by the GA-1 back end")
    return (int(dims[0]), int(dims[1]))

@dataclass
class Op:
    res: str; op: str; args: list; attrs: dict; shape: tuple

def parse(text):
    funcs = {}; cur = None
    for raw in text.split("\n"):
        line = raw.strip()
        m = re.match(r"func\.func @(\w+)\((.*?)\) -> .*\{$", line)
        if m:
            cur = dict(params=[(a, shape_of(t)) for a, t in re.findall(r"%(\w+): (tensor<[\dx?]+xf64>)", m.group(2))], ops=[], prints=[], ret=None); funcs[m.group(1)] = cur; continue
        if cur is None or line in ("}", ""): continue
        m = re.match(r"%(\w+) = mg\.constant dense<(.*)> : (tensor<.*>)$", line)
        if m:
            data = eval(m.group(2).replace("e", "e")); cur["ops"].append(Op(m.group(1), "constant", [], {"data": data}, shape_of(m.group(3)))); continue
        m = re.match(r"%(\w+) = func\.call @(\w+)\((.*?)\) : .*-> (tensor<.*>)$", line)
        if m: cur["ops"].append(Op(m.group(1), "call", [a.strip()[1:] for a in m.group(3).split(",") if a.strip()], {"callee": m.group(2)}, shape_of(m.group(4)))); continue
        m = re.match(r"mg\.print %(\w+) :", line)
        if m: cur["ops"].append(Op("", "print", [m.group(1)], {}, (0, 0))); continue
        m = re.match(r"func\.return %(\w+) :", line)
        if m and cur is not None: cur["ret"] = m.group(1); continue
        m = re.match(r"%(\w+) = mg\.(\w+) ([^{:]*?)\s*(\{.*\})?\s*: (.*)$", line)
        if m:
            attrs = {}
            for k, v in re.findall(r'(\w+) = ("[^"]*"|[^,}]+?)(?= :|,|\})', m.group(4) or ""): attrs[k] = v.strip().strip('"')
            cur["ops"].append(Op(m.group(1), m.group(2), [a.strip()[1:] for a in m.group(3).split(",") if a.strip()], attrs, shape_of(m.group(5)))); continue
    return funcs

def flatten(funcs, entry="main"):
    """Inline every call: one flat list of ops with globally unique names, the constants, and the printed values (in order)."""
    ops = []; counter = [0]
    def inline(fname, argmap, prefix):
        env = dict(argmap)
        for o in funcs[fname]["ops"]:
            if o.op == "call":
                counter[0] += 1; sub = f"{prefix}c{counter[0]}_"
                ret = inline(o.attrs["callee"], {p[0]: env[a] for p, a in zip(funcs[o.attrs["callee"]]["params"], o.args)}, sub); env[o.res] = ret
            elif o.op == "print": ops.append(Op("", "print", [env[o.args[0]]], {}, (0, 0)))
            else:
                name = prefix + o.res; env[o.res] = name; ops.append(Op(name, o.op, [env[a] for a in o.args], dict(o.attrs), o.shape))
        return env.get(funcs[fname]["ret"])
    inline(entry, {}, "")
    return ops

# ------------------------------------------------------------------ middle: cse, dead code, fusion groups
PURE_EW = {"add", "sub", "mul", "div", "ge", "neg", "relu", "exp", "sqrt", "log", "scalar", "broadcast"}
def cse(ops):
    seen = {}; alias = {}; out = []
    for o in ops:
        o.args = [alias.get(a, a) for a in o.args]
        if o.op in ("constant", "print"): out.append(o); continue
        key = (o.op, tuple(o.args), tuple(sorted(o.attrs.items())), o.shape)
        if key in seen: alias[o.res] = seen[key]
        else: seen[key] = o.res; out.append(o)
    return out
def dce(ops):
    live = {a for o in ops if o.op == "print" for a in o.args}; keep = []
    for o in reversed(ops):
        if o.op == "print" or o.res in live: keep.append(o); live.update(o.args)
    return list(reversed(keep))

# ------------------------------------------------------------------ code generation
class Compiler:
    def __init__(self, machine, cse_on=True, fuse=True, double_buffer=True, matmul_block=None):
        self.m = machine; self.cse_on, self.fuse, self.db, self.matmul_block = cse_on, fuse, double_buffer, matmul_block
        self.prog = []; self.dram = {}; self.kernels = []; self.shapes = {}; self.useful_macs = 0     # useful_macs: the multiply-adds the PROGRAM asks for (padding excluded)
    def tiles(self, n): return math.ceil(n / self.m.T)
    # ---- entry point
    def compile(self, ops):
        if self.cse_on: ops = cse(ops)
        ops = dce(ops); self.ops = ops
        for o in ops:
            if o.op == "constant": self.dram[o.res] = [[float(x) for x in row] for row in o.attrs["data"]]; self.shapes[o.res] = o.shape
            elif o.op != "print": self.shapes[o.res] = o.shape
        uses = {}
        for k, o in enumerate(ops):
            for a in o.args: uses.setdefault(a, []).append(k)
        groups = []; cur = []
        def flush():
            nonlocal cur
            if cur: groups.append(("ew", cur)); cur = []
        for k, o in enumerate(ops):
            if o.op == "constant": continue
            if o.op in PURE_EW and self.fuse:
                inside = {c.res for c in cur}
                if cur and (o.shape != cur[-1].shape or (o.op == "broadcast" and o.args[0] in inside)): flush()
                cur.append(o)
            elif o.op in PURE_EW: flush(); groups.append(("ew", [o]))
            elif o.op == "print": flush()
            else: flush(); groups.append((o.op, [o]))
        flush()
        position = {o.res: k for k, o in enumerate(ops)}
        for kind, grp in groups:
            start = len(self.prog); first_op = grp[0]
            if kind == "ew": self.emit_group(grp, uses, position)
            elif kind == "matmul": self.emit_matmul(grp[0])
            elif kind == "reduce": self.emit_reduce(grp[0])
            elif kind == "transpose": self.emit_transpose(grp[0])
            else: raise Unsupported(f"operation mg.{kind} is not supported by the GA-1 back end")
            self.kernels.append(dict(kind=kind if kind != "ew" else "elementwise x%d" % len(grp), shape=grp[-1].shape, start=start, end=len(self.prog), ops=[g.op for g in grp]))
        self.prints = [o.args[0] for o in ops if o.op == "print"]
        return self
    def alloc(self, name, shape): self.dram[name] = [[0.0] * shape[1] for _ in range(shape[0])]
    # ---- matmul
    def choose_block(self, M, K, N):
        if self.matmul_block: return self.matmul_block
        Mt, Nt = self.tiles(M), self.tiles(N); best = None
        for bi in range(1, min(Mt, 8) + 1):
            for bj in range(1, min(Nt, 16) + 1):
                for db in ((True, False) if self.db else (False,)):
                    if bi * bj + (2 if db else 1) * (bi + bj) > self.m.slots: continue
                    est = S.predicted_double_buffered_cycles(M, K, N, self.m, bi, bj) if db else S.predicted_double_buffered_cycles(M, K, N, self.m, bi, bj) * 1.3
                    if best is None or est < best[0] - 1e-9: best = (est, bi, bj, db)
        if best is None: raise Unsupported("scratchpad too small for any matmul block")
        return best[1:]
    def emit_matmul(self, o):
        a, b = o.args; (M, K), (K2, N) = self.shapes[a], self.shapes[b]; bi, bj, db = self.choose_block(M, K, N); self.alloc(o.res, o.shape); self.useful_macs += M * K * N
        self.prog += S.blocked(M, K, N, self.m, bi, bj, db, a=a, b=b, c=o.res); self.last_block = (bi, bj, db)
        o.attrs["block"] = f"{bi}x{bj}{', double buffered' if db else ''}"
    # ---- software pipelining helper for per-tile loops: items = list of (loads, compute, stores) instruction lists using pool p
    def pipeline(self, steps):
        """steps: one function per tile, pool index -> (loads, compute, stores). Without double buffering: loads, compute, stores of each tile in turn. With it, the loads of tile
        t + 1 are emitted BEFORE the compute of tile t, into the other pool of slots, so the DMA queue can run ahead of the vector unit."""
        n = len(steps)
        if not n: return
        if not self.db:
            for f in steps: l, c, st = f(0); self.prog += l + c + st
            return
        parts = [None] * n; parts[0] = steps[0](0); self.prog += parts[0][0]
        for t in range(n):
            if t + 1 < n: parts[t + 1] = steps[t + 1]((t + 1) % 2); self.prog += parts[t + 1][0]
            self.prog += parts[t][1] + parts[t][2]
    # ---- elementwise group
    def emit_group(self, grp, uses, position):
        R, C = grp[-1].shape; Mt, Nt = self.tiles(R), self.tiles(C); names = {o.res for o in grp}; last = position[grp[-1].res]
        ext = []
        for o in grp:
            for a in o.args:
                if a not in names and a not in ext: ext.append(a)
        outs = [o.res for o in grp if any(self.ops[u].res not in names for u in uses.get(o.res, [])) or not uses.get(o.res)]
        for r in outs: self.alloc(r, self.shapes[r])
        # ---- liveness and slot assignment, once per group (relative slot numbers)
        lastuse = {}
        for k, o in enumerate(grp):
            for a in o.args: lastuse[a] = k
        for r in outs: lastuse[r] = len(grp)
        free = []; nxt = [0]; slot = {}
        def take():
            if free: return free.pop()
            nxt[0] += 1; return nxt[0] - 1
        for a in ext: slot[a] = take()
        body = []                                              # (kind, ...) with relative slots
        for k, o in enumerate(grp):
            srcs = [slot[a] for a in o.args]
            d = take(); slot[o.res] = d; body.append((o, srcs, d))
            for a in set(o.args):
                if lastuse.get(a) == k and a not in outs: free.append(slot[a])
            if lastuse.get(o.res, -1) <= k and o.res not in outs: free.append(d)
        K = nxt[0]
        if self.db and 2 * K > self.m.slots: raise Unsupported("scratchpad too small to double buffer an elementwise kernel")
        if K > self.m.slots: raise Unsupported("scratchpad too small for an elementwise kernel")
        def step(i, j):
            def f(pool):
                base = pool * K
                coords = lambda v: (i if self.shapes[v][0] == R else 0, j if self.shapes[v][1] == C else 0)      # a size-1 axis is broadcast: its only tile is tile 0
                loads = [("load", base + slot[a], a, *coords(a)) for a in ext]; comp = []
                for o, srcs, d in body:
                    s = [base + x for x in srcs]; d = base + d; kind = o.op
                    if kind in ("add", "sub", "mul", "div"): comp.append(("vop", kind, d, s[0], s[1]))
                    elif kind == "ge": comp.append(("vop2", "ge", d, s[0], s[1]))
                    elif kind in ("neg", "relu"): comp.append(("vop", kind, d, s[0], None))
                    elif kind in ("exp", "sqrt", "log"): comp.append((kind, d, s[0]))
                    elif kind == "scalar": comp.append(("vscalar", o.attrs["op"], d, s[0], float(o.attrs["value"].split()[0]), o.attrs["reversed"] == "true"))
                    elif kind == "broadcast":
                        sr, sc = self.shapes[o.args[0]]
                        if (sr, sc) == o.shape: comp.append(("vscalar", "mul", d, s[0], 1.0, False))
                        elif sc == 1 and sr == o.shape[0]: comp.append(("bcastcol", d, s[0]))
                        elif sr == 1 and sc == o.shape[1]: comp.append(("bcastrow", d, s[0]))
                        elif (sr, sc) == (1, 1): comp.append(("bcastall", d, s[0]))
                        else: raise Unsupported(f"broadcast {self.shapes[o.args[0]]} -> {o.shape}")
                    else: raise Unsupported(kind)
                stores = [("store", base + slot[r], r, i, j) for r in outs]
                return loads, comp, stores
            return f
        self.pipeline([step(i, j) for i in range(Mt) for j in range(Nt)])
    # ---- reductions
    def emit_reduce(self, o):
        (R, C) = self.shapes[o.args[0]]; axis = int(o.attrs["axis"].split()[0]); kind = o.attrs["kind"]; src = o.args[0]; self.alloc(o.res, o.shape)
        fill = 0.0 if kind == "sum" else float("-inf"); T = self.m.T; Mt, Nt = self.tiles(R), self.tiles(C)
        outer, inner = (Mt, Nt) if axis == 1 else (Nt, Mt)
        if (4 if self.db else 3) * 2 > self.m.slots: raise Unsupported("scratchpad too small")
        for x in range(outer):
            def steps():
                out = []
                for y in range(inner):
                    ti, tj = (x, y) if axis == 1 else (y, x)
                    def f(pool, ti=ti, tj=tj, y=y):
                        base = 2 + pool * 2; t, r = base, base + 1; loads = [("load", t, src, ti, tj)]; comp = []
                        vr, vc = min(T, R - ti * T), min(T, C - tj * T)
                        if vr < T or vc < T: comp.append(("masklen", t, vr, vc, fill))
                        comp.append(("rowred" if axis == 1 else "colred", kind, r, t))
                        comp.append(("vop", "add" if kind == "sum" else "max", 0, 0, r) if y else ("vscalar", "mul", 0, r, 1.0, False))
                        return loads, comp, []
                    out.append(f)
                return out
            st = steps(); self.pipeline(st)
            self.prog.append(("store", 0, o.res, x if axis == 1 else 0, 0 if axis == 1 else x))
    def emit_transpose(self, o):
        src = o.args[0]; R, C = o.shape; Mt, Nt = self.tiles(R), self.tiles(C); self.alloc(o.res, o.shape)
        steps = []
        for i in range(Mt):
            for j in range(Nt):
                def f(pool, i=i, j=j):
                    a, b = pool * 2, pool * 2 + 1
                    return [("load", a, src, j, i)], [("trans", b, a)], [("store", b, o.res, i, j)]
                steps.append(f)
        self.pipeline(steps)

# ------------------------------------------------------------------ running
def compile_text(text, machine, **options):
    return Compiler(machine, **options).compile(flatten(parse(text)))
def run_compiled(c, machine):
    sim = Sim2(machine, c.dram).run(c.prog); sim.useful_utilization = c.useful_macs / (machine.peak_macs_per_cycle * max(sim.cycles, 1)); return sim, [c.dram[p] for p in c.prints]

def kernel_table(c, sim):
    """Per kernel: its instructions, the cycles each engine was busy for it, and the span from its first start to its last finish (kernels overlap through the queues)."""
    rows = []
    for k in c.kernels:
        t = [x for x in sim.trace if k["start"] <= x[0] < k["end"]]
        if not t: continue
        busy = {e: sum(x[3] - x[2] for x in t if x[1] == e) for e in ("dma", "mxu", "vpu")}
        rows.append(dict(kind=k["kind"], shape=k["shape"], instructions=len(t), dma=busy["dma"], mxu=busy["mxu"], vpu=busy["vpu"], span=max(x[3] for x in t) - min(x[2] for x in t)))
    return rows

The simulator needed more operations than Chapter 40 gave it (ge, sqrt, log, broadcasts along either axis, column reductions, a transposer, and exp that overflows to infinity the way IEEE does instead of raising); they live in a subclass:

#!/usr/bin/env python3
"""GA-1 plus the vector-unit instructions the back end of Chapter 41 needs (Chapter 40's simulator is imported unchanged and extended by subclassing):

  ("vscalar", op, dst, a, k, reversed)   dst := a (op) k, or k (op) a when reversed;  op in add sub mul div
  ("vop2", op, dst, a, b)                 elementwise compare / min: op in ge min  (ge gives 1.0 or 0.0)
  ("exp", dst, a) ("sqrt", dst, a) ("log", dst, a)   on the special-function path (exp overrides Chapter 40's so that it overflows to infinity, as IEEE arithmetic does)
  ("bcastrow", dst, a)                    dst[i][j] := a[0][j]
  ("bcastall", dst, a)                    dst[i][j] := a[0][0]
  ("trans", dst, a)                       dst := transpose(a)
  ("colred", op, dst, a)                  reduce each column of a with op in sum max, broadcast down the column
  ("masklen", slot, rows, cols, fill)     slot entries outside the first rows x cols become `fill` (0 for a sum, -inf for a max): hides the padding of a partial tile from a reduction
"""
import math
from ga import Machine, Sim

class Sim2(Sim):
    def run(self, program):
        m = self.m; T = m.T; S = self.spad
        for ins in program:
            op = ins[0]; vec = math.ceil(m.tile_words / m.vpu_lanes) + 1; sfu = math.ceil(m.tile_words / m.sfu_lanes) + 1
            if op == "vscalar":
                _, f, d, a, k, rev = ins; self._run("vpu", vec, (a,), (d,)); A = S[a]
                g = {"add": lambda x, y: x + y, "sub": lambda x, y: x - y, "mul": lambda x, y: x * y, "div": lambda x, y: x / y if y != 0 else float("inf") if x > 0 else float("-inf") if x < 0 else float("nan")}[f]
                S[d] = [[g(k, x) if rev else g(x, k) for x in row] for row in A]
            elif op == "vop2":
                _, f, d, a, b = ins; self._run("vpu", vec, (a, b), (d,)); A, B = S[a], S[b]
                S[d] = [[(1.0 if A[i][j] >= B[i][j] else 0.0) if f == "ge" else min(A[i][j], B[i][j]) for j in range(T)] for i in range(T)]
            elif op == "exp":                                         # IEEE behaviour on overflow (Python's math.exp raises): exp of a large number is infinity, as on the CPU path
                _, d, a = ins; self._run("vpu", sfu, (a,), (d,))
                def safe_exp(x):
                    try: return math.exp(x)
                    except OverflowError: return float("inf")
                S[d] = [[safe_exp(x) for x in row] for row in S[a]]
            elif op in ("sqrt", "log"):
                _, d, a = ins; self._run("vpu", sfu, (a,), (d,)); fn = (lambda x: math.sqrt(x) if x >= 0 else float("nan")) if op == "sqrt" else (lambda x: math.log(x) if x > 0 else float("-inf") if x == 0 else float("nan"))
                S[d] = [[fn(x) for x in row] for row in S[a]]
            elif op == "bcastrow":
                _, d, a = ins; self._run("vpu", vec, (a,), (d,)); S[d] = [list(S[a][0]) for _ in range(T)]
            elif op == "bcastall":
                _, d, a = ins; self._run("vpu", vec, (a,), (d,)); S[d] = [[S[a][0][0]] * T for _ in range(T)]
            elif op == "trans":
                _, d, a = ins; self._run("vpu", vec, (a,), (d,)); S[d] = [[S[a][j][i] for j in range(T)] for i in range(T)]
            elif op == "colred":
                _, f, d, a = ins; self._run("vpu", vec + 2, (a,), (d,)); fn = sum if f == "sum" else max
                col = [fn(S[a][i][j] for i in range(T)) for j in range(T)]; S[d] = [list(col) for _ in range(T)]
            elif op == "masklen":
                _, s, r, c, fill = ins; self._run("vpu", 1, (s,), (s,)); S[s] = [[S[s][i][j] if i < r and j < c else fill for j in range(T)] for i in range(T)]
            else: super().run([ins])
        return self

One small program, kernel by kernel

Chapter 29's stable softmax (exp(x - row_max(x)) / row_sum(exp(x - row_max(x))), applied three times), with nothing turned on and then with CSE and fusion:

== no options: 34 kernels, 1854 cycles, 129 instructions
   reduce           (2, 1)   reduce
   elementwise x1   (2, 3)   broadcast
   elementwise x1   (2, 3)   sub
   elementwise x1   (2, 3)   exp
   reduce           (2, 1)   reduce
   elementwise x1   (2, 3)   broadcast
   elementwise x1   (2, 3)   sub
   elementwise x1   (2, 3)   exp
   reduce           (2, 1)   reduce
   elementwise x1   (2, 3)   broadcast
   elementwise x1   (2, 3)   div
   elementwise x1   (2, 3)   scalar
   reduce           (2, 1)   reduce
   elementwise x1   (2, 3)   broadcast
   elementwise x1   (2, 3)   sub
   elementwise x1   (2, 3)   exp
   reduce           (2, 1)   reduce
   elementwise x1   (2, 3)   broadcast
   elementwise x1   (2, 3)   sub
   elementwise x1   (2, 3)   exp
   reduce           (2, 1)   reduce
   elementwise x1   (2, 3)   broadcast
   elementwise x1   (2, 3)   div
   reduce           (2, 1)   reduce
   elementwise x1   (2, 3)   broadcast
   elementwise x1   (2, 3)   sub
   elementwise x1   (2, 3)   exp
   reduce           (2, 1)   reduce
   elementwise x1   (2, 3)   broadcast
   elementwise x1   (2, 3)   sub
   elementwise x1   (2, 3)   exp
   reduce           (2, 1)   reduce
   elementwise x1   (2, 3)   broadcast
   elementwise x1   (2, 3)   div
== cse + fusion: 13 kernels, 804 cycles, 66 instructions
   reduce           (2, 1)   reduce
   elementwise x3   (2, 3)   broadcast sub exp
   reduce           (2, 1)   reduce
   elementwise x2   (2, 3)   broadcast div
   elementwise x1   (2, 3)   scalar
   reduce           (2, 1)   reduce
   elementwise x3   (2, 3)   broadcast sub exp
   reduce           (2, 1)   reduce
   elementwise x2   (2, 3)   broadcast div
   reduce           (2, 1)   reduce
   elementwise x3   (2, 3)   broadcast sub exp
   reduce           (2, 1)   reduce
   elementwise x2   (2, 3)   broadcast div
-- the first fused kernel (broadcast sub exp) is 6 instructions (the 2x3 matrix is one tile):
    ('load', 0, 'c1_t1', 0, 0)
    ('load', 1, 't1', 0, 0)
    ('bcastcol', 2, 0)
    ('vop', 'sub', 0, 1, 2)
    ('exp', 2, 0)
    ('store', 2, 'c1_t4', 0, 0)

Without options the operations come out as 34 kernels and 1,854 cycles; with CSE and fusion, 13 kernels and 804. Each softmax is now reduce, a fused broadcast sub exp, reduce, a fused broadcast div. The fused kernel stores c1_t4 (the exponentials) because later kernels need them, and nothing else: its subtraction result lives and dies in a scratchpad slot.

Does it give the right answers?

compare_with_cpu.py runs every Mountain Goat example of Chapters 20 to 32 and the tour (excluding the large training programs, which are too slow for a Python simulator) on GA-1 and compares each printed matrix with mgc run's, to a tolerance of 1e-5 (relative, or absolute below 1; mgc run prints only six digits):

part20/code/examples/01_hello.mg                               same             0 cycles, useful MXU 0%
part20/code/examples/02_add_transpose.mg                       same             110 cycles, useful MXU 0%
part20/code/examples/03_dynamic.mg                             unsupported      dynamic shapes are not supported by the GA-1 back end
part20/code/examples/04_mismatch.mg                            unsupported      dynamic shapes are not supported by the GA-1 back end
part20/code/examples/05_gpu.mg                                 prints nothing   
part20/code/examples/06_transpose_twice.mg                     same             90 cycles, useful MXU 0%
part20/code/examples/07_expressions.mg                         same             240 cycles, useful MXU 0%
part20/code/examples/08_functions.mg                           same             90 cycles, useful MXU 0%
part20/code/examples/09_vectors.mg                             unsupported      dynamic shapes are not supported by the GA-1 back end
part20/code/examples/10_mixed.mg                               unsupported      dynamic shapes are not supported by the GA-1 back end
part21/code/examples/01_subtract.mg                            same             130 cycles, useful MXU 0%
part21/code/examples/02_hadamard.mg                            same             65 cycles, useful MXU 0%
part21/code/examples/03_divide.mg                              same             65 cycles, useful MXU 0%
part21/code/examples/04_matmul.mg                              same             152 cycles, useful MXU 0%
part21/code/examples/05_scalar_ops.mg                          same             270 cycles, useful MXU 0%
part21/code/examples/06_precedence.mg                          same             336 cycles, useful MXU 0%
part21/code/examples/07_linear_layer.mg                        same             141 cycles, useful MXU 0%
part21/code/examples/08_gram.mg                                same             197 cycles, useful MXU 0%
part21/code/examples/09_identity.mg                            same             216 cycles, useful MXU 0%
part21/code/examples/10_dynamic_matmul.mg                      unsupported      dynamic shapes are not supported by the GA-1 back end
part21/code/examples/11_dynamic_scalar.mg                      unsupported      dynamic shapes are not supported by the GA-1 back end
part21/code/examples/12_matmul_mismatch.mg                     unsupported      dynamic shapes are not supported by the GA-1 back end
part21/code/examples/13_gpu_matmul.mg                          prints nothing   
part21/code/examples/14_gpu_ops.mg                             prints nothing   
part21/code/examples/15_elementwise_mismatch.mg                unsupported      dynamic shapes are not supported by the GA-1 back end
part21/code/examples/15b_mul_mismatch.mg                       unsupported      dynamic shapes are not supported by the GA-1 back end
part21/code/examples/15c_div_mismatch.mg                       unsupported      dynamic shapes are not supported by the GA-1 back end
part21/code/examples/16_negate.mg                              unsupported      dynamic shapes are not supported by the GA-1 back end
part21/code/examples/17_dynamic_divide.mg                      unsupported      dynamic shapes are not supported by the GA-1 back end
part22/code/examples/01_relu.mg                                same             45 cycles, useful MXU 0%
part22/code/examples/02_row_and_col_sums.mg                    same             106 cycles, useful MXU 0%
part22/code/examples/03_max_and_mean.mg                        same             302 cycles, useful MXU 0%
part22/code/examples/04_broadcast_bias.mg                      same             195 cycles, useful MXU 0%
part22/code/examples/05_center_columns.mg                      same             216 cycles, useful MXU 0%
part22/code/examples/06_normalize_rows.mg                      same             171 cycles, useful MXU 0%
part22/code/examples/07_mlp_layer.mg                           same             146 cycles, useful MXU 0%
part22/code/examples/08_two_layers.mg                          same             421 cycles, useful MXU 0%
part22/code/examples/09_dynamic_reductions.mg                  unsupported      dynamic shapes are not supported by the GA-1 back end
part22/code/examples/10_gpu_layer.mg                           prints nothing   
part22/code/examples/11_dynamic_no_broadcast.mg                unsupported      dynamic shapes are not supported by the GA-1 back end
part22/code/examples/12_broadcast_left.mg                      same             275 cycles, useful MXU 0%
part22/code/examples/13_negative_max.mg                        same             106 cycles, useful MXU 0%
part23/code/examples/matmul_odd.mg                             unsupported      dynamic shapes are not supported by the GA-1 back end
part23/code/examples/matmul_static.mg                          same             76 cycles, useful MXU 1%
part24/code/examples/matmul_odd.mg                             unsupported      dynamic shapes are not supported by the GA-1 back end
part24/code/examples/matmul_static.mg                          same             76 cycles, useful MXU 1%
part28/code/examples/01_reshape_and_permute.mg                 unsupported      rank 3 tensors are not supported by the GA-1 back end
part28/code/examples/02_matmul_is_a_contraction.mg             same             76 cycles, useful MXU 0%
part28/code/examples/03_double_contraction.mg                  unsupported      rank 3 tensors are not supported by the GA-1 back end
part28/code/examples/04_axis_at_the_front.mg                   unsupported      rank 3 tensors are not supported by the GA-1 back end
part28/code/examples/05_axes_in_reverse_order.mg               unsupported      rank 3 tensors are not supported by the GA-1 back end
part28/code/examples/06_vectors_dot_and_outer.mg               unsupported      rank 1 tensors are not supported by the GA-1 back end
part28/code/examples/07_rank4_result.mg                        unsupported      rank 3 tensors are not supported by the GA-1 back end
part28/code/examples/08_in_a_function.mg                       unsupported      rank 3 tensors are not supported by the GA-1 back end
part29/code/examples/01_exp.mg                                 same             221 cycles, useful MXU 0%
part29/code/examples/02_softmax_naive.mg                       same             228 cycles, useful MXU 0%
part29/code/examples/03_softmax_overflow.mg                    same             232 cycles, useful MXU 0%
part29/code/examples/04_softmax_stable.mg                      same             804 cycles, useful MXU 0%
part29/code/examples/05_temperature.mg                         same             849 cycles, useful MXU 0%
part29/code/examples/06_attention.mg                           same             495 cycles, useful MXU 0%
part29/code/examples/07_causal_attention.mg                    same             515 cycles, useful MXU 0%
part29/code/examples/08_self_attention_block.mg                same             864 cycles, useful MXU 0%
part30/code/examples/01_layer_norm.mg                          same             513 cycles, useful MXU 0%
part30/code/examples/02_feed_forward.mg                        same             513 cycles, useful MXU 1%
part30/code/examples/03_tiny_transformer.mg                    same             6595 cycles, useful MXU 0%
part30/code/examples/04_last_token_changed.mg                  same             6595 cycles, useful MXU 0%
part31/code/examples/01_log.mg                                 same             409 cycles, useful MXU 0%
part31/code/examples/02_cross_entropy.mg                       same             261 cycles, useful MXU 0%
part31/code/examples/03_one_gradient_step.mg                   same             1431 cycles, useful MXU 0%
part32/code/examples/01_ge.mg                                  same             431 cycles, useful MXU 0%
tour/code/01_first_program.mg                                  same             0 cycles, useful MXU 0%
tour/code/02_literals.mg                                       same             45 cycles, useful MXU 0%
tour/code/03_let_and_shadowing.mg                              same             50 cycles, useful MXU 0%
tour/code/04_functions.mg                                      same             50 cycles, useful MXU 0%
tour/code/05_dynamic_types.mg                                  unsupported      dynamic shapes are not supported by the GA-1 back end
tour/code/06_operators.mg                                      same             476 cycles, useful MXU 0%
tour/code/07_shapes_flow.mg                                    same             227 cycles, useful MXU 0%
tour/code/08_column_statistics.mg                              same             306 cycles, useful MXU 0%
tour/code/09_polynomial.mg                                     same             76 cycles, useful MXU 0%
tour/code/10_rotation.mg                                       same             273 cycles, useful MXU 0%
tour/code/11_markov_chain.mg                                   same             456 cycles, useful MXU 0%
tour/code/12_graph_walks.mg                                    same             152 cycles, useful MXU 1%
tour/code/13_blur.mg                                           same             76 cycles, useful MXU 2%
tour/code/14_affine_layers.mg                                  same             282 cycles, useful MXU 0%
tour/code/15_batches.mg                                        unsupported      dynamic shapes are not supported by the GA-1 back end
tour/code/16_dynamic_bias_gotcha.mg                            unsupported      dynamic shapes are not supported by the GA-1 back end
tour/code/17_more_literals.mg                                  same             90 cycles, useful MXU 0%
tour/code/18_let_chains.mg                                     same             326 cycles, useful MXU 0%
tour/code/19_function_composition.mg                           same             483 cycles, useful MXU 0%
tour/code/20_one_function_many_shapes.mg                       unsupported      dynamic shapes are not supported by the GA-1 back end
tour/code/21_scalars_on_either_side.mg                         same             280 cycles, useful MXU 0%
tour/code/22_precedence.mg                                     same             608 cycles, useful MXU 0%

prints nothing: 4, same: 61, unsupported: 27

(The listing is long; the last line is the total.) 61 programs print the same matrices; 4 print nothing (they are the four GPU-path examples, which are compile-only on this machine and have nothing to run); 27 are refused, each for dynamic shapes or a rank other than 2. No program differs. "Same" means numerically close, not bit-identical: GA-1 sums reductions in tile order, and the CPU path sums in row order.

What the options are worth

studies.py compiles the same programs with each combination of options and runs them on GA-1 with 32 scratchpad slots. Every cell is a run whose result was compared with the CPU's:

STUDY 1. compiler options (32 tile slots). cycles for each combination; every program matches the CPU result
program                                              nothing                     + cse            + cse + fusion  + cse + double buffering                 all three
part29 04_softmax_stable                    1854                      1194 (1.55x)             804 (2.31x)            1194 (1.55x)             804 (2.31x)  
part29 08_self_attention_block              1214                       994 (1.22x)             864 (1.41x)             994 (1.22x)             864 (1.41x)  
part30 03_tiny_transformer                 12995                      9715 (1.34x)            6595 (1.97x)            9715 (1.34x)            6595 (1.97x)  
block 8 tokens x 16                         7484                      5809 (1.29x)            3794 (1.97x)            5419 (1.38x)            3564 (2.10x)  
block 32 tokens x 32                       60336                     46604 (1.29x)           29724 (2.03x)           41408 (1.46x)           26078 (2.31x)  

STUDY 2. one feed-forward network, relu(x @ w1 + b1) @ w2 + b2, width 64 -> 256 -> 64, for T tokens at once (64 tile slots, all options on)
tokens    cycles  cycles/token  useful MXU  DMA busy  weight tiles moved
     1     14400         14400       3.6%      100%                 720
     2     14400          7200       7.1%      100%                 720
     4     14400          3600      14.2%      100%                 720
     8     14400          1800      28.4%      100%                 720
    16     18880          1180      43.4%      100%                 944
    32     28904           903      56.7%       97%                1408
    64     57808           903      56.7%       97%                2816

STUDY 3. one transformer block, 64 tokens, width 64, feed-forward 128 (32 tile slots, all options on): the kernels, in program order
total 127006 cycles; useful MXU utilization 32%; DMA busy 100%, MXU busy 32%, VPU busy 14%
kernel kind    count  DMA busy  MXU busy  VPU busy  sum of spans
matmul             8     62720     40960       576         63765
elementwise       14     52480         0     12480         52510
reduce             6      8640         0      4608          9216
transpose          1      2560         0       320          2560

real    0m49.517s
user    0m44.102s
sys 0m4.740s

Study 1, the options (the first table):

  • CSE alone is worth 1.2× to 1.55×: the front end evaluates repeated subexpressions repeatedly, and on GA-1 every repetition is a set of loads, vector instructions and stores.
  • Fusion on top takes the stable softmax from 1,194 to 804 cycles and the 8-token block from 5,809 to 3,794: intermediates no longer make a round trip through DRAM.
  • Double buffering alone, after CSE, gives nothing for the softmax (every kernel is a single tile, so there is nothing to overlap) and 5,809 to 5,419 for the block. With all three options the block reaches 2.10×, and the 32-token, width-32 block 2.31×.

Decode and prefill

Study 2 is the point of the chapter for anyone interested in inference hardware. A feed-forward network (64 to 256 to 64) is compiled for 1, 2, 4, ... 64 tokens at once:

  • 1, 2, 4 and 8 tokens all take 14,400 cycles. The weights are 720 tiles that must stream through the DMA whatever the number of tokens, and the DMA is 100% busy; the extra tokens ride along in tile rows that were being padded anyway. So the cycles per token fall 8× and the useful matrix-unit utilization (padding excluded) is exactly proportional to the tokens: 3.6%, 7.1%, 14.2%, 28.4%.
  • Past 8 tokens the extra tokens need extra tile rows, so the cycles start to grow (18,880 for 16 tokens; 28,904 for 32) and the utilization rises more slowly (43%, then 57%; 64 tokens gives 57% again). I did not break down what limits the schedule at 32 and 64 tokens; the DMA is still 97% busy there.

This is the decode-versus-prefill picture in miniature: producing one token at a time is limited by moving the weights, and the cure (batching, or processing a long prompt at once) works because the weights are reused. It is a statement about GA-1's ratios; a real device's balance point would be elsewhere.

Where a block spends its time

Study 3 compiles one 64-token transformer block (width 64, feed-forward 128) and splits the busy cycles by kernel type:

  • Total 127,006 cycles; useful MXU utilization 32%; the DMA is busy 100% of the time, the MXU 32%, the vector unit 14%.
  • The 8 matmuls take 62,720 DMA cycles and 40,960 matrix-unit cycles; the 14 elementwise kernels take 52,480 DMA cycles (and only 12,480 of vector work). About 41% of the DMA time is elementwise kernels that do almost no arithmetic. That is what the fusion of this chapter shrinks and what a back end with cross-kernel on-chip fusion (flash-attention-style) would shrink further; this back end does not do that, so the number is an upper bound on what remains to be gained, not a prediction.

Tests

One new lit group test/ga41/ (the suite is now 153 tests). check_backend.py checks, in order:

  1. Agreement with the CPU: the 61 programs print the same matrices; the 27 refusals are for the two stated reasons only.
  2. All option combinations: each of the 8 combinations of CSE / fusion / double buffering gives the right answer on three programs.
  3. Other machines: tiles of 4, 8 and 16 words and scratchpads of 24 and 64 slots give the right answer, so nothing depends on the tile being 8.
  4. What the options do: pinned cycle counts (1,854 → 1,194 → 804; 7,484 → 3,564); fusion turns 22 kernels into 13 (the 9-to-4 figure one might first guess is wrong: the stored values the later code needs stay as kernels); double buffering changes nothing for one-tile kernels and helps the block.
  5. Decode and prefill: 14,400 cycles for 1, 2, 4 and 8 tokens; utilization proportional to the tokens.
  6. Unsupported: dynamic shapes and rank 3 are refused with a message.
-- 1. agreement with the CPU path on the book's examples
  ok   61 programs print the same matrices as the CPU path, 0 differ
  ok   the 27 rejected programs are rejected for dynamic shapes or rank, nothing else
-- 2. all option combinations
  ok   all 8 combinations of cse / fusion / double buffering give the CPU's answer on 3 programs
-- 3. other machines
  ok   tiles of 4, 8 and 16 and scratchpads of 24 and 64 slots all give the CPU's answer
-- 4. what the options do
  ok   softmax (stable): 1854 cycles with nothing, 1194 with cse, 804 with cse and fusion
  ok   8-token block: 7484 cycles with nothing, 3564 with all three options
  ok   fusion turns 22 kernels into 13; each softmax becomes reduce, elementwise x3, reduce, elementwise x2
  ok   double buffering changes nothing for one-tile kernels (softmax) and helps the block
-- 5. decode and prefill
  ok   1, 2, 4 and 8 tokens all take 14400 cycles (the weights are streamed once for all of them)
  ok   useful matrix-unit utilization is proportional to the tokens: 1: 3.6%, 2: 7.1%, 4: 14.2%, 8: 28.4%
-- 6. unsupported programs
  ok   dynamic shapes and rank 3 are refused with a message: dynamic shapes are not supported by the GA-1 back end; rank 3 tensors are not supported by the GA-1 back end
all checks pass

Are the checks good enough? Fourteen wrong back ends.

NOTE: this script deliberately breaks the back end. 'caught' lines are EXPECTED: they show the checker can detect the mistake.
baseline (nothing broken): checker passes (expected)
a broadcast operand is fetched at the output tile's own coordinates (not projected onto its size-1 axis): caught by 2 failing check(s) (the option combinations, other machines)
a max-reduction pads a partial tile with 0 instead of -infinity: caught by 2 failing check(s) (agreement with the CPU on the examples, the option combinations)
common-subexpression elimination ignores the operation's attributes (two scalar ops with different constants look equal): caught by 1 failing check(s) (agreement with the CPU on the examples)
transpose loads the tile at the output's coordinates (the tiles are not swapped): caught by 2 failing check(s) (the option combinations, other machines)
a scalar operation ignores 'reversed' (k - x computed as x - k): caught by 1 failing check(s) (agreement with the CPU on the examples)
matmul is compiled with its operands swapped: caught by 3 failing check(s) (agreement with the CPU on the examples, the option combinations, other machines)
a fused kernel stores only its last value (values used later by other kernels are lost): caught (the checker did not finish: KeyError: 't3')
double buffering uses the same slots for both tiles (the next tile's loads overwrite the current tile): caught by 3 failing check(s) (the option combinations, other machines, block cycles)
the matmul block is the first that fits, not the best by the model: caught by 2 failing check(s) (block cycles, decode cycles)
useful utilization counts the padding: caught by 1 failing check(s) (utilization)
a column vector is broadcast as a row (bcastrow where bcastcol belongs): caught by 3 failing check(s) (agreement with the CPU on the examples, the option combinations, other machines)
masklen in the simulator ignores its fill value: caught by 2 failing check(s) (agreement with the CPU on the examples, the option combinations)
the simulator's vscalar swaps its operands when reversed is false: caught by 3 failing check(s) (agreement with the CPU on the examples, the option combinations, other machines)
exp overflow raises instead of giving infinity: caught (the checker did not finish: OverflowError: math range error)

restored: ga_backend.py and ga_ext.py are back to their originals

All 14 are caught. Two are caught by a crash rather than by a failing comparison (a fused kernel that stores only its last value loses a tile a later kernel loads; an exp that raises on overflow stops the run). A crash is a valid detection but a weaker one, and the page says so rather than counting it as a clean mismatch. The mutants that need a particular program to show up (the max-reduction padded with 0 instead of minus infinity; CSE ignoring attributes) are caught only because the book's own examples happen to contain them; a program whose partial tile holds all negative values for a max is the kind the examples might not have included.

Limits and what is not established

  • A model, not a chip. Cycle counts are GA-1's; no real device was consulted.
  • Not an MLIR pass. The back end parses printed IR with regular expressions. It would not survive a change in the printing format (the check would notice), and it shares nothing with the passes of Chapters 14 to 38.
  • No dynamic shapes, no rank other than 2 (27 of the book's examples are refused for this).
  • No cross-kernel on-chip fusion. Attention's scores are written to DRAM and read back; a back end that kept them in the scratchpad (as flash attention does) would remove the traffic, and this one cannot.
  • Tiny sizes. The blocks are 8 to 64 tokens wide; the simulator is in Python and large programs take minutes. Trends at real model sizes (thousands of wide matrices, scratchpads of megabytes) are not established.
  • Partial tiles cost full tiles. A 2 × 3 matrix takes a whole 8 × 8 tile of work; that is why the "useful utilization" excludes padding and why decode looks so poor.
  • The block choice rests on Chapter 40's analytic model, which was within 4% to 6% of the simulator on the sweep there, but was not re-validated on every matmul here.
  • Numerics: results agree to 1e-5 (the CPU path prints six digits), not bit for bit; GA-1 sums reductions in tile order.
  • The large training programs were not compiled: the comparison skips every example of 400 lines or more, because simulating them in Python takes too long.

Reproducing

cd docs/part41/code
./show_program.py > show_program_out.txt && ./compare_with_cpu.py > compare_out.txt && ./studies.py > studies_out.txt
./check_backend.py && ./backend_mutation.py > backend_mutation_out.txt
cd ../../part15/code && ./run_lit.sh          # 153 tests

Chapter summary

  • A Python back end over the mg dialect text compiles Mountain Goat to GA-1: matmul to Chapter 40's blocked schedule with a model-chosen block, runs of elementwise operations to fused kernels, reductions with masks, transposes on the vector unit.
  • 61 of the book's examples give the CPU's matrices; 27 are refused for dynamic shapes or rank; 4 print nothing; none differ.
  • CSE, fusion and double buffering together are worth 2.1× on an 8-token block; fusion alone cuts the stable softmax from 1,194 to 804 cycles.
  • The feed-forward network takes the same 14,400 cycles for 1 to 8 tokens: decode is limited by streaming the weights, and batching is nearly free until the matrix unit fills.
  • About 41% of a block's DMA time goes to elementwise kernels; what remains to fuse would need on-chip fusion across kernels, which this back end does not have.
  • Fourteen broken back ends are caught (two only by a crash).

Self-check questions

Each answer is collapsed; try the question first.

  1. Why does the front end's output contain the same exp(x - row_max(x)) several times, and what does CSE do on GA-1?

    Answer

    The front end expands a call to the softmax definition each time it is used, and the definition mentions exp(x - row_max(x)) twice. The IR therefore holds identical operations. On GA-1 each repetition is loads, vector instructions and stores; merging them takes the stable softmax from 1,854 to 1,194 cycles.

  2. What does fusion save, and what does it not?

    Answer

    It saves the stores and loads of intermediates that nothing else needs: the result of the subtraction stays in a scratchpad slot. It does not save stores of values that later kernels read (the exponentials, needed by the sum and the division), which is why 22 kernels become 13 rather than fewer.

  3. Why do 1, 2, 4 and 8 tokens take the same number of cycles through the feed-forward network?

    Answer

    The weight tiles (720 of them) dominate and are streamed once for any number of tokens up to a tile's 8 rows; the matrix unit's rows for the extra tokens were being padded anyway. The DMA is 100% busy, so more rows cost nothing until they exceed 8.

  4. The tool reports "useful" utilization separately from utilization. Why?

    Answer

    A padded tile multiplies zeros at full speed, so counting padded multiply-adds makes a one-token matmul look as busy as an eight-token one. Useful utilization counts only the multiply-adds the program asked for, so it rises in proportion to the tokens: 3.6%, 7.1%, 14.2%, 28.4%. The mutant that counts padding is caught.

  5. Two of the fourteen mutants are "caught" by a crash. Why does the page call that weaker?

    Answer

    A crash proves the checker noticed something was badly wrong, but not that it would notice a subtle wrong answer of the same kind. A mismatch against the CPU result is evidence the checker compares numbers.

  6. What would a back end need in order to remove the 41% of DMA time spent on elementwise kernels?

    Answer

    Fusion across kernels with matmuls and reductions: keeping an attention score tile in the scratchpad from the matmul through the softmax into the next matmul. This back end ends every group at a matmul or a reduction and stores the result to DRAM, so it cannot.

  7. Name three limits of the claims on this page.

    Answer

    Any of: the machine is invented, so cycle counts are not a prediction for a real device; the sizes are tiny; the back end is not an MLIR pass and parses printed text; nothing with dynamic shapes or rank other than 2 is accepted; the tile-block choice trusts Chapter 40's analytic model; results are close to, not identical to, the CPU's.