gstoner · gstoner · Jul 10, 2026 · Jul 10, 2026 · Jul 10, 2026
diff --git a/benchmarks/baselines/rocm_gfx1151_hot_paths.json b/benchmarks/baselines/rocm_gfx1151_hot_paths.json
@@ -7,82 +7,106 @@
       "shape": "512x512x512",
       "dtype": "f16",
       "mode": "wmma",
-      "median_ms": 0.6322,
-      "max_latency_ms": 1.2643,
-      "achieved_tflops": 0.4246,
-      "pct_peak": 0.00715,
-      "attainment_floor": 0.00358
+      "median_ms": 0.6391,
+      "max_latency_ms": 1.2781,
+      "achieved_tflops": 0.42,
+      "pct_peak": 0.00707,
+      "attainment_floor": 0.00353
     },
     {
       "op": "matmul",
       "shape": "1024x1024x1024",
       "dtype": "f16",
       "mode": "wmma",
-      "median_ms": 5.3514,
-      "max_latency_ms": 10.7028,
-      "achieved_tflops": 0.4013,
-      "pct_peak": 0.00676,
-      "attainment_floor": 0.00338
+      "median_ms": 5.6278,
+      "max_latency_ms": 11.2557,
+      "achieved_tflops": 0.3816,
+      "pct_peak": 0.00642,
+      "attainment_floor": 0.00321
     },
     {
       "op": "matmul",
       "shape": "2048x2048x2048",
       "dtype": "f16",
       "mode": "wmma",
-      "median_ms": 9.4433,
-      "max_latency_ms": 18.8867,
-      "achieved_tflops": 1.8193,
-      "pct_peak": 0.03063,
-      "attainment_floor": 0.01532
+      "median_ms": 10.8858,
+      "max_latency_ms": 21.7717,
+      "achieved_tflops": 1.5782,
+      "pct_peak": 0.02657,
+      "attainment_floor": 0.01328
     },
     {
       "op": "flash_attn",
       "shape": "1x8x512x64",
       "dtype": "f16",
       "mode": "flash_attn",
-      "median_ms": 2.9367,
-      "max_latency_ms": 5.8733,
-      "achieved_tflops": 0.1828,
-      "pct_peak": 0.00308,
-      "attainment_floor": 0.00154
+      "median_ms": 2.9053,
+      "max_latency_ms": 5.8106,
+      "achieved_tflops": 0.1848,
+      "pct_peak": 0.00311,
+      "attainment_floor": 0.00155
     },
     {
       "op": "flash_attn",
       "shape": "1x8x1024x64",
       "dtype": "f16",
       "mode": "flash_attn",
-      "median_ms": 4.6006,
-      "max_latency_ms": 9.2011,
-      "achieved_tflops": 0.4668,
-      "pct_peak": 0.00786,
-      "attainment_floor": 0.00393
+      "median_ms": 4.6324,
+      "max_latency_ms": 9.2649,
+      "achieved_tflops": 0.4636,
+      "pct_peak": 0.0078,
+      "attainment_floor": 0.0039
     },
     {
       "op": "flash_attn",
       "shape": "1x16x1024x128",
       "dtype": "f16",
       "mode": "flash_attn",
-      "median_ms": 11.2024,
-      "max_latency_ms": 22.4049,
-      "achieved_tflops": 0.7668,
-      "pct_peak": 0.01291,
-      "attainment_floor": 0.00645
+      "median_ms": 11.7242,
+      "max_latency_ms": 23.4484,
+      "achieved_tflops": 0.7327,
+      "pct_peak": 0.01233,
+      "attainment_floor": 0.00617
     },
     {
       "op": "flash_attn_bwd",
       "shape": "1x8x512x64",
       "dtype": "f16",
       "mode": "flash_attn_bwd",
-      "median_ms": 44.0716,
-      "max_latency_ms": 88.1433
+      "median_ms": 43.5883,
+      "max_latency_ms": 87.1767
     },
     {
       "op": "flash_attn_bwd",
       "shape": "1x16x1024x128",
       "dtype": "f16",
       "mode": "flash_attn_bwd",
-      "median_ms": 87.3895,
-      "max_latency_ms": 174.779
+      "median_ms": 88.1148,
+      "max_latency_ms": 176.2297
+    },
+    {
+      "op": "gemm_f32",
+      "shape": "256x256x256",
+      "dtype": "f32",
+      "mode": "gemm_f32",
+      "median_ms": 1.9705,
+      "max_latency_ms": 3.941
+    },
+    {
+      "op": "gemm_f32",
+      "shape": "512x512x512",
+      "dtype": "f32",
+      "mode": "gemm_f32",
+      "median_ms": 2.8024,
+      "max_latency_ms": 5.6048
+    },
+    {
+      "op": "gemm_f32",
+      "shape": "1024x1024x1024",
+      "dtype": "f32",
+      "mode": "gemm_f32",
+      "median_ms": 7.5255,
+      "max_latency_ms": 15.051
     }
   ]
 }
