[Fix] Restore the fx.maxnumf contract, and keep ninf off the score reductions

[JohnQinAMD]

Follow-ups to #1035, both from @jhinpan's post-merge review.

1. fx.maxnumf result contract

fx.maxnumf is in the stable catalog and documented to return the DSL type of a. The bare @dsl_math_wrap_result it moved to does not do that:

call	before #1035	on main
maxnumf(a=x, b=x), Float32	Float32	raw ArithValue
keyword Vector shape (2, 2)	(2, 2)	flat (4,)
positional Vector(..., dtype=Float)	Float	Float32

exemplar="a" fixes the first two but not the third, and preserve_numeric_type=True changes generic-scalar behaviour, so this hand-wraps it again and resolves the ambient inside — exact parity plus the flag. Three regressions added for keyword scalar, keyword vector and generic vector dtype; they fail on main.

2. ninf on the score reductions

The 16 maxnumf sites in flash_attn_utils.py inherit the ambient fast, which carries ninf. But _score_pair_max is seeded with literal -inf and masked scores are -inf, and an infinite operand under ninf is poison. pa_decode_tile.py already records the invariant:

nnan lets maxnum lower to a bare v_max […] (ninf must NOT be set: -inf is load-bearing.)

So they carry nnan explicitly now. swa_gfx950.py masks with -inf too and was reaching for the ambient through a current_fastmath() workaround whose comment no longer held; it uses the same flag.

The speed was in nnan. Two rounds, cache cold, B=2 H=32 D=128 bf16:

S	main (ambient fast)	this PR (nnan)
61,380	1210.3 / 1212.6	1213.5 / 1213.4
180,180	1208.3 / 1207.0	1207.2 / 1206.7

Within noise, so #1035's recovery is untouched.

Testing

• test_fastmath_context.py — 25 passed, no GPU. Includes an IR guard that the reduction flag emits fastmath<nnan> with no ninf.
• test_flash_attn_fwd.py — 94 passed. Includes a check that the flag excludes ninf and that no bare fx.maxnumf is left in that module.
• test_swa_gfx950.py — 5 passed.
• Benchmarks — table above.
• No new third-party dependencies added

Breaking Changes

None. The contract change restores the pre-#1035 behaviour.

[Copilot]

Pull request overview

Restores fx.maxnumf’s stable return-type/shape contract and updates attention score-reduction call sites to use a fastmath flag that preserves -inf mask semantics (i.e., keep ninf off score reductions while still allowing nnan-driven lowering).

Changes:

• Re-introduce a hand-wrapped fx.maxnumf implementation so it returns the DSL type/shape of operand a (including logical vector shape and generic vector dtype).
• Introduce _MAX_FASTMATH = FastMathFlags.nnan for attention score reductions and thread it through all fx.maxnumf reduction sites (and SWA gfx950) to avoid ninf poisoning with -inf sentinels.
• Add regression tests covering the fx.maxnumf contract and guarding that score-reduction fastmath excludes ninf.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
tests/unit/test_fastmath_context.py	Adds contract + fastmath-flag regression tests for fx.maxnumf.
tests/kernels/test_flash_attn_fwd.py	Adds a guard asserting score-reduction fastmath includes nnan but excludes ninf, and that no bare fx.maxnumf remains in flash_attn_utils.py.
python/flydsl/expr/arith.py	Re-wraps maxnumf to preserve the stable return-type/shape contract and resolve ambient fastmath explicitly.
kernels/attention/swa_gfx950.py	Switches SWA masking max to use _MAX_FASTMATH (nnan-only) instead of inheriting ambient flags (which may include ninf).
kernels/attention/flash_attn_utils.py	Defines _MAX_FASTMATH and applies it to all score-reduction fx.maxnumf sites to prevent ninf interaction with -inf masks.

---

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

[JohnQinAMD]

Merge-order note: this PR's guard asserts no bare fx.maxnumf is left in flash_attn_utils.py, and #1033 adds one — m_new = fx.maxnumf(m_row, m_tile_max) in _lazy_rescale_o_rescale, deliberately bare after review there asked for the explicit flag to be dropped once ambient resolution landed.

Either PR is fine on its own; whichever merges second needs that one call to carry _MAX_FASTMATH. Flagging it so it does not turn up as a surprise CI failure. Happy to fold the change into whichever lands first if you tell me the order you want.