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(