gqa(adr-0065): N2 — softmax_merge recipe MATH ops compute in data mode

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) <noreply@anthropic.com>

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(
+    operands = list(op.operands.values())
-        StageType.MATH, f"{pe_prefix}.pe_math",
+    params: dict = {
-        {"op_kind": op.kind, "num_elements": num_elements, "scope": "kernel"},
+        "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(

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:

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

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()))

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"