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

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