ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

gemm(perf): fix back-to-back DMA pipelining + tighten analytic model

Browse Source 
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 in
  92b9221 / e45626c but §5 still cites them by name).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yangwook Kang
2026-06-17 09:56:05 -07:00
parent 92b9221533
commit 23992548f7
15 changed files with 218 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kernbench/components/builtin/pe_dma.py
+25 -16
View File
@@ -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.