[Fix] Let fx.maxnumf inherit the ambient fastmath

kernels/attention/mla_fwd_decode_m16x8_fp8_fp8.py

@@ -735,7 +735,9 @@ def _warp_reduce_max_16(val):
         """Butterfly max reduce across MFMA column groups (strides 32, 16)."""
         w = _f32(val)
         for sh in [32, 16]:
-            w = fx.maxnumf(w, _shfl_xor_f32(w, sh))
+            # Pinned: fx.maxnumf now inherits this launcher's fast_fp_math, and
+            # test_mla_decode.py aborts here so that widening is unverified.
+            w = fx.maxnumf(w, _shfl_xor_f32(w, sh), fastmath=arith.FastMathFlags.contract)
         return w
 
     def _warp_reduce_add_16(val):
@@ -895,7 +897,7 @@ def _softmax(
         # Local max
         local_max = scaled[0]
         for i in range_constexpr(1, P_VALS_PER_THR):
-            local_max = fx.maxnumf(local_max, scaled[i])
+            local_max = fx.maxnumf(local_max, scaled[i], fastmath=arith.FastMathFlags.contract)
 
         # Warp reduce max (within 16-lane groups)
         local_max = _warp_reduce_max_16(local_max)
@@ -905,7 +907,7 @@ def _softmax(
             new_row_max = local_max
             rescale = c_one_f32
         else:
-            new_row_max = fx.maxnumf(local_max, row_max_old)
+            new_row_max = fx.maxnumf(local_max, row_max_old, fastmath=arith.FastMathFlags.contract)
             # rescale = exp2((old_max - new_max) * log2e)
             diff = row_max_old - new_row_max
             rescale = _fast_exp2(diff * c_log2e)

python/flydsl/expr/arith.py

@@ -290,17 +290,10 @@ def ceildiv(lhs, rhs, **kwargs):
 
 
 @dsl_loc_tracing
-def maxnumf(a, b, **kwargs):
+@dsl_math_wrap_result
+def maxnumf(a, b, *, fastmath=None, **kwargs):
     """Floating-point maximum, returning the non-NaN operand when one input is NaN (libm ``fmax``)."""
-    from .numeric import Numeric
-    from .typing import Vector
-
-    result = arith.maxnumf(as_ir_value(a), as_ir_value(b), **kwargs)
-    if isinstance(a, Vector):
-        return Vector(result, a.shape, a.dtype)
-    if isinstance(a, Numeric):
-        return Numeric.from_ir_type(result.type)(result)
-    return result
+    return arith.maxnumf(as_ir_value(a), as_ir_value(b), fastmath=fastmath, **kwargs)
 
 
 @dsl_loc_tracing

tests/unit/test_fastmath_context.py

@@ -301,3 +301,42 @@ def test_hint_changes_cache_key():
 
     assert key_none != key_fast
     assert key_fast != key_contract
+
+
+@pytest.mark.l0_backend_agnostic
+@pytest.mark.parametrize("op_name", ["maxnumf", "maximumf", "minimumf"])
+def test_minmax_inherits_ambient(op_name):
+    """These reach the raw MLIR builder, so they need the ambient lookup themselves."""
+    op = getattr(fx, op_name)
+
+    def build(a, b):
+        fa, fb = Float32(a), Float32(b)
+        with fx.fastmath(fx.FastMathFlags.fast):
+            op(fa, fb)
+
+    ir_text = _build(build, [ir.F32Type.get, ir.F32Type.get])
+    assert f"arith.{op_name}" in ir_text
+    assert "fastmath<fast>" in ir_text, f"{op_name} did not inherit the ambient flags:\n{ir_text}"
+
+
+@pytest.mark.l0_backend_agnostic
+@pytest.mark.parametrize("op_name", ["maxnumf", "maximumf", "minimumf"])
+def test_minmax_explicit_arg_overrides_ambient(op_name):
+    def build(a, b):
+        fa, fb = Float32(a), Float32(b)
+        with fx.fastmath(fx.FastMathFlags.fast):
+            getattr(fx, op_name)(fa, fb, fastmath="contract")
+
+    ir_text = _build(build, [ir.F32Type.get, ir.F32Type.get])
+    assert "fastmath<contract>" in ir_text and "fastmath<fast>" not in ir_text
+
+
+@pytest.mark.l0_backend_agnostic
+@pytest.mark.parametrize("op_name", ["maxnumf", "maximumf", "minimumf"])
+def test_minmax_unflagged_outside_block(op_name):
+    def build(a, b):
+        fa, fb = Float32(a), Float32(b)
+        getattr(fx, op_name)(fa, fb)
+
+    ir_text = _build(build, [ir.F32Type.get, ir.F32Type.get])
+    assert f"arith.{op_name}" in ir_text and "fastmath" not in ir_text