[Fix] Let fx.maxnumf inherit the ambient fastmath

[JohnQinAMD]

Summary

fx.maxnumf cannot see the ambient fastmath. It forwards straight to the raw MLIR builder, so with no explicit argument the op comes out fastmath<none> whatever the enclosing kernel asked for. #930 moved the kernels onto this function and the flag was dropped at every site it touched — 29 of them across attention, paged decode, MLA and moe. On the gfx950 dense bf16 path that costs 11.1–11.9% at long sequence lengths, and nothing failed.

Root cause

fx.maximumf and fx.minimumf carry @dsl_math_wrap_result and a fastmath parameter, and the decorator fills that parameter from the ambient context when the caller leaves it unset:

# expr/math.py
accepts_fastmath = "fastmath" in sig.parameters
...
if accepts_fastmath and kwargs.get("fastmath") is None:
    ambient = current_fastmath()
    if ambient is not None:
        kwargs["fastmath"] = ambient

fx.maxnumf had neither. It took **kwargs, hand-rolled its own Vector/Numeric wrapping, and passed straight to the raw builder, so the op came out fastmath<none> whatever the enclosing kernel asked for. #930 moved the attention kernels onto it and the flag was dropped at all 29 sites it touched.

This gives it the same shape as its neighbours; the hand-written wrapping goes away with it.

Result

MI355X, B=2 H=32 D=128 bf16 non-causal, stock kwargs, two rounds each, cache cold (FLYDSL_RUNTIME_ENABLE_CACHE=0), aiter measured in the same process as a control:

S	main	with this PR	change
61,380	1229.5	1226.4	−0.3%
180,180	1080.2	1171.1	+8.4%
239,580	1057.9	1153.4	+9.0%

Six timed rounds per point, alternating arms, idle machine, spread under 0.6% within each arm.

The step lands on #930 in a bisect: 47009a3 (8/14) measures 1188.6 / 1183.4 at the two long lengths, daee90e (#930, 8/17) measures 1077.0 / 1060.9, and everything between then and now is flat.

What each path inherits, and how it was checked

path	sites	inherits	verification
flash attention	16	fast	bit-identical over four input distributions including a planted NaN — ndiff=0/67108864, matching NaN masks, rel L2 identical to seven digits. 87 tests pass.
paged decode	22	contract	2304 tests, identical results before and after — including the three test_sliding_window_accuracy[...fp8] cases that already fail on main. contract permits only FMA fusion, which is inert for a max.
moe	2	none	unchanged; there is no hint on those launchers to inherit. 9 passed, 848 skipped, 22 xfailed, same as before.
MLA fp8	3	fast	pinned to contract in this PR rather than widened.

The MLA sites are pinned because tests/kernels/test_mla_decode.py aborts on this machine (Fatal Python error: Aborted, identically with and without this change), so the widening cannot be verified there. Pinning keeps them exactly where they are today; contract is inert for a max, and it matches what #930 did for the paged path. The argument should be dropped once that test runs.

Worth noting for whoever picks that up: test_pa.py and test_mla_decode.py skip silently when the installed aiter does not match flydsl (cannot import name 'fly_values' from 'flydsl.compiler.protocol'), collecting zero of 2305 tests and exiting 0. The runs above went through a shim that stubs aiter.ops.flydsl so the suites actually execute.

Guard

Three tests next to the existing ambient-fastmath ones, parameterised over maxnumf, maximumf and minimumf: each inherits the ambient flags, an explicit argument overrides them, and no flag appears outside the block. On main the maxnumf case fails and the other two pass, which is the asymmetry this fixes. IR-string checks marked l0_backend_agnostic, so CI runs them without a GPU.

Testing

python3 -m pytest tests/unit/test_fastmath_context.py -q                       # 17 passed, no GPU
FLYDSL_RUNTIME_ENABLE_CACHE=0 python3 -m pytest tests/kernels/test_flash_attn_fwd.py -q   # 87 passed

• Unit tests added — three, above.
• Existing tests pass — flash attention 87, paged decode 2304 (identical to main), moe unchanged, DSL unit file 17.
• Performance benchmarks run — table above.
• Tested on MI355X (gfx950).
• No new third-party dependencies added

Breaking Changes

None, on everything that could be measured: bit-identical on the attention path, identical test results on paged decode and moe, and MLA held at its current flags.

[Copilot]

Pull request overview

