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:
+132
-30
@@ -313,19 +313,26 @@ def test_pipeline_overlap_within_command():
|
||||
|
||||
|
||||
def test_pe_dma_record_start_after_channel_acquire():
|
||||
"""Three back-to-back DMA_READs serialise on pe_dma.cap=1.
|
||||
"""Three back-to-back DMA_READs at the same PE_DMA.
|
||||
|
||||
With ``_DEFER_RECORD_START = True`` on PeDmaComponent, each op's
|
||||
``t_start`` is captured right after ``yield req`` succeeds. Result:
|
||||
``t_start`` is captured right after ``yield req`` succeeds (Option B:
|
||||
record_start fires post-channel-acquire, not at command-entry).
|
||||
The DMA channel is held only for the **issue path** (ADR-0014 D4),
|
||||
not for the full HBM round-trip — HBM-level serialization is the
|
||||
HBM_CTRL's responsibility (per-PC `available_at`).
|
||||
|
||||
- op N's ``(t_end - t_start)`` is the *actual transfer time* — same
|
||||
across all three ops (no queueing inflation).
|
||||
- op N+1's ``t_start`` >= op N's ``t_end - epsilon`` (waited for the
|
||||
previous holder to release the channel before being recorded).
|
||||
Under that model, three back-to-back DMAs at t=0 give:
|
||||
|
||||
Counter-example (the bug this fix addresses): if ``record_start`` fired
|
||||
on command entry, all three ops would share ``t_start == 0`` and the
|
||||
second/third would show inflated ``t_end - t_start``.
|
||||
- ``t_start`` values monotonic and tightly clustered (channel
|
||||
release is fast → each subsequent op's ``t_start`` is at or
|
||||
immediately after the previous one's, NOT after the previous
|
||||
one's HBM round-trip completes).
|
||||
- ``t_end`` values reflect HBM_CTRL serialization (later ops can
|
||||
have longer duration as they queue at the per-PC).
|
||||
|
||||
Counter-example (the bug this rules out): if ``record_start`` fired
|
||||
on command entry, all three ops would share ``t_start == 0``.
|
||||
"""
|
||||
from pathlib import Path
|
||||
|
||||
@@ -380,28 +387,123 @@ def test_pe_dma_record_start_after_channel_acquire():
|
||||
)
|
||||
|
||||
durations = [r.t_end - r.t_start for r in dma_records]
|
||||
# All three should have similar transfer time. Under the PC striping
|
||||
# model (ADR-0033 D1), per-PC `available_at` state introduces small
|
||||
# timing differences between consecutive same-direction reads to the
|
||||
# same PC set (the second read's start = max(eff_start, pc_avail[pc])).
|
||||
# Tolerance widened from ±1ns to ±3ns to absorb this variance without
|
||||
# weakening the invariant that queue wait is excluded from the recorded
|
||||
# interval (still validated by the t_start >= prev_end check below).
|
||||
base = durations[0]
|
||||
assert base > 0, f"first dma duration must be positive, got {base}"
|
||||
for i, d in enumerate(durations):
|
||||
assert abs(d - base) <= 3.0, (
|
||||
f"op {i} duration {d} differs from baseline {base} by >3 ns "
|
||||
f"— record_start may still be including queue wait"
|
||||
)
|
||||
|
||||
# Each subsequent op's t_start must be at or after the previous op's
|
||||
# t_end (modulo a few ns of scheduler overhead) — i.e. the wait is
|
||||
# *excluded* from the recorded interval, not folded into it.
|
||||
for i in range(1, len(dma_records)):
|
||||
prev_end = dma_records[i - 1].t_end
|
||||
cur_start = dma_records[i].t_start
|
||||
assert cur_start >= prev_end - 1.0, (
|
||||
f"op {i} t_start={cur_start} began before op {i-1} t_end={prev_end} "
|
||||
f"— channel was not actually held, fix is incorrect"
|
||||
# t_start values must be monotonic and tightly clustered. The window
|
||||
# over all three t_starts is bounded by the per-op channel-issue cost,
|
||||
# NOT by the HBM round-trip — i.e. record_start fires post channel
|
||||
# acquire AND channel release is fast.
|
||||
t_starts = [r.t_start for r in dma_records]
|
||||
for i in range(1, len(t_starts)):
|
||||
assert t_starts[i] >= t_starts[i - 1], (
|
||||
f"op {i} t_start={t_starts[i]} regressed before "
|
||||
f"op {i-1} t_start={t_starts[i-1]} — record order broken"
|
||||
)
|
||||
t_start_window = t_starts[-1] - t_starts[0]
|
||||
assert t_start_window < base, (
|
||||
f"t_start window {t_start_window:.1f}ns across 3 back-to-back "
|
||||
f"ops is >= a single op's duration {base:.1f}ns — PE_DMA "
|
||||
f"channel is being held beyond the issue path (ADR-0014 D4)."
|
||||
)
|
||||
|
||||
|
||||
# ── 7. Back-to-back DMA pipelining (head amortization) ───────────────
|
||||
|
||||
|
||||
def test_pe_dma_back_to_back_pipelining():
|
||||
"""Back-to-back DMA_READs should pipeline so that per-tile head latency
|
||||
is paid once, not N times.
|
||||
|
||||
Setup: 8 back-to-back 4 KiB DMA_READ commands to the same PE_DMA's
|
||||
inbox at t=0. HBM_CTRL serializes via its per-PC availability
|
||||
timestamps; PE_DMA should release its channel after issuing the
|
||||
request, NOT hold it through the full HBM round-trip.
|
||||
|
||||
Expected (post-fix model):
|
||||
- First op pays full round-trip: head + transmission.
|
||||
- Subsequent ops pay only transmission time (head amortized): each
|
||||
op's t_end advances by ~the per-request HBM service time, NOT by
|
||||
the full round-trip.
|
||||
- Total wall (last t_end - first t_start) is close to the
|
||||
transmission-bound limit: ``head_once + N * service``, not
|
||||
``N * (head + service)``.
|
||||
|
||||
Anti-counter-example: under the current PE_DMA round-trip-hold model
|
||||
each op N+1's t_start equals op N's t_end, so the wall is exactly
|
||||
``N * (head + service)``. That is what this test rules out.
|
||||
|
||||
Concrete numbers (32 KiB probe = 141 ns; 4 KiB single-request inferred
|
||||
at ~29 ns): the current model gives wall ~232 ns for 8 ops; the
|
||||
fixed model should give wall ≲ 160 ns (head ~13 + 8 * ~16 ns).
|
||||
"""
|
||||
from pathlib import Path
|
||||
|
||||
from kernbench.common.pe_commands import DmaReadCmd, PeInternalTxn, TensorHandle
|
||||
from kernbench.policy.address.phyaddr import PhysAddr
|
||||
from kernbench.sim_engine.engine import GraphEngine
|
||||
from kernbench.topology.builder import load_topology
|
||||
|
||||
TOPOLOGY_PATH = Path(__file__).parent.parent / "topology.yaml"
|
||||
|
||||
def _hbm_pa(offset: int) -> int:
|
||||
slice_bytes = 48 * (1 << 30) // 8
|
||||
pa = PhysAddr.pe_hbm_addr(
|
||||
sip_id=0, die_id=0, pe_id=0,
|
||||
pe_local_hbm_offset=offset, slice_size_bytes=slice_bytes,
|
||||
)
|
||||
return pa.encode()
|
||||
|
||||
engine = GraphEngine(load_topology(TOPOLOGY_PATH), enable_data=True)
|
||||
pe_dma_id = "sip0.cube0.pe0.pe_dma"
|
||||
pe_dma = engine._components[pe_dma_id]
|
||||
env = engine._env
|
||||
|
||||
N_OPS = 8
|
||||
BYTES_PER_OP = 4096
|
||||
|
||||
handles = [
|
||||
TensorHandle(id=f"r{i}", addr=0x1000 + i * 0x1000,
|
||||
shape=(32, 64), dtype="f16", nbytes=BYTES_PER_OP)
|
||||
for i in range(N_OPS)
|
||||
]
|
||||
# Spread across distinct cache lines so HBM_CTRL hits different PCs
|
||||
# and we measure the BW-only floor, not per-PC bank conflict.
|
||||
cmds = [
|
||||
DmaReadCmd(handle=h, src_addr=_hbm_pa(0x1000 * (1 + i)),
|
||||
nbytes=BYTES_PER_OP)
|
||||
for i, h in enumerate(handles)
|
||||
]
|
||||
txns = [PeInternalTxn(command=c, done=env.event()) for c in cmds]
|
||||
|
||||
def submit_all():
|
||||
for txn in txns:
|
||||
yield pe_dma._inbox.put(txn)
|
||||
|
||||
env.process(submit_all())
|
||||
env.run()
|
||||
|
||||
dma_records = [
|
||||
r for r in engine.op_log
|
||||
if r.op_name == "dma_read" and r.component_id == pe_dma_id
|
||||
]
|
||||
assert len(dma_records) == N_OPS, (
|
||||
f"expected {N_OPS} dma_read records, got {len(dma_records)}"
|
||||
)
|
||||
dma_records.sort(key=lambda r: r.t_start)
|
||||
|
||||
single_op_ns = dma_records[0].t_end - dma_records[0].t_start
|
||||
wall = dma_records[-1].t_end - dma_records[0].t_start
|
||||
|
||||
# Under the round-trip-hold model: wall == N_OPS * single_op_ns.
|
||||
# Under the issue-only-hold model: wall ≈ single_op_ns + (N-1)*service
|
||||
# where service ≈ BYTES_PER_OP / hbm_bw_per_pe ≈ 16 ns.
|
||||
# We assert wall < 0.75 * (N_OPS * single_op_ns) — i.e. amortization
|
||||
# of head must save at least 25% vs strict serialization.
|
||||
strict_serialized = N_OPS * single_op_ns
|
||||
assert wall < 0.75 * strict_serialized, (
|
||||
f"wall={wall:.1f}ns vs strict-serialized={strict_serialized:.1f}ns "
|
||||
f"(N_OPS={N_OPS}, single_op={single_op_ns:.1f}ns). "
|
||||
f"PE_DMA is still holding channel through full round-trip — "
|
||||
f"head latency is not amortized across back-to-back DMAs."
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user