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>

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)