gemm(bench): composite-vs-async dispatch comparison harness + single-op-model sweep

Adds the user-orchestrated async GEMM variants used in paper §3.4 to
contrast with the composite (load_ref) path:
  - matmul-async              : naive (tl.load full A+B, one tl.dot)
  - matmul-async-chunked      : depth-inf prefetch (all B-tiles queued)
  - matmul-async-chunked-db   : depth-2 double-buffer (TCM-bounded)
plus scripts/paper/paper_plot_gemm_async_vs_composite.py (4-way per-PE
TFLOPS sweep + figure) and the regenerated outputs under
1H_milestone_output/gemm/.

Sweep regenerated under the ADR-0064 D8 single-op cost model (commit
2d8271c). At K=3072 the composite-vs-async-tiled gap narrows from the
pre-D8 ~6.3x to ~2.8x (composite 7.18 / async-tiled 2.54 TFLOPS): D8
makes the async-tiled kernel's ~191 single-op dispatches 5x cheaper, so
its dispatch overhead drops to ~1.5us. Qualitative order persists:
composite > async-full (3.91) > async-tiled (2.54).

test_bench_registry.py: register matmul-async / -chunked / -chunked-db
(also picks up the earlier milestone-gqa-headline -> milestone-1h-gqa
rename already present in the tree).

--- Remaining work (resume here if interrupted) ---
- Paper §3.4 (03-gemm.tex sec:gemm-vs-async): finish naive->async-full /
  chunked->async-tiled rename AND reframe "FIXED_DMA=8 for DMA descriptors"
  to single-op(8) vs composite(40). Figure already copied to
  docs/report/1H-codesign-paper/figures/gemm_composite_vs_async_tflops.png.
- Paper §3.4 K=3072 para + mechanism #3: update dispatch breakdown to new
  model (191 single-op * 8c ~= 1.5us, was 4.6us) and headline gap
  ~6.3x -> ~2.8x.
- Rebuild build/main.pdf (tectonic) + verify via pdftotext.
- Then commit Group 3 (paper + diagrams + PDF).
- Table 2 / §2 prose / ADR-0064 D8 already done in 2d8271c.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

scripts/paper/paper_plot_gemm_async_vs_composite.py

