[hephaex]

Summary

GeluActivation::TanhF32Compiled reaches Metal as a mlx::core::compile region that expands to
roughly eleven eager elementwise passes: upcast to f32, square, cube, the
0.7978845608028654 * (x + 0.044715 x^3) polynomial, tanh, 0.5 * (1 + ...), multiply, cast
back. This makes it one Metal kernel.

This is the Metal counterpart to #2780, which reaches the same conclusion for CUDA and puts
Metal explicitly out of scope ("Metal and CPU keep today's paths", "Out of scope: ... Metal and
ROCm kernels"). Closes #2857.

Design

• Kernel source in src/lib/mlxcelverse/kernels/gelu_tanh_f32.{h,cpp}, one thread per lane,
  threadgroup 256, metal::precise::tanh. Same formula and op order as the eager chain, so a
  lane may differ only by f32 rounding where the fused body contracts a multiply-add into an FMA.
• A kernel port with .metal = ..., .cuda = nullptr, .rocm = nullptr — the perf(rocm): allocator footprint and ROCm ports of mlxcel fused kernels #1814 convention
  for a backend-specific port.
• fused_gelu_tanh_f32_available() reads the same table the launcher reads and requires the GPU
  stream, so MLXCEL_DEVICE=cpu keeps the eager chain instead of throwing (mirrors
  fused_xielu_kernel_available, and fused_xielu is the pattern this follows throughout).
• GeluActivation::TanhF32Compiled goes through gelu_tanh_f32_compiled_where_available, falling
  back to compiled_gelu_tanh_f32 when the port is null for the backend.

One deliberate deviation from #2780: plain TanhF32 stays eager. in_f32() has no production
caller — Tanh reaches the GPU through in_f32_compiled(), which already folds to
TanhF32Compiled — and TanhF32 is pinned bit-for-bit to the private copies it replaced in
utils_gelu_tests.rs. Routing it would buy nothing measurable and would break that pin. #2780
contemplates reversing its own two-variant routing and updating those doc comments together; if
you want that here, it is a small follow-up on top of this.

Type of change

• perf — performance improvement (before/after numbers below)

Before / after

M1 Pro, bf16, microseconds per call, minimum of 7 interleaved trials, run under caffeinate -i.
Harness: examples/fused_gelu_tanh_microbench.rs. Lower is better.

Shape	eager	compiled	fused	fused / compiled	fused / eager
tower prefill [64, 4304]	587.29	342.31	279.96	0.82x	0.48x
decode row [1, 4304]	314.52	252.87	220.77	0.87x	0.70x
large prefill [1024, 4096]	2497.79	959.47	738.92	0.77x	0.30x

Variance, disclosed: across three sweeps the ordering held (fused fastest, eager slowest) but
large-prefill fused/compiled ranged 0.71x–0.85x and run-to-run absolute figures moved by up to
~30%, so read that row as a range rather than a point. Host load was not sampled; caffeinate -i
held throughout, which is why the harness takes the minimum of interleaved trials: interference
only ever adds time, so the minimum is the estimator least biased by it, and interleaving keeps
the three arms under identical conditions. These are warm-cache op-level figures, not production
latency.

Traced, per "Judging a change that moves the numbers"

• Where the disagreement lands: at decided positions only. fused_tanh_f32_matches_the_eager_helper_where_the_kernel_runs
  reports lane-level differing / max_abs / max_rel; on the tower shape [64, 4304] that is
  1941 of 275456 lanes, max abs 2.3841858e-7, i.e. f32 rounding, not a different function.
• Width and context traced at: bf16 activations, widths 4304 (tower and decode) and 4096
  (large prefill), five-context prefill at [64, 4304]. Magnitudes reach ~170, past the vision
  projection this targets, so the tolerance is not fitted to the easy case.
• Not vacuous: substituting a no-op body (out[i] = (T)xv;) fails 3 of the 27 gelu tests,
  including the dtype-preservation test that reads values back rather than checking dtype alone.

Test plan

• cargo test -p mlxcel-core gelu → 27 passed, 0 failed
• make verify-fmt → exit 0
• make verify-clippy (--workspace --all-targets -D warnings) → exit 0
• make verify-kernel-port-dispatch → exit 0 (13 launchers, 0 exempt)
• make verify-kernel-dtype-keys → exit 0

New tests: fused_kernel_returns_the_input_dtype (value-level, per-dtype tolerance),
fused_tanh_f32_matches_the_eager_helper_where_the_kernel_runs,
compiled_dispatch_falls_back_to_the_compiled_chain_on_the_cpu (holds
streams::lock_default_device() so it cannot leave the process on CPU for gated tests),
empty_input_takes_the_same_route_on_both_paths, plus dual tolerance for
compiled_tanh_f32_matches_the_eager_helper_on_tower_scale_bf16 where 1 + tanh(...) cancels to
zero in the far-negative tail and a relative-only tolerance is meaningless.

Two verify-test failures (a Xiaomi entry in family_order_is_exhaustive, and a 49-vs-50 roster
count in create_causal_mask_roster_count_matches_repo) reproduce on a clean 630aaeba, so they
are pre-existing and not from this branch.

Checklist

• make verify passes apart from the two pre-existing failures above
• No #[allow(dead_code)] added; doc comments describe the real two-stage dispatch
• Benchmarks show differences, and the kernel is registered in docs/benchmarks.md

[cla-assistant]

All committers have signed the CLA.