diff --git a/benchmarks/rocm/record_hot_path_baseline.py b/benchmarks/rocm/record_hot_path_baseline.py
@@ -54,6 +54,13 @@
 # floor, not an MFU claim (pct_peak is expectedly tiny — repo Decision #26).
 FLASH_ATTN_BWD_SHAPES = [(1, 8, 512, 64), (1, 16, 1024, 128)]
 
+# Register-blocked f32 GEMM ladder (M, N, K) — the plain-VALU f32 kernel
+# (generate-rocm-gemm-f32-kernel) grouped-SwiGLU rides. TM×TN=4×4 output-tile
+# register blocking is a ~1.6x win over one-thread-per-output at 1024³
+# (STRIX_HALO Stage F: register-budget tiling is the lever). Correctness-first
+# still — the ratchet rows are a regression floor, not an MFU claim.
+GEMM_F32_SIZES = [(256, 256, 256), (512, 512, 512), (1024, 1024, 1024)]
+
 
 def _median_ms(fn, reps: int = 20, warmup: int = 3) -> float:
     for _ in range(warmup):
@@ -130,6 +137,14 @@ def _run():
                 raise RuntimeError(res.get("reason"))
         return _run
 
+    def _make_gemm_f32(m, n, k):
+        a = rng.standard_normal((m, k)).astype(np.float32)
+        bb = rng.standard_normal((k, n)).astype(np.float32)
+
+        def _run():
+            rt._rocm_f32_gemm(a, bb, np)
+        return _run
+
     cases = []
     for (m, n, k) in HOT_PATH_SIZES:
         cases.append(("matmul", f"{m}x{n}x{k}", "f16", "wmma", _make(m, n, k)))
@@ -140,6 +155,13 @@ def _run():
         for (b, h, s, d) in FLASH_ATTN_BWD_SHAPES:
             cases.append(("flash_attn_bwd", f"{b}x{h}x{s}x{d}", "f16",
                           "flash_attn_bwd", _make_fa_bwd(b, h, s, d)))
+    # f32 GEMM gates on its OWN probe (generate-rocm-gemm-f32-kernel is a
+    # DIFFERENT pass from the flash lane), so a host where the f32 GEMM works but
+    # the flash lane is missing/broken still records its gemm_f32 ratchet rows.
+    if rt._rocm_compiled_gemm_f32_available():
+        for (m, n, k) in GEMM_F32_SIZES:
+            cases.append(("gemm_f32", f"{m}x{n}x{k}", "f32", "gemm_f32",
+                          _make_gemm_f32(m, n, k)))
     return cases
 
 

diff --git a/docs/rocm_mfma_kernel_inventory.md b/docs/rocm_mfma_kernel_inventory.md
@@ -492,10 +492,15 @@ promotes them to `compileable`. See §9 for the concrete done / open / blocked s
 - (flash_attn backward is now the **full forward variant surface** — MHA +
   GQA/MQA + attn_bias + sliding-window + logit-softcap — all runtime-wired)
 - (`grad_clip_norm` now executes via `rocm_grad_clip_compiled`, x86 parity)
