perf(metal): fuse the tanh-GELU f32 chain into one kernel

docs/benchmark_results/fused-gelu-tanh-f32-m1pro-2026-11-03.md

@@ -0,0 +1,103 @@
+# Fused tanh-GELU f32: measure-then-keep evidence
+
+Op-level measurement for the fused `gelu_tanh_f32` Metal kernel against the two
+chains it replaces. This follows the same measure-then-keep policy as the
+`#905` fused decode kernels: the wiring keeps its place only if it beats the
+graph it replaced at op level.
+
+`examples/fused_gelu_tanh_microbench.rs` produces this evidence. It loads no
+model and sweeps the three shapes the tanh-switched encoders feed.
+
+## Environment
+
+| Field | Value |
+|---|---|
+| Hardware | Apple M1 Pro, 10 cores, 32 GB unified memory |
+| OS | macOS 26.6.2 (Darwin 25.6.0) |
+| Backend | Metal, activations bf16 |
+| mlxcel | 0.8.1, worktree `mlxcel-2423`, base `630aaeba` |
+| Harness | `examples/fused_gelu_tanh_microbench.rs`, warmup 64, iters 512, trials 7, min per arm |
+| Host load | not sampled; `caffeinate -i` held, min of 7 interleaved trials per arm |
+
+The host was **not** quiet. An earlier single non-interleaved pass on this host
+inverted the ordering (`compiled` measured `+96%` against `eager`), which is why
+the harness interleaves its trials and reports the minimum: background load and
+throttling only ever add time. The three sweeps below keep the ordering but do
+**not** agree closely on absolute microseconds: run-to-run figures moved by up
+to ~30%, and the fused/compiled ratio on large prefill ranged 0.71x-0.85x.
+Treat an absolute microsecond as a warm-cache op-level figure on a loaded host,
+not as a production latency, and read the large-prefill ratio as a range.
+
+No CUDA or ROCm measurement was taken. The port table sets `.cuda = nullptr`
+and `.rocm = nullptr` for this kernel, so those backends keep the
+`compile_shapeless` path by construction.
+
+## Method
+
+Dtype bf16, because that is what the tanh-switched encoders hand the activation
+after the f32 GEMM: the tower runs the f32 formula and the result is stored
+back in the input dtype. Input is a deterministic ramp —
+`sin(i * 0.01) * 4.0` — which spans the saturated negative tail, the linear
+region and the saturated positive tail without an RNG.
+
+Every arm allocates its input once and reuses it on each timed iteration, so
+this is a warm last-level-cache measurement; see "Warm vs cold last-level cache"
+in `docs/benchmarks.md`. Each call is `eval`-ed and the timed region is
+bracketed by `synchronize_default()`.
+
+The sweep ran three times rather than once, for the reason given above.
+
+```bash
+caffeinate -i cargo run --release --features metal,accelerate \
+    --example fused_gelu_tanh_microbench
+```
+
+## Results
+
+Per-call microseconds, minimum of 7 interleaved trials; ratios are
+`fused / baseline`, so **below 1.00 is the fused kernel winning**.
+
+| 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** |
+
+Sweeps 2 and 3, same binary, same host, to show the spread:
+
+| Shape | fused / compiled (s2) | fused / compiled (s3) | fused / eager (s2) | fused / eager (s3) |
+|---|---|---|---|---|
+| tower prefill `[64, 4304]` | 0.81x | 0.76x | 0.45x | 0.48x |
+| decode row `[1, 4304]` | 0.96x | 0.97x | 0.77x | 0.73x |
+| large prefill `[1024, 4096]` | 0.71x | 0.85x | 0.32x | 0.32x |
+
+Raw CSV for the first sweep:
+
+```csv
+shape,eager_us,compiled_us,fused_us,fused_vs_compiled,fused_vs_eager
+"tower prefill [64, 4304]",587.29,342.31,279.96,0.8178,0.4767
+"decode row [1, 4304]",314.52,252.87,220.77,0.8731,0.7019
+"large prefill [1024, 4096]",2497.79,959.47,738.92,0.7701,0.2958
+```
+
+## What this does and does not establish
+
+The fused kernel wins on all three shapes against both baselines, in all three
+sweeps, and the margin is widest where the work is largest — which is the image
+prefill the fusion exists for. The narrowest result is the decode row against
+the compiled chain at 0.87x-0.97x, so a single token still gains, but by the
+least and with the widest spread;
+that is expected, since a batch-1 row is launch-bound rather than
+bandwidth-bound.
+
+The op-level ratio is a **lower bound** on the end-to-end effect. The fused
+kernel also removes the whole full-width intermediate from the MLX graph, which
+shows up as allocator and dependency-tracking pressure rather than as kernel
+time.
+
+What this does not establish: no end-to-end prefill or decode number was taken
+on a real checkpoint, because no tanh-switched vision checkpoint was loaded for
+this measurement. The decision recorded here is the op-level one the
+measure-then-keep policy asks for, and the wiring stays behind
+`fused_gelu_tanh_f32_available()` so `MLXCEL_DEVICE=cpu` keeps the compiled
+fallback.

