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