-- **Perf tuning** (LDS/blocked kernels) — the compiled lanes now carry *measured*
-  ratchet baselines (matmul + flash_attn fwd/bwd in `rocm_gfx1151_hot_paths.json`)
-  but stay correctness-first; a real MFU ladder (register/LDS blocking) is future
-  and is where ROCm's lead-performance-target status (Decision #28) gets earned
+- **Perf tuning** — the WMMA GEMM ladder is done (rung-1 register blocking is
+  production; LDS/pipelining are wash-to-regression on this unified-memory APU —
+  `STRIX_HALO_EXECUTION_PLAN.md` Stage F/H). The **f32 GEMM** (`gemm_f32`, grouped-
+  SwiGLU's kernel) is now **register-blocked** (TM×TN=4×4 output tile per thread,
+  ~1.6× over one-thread-per-output at 1024³) — applying the same proven lever;
+  measured `gemm_f32` ratchet rows are in `rocm_gfx1151_hot_paths.json`. The lanes
+  still stay correctness-first overall; the deeper WMMA occupancy/dual-issue lever
+  (Stage F "next lever") is where ROCm's lead-performance status (Decision #28)
+  gets fully earned, and remains future
 - **Fused paged-attention** — the §5.6 movement core (`kv_cache_append/read/prune`)
   now executes via `rocm_kv_cache_compiled` (scatter/gather compose,
   execute-compare vs `KVCacheHandle`); a single fused gather→attention paged

diff --git a/python/tessera/runtime.py b/python/tessera/runtime.py
@@ -1561,6 +1561,9 @@ def _execute_rocm_moe_transport(artifact: RuntimeArtifact, args: Any) -> Any:
 #: Cached probe for the compiler-generated ROCm WMMA flash_attn forward lane
 #: (separate from the GEMM lane — it shells to ``tessera-opt`` + a live GPU).
 _rocm_compiled_flash_attn_probe_ok: bool | None = None