docs/benchmarks.md

@@ -431,6 +431,26 @@ end-to-end decode sweep (one dense model and one MoE model, batch 1 and 4, with
 each kill switch flipped for the A/B) is the deciding measurement. Record both
 in `docs/benchmark_results/fused-norm-rope-<hw>-<date>.md`.
 
+## Fused tanh-GELU f32: the same measure-then-keep gate
+
+The fused `gelu_tanh_f32` Metal kernel from issue #2780's family lands under the
+same policy: it keeps its wiring only if it beats the graph it replaced at op
+level. `examples/fused_gelu_tanh_microbench.rs` produces that evidence. It loads
+no model and times the eager chain, the `compile_shapeless` chain and the fused
+kernel on the three shapes the tanh-switched encoders feed.
+
+```bash
+caffeinate -i cargo run --release --features metal,accelerate \
+    --example fused_gelu_tanh_microbench
+```
+
+The harness interleaves its trials and reports the per-arm minimum, because
+background load only ever adds time and a single non-interleaved pass on a
+loaded host inverted the ordering. It prints why and exits when
+`fused_gelu_tanh_f32_available()` is false, which is the case with
+`MLXCEL_DEVICE=cpu`. Record the result in
+`docs/benchmark_results/fused-gelu-tanh-f32-<hw>-<date>.md`.
+
 ## Warm vs cold last-level cache (issue #906)
 
 An op-level microbenchmark that allocates its inputs once and reuses them on

examples/fused_gelu_tanh_microbench.rs