@@ -0,0 +1,261 @@
+"""Composite vs async-loading GEMM TFLOPS comparison.
+
+For each shape in the milestone sweep, run both benches:
+  - matmul-composite (load_ref variant — A pre-staged, B streamed by
+    scheduler inside one composite command)
+  - matmul-async (A and B both async-loaded via tl.load, then a single
+    tl.dot — no per-tile overlap of streaming-B with GEMM)
+
+Compute per-PE TFLOPS = 2*M*K*N / pe_window_ns for each kernel and emit a
+side-by-side bar chart PNG to
+  src/kernbench/benches/1H_milestone_output/gemm/gemm_composite_vs_async_tflops.png
+
+Run from repo root:
+  python scripts/paper/paper_plot_gemm_async_vs_composite.py
+"""
+from __future__ import annotations
+
+import json
+import os
+import sys
+from pathlib import Path
+
+REPO = Path(__file__).resolve().parents[2]
+OUT_DIR = REPO / "src" / "kernbench" / "benches" / "1H_milestone_output" / "gemm"
+OUT_PNG = OUT_DIR / "gemm_composite_vs_async_tflops.png"
+OUT_JSON = OUT_DIR / "gemm_composite_vs_async.json"
+
+TOPO_PATH = REPO / "topology.yaml"
+
+# Same 7 shapes as the milestone composite sweep (skip square=512: B alone
+# is 512 KiB; async needs all of A+B+out in TCM scratch and the 512×512
+# square pushes total scratch use past the 1 MiB cap).
+SHAPES: list[tuple[int, int, int]] = [
+    (32, 32, 32),     # 1 tile, K<TILE_K (under-tile in K)
+    (32, 64, 32),     # 1 tile
+    (32, 128, 32),    # 2 tiles
+    (32, 128, 128),   # 8 tiles
+    (32, 3072, 32),   # 48 tiles (deep-K)
+    (8, 128, 128),    # under-tile in M
+    (128, 8, 128),    # under-tile in K
+]
+
+ENGINES = ("pe_dma", "pe_fetch_store", "pe_gemm", "pe_math")
+STAGES = ("DMA_READ", "DMA_WRITE", "FETCH", "STORE", "GEMM", "MATH")
+
+
+def _pe_records(op_log):
+    return [r for r in op_log
+            if any(r.component_id.endswith("." + e) for e in ENGINES)]
+
+
+def _pe_window_ns(op_log) -> float:
+    pe = _pe_records(op_log)
+    if not pe:
+        return 0.0
+    return max(r.t_end for r in pe) - min(r.t_start for r in pe)
+
+
+def _composite_window_ns(op_log) -> float:
+    """For the composite kernel: window of records carrying a stage_type
+    set by the composite plan (DMA_READ, FETCH, GEMM, STORE, DMA_WRITE).
+    Excludes the initial up-front tl.load(A) because that record is an
+    atomic DmaReadCmd with no stage_type. Matches the existing
+    milestone_1h_gemm.py / gemm_per_pe_tflops.png methodology.
+    """
+    stage_records = [r for r in op_log
+                     if r.params.get("stage_type") in STAGES]
+    if not stage_records:
+        return 0.0
+    return max(r.t_end for r in stage_records) \
+           - min(r.t_start for r in stage_records)
+
+
+def _async_engine_window_ns(op_log) -> float:
+    """For the async kernel: engine pipeline window that excludes the
+    initial tl.load(A), paralleling composite_window's exclusion of the
+    up-front A pre-stage. The first DMA_READ record on pe_dma is the
+    tl.load(A); the window starts at the SECOND pe_dma record's t_start
+    (= tl.load(B)) and ends at the last engine record's t_end.
+    """
+    pe = _pe_records(op_log)
+    if not pe:
+        return 0.0
+    dma = sorted(
+        (r for r in op_log if r.component_id.endswith(".pe_dma")),
+        key=lambda r: r.t_start,
+    )
+    if len(dma) < 2:
+        return _pe_window_ns(op_log)
+    window_start = dma[1].t_start
+    window_end = max(r.t_end for r in pe)
+    return window_end - window_start
+
+
+def _run_one(bench_name: str, variant: str | None, M: int, K: int, N: int) -> dict:
+    os.environ["MATMUL_M"] = str(M)
+    os.environ["MATMUL_K"] = str(K)
+    os.environ["MATMUL_N"] = str(N)
+    if variant is not None:
+        os.environ["MATMUL_VARIANT"] = variant
+    elif "MATMUL_VARIANT" in os.environ:
+        del os.environ["MATMUL_VARIANT"]
+
+    from kernbench.benches.registry import resolve as resolve_bench
+    from kernbench.runtime_api.bench_runner import run_bench
+    from kernbench.runtime_api.types import resolve_device
+    from kernbench.sim_engine.engine import GraphEngine
+    from kernbench.topology.builder import resolve_topology
+
+    # Chunked async loads A in TILE_K-sized blocks at sequential offsets;
+    # those offsets do not match the row-major (M, K) layout of A, so the
+    # DataExecutor would fail on the resulting mismatched read. The
+    # simulator's *timing* path doesn't need correct data — only the
+    # number of bytes per DMA matters — so bypass DataExecutor for the
+    # chunked kernel. Composite/naive-async loads use the full (M, K)
+    # shape and remain data-correct.
+    if bench_name in ("matmul-async-chunked", "matmul-async-chunked-db"):
+        GraphEngine._flush_data_phase = lambda self: None
+
+    topo = resolve_topology(str(TOPO_PATH))
+    bench = resolve_bench(bench_name).run
+    device = resolve_device(None)
+    result = run_bench(
+        topology=topo, bench_fn=bench, device=device,
+        engine_factory=lambda t, d: GraphEngine(
+            getattr(t, "topology_obj", t), enable_data=True,
+        ),
+    )
+    if not result.completion.ok:
+        raise RuntimeError(f"{bench_name} failed at {M}x{K}x{N}: {result.completion}")
+    log = result.engine.op_log
+    pe_window = _pe_window_ns(log)
+    if bench_name == "matmul-composite":
+        engine_window = _composite_window_ns(log)
+    elif bench_name in ("matmul-async-chunked", "matmul-async-chunked-db"):
+        # First N_chunks pe_dma records are A pre-stage; engine window
+        # starts at the (N_chunks+1)-th pe_dma record (= first B-chunk
+        # load). Fall back to the naive analog for K <= TILE_K (kernel
+        # collapses to single load+dot+store).
+        TILE_K = 64
+        n_chunks = max(K // TILE_K, 1)
+        if n_chunks <= 1:
+            engine_window = _async_engine_window_ns(log)
+        else:
+            dma = sorted(
+                (r for r in log if r.component_id.endswith(".pe_dma")),
+                key=lambda r: r.t_start,
+            )
+            pe = _pe_records(log)
+            if len(dma) > n_chunks and pe:
+                engine_window = max(r.t_end for r in pe) - dma[n_chunks].t_start
+            else:
+                engine_window = _async_engine_window_ns(log)
+    else:
+        engine_window = _async_engine_window_ns(log)
+    flops = 2 * M * K * N
+    # flops / ns * 1e-3 = TFLOP/s  (since 1 flop/ns = 1 GFLOP/s)
+    return {
+        "M": M, "K": K, "N": N,
+        "bench": bench_name, "variant": variant,
+        "pe_window_ns": pe_window,
+        "engine_window_ns": engine_window,
+        "flops": flops,
+        "tflops": (flops / engine_window / 1000.0) if engine_window > 0 else 0.0,
+        "n_records": len(log),
+    }
+
+
+def collect() -> list[dict]:
+    rows: list[dict] = []
+    for M, K, N in SHAPES:
+        print(f"  shape M={M:4d} K={K:5d} N={N:4d} ...", flush=True)
+        comp = _run_one("matmul-composite", "load_ref", M, K, N)
+        asyn = _run_one("matmul-async", None, M, K, N)
+        chnk = _run_one("matmul-async-chunked", None, M, K, N)
+        chnkdb = _run_one("matmul-async-chunked-db", None, M, K, N)
+        print(f"    composite:    {comp['engine_window_ns']:8.1f} ns  "
+              f"{comp['tflops']:6.3f} TFLOPS")
+        print(f"    async naive:  {asyn['engine_window_ns']:8.1f} ns  "
+              f"{asyn['tflops']:6.3f} TFLOPS")
+        print(f"    chunked all:  {chnk['engine_window_ns']:8.1f} ns  "
+              f"{chnk['tflops']:6.3f} TFLOPS")
+        print(f"    chunked db=2: {chnkdb['engine_window_ns']:8.1f} ns  "
+              f"{chnkdb['tflops']:6.3f} TFLOPS")
+        rows.append(comp)
+        rows.append(asyn)
+        rows.append(chnk)
+        rows.append(chnkdb)
+    return rows
+
+
+def plot(rows: list[dict]) -> None:
+    import matplotlib.pyplot as plt
+    import numpy as np
+
+    shape_keys = []
+    for M, K, N in SHAPES:
+        shape_keys.append((M, K, N))
+
+    by_key: dict[tuple[int, int, int], dict[str, dict]] = {}
+    for r in rows:
+        key = (r["M"], r["K"], r["N"])
+        by_key.setdefault(key, {})[r["bench"]] = r
+
+    labels = [f"M={M}\nK={K}\nN={N}" for (M, K, N) in shape_keys]
+    comp_tflops = [by_key[k]["matmul-composite"]["tflops"] for k in shape_keys]
+    asyn_tflops = [by_key[k]["matmul-async"]["tflops"] for k in shape_keys]
+    chnk_tflops = [by_key[k]["matmul-async-chunked"]["tflops"] for k in shape_keys]
+    chnkdb_tflops = [by_key[k]["matmul-async-chunked-db"]["tflops"]
+                     for k in shape_keys]
+
+    x = np.arange(len(shape_keys))
+    width = 0.20
+
+    fig, ax = plt.subplots(figsize=(14, 5.5))
+    b_c = ax.bar(x - 1.5*width, comp_tflops, width,
+                 label="Composite (load_ref)", color="#10b981")
+    b_a = ax.bar(x - 0.5*width, asyn_tflops, width,
+                 label="Async naive (tl.load full + tl.dot)",
+                 color="#f59e0b")
+    b_db = ax.bar(x + 0.5*width, chnkdb_tflops, width,
+                  label="Async chunked-prefetch, depth=2 "
+                        "(TCM-bounded)", color="#a855f7")
+    b_k = ax.bar(x + 1.5*width, chnk_tflops, width,
+                 label="Async chunked-prefetch, depth=$\\infty$ "
+                       "(all B-tiles queued up front)", color="#3b82f6")
+
+    ax.axhline(8.0, linestyle="--", color="#94a3b8", linewidth=0.8,
+               label="Per-PE GEMM peak (8 TFLOP/s)")
+    ax.set_ylabel("Per-PE achieved TFLOP/s")
+    ax.set_title("Composite vs async-loading GEMM — per-PE throughput "
+                 "(engine pipeline window, A pre-stage excluded)")
+    ax.set_xticks(x)
+    ax.set_xticklabels(labels, fontsize=8)
+    ax.grid(axis="y", alpha=0.25)
+    ax.legend(loc="upper left", fontsize=9)
+
+    for bar in (*b_c, *b_a, *b_db, *b_k):
+        h = bar.get_height()
+        ax.annotate(f"{h:.2f}", xy=(bar.get_x() + bar.get_width()/2, h),
+                    xytext=(0, 2), textcoords="offset points",
+                    ha="center", va="bottom", fontsize=7, color="#475569")
+
+    fig.tight_layout()
+    OUT_DIR.mkdir(parents=True, exist_ok=True)
+    fig.savefig(OUT_PNG, dpi=150)
+    print(f"wrote {OUT_PNG.relative_to(REPO)}")
+    plt.close(fig)
+
+
+def main() -> int:
+    rows = collect()
+    OUT_DIR.mkdir(parents=True, exist_ok=True)
+    OUT_JSON.write_text(json.dumps(rows, indent=2))
+    print(f"wrote {OUT_JSON.relative_to(REPO)}")
+    plot(rows)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

src/kernbench/benches/1H_milestone_output/gemm/gemm_composite_vs_async.json

@@ -0,0 +1,338 @@
+[
+  {
+    "M": 32,
+    "K": 32,
+    "N": 32,
+    "bench": "matmul-composite",
+    "variant": "load_ref",
+    "pe_window_ns": 171.394,
+    "engine_window_ns": 106.38400000000001,
+    "flops": 65536,
+    "tflops": 0.6160324860881335,
+    "n_records": 7
+  },
+  {
+    "M": 32,
+    "K": 32,
+    "N": 32,
+    "bench": "matmul-async",
+    "variant": null,
+    "pe_window_ns": 103.22199999999998,
+    "engine_window_ns": 71.71199999999999,
+    "flops": 65536,
+    "tflops": 0.9138777331548418,
+    "n_records": 4
+  },
+  {
+    "M": 32,
+    "K": 32,
+    "N": 32,
+    "bench": "matmul-async-chunked",
+    "variant": null,
+    "pe_window_ns": 103.22199999999998,
+    "engine_window_ns": 71.71199999999999,
+    "flops": 65536,
+    "tflops": 0.9138777331548418,
+    "n_records": 4
+  },
+  {
+    "M": 32,
+    "K": 32,
+    "N": 32,
+    "bench": "matmul-async-chunked-db",
+    "variant": null,
+    "pe_window_ns": 103.22199999999998,
+    "engine_window_ns": 71.71199999999999,
+    "flops": 65536,
+    "tflops": 0.9138777331548418,
+    "n_records": 4
+  },
+  {
+    "M": 32,
+    "K": 64,
+    "N": 32,
+    "bench": "matmul-composite",
+    "variant": "load_ref",
+    "pe_window_ns": 179.394,
+    "engine_window_ns": 106.38400000000001,
+    "flops": 131072,
+    "tflops": 1.232064972176267,
+    "n_records": 7
+  },
+  {
+    "M": 32,
+    "K": 64,
+    "N": 32,
+    "bench": "matmul-async",
+    "variant": null,
+    "pe_window_ns": 127.41399999999999,
+    "engine_window_ns": 87.904,
+    "flops": 131072,
+    "tflops": 1.4910811794685113,
+    "n_records": 4
+  },
+  {
+    "M": 32,
+    "K": 64,
+    "N": 32,
+    "bench": "matmul-async-chunked",
+    "variant": null,
+    "pe_window_ns": 127.41399999999999,
+    "engine_window_ns": 87.904,
+    "flops": 131072,
+    "tflops": 1.4910811794685113,
+    "n_records": 4
+  },
+  {
+    "M": 32,
+    "K": 64,
+    "N": 32,
+    "bench": "matmul-async-chunked-db",
+    "variant": null,
+    "pe_window_ns": 127.41399999999999,
+    "engine_window_ns": 87.904,
+    "flops": 131072,
+    "tflops": 1.4910811794685113,
+    "n_records": 4
+  },
+  {
+    "M": 32,
+    "K": 128,
+    "N": 32,
+    "bench": "matmul-composite",
+    "variant": "load_ref",
+    "pe_window_ns": 211.77800000000002,
+    "engine_window_ns": 122.76800000000003,
+    "flops": 262144,
+    "tflops": 2.135279551674703,
+    "n_records": 10
+  },
+  {
+    "M": 32,
+    "K": 128,
+    "N": 32,
+    "bench": "matmul-async",
+    "variant": null,
+    "pe_window_ns": 175.798,
+    "engine_window_ns": 120.28800000000001,
+    "flops": 262144,
+    "tflops": 2.1793030061186482,
+    "n_records": 4
+  },
+  {
+    "M": 32,
+    "K": 128,
+    "N": 32,
+    "bench": "matmul-async-chunked",
+    "variant": null,
+    "pe_window_ns": 185.81799999999998,
+    "engine_window_ns": 130.308,
+    "flops": 262144,
+    "tflops": 2.011726064401265,
+    "n_records": 8
+  },
+  {
+    "M": 32,
+    "K": 128,
+    "N": 32,
+    "bench": "matmul-async-chunked-db",
+    "variant": null,
+    "pe_window_ns": 185.81799999999998,
+    "engine_window_ns": 130.308,
+    "flops": 262144,
+    "tflops": 2.011726064401265,
+    "n_records": 8
+  },
+  {
+    "M": 32,
+    "K": 128,
+    "N": 128,
+    "bench": "matmul-composite",
+    "variant": "load_ref",
+    "pe_window_ns": 315.394,
+    "engine_window_ns": 226.38400000000001,
+    "flops": 1048576,
+    "tflops": 4.631846773623577,
+    "n_records": 34
+  },
+  {
+    "M": 32,
+    "K": 128,
+    "N": 128,
+    "bench": "matmul-async",
+    "variant": null,
+    "pe_window_ns": 394.102,
+    "engine_window_ns": 338.592,
+    "flops": 1048576,
+    "tflops": 3.096871751252245,
+    "n_records": 4
+  },
+  {
+    "M": 32,
+    "K": 128,
+    "N": 128,
+    "bench": "matmul-async-chunked",
+    "variant": null,
+    "pe_window_ns": 368.12200000000007,
+    "engine_window_ns": 312.6120000000001,
+    "flops": 1048576,
+    "tflops": 3.354241040011259,
+    "n_records": 8
+  },
+  {
+    "M": 32,
+    "K": 128,
+    "N": 128,
+    "bench": "matmul-async-chunked-db",
+    "variant": null,
+    "pe_window_ns": 368.12200000000007,
+    "engine_window_ns": 312.6120000000001,
+    "flops": 1048576,
+    "tflops": 3.354241040011259,
+    "n_records": 8
+  },
+  {
+    "M": 32,
+    "K": 3072,
+    "N": 32,
+    "bench": "matmul-composite",
+    "variant": "load_ref",
+    "pe_window_ns": 1701.4420000000007,
+    "engine_window_ns": 876.4320000000005,
+    "flops": 6291456,
+    "tflops": 7.178487321320988,
+    "n_records": 148
+  },
+  {
+    "M": 32,
+    "K": 3072,
+    "N": 32,
+    "bench": "matmul-async",
+    "variant": null,
+    "pe_window_ns": 2401.4620000000004,
+    "engine_window_ns": 1609.952,
+    "flops": 6291456,
+    "tflops": 3.9078531533859397,
+    "n_records": 4
+  },
+  {
+    "M": 32,
+    "K": 3072,
+    "N": 32,
+    "bench": "matmul-async-chunked",
+    "variant": null,
+    "pe_window_ns": 3264.8920000000217,
+    "engine_window_ns": 2473.3820000000214,
+    "flops": 6291456,
+    "tflops": 2.5436653133240017,
+    "n_records": 192
+  },
+  {
+    "M": 32,
+    "K": 3072,
+    "N": 32,
+    "bench": "matmul-async-chunked-db",
+    "variant": null,
+    "pe_window_ns": 3286.402000000022,
+    "engine_window_ns": 2494.8920000000217,
+    "flops": 6291456,
+    "tflops": 2.5217348085608298,
+    "n_records": 192
+  },
+  {
+    "M": 8,
+    "K": 128,
+    "N": 128,
+    "bench": "matmul-composite",
+    "variant": "load_ref",
+    "pe_window_ns": 291.394,
+    "engine_window_ns": 226.38400000000001,
+    "flops": 262144,
+    "tflops": 1.1579616934058943,
+    "n_records": 34
+  },
+  {
+    "M": 8,
+    "K": 128,
+    "N": 128,
+    "bench": "matmul-async",
+    "variant": null,
+    "pe_window_ns": 247.798,
+    "engine_window_ns": 216.288,
+    "flops": 262144,
+    "tflops": 1.212013611480988,
+    "n_records": 4
+  },
+  {
+    "M": 8,
+    "K": 128,
+    "N": 128,
+    "bench": "matmul-async-chunked",
+    "variant": null,
+    "pe_window_ns": 251.42399999999998,
+    "engine_window_ns": 214.914,
+    "flops": 262144,
+    "tflops": 1.2197623235340647,
+    "n_records": 8
+  },
+  {
+    "M": 8,
+    "K": 128,
+    "N": 128,
+    "bench": "matmul-async-chunked-db",
+    "variant": null,
+    "pe_window_ns": 251.42399999999998,
+    "engine_window_ns": 214.914,
+    "flops": 262144,
+    "tflops": 1.2197623235340647,
+    "n_records": 8
+  },
+  {
+    "M": 128,
+    "K": 8,
+    "N": 128,
+    "bench": "matmul-composite",
+    "variant": "load_ref",
+    "pe_window_ns": 462.01,
+    "engine_window_ns": 397.0,
+    "flops": 262144,
+    "tflops": 0.6603123425692695,
+    "n_records": 82
+  },
+  {
+    "M": 128,
+    "K": 8,
+    "N": 128,
+    "bench": "matmul-async",
+    "variant": null,
+    "pe_window_ns": 247.798,
+    "engine_window_ns": 216.288,
+    "flops": 262144,
+    "tflops": 1.212013611480988,
+    "n_records": 4
+  },
+  {
+    "M": 128,
+    "K": 8,
+    "N": 128,
+    "bench": "matmul-async-chunked",
+    "variant": null,
+    "pe_window_ns": 247.798,
+    "engine_window_ns": 216.288,
+    "flops": 262144,
+    "tflops": 1.212013611480988,
+    "n_records": 4
+  },
+  {
+    "M": 128,
+    "K": 8,
+    "N": 128,
+    "bench": "matmul-async-chunked-db",
+    "variant": null,
+    "pe_window_ns": 247.798,
+    "engine_window_ns": 216.288,
+    "flops": 262144,
+    "tflops": 1.212013611480988,
+    "n_records": 4
+  }
+]

src/kernbench/benches/matmul_async.py

@@ -0,0 +1,57 @@
+"""Single-PE async-loading GEMM (tl.load + tl.dot) for composite-vs-async
+dispatch comparison.
+
+Pairs with ``matmul_composite``: both run the same (M,K,N) sweep, but
+this kernel uses atomic primitives:
+
+  a = tl.load(a_ptr, ...)   # async DMA into TCM (1 PE_CPU cmd)
+  b = tl.load(b_ptr, ...)   # async DMA into TCM (1 PE_CPU cmd), overlaps with a
+  out = tl.dot(a, b)        # 1 PE_CPU cmd; GEMM engine starts only after
+                            # both loads complete (tl.dot _await_pending)
+  tl.store(out_ptr, out)    # 1 PE_CPU cmd
+
+Compared to composite:
+  - dispatch count grows to 4 instead of ~2
+  - GEMM cannot start until *all* of B is in TCM (no per-tile overlap of
+    streaming B-load with GEMM compute)
+
+Env vars: MATMUL_M, MATMUL_K, MATMUL_N, MATMUL_DTYPE (mirror matmul_composite).
+"""
+import os
+
+from kernbench.benches.registry import bench
+from kernbench.policy.placement.dp import DPPolicy
+
+
+def _kernel_async(a_ptr, b_ptr, out_ptr, M, K, N, tl, DTYPE="f16"):
+    M, K, N = int(M), int(K), int(N)
+    # Pre-stage A: load A and block the user kernel until its DMA finishes
+    # before dispatching load(B). Without this barrier, load(A) and load(B)
+    # run concurrently on the DMA engine and share HBM bandwidth, so the
+    # "exclude load(A) from the engine window" methodology does not
+    # actually subtract A's time — A's traffic still shows up as halved-BW
+    # contention with B. Calling _await_pending makes the load(A) the
+    # explicit pre-staging step, matching the composite load_ref semantic
+    # (where the composite waits for A before its first tile fetch).
+    a = tl.load(int(a_ptr), shape=(M, K), dtype=DTYPE)
+    tl._await_pending(a)
+    b = tl.load(int(b_ptr), shape=(K, N), dtype=DTYPE)
+    out = tl.dot(a, b)
+    tl.store(int(out_ptr), out)
+
+
+@bench(
+    name="matmul-async",
+    description="Single-PE async-loading GEMM (tl.load + tl.dot) for "
+                "composite-vs-async dispatch comparison.",
+)
+def run(torch):
+    M = int(os.environ.get("MATMUL_M", "256"))
+    K = int(os.environ.get("MATMUL_K", "256"))
+    N = int(os.environ.get("MATMUL_N", "256"))
+    DTYPE = os.environ.get("MATMUL_DTYPE", "f16")
+    dp = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=1)
+    a = torch.empty((M, K), dtype=DTYPE, dp=dp, name="a")
+    b = torch.empty((K, N), dtype=DTYPE, dp=dp, name="b")
+    out = torch.empty((M, N), dtype=DTYPE, dp=dp, name="out")
+    torch.launch("matmul_async", _kernel_async, a, b, out, M, K, N)

