benches: package as kernbench.benches, add @bench registry + list subcommand

Move benches/ -> src/kernbench/benches/ and src/kernbench/cli/probe.py -> src/kernbench/probes/probe.py. Each bench self-registers via @bench(name=..., description=...); kernbench list enumerates benches with auto-assigned indices, --bench accepts kebab-case name or numeric index. Audit at package-import time fails if any non-underscore module forgets the decorator. ADR-0010 (EN + KO) updated to reflect the new resolver path, list subcommand, and probes package separation. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-20 14:42:10 -07:00
parent 168b0c89f0
commit 049e3d8bb3
28 changed files with 398 additions and 79 deletions
@@ -0,0 +1,46 @@
+"""QKV GEMM benchmark: Q*K^T projection on a single PE.
+
+Demonstrates the full host-to-PE kernel launch pipeline:
+  Host → PCIE_EP → IO_CPU → M_CPU → NOC → PE_CPU → PE_SCHEDULER → engines
+
+Kernel: tl.load(a) + tl.ref(b) + tl.composite(gemm) + tl.wait()
+  - Tensor a is loaded into TCM via DMA
+  - Tensor b stays in HBM; PE_SCHEDULER streams it per-tile (32x64x32)
+"""
+from kernbench.benches.registry import bench
+from kernbench.policy.placement.dp import DPPolicy
+
+# GEMM dimensions: (M, K) x (K, N) → (M, N)
+# Small dims (1 tile) for fast regression. The test verifies the full
+# host→PE pipeline, not large-matrix throughput.
+M, K, N = 32, 64, 32
+DTYPE = "f16"
+
+
+def _gemm_kernel(a_ptr, b_ptr, out_ptr, M, K, N, tl, DTYPE="f16"):
+    """QKV GEMM kernel: out = a @ b.
+
+    a is loaded into TCM (DMA_READ).
+    b is referenced in HBM (tl.ref, no DMA — scheduler streams per-tile).
+    """
+    a = tl.load(a_ptr, shape=(M, K), dtype=DTYPE)
+    b = tl.ref(b_ptr, shape=(K, N), dtype=DTYPE)
+    handle = tl.composite(op="gemm", a=a, b=b, out_ptr=out_ptr)
+    tl.wait(handle)
+
+
+@bench(
+    name="qkv-gemm",
+    description="QKV GEMM (Q*K^T) on a single PE — full host-to-PE pipeline.",
+)
+def run(torch):
+    """Run the QKV GEMM benchmark."""
+    # DP placement: a=replicate (cube-level), b/out=column_wise (N-axis, single PE)
+    a = torch.zeros((M, K), dtype=DTYPE, dp=DPPolicy(cube="replicate", pe="replicate"), name="a")
+    b = torch.zeros((K, N), dtype=DTYPE, dp=DPPolicy(cube="replicate", pe="column_wise"), name="b")
+    out = torch.empty(
+        (M, N), dtype=DTYPE, dp=DPPolicy(cube="replicate", pe="column_wise"), name="out",
+    )
+
+    # Launch GEMM kernel
+    torch.launch("qkv_gemm", _gemm_kernel, a, b, out, M, K, N)