@@ -0,0 +1,171 @@
+// Copyright 2025-2026 Lablup Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Op-level measurement of the fused tanh-GELU f32 Metal kernel against the two
+//! chains it replaces.
+//!
+//! `gelu_tanh_f32` applies one eager elementwise formula and passes about a
+//! dozen times over a full-width intermediate. `compiled_gelu_tanh_f32` folds
+//! that formula into one `mlx::core::compile` region. The fused kernel is a
+//! single dispatch, and this harness prints all three for the shapes the
+//! tanh-switched vision encoders actually feed, so the wiring decision has
+//! evidence rather than a guess behind it.
+//!
+//! # Method
+//!
+//! Every arm allocates its input once and reuses it, so this measures a warm
+//! last-level cache; see "Warm vs cold last-level cache" in
+//! `docs/benchmarks.md` for what that does and does not say about production.
+//! Each `body` is evaluated per iteration so the lazy graph actually runs, with
+//! `synchronize_default()` bracketing the timed region.
+//!
+//! Trials are interleaved and the per-arm minimum is reported. Background load
+//! and host throttling only ever add time, so the minimum is the estimator least
+//! biased by them, and interleaving puts all three arms under the same
+//! conditions. A single non-interleaved pass proved noisy enough to invert the
+//! ordering on a loaded host, so do not read one trial.
+//!
+//! The absolute microseconds are not the claim. The op-level ratio is also a
+//! lower bound on the end-to-end effect: the fused kernel removes the whole
+//! intermediate from the graph, which shows up as allocator and
+//! dependency-tracking pressure rather than as kernel time.
+//!
+//! Reproduce (use `caffeinate -i` so the host does not idle-throttle the GPU
+//! mid-run, and let the machine run cool between sweeps):
+//!
+//! ```text
+//! caffeinate -i cargo run --release --features metal,accelerate \
+//!     --example fused_gelu_tanh_microbench
+//! ```
+//!
+//! `fused_gelu_tanh_f32_available()` is false off the GPU stream, so with
+//! `MLXCEL_DEVICE=cpu` the harness prints why and exits instead of timing the
+//! eager path three times.
+
+use std::time::{Duration, Instant};
+
+use mlxcel_core::utils::{compiled_gelu_tanh_f32, fused_gelu_tanh_f32, gelu_tanh_f32};
+use mlxcel_core::{
+    MlxArray, UniquePtr, astype, dtype, eval, from_slice_f32, fused_gelu_tanh_f32_available,
+    synchronize_default,
+};
+
+const WARMUP: usize = 64;
+const ITERS: usize = 512;
+const TRIALS: usize = 7;
+
+/// Time `iters` calls of `body`, evaluating each result so the lazy graph is
+/// actually executed, with `synchronize_default()` bracketing the timed region
+/// so queued GPU work is not attributed to the next arm.
+fn time_body<F>(warmup: usize, iters: usize, mut body: F) -> Duration
+where
+    F: FnMut() -> UniquePtr<MlxArray>,
+{
+    for _ in 0..warmup {
+        let out = body();
+        eval(&out);
+    }
+    synchronize_default();
+    let start = Instant::now();
+    for _ in 0..iters {
+        let out = body();
+        eval(&out);
+    }
+    synchronize_default();
+    start.elapsed()
+}
+
+fn per_call_us(total: Duration, iters: usize) -> f64 {
+    total.as_secs_f64() * 1e6 / iters as f64
+}
+
+/// Interleaved trials, minimum per arm. Returns `(eager, compiled, fused)`.
+fn sweep(x: &MlxArray) -> (f64, f64, f64) {
+    let (mut eager, mut compiled, mut fused) = (f64::INFINITY, f64::INFINITY, f64::INFINITY);
+    for _ in 0..TRIALS {
+        eager = eager.min(per_call_us(
+            time_body(WARMUP, ITERS, || gelu_tanh_f32(x)),
+            ITERS,
+        ));
+        compiled = compiled.min(per_call_us(
+            time_body(WARMUP, ITERS, || compiled_gelu_tanh_f32(x)),
+            ITERS,
+        ));
+        fused = fused.min(per_call_us(
+            time_body(WARMUP, ITERS, || fused_gelu_tanh_f32(x)),
+            ITERS,
+        ));
+    }
+    (eager, compiled, fused)
+}
+
+fn main() {
+    if !fused_gelu_tanh_f32_available() {
+        println!(
+            "fused_gelu_tanh_f32 has no port for this backend; run with MLXCEL_DEVICE=gpu \
+             on a Metal build"
+        );
+        return;
+    }
+
+    // The tanh-switched encoders feed [batch, hidden] f32 GEMM output in bf16,
+    // so that is the dtype under test. The tower row is the shape the GELU
+    // tests use; the decode row is one token's worth; the large prefill row is
+    // an image chunk the tower can actually hand over.
+    let shapes: [(&str, [i32; 2]); 3] = [
+        ("tower prefill", [64, 4304]),
+        ("decode row", [1, 4304]),
+        ("large prefill", [1024, 4096]),
+    ];
+
+    println!("dtype bf16, warmup {WARMUP}, iters {ITERS}, trials {TRIALS}, min per arm");
+    println!();
+    println!(
+        "{:<16} {:>10} {:>10} {:>10} {:>10} {:>10}",
+        "shape", "eager us", "compiled", "fused", "f/compiled", "f/eager"
+    );
+
+    let mut csv =
+        String::from("shape,eager_us,compiled_us,fused_us,fused_vs_compiled,fused_vs_eager\n");
+    for (label, shape) in shapes {
+        let n: i32 = shape.iter().product();
+        // A deterministic ramp rather than random: the tanh CDF is monotone, so
+        // this spans the saturated negative tail, the linear region and the
+        // saturated positive tail without depending on an RNG.
+        let data: Vec<f32> = (0..n).map(|i| ((i as f32) * 0.01).sin() * 4.0).collect();
+        let x = from_slice_f32(&data, &shape);
+        let x = astype(&x, dtype::BFLOAT16);
+        eval(&x);
+
+        let (eager, compiled, fused) = sweep(&x);
+        let vs_compiled = fused / compiled;
+        let vs_eager = fused / eager;
+
+        let shape_name = format!("{label} {:?}", shape);
+        println!(
+            "{:<16} {eager:>10.2} {compiled:>10.2} {fused:>10.2} {vs_compiled:>9.2}x {vs_eager:>9.2}x",
+            shape_name
+        );
+        csv.push_str(&format!(
+            "\"{shape_name}\",{eager:.2},{compiled:.2},{fused:.2},{vs_compiled:.4},{vs_eager:.4}\n"
+        ));
+    }
+
+    println!();
+    println!("A ratio below 1.00 is the fused kernel winning.");
+    println!();
+    println!("```csv");
+    print!("{csv}");
+    println!("```");
+}

