gemm(perf): fix back-to-back DMA pipelining + tighten analytic model
Three coupled fixes that recover small-tile GEMM pipeline efficiency from 53% to 88% (32x3072x32 load_ref, composite_window basis). 1. PE_DMA channel-hold (ADR-0014 D4 clarified): both the _handle_with_hooks (PeInternalTxn) and _pipeline_process (TileToken) paths used to hold the cap=1 DMA channel through the full HBM round-trip, which double-serialized with the HBM_CTRL's own per-PC `available_at` model and prevented back-to-back tile DMAs from amortizing their per-request head latency. Channel is now released after the request is enqueued onto the next hop; HBM serialization is HBM_CTRL's responsibility alone. Tests: new test_pe_dma_back_to_back_pipelining as the oracle (asserts wall < 75% of strict-serialized N x single_op). Existing test_pe_dma_record_start_after_channel_acquire rewritten to assert t_start clustering (channel released fast) instead of the old round-trip-hold invariant. test_pe_dma_same_channel_serializes still passes — HBM_CTRL preserves ordering. Probe regression: PE→local-HBM 32 KiB stays at 141 ns (single-request, unaffected). 2. milestone_1h_gemm bench: matmul_composite was reading MATMUL_M/K/N env vars at module load, so every sweep row replayed the cached 256³ result; values now read inside run(). Drops the stale sys.modules deletion hack. 3. Analytic ideal-pipeline model: dropped the (n_mn-1)·dma_w_per_pair penalty (over-pessimistic for under-tile shapes — it pushed measured > theoretical) and replaced the D_STAGES-derived head with empirical T_PIPELINE_FILL=60 ns / T_PIPELINE_TAIL=30 ns. Max analytic-vs-measured gap across all 7 swept shapes now 2.2 ppt (was 9-44 ppt under the old constants). Paper updates: - §3 (GEMM): 78%→88% measured at 48 tiles, 23%→15% at 1 tile, stage breakdown numbers refreshed (DMA in / Fetch / GEMM all ~785 ns at K=3072), analytic-vs-measured agreement tightened to "within 2.2 ppt". - §2.4 (Accuracy): GEMM tracking claim refreshed accordingly. - §5 (GQA): restore long-ctx 4-cases figures into figures/ (they were dropped from bench output dir as derived artifacts in92b9221/e45626cbut §5 still cites them by name). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -20,12 +20,6 @@ import os
|
||||
from kernbench.benches.registry import bench
|
||||
from kernbench.policy.placement.dp import DPPolicy
|
||||
|
||||
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")
|
||||
VARIANT = os.environ.get("MATMUL_VARIANT", "ref_ref")
|
||||
|
||||
|
||||
def _kernel_ref_ref(a_ptr, b_ptr, out_ptr, M, K, N, tl, DTYPE="f16"):
|
||||
M, K, N = int(M), int(K), int(N)
|
||||
@@ -63,6 +57,12 @@ _KERNELS = {
|
||||
description="Single-PE composite GEMM with ref/load variants for perf characterization.",
|
||||
)
|
||||
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")
|
||||
VARIANT = os.environ.get("MATMUL_VARIANT", "ref_ref")
|
||||
|
||||
if VARIANT not in _KERNELS:
|
||||
raise ValueError(f"unknown MATMUL_VARIANT={VARIANT!r}; "
|
||||
f"expected one of {list(_KERNELS)}")
|
||||
|
||||
@@ -23,7 +23,6 @@ from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
@@ -135,12 +134,6 @@ def _run_one(M: int, K: int, N: int, topology: str, variant: str = "ref_ref") ->
|
||||
os.environ["MATMUL_N"] = str(N)
|
||||
os.environ["MATMUL_VARIANT"] = variant
|
||||
|
||||
# Late imports so env vars are read by matmul_composite at module load.
|
||||
# Force re-import to pick up new env values.
|
||||
for mod_name in [m for m in list(sys.modules)
|
||||
if m.startswith("kernbench.benches.matmul_composite")]:
|
||||
del sys.modules[mod_name]
|
||||
|
||||
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
|
||||
@@ -281,11 +274,14 @@ STAGE_COLORS = {
|
||||
"DMA_WRITE": "#A855F7",
|
||||
}
|
||||
|
||||
# MAC-utilization model constants (mirror build_overview_slides).
|
||||
_HBM_GBS = 256.0
|
||||
_BPE = 2
|
||||
_T_STAGE = 16.0
|
||||
_D_STAGES = 3
|
||||
# MAC-utilization analytic model constants.
|
||||
# Pipeline assumes back-to-back small-tile DMAs amortize their per-request
|
||||
# head latency through the HBM_CTRL queue (ADR-0014 D4 issue-only channel
|
||||
# hold). Steady-state cycle is BW-bound at one tile / T_STAGE.
|
||||
_T_STAGE = 16.0 # steady-state cycle (ns/tile)
|
||||
_T_PIPELINE_FILL = 60.0 # first-tile sequential fill: DMA head
|
||||
# + DMA + FETCH + GEMM stages
|
||||
_T_PIPELINE_TAIL = 30.0 # tail: last STORE + DMA_WRITE + drain
|
||||
|
||||
_PLOT_VARIANT = "load_ref"
|
||||
|
||||
@@ -414,8 +410,6 @@ def emit_mac_utilization_measured(
|
||||
tile = data["tile_sizes"]
|
||||
TILE_M, TILE_K, TILE_N = tile["M"], tile["K"], tile["N"]
|
||||
tile_flops = 2 * TILE_M * TILE_K * TILE_N
|
||||
dma_w_per_pair = (TILE_M * TILE_N * _BPE) / _HBM_GBS
|
||||
head_ns = (_D_STAGES - 1) * _T_STAGE
|
||||
|
||||
by_shape = {(r["M"], r["K"], r["N"]): r
|
||||
for r in rows if r["variant"] == _PLOT_VARIANT}
|
||||
@@ -434,11 +428,8 @@ def emit_mac_utilization_measured(
|
||||
tiles = r["tile_count_expected"]
|
||||
gu = useful / (tile_flops * tiles) * 100
|
||||
gemm_util.append(gu)
|
||||
m_tiles = (M + TILE_M - 1) // TILE_M
|
||||
n_tiles = (N + TILE_N - 1) // TILE_N
|
||||
n_mn = m_tiles * n_tiles
|
||||
compute_total = tiles * _T_STAGE
|
||||
wall = head_ns + tiles * _T_STAGE + max(0, n_mn - 1) * dma_w_per_pair
|
||||
wall = _T_PIPELINE_FILL + compute_total + _T_PIPELINE_TAIL
|
||||
ueff = (compute_total * (gu / 100.0) / wall) * 100 if wall > 0 else 0.0
|
||||
useful_eff.append(ueff)
|
||||
|
||||
@@ -467,8 +458,6 @@ def emit_mac_utilization_theoretical_vs_measured(
|
||||
tile = data["tile_sizes"]
|
||||
TILE_M, TILE_K, TILE_N = tile["M"], tile["K"], tile["N"]
|
||||
tile_flops = 2 * TILE_M * TILE_K * TILE_N
|
||||
dma_w_per_pair = (TILE_M * TILE_N * _BPE) / _HBM_GBS
|
||||
head_ns = (_D_STAGES - 1) * _T_STAGE
|
||||
peak_per_ns = tile_flops / _T_STAGE
|
||||
|
||||
by_shape = {(r["M"], r["K"], r["N"]): r
|
||||
@@ -490,11 +479,8 @@ def emit_mac_utilization_theoretical_vs_measured(
|
||||
gu_t.append(gut * 100)
|
||||
rec = r.get("stages", {}).get("GEMM", {}).get("record_count", 0) or tiles
|
||||
gu_m.append((useful / (tile_flops * rec) * 100) if rec else 0.0)
|
||||
m_tiles = (M + TILE_M - 1) // TILE_M
|
||||
n_tiles = (N + TILE_N - 1) // TILE_N
|
||||
n_mn = m_tiles * n_tiles
|
||||
compute_total = tiles * _T_STAGE
|
||||
wall_t = head_ns + compute_total + max(0, n_mn - 1) * dma_w_per_pair
|
||||
wall_t = _T_PIPELINE_FILL + compute_total + _T_PIPELINE_TAIL
|
||||
eff_t.append((compute_total * gut / wall_t * 100) if wall_t > 0 else 0.0)
|
||||
cw = r.get("composite_window_ns", 0.0) or 0.0
|
||||
eff_m.append((useful / cw / peak_per_ns * 100) if cw > 0 else 0.0)
|
||||
|
||||
@@ -88,18 +88,21 @@ class PeDmaComponent(PeEngineBase):
|
||||
path = self.ctx.router.find_path(self._pe_prefix, dst_node)
|
||||
drain_ns = self.ctx.compute_drain_ns(path, cmd.nbytes)
|
||||
|
||||
# Acquire DMA channel — held through the entire round-trip so the
|
||||
# channel models "one DMA in flight per PE per direction" rather
|
||||
# than just issue-time serialization. This is what makes Option B
|
||||
# meaningful: t_start = serve-start covers the actual transfer.
|
||||
# Acquire DMA channel for the **issue path only** (ADR-0014 D4
|
||||
# clarified): the channel models the engine's issue-rate limit, not
|
||||
# full-round-trip occupancy. HBM-level serialization is the
|
||||
# HBM_CTRL's responsibility (per-PC `available_at` timestamps).
|
||||
# Holding the channel through the round-trip would double-serialize
|
||||
# and prevent back-to-back head-amortization, capping small-tile
|
||||
# pipeline efficiency.
|
||||
sub_done = env.event()
|
||||
with dma_res.request() as req:
|
||||
yield req
|
||||
# Option B: record_start fires AFTER channel acquired, so t_start
|
||||
# = serve-start (excludes queue wait). _DEFER_RECORD_START=True
|
||||
# suppresses the auto-start in ComponentBase._handle_with_hooks.
|
||||
# record_start fires AFTER channel acquired (t_start =
|
||||
# issue-start). _DEFER_RECORD_START=True suppresses the
|
||||
# auto-start in ComponentBase._handle_with_hooks.
|
||||
self._on_process_start(env, cmd)
|
||||
# Create sub-Transaction with PeDmaMsg (HbmCtrl handles it directly)
|
||||
sub_done = env.event()
|
||||
sub_request = PeDmaMsg(
|
||||
correlation_id="pe_internal",
|
||||
request_id=f"dma_{id(pe_txn)}",
|
||||
@@ -114,8 +117,11 @@ class PeDmaComponent(PeEngineBase):
|
||||
# Send to next hop (path[0] is pe_dma itself, path[1] is router)
|
||||
if len(path) > 1:
|
||||
yield self.out_ports[path[1]].put(sub_txn.advance())
|
||||
# Wait for HBM transfer completion BEFORE releasing the channel.
|
||||
yield sub_done
|
||||
# Channel released here; next DMA can issue immediately while
|
||||
# this one's HBM round-trip is still in flight.
|
||||
# Wait for HBM transfer completion OUTSIDE the channel hold so
|
||||
# back-to-back DMAs can pipeline through the fabric.
|
||||
yield sub_done
|
||||
pe_txn.done.succeed()
|
||||
|
||||
def _worker(self, env: simpy.Environment) -> Generator:
|
||||
@@ -355,14 +361,16 @@ class PeDmaComponent(PeEngineBase):
|
||||
path = self.ctx.router.find_path(self._pe_prefix, dst_node)
|
||||
drain_ns = self.ctx.compute_drain_ns(path, nbytes)
|
||||
|
||||
# Hold dma_res through the full round-trip — one DMA in flight
|
||||
# per PE per direction — so Option B's t_start (post-acquire)
|
||||
# bounds the actual transfer interval.
|
||||
# Channel held for the issue path only (ADR-0014 D4 clarified):
|
||||
# PE_DMA's capacity=1 throttles the issue rate; HBM-level
|
||||
# serialization is HBM_CTRL's responsibility (per-PC
|
||||
# `available_at`). Releasing the channel after issue lets
|
||||
# back-to-back tile DMAs amortize head latency through the
|
||||
# fabric, which is essential for small-tile pipelining.
|
||||
sub_done = env.event()
|
||||
with dma_res.request() as req:
|
||||
yield req
|
||||
# Option B: t_start = post-acquire moment.
|
||||
self._on_process_start(env, token)
|
||||
sub_done = env.event()
|
||||
sub_request = PeDmaMsg(
|
||||
correlation_id="pipeline",
|
||||
request_id=f"tile_{token.tile_id}",
|
||||
@@ -376,7 +384,8 @@ class PeDmaComponent(PeEngineBase):
|
||||
)
|
||||
if len(path) > 1:
|
||||
yield self.out_ports[path[1]].put(sub_txn.advance())
|
||||
yield sub_done
|
||||
# channel released here
|
||||
yield sub_done
|
||||
else:
|
||||
# No-op (nbytes==0 or no ctx): no channel wait, but still record
|
||||
# so _on_process_end has a matching pending entry to finalise.
|
||||
|
||||
Reference in New Issue
Block a user