This PR fixes a performance regression introduced when kernels were migrated to fx.maxnumf by ensuring fx.maxnumf inherits the ambient fastmath flags from tracing_context (unless an explicit fastmath= is provided), matching the behavior of other DSL operators.

Changes:

• Update fx.maxnumf to resolve fastmath via resolve_fastmath() and forward it to the raw MLIR builder only when effective flags are present.
• Add backend-agnostic IR-string unit tests to guard ambient fastmath inheritance and explicit override behavior for fx.maxnumf.
• Pin a few MLA fx.maxnumf call sites to contract fastmath explicitly to avoid widening behavior in a path that can’t currently be validated on the author’s machine.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
tests/unit/test_fastmath_context.py	Adds IR-level unit tests verifying fx.maxnumf inherits ambient fastmath, explicit overrides work, and flags don’t leak outside a block.
python/flydsl/expr/arith.py	Makes maxnumf() resolve ambient fastmath (when unset) and forward the effective flag into arith.maxnumf.
kernels/attention/mla_fwd_decode_m16x8_fp8_fp8.py	Pins fx.maxnumf(..., fastmath=contract) at key reduction/max sites pending validation of widening.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[JohnQinAMD]

Right — it was already in the import list; the added entry was redundant and out of alphabetical order. Removed. Pushed.

[yanguahe]

LGTM

[jhinpan]

Post-merge follow-up from a three-pass review. Two items should be fixed before the next 0.3.x release.

1. Restore the stable fx.maxnumf result contract

Replacing the hand-written wrapper with bare @dsl_math_wrap_result changed valid result forms:

• maxnumf(a=x, b=x) with Float32 returned Float32 at the merge base and now returns raw ArithValue.
• A keyword Vector with logical shape (2, 2) now comes back as the flat (4,) shape.
• A positional Vector(..., dtype=Float) now comes back with concrete Float32 dtype instead of preserving the operand's generic dtype.

maxnumf is in the stable catalog, and the API documentation promises to preserve the DSL type of a, so this is a §4 result-type/layout regression. Existing tests cover only positional, concrete-dtype operands. exemplar="a" fixes keyword binding and logical shape but does not restore the generic vector dtype; preserve_numeric_type=True changes the existing generic-scalar behavior. The exact-parity fix is to retain the former hand wrapping while passing fastmath=resolve_fastmath(fastmath) to the raw builder. Please add keyword scalar/vector and generic-vector regression cases.

2. Exclude ninf where -inf is load-bearing

The 16 maxnumf sites in flash_attn_utils.py now inherit ambient fast, which includes ninf, while _score_pair_max is unconditionally seeded with literal -inf and masked scores also use -inf. LLVM defines an infinite operand under ninf as poison. The repository already records the corresponding invariant in pa_decode_tile.py: “ninf must NOT be set: -inf is load-bearing,” using nnan for the max optimization instead.

Current gfx950 MD5 comparisons and targeted masked/inf tests did not show a miscompile, so this is a structural poison risk rather than an observed numerical failure. Please carry a max-specific flag that excludes ninf at these reductions (the paged-decode nnan precedent is a good starting point), add an IR guard, and re-run the long-sequence A/B benchmark to confirm the performance recovery remains.

Smaller cleanup: swa_gfx950.py still contains the now-redundant current_fastmath() workaround and a comment that is no longer true. The “29 sites” figure also does not match the PR's own table or the source census.

[JohnQinAMD]

Both fixed in #1050.

Contract — hand-wrapped again with resolve_fastmath inside, which is the exact-parity option you named. Regressions for keyword scalar, keyword vector and generic vector dtype; they fail on main.

ninf — the 16 sites carry nnan explicitly now, same flag and same reasoning as pa_decode_tile.py. swa_gfx950.py masks with -inf too, so it uses the flag instead of the current_fastmath() workaround, whose comment was stale as you said.

Worth noting for the perf question: the speed was entirely in nnan. 1213.5 / 1207.2 TFLOP/s against main's 1210.3 / 1208.3 at S=61,380 / 180,180, two rounds — nothing lost by dropping the rest of fast.

Guards added: an IR check that the flag emits fastmath<nnan> with no ninf, and a check that no bare fx.maxnumf is left in that module.

You are right about the census — 16 in flash_attn_utils.py, not 29. The 29 counted every unflagged site in the repo, including paged decode and moe, which that PR never touched. Fixed in its body.