src/lib/mlxcel-core/build.rs

@@ -155,6 +155,12 @@ fn main() {
         // The one place a fused kernel chooses its per-backend port (#1801).
         .file("../mlxcelverse/kernels/kernel_port.cpp")
         .file("../mlxcelverse/kernels/sparse_v_sdpa.cpp")
+        // Fused tanh-GELU real Metal port launcher (#2423). Lives under
+        // `src/lib/mlxcelverse/kernels/` so the MLX-upstream-commit upgrade
+        // checklist ("Bumping the MLX upstream pin" in CONTRIBUTING.md)
+        // treats this directory as in-scope. Forwards
+        // `mlxcel::turbo::fused_gelu_tanh_f32` into the cxx-bridge ABI.
+        .file("../mlxcelverse/kernels/gelu_tanh_f32.cpp")
         // Fused Turbo4Delegated cold-V weighted-sum kernel
         // launcher. Reads the packed cold V directly so the dequantised
         // FP16 cold body never materialises in global memory; the host
@@ -469,6 +475,9 @@ fn main() {
     println!("cargo:rerun-if-changed=../mlxcelverse/kernels/fused_rope_append_hip.h");
     println!("cargo:rerun-if-changed=../mlxcelverse/kernels/kv_inplace_write.h");
     println!("cargo:rerun-if-changed=../mlxcelverse/kernels/kv_inplace_write.cpp");
+    // Fused tanh-GELU real Metal port (#2423).
+    println!("cargo:rerun-if-changed=../mlxcelverse/kernels/gelu_tanh_f32.h");
+    println!("cargo:rerun-if-changed=../mlxcelverse/kernels/gelu_tanh_f32.cpp");
     println!("cargo:rerun-if-env-changed=MLX_CUDA_ARCHITECTURES");
     println!("cargo:rerun-if-env-changed=MLX_ROCM_ARCHITECTURES");
     println!("cargo:rerun-if-env-changed=ROCM_PATH");

src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp

@@ -3,6 +3,7 @@
 
 #include "mlx_cxx_internal.h"
 #include "../../mlxcelverse/kernels/gpu_backend.h"
+#include "../../mlxcelverse/kernels/gelu_tanh_f32.h"
 
 #include "mlx/primitives.h"
 // Backend-agnostic CUDA availability probe; non-CUDA builds link the stub.
@@ -1975,6 +1976,20 @@ std::unique_ptr<MlxArray> compiled_gelu_tanh_f32(const MlxArray& x) {
     return std::make_unique<MlxArray>(std::move(result[0]));
 }
 
+// Real Metal custom kernel for `fused_gelu_tanh_f32` (#2423, Option A): the
+// tanh-GELU as one 1-D elementwise JIT Metal kernel instead of the compiled
+// graph above. Flattens any shape to 1-D for the launch and reshapes back;
+// f32, bfloat16 and float16 are upcast inside the kernel and rounded back at
+// the store. Throws when there is no port on this backend — Rust-side callers
+// gate on `fused_gelu_tanh_f32_available()` first or route to the eager chain.
+std::unique_ptr<MlxArray> fused_gelu_tanh_f32(const MlxArray& x) {
+    return std::make_unique<MlxArray>(mlxcel::turbo::fused_gelu_tanh_f32(x.inner));
+}
+
+bool fused_gelu_tanh_f32_available() {
+    return mlxcel::turbo::fused_gelu_tanh_f32_available();
+}
+
 // Compiled GeGLU: gelu(gate) * x
 // Used by: legacy/tests for precise GeGLU
 namespace {

src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h

@@ -591,6 +591,17 @@ std::unique_ptr<MlxArray> compiled_gelu_approx(const MlxArray& x);
 // Used by: Rust `GeluActivation::TanhF32Compiled` (#2423)
 std::unique_ptr<MlxArray> compiled_gelu_tanh_f32(const MlxArray& x);
 
+// Whether the real Metal tanh-GELU kernel runs on this backend (#2423, Option A):
+// a GPU default device with a Metal port in the table. False with
+// MLXCEL_DEVICE=cpu — callers take the eager chain. Read from the same
+// table the launcher dispatches through.
+bool fused_gelu_tanh_f32_available();
+
+// Real Metal custom kernel for `fused_gelu_tanh_f32` (#2423, Option A): the
+// tanh-GELU as one 1-D elementwise JIT Metal kernel. Rust-side callers gate on
+// `fused_gelu_tanh_f32_available()` first or route to the eager chain.
+std::unique_ptr<MlxArray> fused_gelu_tanh_f32(const MlxArray& x);
+
 // Compiled gelu_topk: sparse GELU with dynamic threshold — single fused kernel
 // gelu_tanh(max(0, x - (mean + std * multiplier))) in the input dtype (#2548)
 // Used by: Gemma3n MLP layers with activation_sparsity > 0

src/lib/mlxcel-core/src/lib.rs

@@ -808,6 +808,19 @@ mod ffi {
         /// Used by: `GeluActivation::TanhF32Compiled` (#2423).
         fn compiled_gelu_tanh_f32(x: &MlxArray) -> UniquePtr<MlxArray>;
 
+        /// Whether the fused tanh-GELU Metal kernel can run on this backend
+        /// (issue #2423): the default device is the GPU and the kernel has a
+        /// Metal port. Read from the kernel's port table, the one the dispatch
+        /// reads. False on a CPU-only host, where the eager `gelu_tanh_f32`
+        /// chain runs instead.
+        fn fused_gelu_tanh_f32_available() -> bool;
+
+        /// tanh-approximate GELU as one 1-D elementwise Metal JIT kernel;
+        /// throws when there is no port on this backend, so declare `-> Result`
+        /// so cxx catches the throw at the FFI boundary (#2423).
+        /// Used by: `GeluActivation::TanhF32Compiled`.
+        fn fused_gelu_tanh_f32(x: &MlxArray) -> Result<UniquePtr<MlxArray>>;
+
         /// Compiled GeGLU activation: gelu(gate) * x — single fused kernel
         /// Used by: legacy/tests for precise GeGLU
         fn compiled_geglu_activation(gate: &MlxArray, x: &MlxArray) -> UniquePtr<MlxArray>;