From 3d4a3d43c44359f282c451ada8c4eb5846fa6bac Mon Sep 17 00:00:00 2001 From: Yangwook Date: Wed, 10 Jun 2026 22:56:25 -0700 Subject: [PATCH] =?UTF-8?q?gqa(adr-0065):=20N2=20=E2=80=94=20softmax=5Fmer?= =?UTF-8?q?ge=20recipe=20MATH=20ops=20compute=20in=20data=20mode?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit data_executor._compute_math gains the 5 recipe ops (rmax/rsum as keepdims reductions, max_elem/exp_diff/mul_bcast as binary; numpy broadcasting covers bcast_axis). tiling._math_stage now carries operand+output addrs/shapes/spaces + axis on the recipe prologue MATH stages. op_log.record_end promotes a MATH stage to op_kind='math' ONLY when it carries input_addrs -> the DataExecutor runs it; legacy epilogue MATH stages (bias/relu, no addrs) stay op_kind-opaque -> byte-equal. Fixed a latent P2 lowering bug exposed by data mode: _resolve_recipe_dst gave reductions shape (G,) (axis removed) but max_elem broadcasts them against the running m=(G,1); reductions now keep the reduced axis as size 1 (keepdims) -> (G,1). Verified end-to-end: running the recipe's 8 MATH ops through the DataExecutor produces m, l (fully updated online-softmax) and O (rescaled by corr) matching a numpy reference. Suite 815 pass / 3 pre-existing fail. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/kernbench/components/builtin/tiling.py | 29 +++++-- src/kernbench/sim_engine/data_executor.py | 12 ++- src/kernbench/sim_engine/op_log.py | 6 ++ src/kernbench/triton_emu/tl_context.py | 9 +- tests/test_recipe_data_mode.py | 95 ++++++++++++++++++++++ 5 files changed, 142 insertions(+), 9 deletions(-) create mode 100644 tests/test_recipe_data_mode.py diff --git a/src/kernbench/components/builtin/tiling.py b/src/kernbench/components/builtin/tiling.py index 6b7f8f2..4933b47 100644 --- a/src/kernbench/components/builtin/tiling.py +++ b/src/kernbench/components/builtin/tiling.py @@ -227,13 +227,32 @@ def generate_math_plan( def _math_stage(op: object, pe_prefix: str) -> Stage: - """Single-shot MATH stage for a KERNEL-scope flat op (ADR-0065 D3).""" + """Single-shot MATH stage for a KERNEL-scope flat op (ADR-0065 D3). + + Carries operand + output addresses (ADR-0065 N2) so the DataExecutor can + compute the recipe's MATH op in data mode. Legacy epilogue MATH stages + (generated in ``generate_gemm_plan``) do NOT carry these, so their op_log + records stay op_kind-opaque → byte-equal. + """ out = getattr(op, "out", None) num_elements = prod(out.shape) if out is not None else 1 - return Stage( - StageType.MATH, f"{pe_prefix}.pe_math", - {"op_kind": op.kind, "num_elements": num_elements, "scope": "kernel"}, - ) + operands = list(op.operands.values()) + params: dict = { + "op_kind": op.kind, "num_elements": num_elements, "scope": "kernel", + "op": op.kind, + "input_addrs": [h.addr for h in operands], + "input_shapes": [h.shape for h in operands], + "input_spaces": [getattr(h, "space", "tcm") for h in operands], + "input_dtypes": [h.dtype for h in operands], + } + if out is not None: + params["dst_addr"] = out.addr + params["dst_space"] = getattr(out, "space", "tcm") + params["dtype"] = out.dtype + axis = op.extra.get("reduce_axis") if hasattr(op, "extra") else None + if axis is not None: + params["axis"] = axis + return Stage(StageType.MATH, f"{pe_prefix}.pe_math", params) def generate_plan_from_ops( diff --git a/src/kernbench/sim_engine/data_executor.py b/src/kernbench/sim_engine/data_executor.py index 28c3989..f00e952 100644 --- a/src/kernbench/sim_engine/data_executor.py +++ b/src/kernbench/sim_engine/data_executor.py @@ -270,9 +270,9 @@ def _compute_math(op: str, inputs: list[np.ndarray], axis: int | None) -> np.nda return np.sin(x) # Reduction - if op == "sum": + if op == "sum" or op == "rsum": return np.sum(x, axis=axis, keepdims=True) - if op == "max": + if op == "max" or op == "rmax": return np.max(x, axis=axis, keepdims=True) if op == "min": return np.min(x, axis=axis, keepdims=True) @@ -296,10 +296,16 @@ def _compute_math(op: str, inputs: list[np.ndarray], axis: int | None) -> np.nda return x * y if op == "div": return x / y - if op == "maximum": + if op == "maximum" or op == "max_elem": return np.maximum(x, y) if op == "minimum": return np.minimum(x, y) + # softmax_merge recipe ops (ADR-0065 D5). numpy broadcasting covers + # the recipe's bcast_axis (e.g. (G,1) corr against (G,d) O). + if op == "exp_diff": + return np.exp(x - y) + if op == "mul_bcast": + return x * y # Ternary if len(inputs) >= 3: diff --git a/src/kernbench/sim_engine/op_log.py b/src/kernbench/sim_engine/op_log.py index 60d5662..c2eea79 100644 --- a/src/kernbench/sim_engine/op_log.py +++ b/src/kernbench/sim_engine/op_log.py @@ -90,6 +90,12 @@ class OpLogger: params["stage_type"] = stage_type if op_name == "TileToken": op_name = f"TileToken/{stage_type}" + # ADR-0065 N2: a recipe MATH stage carries operand addresses → + # promote it to a computable math op so the DataExecutor runs it. + # Legacy epilogue MATH stages (no addresses) stay op_kind-opaque + # (latency-only) → byte-equal. + if stage_type == "MATH" and "input_addrs" in params: + op_kind = "math" # Snapshot data at record time so Phase 2 replay sidesteps # downstream mutations of source addrs (e.g. a tl.store that # overwrites HBM after a load handle was sent, or a slot that diff --git a/src/kernbench/triton_emu/tl_context.py b/src/kernbench/triton_emu/tl_context.py index e30e97d..9543290 100644 --- a/src/kernbench/triton_emu/tl_context.py +++ b/src/kernbench/triton_emu/tl_context.py @@ -975,8 +975,15 @@ class TLContext: ref = slots[recipe.primary_out.from_shape] shape, dtype = ref.shape, ref.dtype elif eop.reduce_axis is not None: + # Reductions keep the reduced axis as size 1 (keepdims), so the + # result broadcasts against the running (m, l) — e.g. rmax of + # s=(G,TILE) over axis -1 → (G,1), not (G,) (ADR-0065 N2). ref = next(iter(operands.values())) - shape = ref.shape[:-1] if len(ref.shape) > 1 else ref.shape + dims = list(ref.shape) + if len(dims) > 1: + ax = eop.reduce_axis % len(dims) + dims[ax] = 1 + shape = tuple(dims) dtype = ref.dtype else: ref = next(iter(operands.values())) diff --git a/tests/test_recipe_data_mode.py b/tests/test_recipe_data_mode.py new file mode 100644 index 0000000..7dccf47 --- /dev/null +++ b/tests/test_recipe_data_mode.py @@ -0,0 +1,95 @@ +"""Data-mode numeric tests for the softmax_merge recipe MATH ops (ADR-0065 N2). + +N2 makes the recipe's 8 MATH ops compute in data mode: `_compute_math` gains +rmax/rsum/max_elem/exp_diff/mul_bcast, the recipe MATH stages carry operand +addresses (`_math_stage`), and op_log `record_end` promotes an address- +carrying MATH stage to op_kind="math". The MATH ops fully compute the online- +softmax update of `m` and `l` (and rescale `O`); the `+P·V` of `O` is the +GEMM-accumulate step (N3). +""" +from __future__ import annotations + +import numpy as np + +from kernbench.common.pe_commands import CompositeCmd, TensorHandle +from kernbench.components.builtin.tiling import generate_plan_from_ops +from kernbench.sim_engine.data_executor import DataExecutor, _compute_math +from kernbench.sim_engine.memory_store import MemoryStore +from kernbench.sim_engine.op_log import OpRecord +from kernbench.triton_emu.tl_context import TLContext + +G, TILE, D = 8, 64, 128 + + +def test_compute_math_recipe_ops(): + x = np.random.randn(4, 6).astype(np.float32) + y = np.random.randn(4, 1).astype(np.float32) + assert np.allclose(_compute_math("rmax", [x], -1), x.max(-1, keepdims=True)) + assert np.allclose(_compute_math("rsum", [x], -1), x.sum(-1, keepdims=True)) + assert np.allclose(_compute_math("max_elem", [y, x.max(-1, keepdims=True)], None), + np.maximum(y, x.max(-1, keepdims=True))) + assert np.allclose(_compute_math("exp_diff", [x, y], None), np.exp(x - y)) + assert np.allclose(_compute_math("mul_bcast", [x, y], None), x * y) + + +def _tcm(addr, shape): + return TensorHandle(id=f"h{addr:x}", addr=addr, shape=shape, dtype="f16", + nbytes=2 * int(np.prod(shape)), space="tcm") + + +def test_softmax_merge_computes_m_l_O_rescale(): + """Run the recipe's prologue MATH ops through the DataExecutor and check + m, l (fully updated) and O (rescaled by corr) match a numpy reference.""" + rng = np.random.default_rng(0) + s = rng.standard_normal((G, TILE)).astype(np.float16) + m0 = rng.standard_normal((G, 1)).astype(np.float16) + l0 = np.abs(rng.standard_normal((G, 1))).astype(np.float16) + O0 = rng.standard_normal((G, D)).astype(np.float16) + + sh = _tcm(0x1000, (G, TILE)) + mh = _tcm(0x2000, (G, 1)) + lh = _tcm(0x3000, (G, 1)) + Oh = _tcm(0x4000, (G, D)) + V = TensorHandle(id="V", addr=0x5000, shape=(TILE, D), dtype="f16", + nbytes=2 * TILE * D, space="hbm") + + tl = TLContext(pe_id=0, num_programs=1, scratch_base=0x100000, + scratch_size=1 << 20) + tl.composite( + prologue=[{"op": "softmax_merge", "s": sh, "m": mh, "l": lh, "O": Oh}], + op="gemm", b=V, out=Oh, + ) + cmd = [c for c in tl.commands if isinstance(c, CompositeCmd)][-1] + plan = generate_plan_from_ops(cmd.ops, tile_m=32, tile_k=32, tile_n=32, + bytes_per_element=2, + pe_prefix="sip0.cube0.pe0") + + store = MemoryStore() + store.write("tcm", sh.addr, s) + store.write("tcm", mh.addr, m0) + store.write("tcm", lh.addr, l0) + store.write("tcm", Oh.addr, O0) + + records = [ + OpRecord(t_start=float(i), t_end=float(i), component_id="pe_math", + op_kind="math", op_name=st.params["op"], params=dict(st.params)) + for i, st in enumerate(plan.prologue_stages) + ] + DataExecutor(records, store).run() + + # numpy reference (online-softmax merge of a single new tile into m/l/O). + sf = s.astype(np.float32) + m_loc = sf.max(-1, keepdims=True) + m_new = np.maximum(m0.astype(np.float32), m_loc) + corr = np.exp(m0.astype(np.float32) - m_new) + P = np.exp(sf - m_new) + l_new = l0.astype(np.float32) * corr + P.sum(-1, keepdims=True) + O_resc = O0.astype(np.float32) * corr + + m_got = store.read("tcm", mh.addr, shape=(G, 1), dtype="f16").astype(np.float32) + l_got = store.read("tcm", lh.addr, shape=(G, 1), dtype="f16").astype(np.float32) + O_got = store.read("tcm", Oh.addr, shape=(G, D), dtype="f16").astype(np.float32) + + assert np.allclose(m_got, m_new, atol=2e-2), f"m: {m_got.ravel()[:3]} vs {m_new.ravel()[:3]}" + assert np.allclose(l_got, l_new, rtol=5e-2, atol=5e-2), "l mismatch" + assert np.allclose(O_got, O_resc, rtol=5e-2, atol=5e-2), "O rescale mismatch"