src/kernbench/benches/matmul_async_chunked.py

@@ -0,0 +1,90 @@
+"""Single-PE chunked-prefetching async-loading GEMM for composite-vs-async
+dispatch comparison.
+
+Splits K into hardware-tile-sized chunks (CHUNK_K = TILE_K = 64) so that:
+  - A chunks are loaded first (A pre-staged), then awaited
+  - B chunks are all issued as async tl.load at once → DMA engine queues
+    them FIFO and starts streaming
+  - Per-chunk tl.dot is issued; each dot blocks on its specific b_chunk
+    only, so dot[i] can start as soon as b[i] lands in TCM, overlapping
+    GEMM[i] with DMA of b[i+1..] still in flight
+
+Per-chunk dot outputs are accumulated with tl.add via a running sum
+(``out = out + tl.dot(...)``), which serializes on the MATH engine due
+to accumulator data dependency. This is the closest user-level analog to
+composite's per-tile pipeline that is achievable with the current tl
+DSL (which lacks a GEMM accumulator parameter and slicing primitive).
+
+Env: MATMUL_M, MATMUL_K, MATMUL_N, MATMUL_DTYPE (same as matmul_composite).
+"""
+import os
+
+from kernbench.benches.registry import bench
+from kernbench.policy.placement.dp import DPPolicy
+
+# Mirror PeSchedulerComponent.TILE_K (also TILE_K in milestone_1h_gemm.py).
+TILE_K = 64
+
+
+def _kernel_async_chunked(a_ptr, b_ptr, out_ptr, M, K, N, tl, DTYPE="f16"):
+    M, K, N = int(M), int(K), int(N)
+    bpe = 2  # f16
+
+    if K <= TILE_K:
+        # Under-tile in K: chunking would make each dot smaller, even
+        # worse off. Fall back to a single dot (naive async behaviour).
+        a = tl.load(int(a_ptr), shape=(M, K), dtype=DTYPE)
+        tl._await_pending(a)
+        b = tl.load(int(b_ptr), shape=(K, N), dtype=DTYPE)
+        out = tl.dot(a, b)
+        tl.store(int(out_ptr), out)
+        return
+
+    n_chunks = K // TILE_K
+    a_chunk_bytes = M * TILE_K * bpe
+    b_chunk_bytes = TILE_K * N * bpe
+
+    # ── Phase 1: pre-stage all A chunks (issue + wait) ──
+    a_chunks = [
+        tl.load(int(a_ptr) + i * a_chunk_bytes,
+                shape=(M, TILE_K), dtype=DTYPE)
+        for i in range(n_chunks)
+    ]
+    tl._await_pending(*a_chunks)
+
+    # ── Phase 2: queue all B-chunk async loads up front ──
+    # The DMA engine processes them FIFO; later chunks are in flight
+    # while the early dots run on GEMM engine.
+    b_chunks = [
+        tl.load(int(b_ptr) + i * b_chunk_bytes,
+                shape=(TILE_K, N), dtype=DTYPE)
+        for i in range(n_chunks)
+    ]
+
+    # ── Phase 3: per-chunk dot with running accumulator ──
+    # dot[i] only blocks on b_chunks[i] (tl.dot's _await_pending). While
+    # dot[i] runs on GEMM, b_chunks[i+1..] continue draining off DMA.
+    out = tl.dot(a_chunks[0], b_chunks[0])
+    for i in range(1, n_chunks):
+        out = out + tl.dot(a_chunks[i], b_chunks[i])
+
+    tl.store(int(out_ptr), out)
+
+
+@bench(
+    name="matmul-async-chunked",
+    description="Single-PE chunked-prefetching async-loading GEMM "
+                "(CHUNK_K = TILE_K = 64; A pre-staged, B chunks "
+                "prefetched, per-chunk tl.dot + tl.add accumulator).",
+)
+def run(torch):
+    M = int(os.environ.get("MATMUL_M", "256"))
+    K = int(os.environ.get("MATMUL_K", "256"))
+    N = int(os.environ.get("MATMUL_N", "256"))
+    DTYPE = os.environ.get("MATMUL_DTYPE", "f16")
+    dp = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=1)
+    a = torch.empty((M, K), dtype=DTYPE, dp=dp, name="a")
+    b = torch.empty((K, N), dtype=DTYPE, dp=dp, name="b")
+    out = torch.empty((M, N), dtype=DTYPE, dp=dp, name="out")
+    torch.launch("matmul_async_chunked", _kernel_async_chunked,
+                 a, b, out, M, K, N)

