gqa(adr-0065): N3 — GEMM-epilogue accumulate dataflow (O = O.corr + P.V)
Two fixes make the recipe's P.V GEMM fold into the rescaled accumulator: (1) op_log marks the composite gemm record accumulate=True when an 'add' epilogue targets the GEMM's own output; data_executor._execute_gemm then does O = O_existing + a@b instead of overwriting. (2) PE_SCHEDULER records the composite's numeric op at COMPLETION (in _retire_on_done) instead of at dispatch, so the GEMM's t_start sorts AFTER its prologue MATH ops -> it reads the computed P and the rescaled O. Verified end-to-end (DataExecutor): the full recipe produces O = O_old*corr + P@V matching a numpy flash-attention reference. Suite 816 pass / 3 pre-existing fail. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -165,14 +165,9 @@ class PeSchedulerComponent(ComponentBase):
|
||||
from kernbench.components.builtin.pe_types import PipelineContext
|
||||
|
||||
yield from self._hazard.admit(env, cmd, pe_txn.done)
|
||||
env.process(self._retire_on_done(env, pe_txn.done, cmd.completion.id))
|
||||
|
||||
# Record the composite as a single numeric op for data-mode replay
|
||||
# (ADR-0027): the tiles model latency; this zero-duration record
|
||||
# carries the matmul (a@b → out) the DataExecutor replays. No-op
|
||||
# without an op_logger (latency-only mode), so latency is unchanged.
|
||||
self._on_process_start(env, cmd)
|
||||
self._on_process_end(env, cmd)
|
||||
# Retire from the hazard tracker AND record the composite's numeric
|
||||
# op — both at completion (see _retire_on_done).
|
||||
env.process(self._retire_on_done(env, pe_txn.done, cmd))
|
||||
|
||||
plan = self._generate_plan(cmd)
|
||||
|
||||
@@ -192,11 +187,20 @@ class PeSchedulerComponent(ComponentBase):
|
||||
yield self._pending_feeds.put((plan, ctx))
|
||||
|
||||
def _retire_on_done(
|
||||
self, env: simpy.Environment, done_event: Any, completion_id: str,
|
||||
self, env: simpy.Environment, done_event: Any, cmd: Any,
|
||||
) -> Generator:
|
||||
"""Retire a composite from the RW hazard tracker on completion."""
|
||||
"""On composite completion: retire from the RW hazard tracker AND
|
||||
record its numeric op (ADR-0027 / ADR-0065 N1+N3).
|
||||
|
||||
Recording at completion (not dispatch) gives the composite's GEMM a
|
||||
t_start AFTER its prologue MATH ops, so in data mode it reads their
|
||||
results (P, the rescaled O) instead of stale values. No-op without an
|
||||
op_logger → latency-only mode unchanged.
|
||||
"""
|
||||
yield done_event
|
||||
self._hazard.retire(completion_id)
|
||||
self._hazard.retire(cmd.completion.id)
|
||||
self._on_process_start(env, cmd)
|
||||
self._on_process_end(env, cmd)
|
||||
|
||||
def _feed_loop(self, env: simpy.Environment) -> Generator:
|
||||
"""Single feeder process: FIFO command ordering (ADR-0014 D6).
|
||||
|
||||
@@ -174,9 +174,21 @@ class DataExecutor:
|
||||
dtype_acc = p.get("dtype_acc", "f32")
|
||||
a_f = a.astype(_resolve_dtype(dtype_acc))
|
||||
b_f = b.astype(_resolve_dtype(dtype_acc))
|
||||
result = np.matmul(a_f, b_f).astype(_resolve_dtype(dtype_out))
|
||||
result = np.matmul(a_f, b_f)
|
||||
|
||||
self.store.write(dst_space, p["dst_addr"], result)
|
||||
# Accumulate into the existing output when requested (recipe O +=
|
||||
# P·V — the GEMM folds into the already-rescaled accumulator). The
|
||||
# composite is recorded at completion so the rescale has run first.
|
||||
if p.get("accumulate"):
|
||||
try:
|
||||
existing = self.store.read(dst_space, p["dst_addr"],
|
||||
shape=result.shape, dtype=dtype_out)
|
||||
result = existing.astype(_resolve_dtype(dtype_acc)) + result
|
||||
except (KeyError, ValueError):
|
||||
pass
|
||||
|
||||
self.store.write(dst_space, p["dst_addr"],
|
||||
result.astype(_resolve_dtype(dtype_out)))
|
||||
|
||||
def _execute_math(self, op: OpRecord) -> None:
|
||||
"""Execute math op: unary, binary, or reduction."""
|
||||
|
||||
@@ -262,10 +262,18 @@ def _extract_op_info(msg: Any) -> tuple[str, str, dict[str, Any]]:
|
||||
if gemm is not None:
|
||||
a = gemm.operands.get("a")
|
||||
b = gemm.operands.get("b")
|
||||
# Accumulate (O += a@b) when an `add` epilogue targets the GEMM's
|
||||
# own output (the recipe pattern: GEMM P·V folds into the already-
|
||||
# rescaled O). Default overwrite for a plain GEMM composite.
|
||||
accumulate = any(
|
||||
o.kind == "add" and o.operands.get("other") is gemm.out
|
||||
for o in msg.ops
|
||||
)
|
||||
params: dict[str, Any] = {
|
||||
"op": "gemm",
|
||||
"out_addr": gemm.out.addr,
|
||||
"out_nbytes": gemm.out.nbytes,
|
||||
"accumulate": accumulate,
|
||||
}
|
||||
if a is not None and b is not None:
|
||||
params.update({
|
||||
|
||||
Reference in New Issue
Block a user