+#: Cached probe for the compiler-generated f32 GEMM lane
+#: (generate-rocm-gemm-f32-kernel — a DIFFERENT pass from the flash lane).
+_rocm_compiled_gemm_f32_probe_ok: bool | None = None
 
 #: C-ABI GEMM symbols shipped by libtessera_rocm_gemm.so, keyed by storage dtype.
 _ROCM_GEMM_SYMBOLS = {
@@ -1679,6 +1682,27 @@ def _rocm_compiled_flash_attn_available() -> bool:
     return _rocm_compiled_flash_attn_probe_ok
 
 
+def _rocm_compiled_gemm_f32_available() -> bool:
+    """Cached host probe: True iff the compiler-generated f32 GEMM lane can run —
+    i.e. ``tessera-opt`` is built AND a live gfx1151 executes a tiny
+    ``_rocm_f32_gemm`` without raising. Its own probe (NOT the flash-attn one):
+    ``generate-rocm-gemm-f32-kernel`` is a DIFFERENT compiler pass, so a host
+    where the f32 GEMM works but the flash lane is missing/broken still records
+    (and re-times) its gemm_f32 ratchet rows. Never fabricates (Decision #26)."""
+    global _rocm_compiled_gemm_f32_probe_ok
+    if _rocm_compiled_gemm_f32_probe_ok is not None:
+        return _rocm_compiled_gemm_f32_probe_ok
+    _rocm_compiled_gemm_f32_probe_ok = False
+    try:
+        import numpy as np
+        a = np.zeros((1, 1), np.float32)
+        _rocm_f32_gemm(a, a, np)
+        _rocm_compiled_gemm_f32_probe_ok = True
+    except Exception:
+        _rocm_compiled_gemm_f32_probe_ok = False
+    return _rocm_compiled_gemm_f32_probe_ok
+
+
 def _execute_rocm_wmma_artifact(artifact: RuntimeArtifact, args: Any) -> Any:
     """Execute a ``target="rocm"`` single-matmul artifact on the AMD GPU through
     the shipped ``tessera_rocm_wmma_gemm_{f16,bf16}`` C-ABI symbol (RDNA WMMA,
@@ -10045,7 +10069,9 @@ def _mr(p, size):
     arr = (cv * len(launch_args))()
     for i, val in enumerate(launch_args):
         arr[i] = ctypes.cast(ctypes.byref(val), cv)
-    gx = (n_c + _GRID_BLOCKDIM - 1) // _GRID_BLOCKDIM
+    # The kernel is register-blocked: one thread per TM×TN=4×4 output tile.
+    n_tiles = ((M + 3) // 4) * ((N + 3) // 4)
+    gx = (n_tiles + _GRID_BLOCKDIM - 1) // _GRID_BLOCKDIM
     rc = hip.hipModuleLaunchKernel(fn, max(gx, 1), 1, 1, _GRID_BLOCKDIM, 1, 1,
                                    0, None, arr, None)
     if rc != 0:

diff --git a/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/GenerateROCMGemmF32Kernel.cpp b/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/GenerateROCMGemmF32Kernel.cpp
@@ -27,6 +27,13 @@ using namespace mlir;
 namespace {
 
 static constexpr int64_t BD = 256;
+// Output-tile register blocking: each thread computes a TM×TN tile of C. Per
+// k-step it loads TM A-values + TN B-values from global and reuses them across
+// TM*TN FMAs (each A elt reused TN times, each B elt TM times) — the arithmetic-
+// intensity / register-budget lever that wins on Strix Halo's unified memory
+// (STRIX_HALO_EXECUTION_PLAN Stage F: "register-budget tiling is the lever").
+static constexpr int64_t TM = 4;
+static constexpr int64_t TN = 4;
 
 void emitGemmF32Body(OpBuilder &b, Location loc, gpu::GPUFuncOp f) {
   Type f32 = b.getF32Type();
@@ -35,38 +42,85 @@ void emitGemmF32Body(OpBuilder &b, Location loc, gpu::GPUFuncOp f) {
   Value A = f.getArgument(0), B = f.getArgument(1), C = f.getArgument(2);
   Value M = f.getArgument(3), N = f.getArgument(4), K = f.getArgument(5);
 
+  auto ci = [&](int64_t v) { return b.create<arith::ConstantIndexOp>(loc, v); };
+  Value c0 = ci(0), c1 = ci(1);
   Value bid = b.create<gpu::BlockIdOp>(loc, gpu::Dimension::x);
   Value tid = b.create<gpu::ThreadIdOp>(loc, gpu::Dimension::x);
-  Value cBD = b.create<arith::ConstantIndexOp>(loc, BD);
   Value gid = b.create<arith::AddIOp>(
-      loc, b.create<arith::MulIOp>(loc, bid, cBD), tid);
-  Value total = b.create<arith::MulIOp>(loc, M, N);
+      loc, b.create<arith::MulIOp>(loc, bid, ci(BD)), tid);
+  // One thread per TM×TN output tile: nTilesN = ceil(N/TN), total = ceil(M/TM)*.
+  Value nTilesN = b.create<arith::DivUIOp>(
+      loc, b.create<arith::AddIOp>(loc, N, ci(TN - 1)), ci(TN));
+  Value nTilesM = b.create<arith::DivUIOp>(
+      loc, b.create<arith::AddIOp>(loc, M, ci(TM - 1)), ci(TM));
+  Value total = b.create<arith::MulIOp>(loc, nTilesM, nTilesN);
   Value inb = b.create<arith::CmpIOp>(loc, slt, gid, total);
   auto guard = b.create<scf::IfOp>(loc, inb, /*withElse=*/false);
   b.setInsertionPointToStart(guard.thenBlock());
 
-  Value c0 = b.create<arith::ConstantIndexOp>(loc, 0);
-  Value c1 = b.create<arith::ConstantIndexOp>(loc, 1);
   Value zero = b.create<arith::ConstantOp>(loc, f32, b.getF32FloatAttr(0.0f));
-  Value m = b.create<arith::DivUIOp>(loc, gid, N);   // row
-  Value n = b.create<arith::RemUIOp>(loc, gid, N);   // col
-  Value abase = b.create<arith::MulIOp>(loc, m, K);  // m*K
-  // acc = Σ_k A[m*K + k] · B[k*N + n]
-  auto kl = b.create<scf::ForOp>(loc, c0, K, c1, ValueRange{zero});
+  Value tr = b.create<arith::DivUIOp>(loc, gid, nTilesN);   // tile row
+  Value tc = b.create<arith::RemUIOp>(loc, gid, nTilesN);   // tile col
+  Value m0 = b.create<arith::MulIOp>(loc, tr, ci(TM));
+  Value n0 = b.create<arith::MulIOp>(loc, tc, ci(TN));
+
+  // Per-row (i) and per-col (j) bounds + safe (clamped) indices, hoisted out of
+  // the k-loop. OOB rows/cols contribute 0 (masked load) and are not stored.
+  SmallVector<Value> mi(TM), inbM(TM), aRowBase(TM);
+  for (int64_t i = 0; i < TM; ++i) {
+    mi[i] = b.create<arith::AddIOp>(loc, m0, ci(i));
+    inbM[i] = b.create<arith::CmpIOp>(loc, slt, mi[i], M);
+    Value miSafe = b.create<arith::SelectOp>(loc, inbM[i], mi[i], c0);
+    aRowBase[i] = b.create<arith::MulIOp>(loc, miSafe, K);
+  }
+  SmallVector<Value> nj(TN), inbN(TN), njSafe(TN);
+  for (int64_t j = 0; j < TN; ++j) {
+    nj[j] = b.create<arith::AddIOp>(loc, n0, ci(j));
+    inbN[j] = b.create<arith::CmpIOp>(loc, slt, nj[j], N);
+    njSafe[j] = b.create<arith::SelectOp>(loc, inbN[j], nj[j], c0);
+  }
+
+  // k-loop with TM*TN register accumulators.
+  SmallVector<Value> initAcc(TM * TN, zero);
+  auto kl = b.create<scf::ForOp>(loc, c0, K, c1, initAcc);
   {
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPointToStart(kl.getBody());
     Value k = kl.getInductionVar();
-    Value av = b.create<memref::LoadOp>(
-        loc, A, ValueRange{b.create<arith::AddIOp>(loc, abase, k)});
-    Value boff = b.create<arith::AddIOp>(
-        loc, b.create<arith::MulIOp>(loc, k, N), n);   // k*N + n
-    Value bv = b.create<memref::LoadOp>(loc, B, ValueRange{boff});
-    Value acc = b.create<arith::AddFOp>(loc, kl.getRegionIterArgs()[0],
-                                        b.create<arith::MulFOp>(loc, av, bv));
-    b.create<scf::YieldOp>(loc, ValueRange{acc});
+    SmallVector<Value> acc(kl.getRegionIterArgs().begin(),
+                           kl.getRegionIterArgs().end());
+    SmallVector<Value> a(TM), bcol(TN);
+    for (int64_t i = 0; i < TM; ++i) {
+      Value v = b.create<memref::LoadOp>(
+          loc, A, ValueRange{b.create<arith::AddIOp>(loc, aRowBase[i], k)});
+      a[i] = b.create<arith::SelectOp>(loc, inbM[i], v, zero);
+    }
+    Value kN = b.create<arith::MulIOp>(loc, k, N);
+    for (int64_t j = 0; j < TN; ++j) {
+      Value v = b.create<memref::LoadOp>(
+          loc, B, ValueRange{b.create<arith::AddIOp>(loc, kN, njSafe[j])});
+      bcol[j] = b.create<arith::SelectOp>(loc, inbN[j], v, zero);
+    }
+    SmallVector<Value> newAcc(TM * TN);
+    for (int64_t i = 0; i < TM; ++i)
+      for (int64_t j = 0; j < TN; ++j)
+        newAcc[i * TN + j] = b.create<arith::AddFOp>(
+            loc, acc[i * TN + j],
+            b.create<arith::MulFOp>(loc, a[i], bcol[j]));
+    b.create<scf::YieldOp>(loc, newAcc);
   }
-  b.create<memref::StoreOp>(loc, kl.getResult(0), C, ValueRange{gid});
+  // Store the tile, guarded on both bounds.
+  auto res = kl.getResults();
+  for (int64_t i = 0; i < TM; ++i)
+    for (int64_t j = 0; j < TN; ++j) {
+      Value inBoth = b.create<arith::AndIOp>(loc, inbM[i], inbN[j]);
+      auto st = b.create<scf::IfOp>(loc, inBoth, /*withElse=*/false);
+      OpBuilder::InsertionGuard g(b);
+      b.setInsertionPointToStart(st.thenBlock());
+      Value cidx = b.create<arith::AddIOp>(
+          loc, b.create<arith::MulIOp>(loc, mi[i], N), nj[j]);
+      b.create<memref::StoreOp>(loc, res[i * TN + j], C, ValueRange{cidx});
+    }
   b.setInsertionPointToEnd(&f.getBody().front());
   b.create<gpu::ReturnOp>(loc);
 }
@@ -77,8 +131,8 @@ struct GenerateROCMGemmF32KernelPass
 
   StringRef getArgument() const final { return "generate-rocm-gemm-f32-kernel"; }
   StringRef getDescription() const final {
-    return "Expand a tessera_rocm.gemm_f32 directive into a plain f32 GEMM "
-           "kernel (C=A@B, one thread per output element, scalar f32 k-loop)";
+    return "Expand a tessera_rocm.gemm_f32 directive into an f32 GEMM kernel "
+           "(C=A@B, register-blocked TMxTN output tile per thread, f32 k-loop)";
   }
   void getDependentDialects(DialectRegistry &registry) const final {
     registry.insert<gpu::GPUDialect, scf::SCFDialect, arith::ArithDialect,