src/kernbench/benches/matmul_async_chunked_db.py

@@ -0,0 +1,110 @@
+"""Single-PE chunked-prefetching async GEMM with depth=2 double-buffer.
+
+A TCM-bounded variant of ``matmul_async_chunked``: at most two B-chunk
+loads are in flight at any time, so peak TCM occupancy for B is
+``2 × TILE_K × N × bpe`` rather than the full ``K × N × bpe``.
+
+This is the architecturally required form for any real LLM workload
+where ``K = context_length`` (4 K–32 K tokens): the
+``matmul_async_chunked`` variant queues *all* B-chunk loads up front
+and therefore needs the full B in TCM; that exceeds the 1 MiB kernel
+scratch cap (topology.yaml ``pe_tcm.kernel_scratch_mb``) at
+sub-attention scale. Double-buffer prefetch is the closest
+user-orchestrated kernel pattern to composite's per-tile streaming.
+
+Structure (after pre-staging A):
+  - Issue B_chunk[0] + B_chunk[1] async (2 outstanding)
+  - For each chunk i:
+    - tl.dot(A_chunk[i], B_chunk[i])   -- blocks on b[i], then issues GEMM
+    - tl.add into running accumulator
+    - Issue B_chunk[i+2] async (if available) -- keeps in-flight count = 2
+
+Compared to the full-prefetch variant, depth-2 changes *when* DMA
+descriptors hit the queue, not their total number; the structural
+dispatch cost amortization story therefore stays the same. The
+purpose of the variant is to honour the TCM cap, so the comparison
+to composite reflects a kernel that could actually run on real
+hardware at LLM-scale K.
+
+Env: MATMUL_M, MATMUL_K, MATMUL_N, MATMUL_DTYPE.
+"""
+import os
+
+from kernbench.benches.registry import bench
+from kernbench.policy.placement.dp import DPPolicy
+
+TILE_K = 64
+PREFETCH_DEPTH = 2
+
+
+def _kernel_async_chunked_db(a_ptr, b_ptr, out_ptr, M, K, N, tl, DTYPE="f16"):
+    M, K, N = int(M), int(K), int(N)
+    bpe = 2  # f16
+
+    if K <= TILE_K:
+        # Single-chunk fallback (under-tile or one-tile case): degenerates
+        # to naive async, no prefetch opportunity.
+        a = tl.load(int(a_ptr), shape=(M, K), dtype=DTYPE)
+        tl._await_pending(a)
+        b = tl.load(int(b_ptr), shape=(K, N), dtype=DTYPE)
+        out = tl.dot(a, b)
+        tl.store(int(out_ptr), out)
+        return
+
+    n_chunks = K // TILE_K
+    a_chunk_bytes = M * TILE_K * bpe
+    b_chunk_bytes = TILE_K * N * bpe
+
+    # Pre-stage all A chunks; this is the "activation pre-staged" half of
+    # load_ref. A is loaded once at the start and reused across every
+    # output dot.
+    a_chunks = [
+        tl.load(int(a_ptr) + i * a_chunk_bytes,
+                shape=(M, TILE_K), dtype=DTYPE)
+        for i in range(n_chunks)
+    ]
+    tl._await_pending(*a_chunks)
+
+    # Prime the prefetch pipeline: issue the first PREFETCH_DEPTH B-chunk
+    # loads up front. Anything beyond depth-2 stays in HBM until we get
+    # to its dot.
+    b_chunks: list = [
+        tl.load(int(b_ptr) + i * b_chunk_bytes,
+                shape=(TILE_K, N), dtype=DTYPE)
+        for i in range(min(PREFETCH_DEPTH, n_chunks))
+    ]
+
+    # First dot — blocks on b_chunks[0]; while it computes, b_chunks[1]
+    # is already in flight.
+    out = tl.dot(a_chunks[0], b_chunks[0])
+
+    for i in range(1, n_chunks):
+        # Issue the next prefetch BEFORE the dot that consumes b_chunks[i],
+        # so peak outstanding B loads stays at PREFETCH_DEPTH.
+        if i + PREFETCH_DEPTH - 1 < n_chunks:
+            b_chunks.append(
+                tl.load(int(b_ptr) + (i + PREFETCH_DEPTH - 1) * b_chunk_bytes,
+                        shape=(TILE_K, N), dtype=DTYPE)
+            )
+        out = out + tl.dot(a_chunks[i], b_chunks[i])
+
+    tl.store(int(out_ptr), out)
+
+
+@bench(
+    name="matmul-async-chunked-db",
+    description="Single-PE chunked-prefetching async GEMM with depth=2 "
+                "double-buffer (TCM-bounded streaming analog of composite "
+                "load_ref).",
+)
+def run(torch):
+    M = int(os.environ.get("MATMUL_M", "256"))
+    K = int(os.environ.get("MATMUL_K", "256"))
+    N = int(os.environ.get("MATMUL_N", "256"))
+    DTYPE = os.environ.get("MATMUL_DTYPE", "f16")
+    dp = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=1)
+    a = torch.empty((M, K), dtype=DTYPE, dp=dp, name="a")
+    b = torch.empty((K, N), dtype=DTYPE, dp=dp, name="b")
+    out = torch.empty((M, N), dtype=DTYPE, dp=dp, name="out")
+    torch.launch("matmul_async_chunked_db", _kernel_async_chunked_db,
+                 a, b, out, M, K, N)

tests/test_bench_registry.py

@@ -11,10 +11,13 @@ EXPECTED_NAMES = [
     "gemm-single-pe",
     "gpt3-qkv",
     "ipcq-allreduce",
+    "matmul-async",
+    "matmul-async-chunked",
+    "matmul-async-chunked-db",
     "matmul-composite",
     "milestone-1h-ccl",
     "milestone-1h-gemm",
-    "milestone-gqa-headline",
+    "milestone-1h-gqa",
     "qkv-gemm",
     "qkv-gemm-multi-pe",
     "va-offset-verify",