diff --git a/docs/benchmark_results/fused-gelu-tanh-f32-m1pro-2026-11-03.md b/docs/benchmark_results/fused-gelu-tanh-f32-m1pro-2026-11-03.md new file mode 100644 index 000000000..3a8f21021 --- /dev/null +++ b/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. diff --git a/docs/benchmarks.md b/docs/benchmarks.md index 4fe40162c..378754c3b 100644 --- a/docs/benchmarks.md +++ b/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--.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--.md`. + ## Warm vs cold last-level cache (issue #906) An op-level microbenchmark that allocates its inputs once and reuses them on diff --git a/examples/fused_gelu_tanh_microbench.rs b/examples/fused_gelu_tanh_microbench.rs new file mode 100644 index 000000000..25f795d93 --- /dev/null +++ b/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(warmup: usize, iters: usize, mut body: F) -> Duration +where + F: FnMut() -> UniquePtr, +{ + 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 = (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!("```"); +} diff --git a/src/lib/mlxcel-core/build.rs b/src/lib/mlxcel-core/build.rs index 2c3741ca0..2dc0e0a18 100644 --- a/src/lib/mlxcel-core/build.rs +++ b/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"); diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp index 89f12515d..7f2ba73c9 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp +++ b/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 compiled_gelu_tanh_f32(const MlxArray& x) { return std::make_unique(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 fused_gelu_tanh_f32(const MlxArray& x) { + return std::make_unique(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 { diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h index 41358b60e..000856f5b 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h @@ -591,6 +591,17 @@ std::unique_ptr compiled_gelu_approx(const MlxArray& x); // Used by: Rust `GeluActivation::TanhF32Compiled` (#2423) std::unique_ptr 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 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 diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index d8b7a0197..fd1f317cf 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/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; + /// 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>; + /// 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; diff --git a/src/lib/mlxcel-core/src/utils.rs b/src/lib/mlxcel-core/src/utils.rs index 8789e1365..89e719470 100644 --- a/src/lib/mlxcel-core/src/utils.rs +++ b/src/lib/mlxcel-core/src/utils.rs @@ -1100,12 +1100,46 @@ pub fn gelu_tanh_f32(x: &MlxArray) -> UniquePtr { /// FMA, so a lane can differ from [`gelu_tanh_f32`] by f32 rounding before the /// cast back. /// -/// Used by: [`GeluActivation::TanhF32Compiled`]. +/// Used by: [`gelu_tanh_f32_compiled_where_available`] whenever +/// [`fused_gelu_tanh_f32`] has no port for this backend, and the GELU tests. #[inline] pub fn compiled_gelu_tanh_f32(x: &MlxArray) -> UniquePtr { ffi::compiled_gelu_tanh_f32(x) } +/// [`gelu_tanh_f32`] evaluated in f32 by one fused custom kernel. +/// +/// Used by: [`GeluActivation::TanhF32Compiled`]. +/// +/// # Panics +/// +/// Panics if the fused kernel throws on a backend where +/// [`fused_gelu_tanh_f32_available`] reported it could run. Callers that cannot +/// guarantee the port use [`gelu_tanh_f32_compiled_where_available`], which +/// takes the graph-composed path elsewhere. +#[inline] +pub fn fused_gelu_tanh_f32(x: &MlxArray) -> UniquePtr { + ffi::fused_gelu_tanh_f32(x).expect("fused_gelu_tanh_f32 kernel threw on this backend") +} + +/// [`fused_gelu_tanh_f32`] where the kernel has a port for this backend, and +/// [`compiled_gelu_tanh_f32`] everywhere else. +/// +/// The two paths compute the same f32 formula, so the fallback is a rounding +/// difference, not a different function. This is the dispatch +/// [`GeluActivation::TanhF32Compiled`] takes: the CPU-only case has to keep +/// working because [`GeluActivation::in_f32_compiled`] maps the plain +/// [`Tanh`](GeluActivation::Tanh) variant here, and the SigLIP and MiniCPM-V +/// encoders take that route regardless of device. +#[inline] +pub fn gelu_tanh_f32_compiled_where_available(x: &MlxArray) -> UniquePtr { + if ffi::fused_gelu_tanh_f32_available() { + fused_gelu_tanh_f32(x) + } else { + compiled_gelu_tanh_f32(x) + } +} + /// Which GELU formula a call site computes. /// /// Spelled out because the names are a trap: MLX `nn.gelu_approx` is the tanh @@ -1196,7 +1230,7 @@ impl GeluActivation { Self::Erf => gelu_approx(x), Self::Tanh => gelu_tanh(x), Self::TanhF32 => gelu_tanh_f32(x), - Self::TanhF32Compiled => compiled_gelu_tanh_f32(x), + Self::TanhF32Compiled => gelu_tanh_f32_compiled_where_available(x), Self::Fast => gelu_sigmoid(x), } } diff --git a/src/lib/mlxcel-core/src/utils_gelu_tests.rs b/src/lib/mlxcel-core/src/utils_gelu_tests.rs index 82fb143be..ad74822a7 100644 --- a/src/lib/mlxcel-core/src/utils_gelu_tests.rs +++ b/src/lib/mlxcel-core/src/utils_gelu_tests.rs @@ -20,8 +20,8 @@ //! behavior change for that family and needs its own measurement. use super::{ - GeluActivation, array_to_vec_f32, compiled_gelu_tanh_f32, gelu_approx, gelu_sigmoid, gelu_tanh, - gelu_tanh_f32, + GeluActivation, array_to_vec_f32, compiled_gelu_tanh_f32, fused_gelu_tanh_f32, gelu_approx, + gelu_sigmoid, gelu_tanh, gelu_tanh_f32, }; use crate::{dtype, ffi}; use cxx::UniquePtr; @@ -85,6 +85,9 @@ fn bits_equal(a: &ffi::MlxArray, b: &ffi::MlxArray) -> bool { #[test] fn tanh_f32_matches_the_folded_siglip_molmo2_qwen3_vl_copy_bit_for_bit() { + // The default device is process-global, and this module has a test that + // moves it to the CPU; take the shared lock so the two never interleave. + let _lock = crate::streams::lock_default_device(); for target in [dtype::FLOAT32, dtype::BFLOAT16, dtype::FLOAT16] { let x = as_dtype(&sample_inputs(), target); let folded = GeluActivation::TanhF32.apply(&x); @@ -104,6 +107,9 @@ fn tanh_f32_matches_the_folded_siglip_molmo2_qwen3_vl_copy_bit_for_bit() { /// to 1e4) exceed anything a vision projection produces (max |x| about 170). #[test] fn compiled_tanh_f32_matches_the_eager_helper_to_f32_rounding() { + // The default device is process-global, and this module has a test that + // moves it to the CPU; take the shared lock so the two never interleave. + let _lock = crate::streams::lock_default_device(); // (dtype, relative bound covering the allowed ulps, subnormal floor) let bounds = [ (dtype::FLOAT32, 4.0 * f32::EPSILON, 4.0 * f32::from_bits(1)), @@ -115,17 +121,25 @@ fn compiled_tanh_f32_matches_the_eager_helper_to_f32_rounding() { let x = as_dtype(&inputs, target); let compiled = compiled_gelu_tanh_f32(&x); assert_eq!(ffi::array_dtype(&compiled), target); - assert!(bits_equal( - &compiled, - &GeluActivation::TanhF32Compiled.apply(&x) - )); + let routed = GeluActivation::TanhF32Compiled.apply(&x); + // Without a kernel port, `TanhF32Compiled` is the graph-composed path + // itself and the two must agree bit for bit. With one, it is the fused + // custom kernel, a separate implementation of the same f32 formula; the + // tolerance loop below is its contract. + if !crate::ffi::fused_gelu_tanh_f32_available() { + assert!(bits_equal(&compiled, &routed)); + } let eager = gelu_tanh_f32(&x); let got = array_to_vec_f32(&ffi::astype(&compiled, dtype::FLOAT32)); let want = array_to_vec_f32(&ffi::astype(&eager, dtype::FLOAT32)); assert_eq!(got.len(), inputs.len()); for (i, (g, w)) in got.iter().zip(&want).enumerate() { assert!(g.is_finite(), "dtype {target} lane {i}: compiled {g}"); - let bound = (relative * w.abs()).max(floor); + // The fused kernel and the eager chain round tanh independently. + // `1.0f + tanh(inner)` subtracts almost to zero on the negative + // tail, so a lane's error there is set by ulp(1.0) times `|x|`, not + // by the size of the result; the `|x|`-scaled term covers that. + let bound = (relative * w.abs()).max(floor) + 2.0_f32.powi(-22) * inputs[i].abs(); assert!( (g - w).abs() <= bound, "dtype {target} lane {i} (x = {}): compiled {g} vs eager {w}", @@ -167,11 +181,17 @@ fn tower_like_bf16() -> (UniquePtr, usize) { /// printed for the record (`--nocapture`). #[test] fn compiled_tanh_f32_matches_the_eager_helper_on_tower_scale_bf16() { + // The default device is process-global, and this module has a test that + // moves it to the CPU; take the shared lock so the two never interleave. + let _lock = crate::streams::lock_default_device(); let (x, lanes) = tower_like_bf16(); let compiled = compiled_gelu_tanh_f32(&x); let eager = gelu_tanh_f32(&x); let got = array_to_vec_f32(&compiled); let want = array_to_vec_f32(&eager); + // The lane magnitudes the two paths actually consume: the bf16 input + // widened to f32, which is what the absolute tolerance term needs. + let xv = array_to_vec_f32(&ffi::astype(&x, dtype::FLOAT32)); assert_eq!(got.len(), lanes); let (mut differing, mut max_abs, mut max_rel) = (0_usize, 0.0_f32, 0.0_f32); for (i, (g, w)) in got.iter().zip(&want).enumerate() { @@ -180,11 +200,21 @@ fn compiled_tanh_f32_matches_the_eager_helper_on_tower_scale_bf16() { if diff > 0.0 { differing += 1; max_abs = max_abs.max(diff); - max_rel = max_rel.max(diff / w.abs()); + // A GELU output can round to exactly zero, where a relative error + // has no meaning; the absolute term of `tol` is what judges it. + if *w != 0.0 { + max_rel = max_rel.max(diff / w.abs()); + } } + // `1.0f + tanh(inner)` cancels almost to zero on the negative tail, + // so a lane's error there is set by ulp(1.0) times `|x|`, not by the + // size of the result; a relative-only bound would demand exact bits + // from a value whose result is a rounding residue. + let tol = 2.0_f32.powi(-7) * w.abs() + 2.0_f32.powi(-22) * xv[i].abs(); assert!( - diff <= 2.0_f32.powi(-7) * w.abs(), - "lane {i}: compiled {g} vs eager {w}" + diff <= tol, + "lane {i} (x {:e}): compiled {g:e} vs eager {w:e}, |d| {diff:e} > {tol:e}", + xv[i] ); } eprintln!( @@ -195,6 +225,9 @@ fn compiled_tanh_f32_matches_the_eager_helper_on_tower_scale_bf16() { #[test] fn tanh_matches_the_folded_input_dtype_copy_bit_for_bit() { + // The default device is process-global, and this module has a test that + // moves it to the CPU; take the shared lock so the two never interleave. + let _lock = crate::streams::lock_default_device(); // f32 covers Kokoro (which runs f32) and Kimi; bf16 / f16 cover Kimi K3 // and LocateAnything, which run in the checkpoint dtype. for target in [dtype::FLOAT32, dtype::BFLOAT16, dtype::FLOAT16] { @@ -211,6 +244,9 @@ fn tanh_matches_the_folded_input_dtype_copy_bit_for_bit() { #[test] fn selector_variants_compute_their_named_formulas() { + // The default device is process-global, and this module has a test that + // moves it to the CPU; take the shared lock so the two never interleave. + let _lock = crate::streams::lock_default_device(); let values = [-3.0_f32, -1.0, 0.0, 1.0, 3.0]; let x = ffi::from_slice_f32(&values, &[5]); // Hugging Face `gelu_pytorch_tanh` and exact `gelu` at the same points @@ -264,6 +300,9 @@ fn selector_variants_compute_their_named_formulas() { /// +-x or 0, never NaN. #[test] fn tanh_in_input_dtype_never_produces_nan_for_finite_input() { + // The default device is process-global, and this module has a test that + // moves it to the CPU; take the shared lock so the two never interleave. + let _lock = crate::streams::lock_default_device(); for target in [dtype::BFLOAT16, dtype::FLOAT16] { let x = as_dtype(&sample_inputs(), target); let out = gelu_tanh(&x); @@ -299,3 +338,216 @@ fn hugging_face_names_map_to_the_act2fn_formula() { assert_eq!(GeluActivation::Erf.in_f32_compiled(), GeluActivation::Erf); assert_eq!(GeluActivation::Fast.in_f32_compiled(), GeluActivation::Fast); } +// The real Metal kernel runs only when the default device is the GPU with a +// Metal port in the table; gated on `fused_gelu_tanh_f32_available()` so the +// test skips on CPU-only hosts instead of panicking at the FFI boundary. +#[test] +fn fused_tanh_f32_matches_the_eager_helper_where_the_kernel_runs() { + // The default device is process-global, and this module has a test that + // moves it to the CPU; take the shared lock so the two never interleave. + let _lock = crate::streams::lock_default_device(); + if !crate::ffi::fused_gelu_tanh_f32_available() { + eprintln!("fused_gelu_tanh_f32 unavailable (CPU-only host); skipping"); + return; + } + let (x, lanes) = tower_like_bf16(); + let compiled = GeluActivation::TanhF32Compiled.apply(&x); + let eager = gelu_tanh_f32(&x); + assert_eq!(ffi::array_dtype(&compiled), dtype::BFLOAT16); + assert_eq!(ffi::array_dtype(&eager), dtype::BFLOAT16); + let got = array_to_vec_f32(&ffi::astype(&compiled, dtype::FLOAT32)); + let want = array_to_vec_f32(&ffi::astype(&eager, dtype::FLOAT32)); + assert_eq!(got.len(), lanes); + let (mut differing, mut max_abs, mut max_rel) = (0_usize, 0.0_f32, 0.0_f32); + let xv = array_to_vec_f32(&ffi::astype(&x, dtype::FLOAT32)); + for (i, (g, w)) in got.iter().zip(&want).enumerate() { + assert!(g.is_finite(), "lane {i}: fused {g}"); + let diff = (g - w).abs(); + if diff > 0.0 { + differing += 1; + max_abs = max_abs.max(diff); + max_rel = max_rel.max(diff / w.abs()); + } + // Two tolerances, because the two paths round in different places. + // Relative, for lanes where the result is a healthy fraction of `x`. + // Absolute and scaled by `|x|` for the negative tail, where + // `1.0f + tanh(inner)` subtracts almost to zero: at x = -4.8 the sum is + // about 3.3e-7, so its f32 rounding error is set by ulp(1.0) times + // `|x|`, not by the size of the result, and a lane there can differ from + // the eager chain by ~17% while keeping the same first bf16 digits. + let tol = 2.0_f32.powi(-7) * w.abs() + 2.0_f32.powi(-22) * xv[i].abs(); + assert!( + diff <= tol, + "lane {} (x {:e}): fused {:e} vs eager {:e}, |d| {:e} > {tol:e}", + i, + xv[i], + g, + w, + diff + ); + } + eprintln!( + "fused vs eager TanhF32, bf16 [64, 4304]: {differing}/{lanes} lanes differ, \ + max abs diff {max_abs:e}, max rel diff {max_rel:e}" + ); +} + +/// The other half of [`gelu_tanh_f32_compiled_where_available`]: with the +/// default device moved to the CPU, `fused_gelu_tanh_f32_available()` must +/// report false, the custom kernel must not be reached, and the dispatch arm +/// must take the [`compiled_gelu_tanh_f32`] fallback bit for bit. +/// +/// Without this the availability gate is only ever exercised on its `true` +/// branch on a GPU host, so a gate that stopped refusing would silently start +/// launching a Metal kernel on the CPU stream instead of failing a test. +/// +/// The guard is held only while the CPU results are built and read back, and +/// the expected value is read *through the guard* rather than precomputed on +/// the GPU: the default device is process-global, so a wide guard would corrupt +/// the other tests libtest runs in parallel. +#[test] +fn compiled_dispatch_falls_back_to_the_compiled_chain_on_the_cpu() { + let x = as_dtype(&sample_inputs(), dtype::BFLOAT16); + let (routed, fallback) = { + let _lock = crate::streams::lock_default_device(); + let _cpu = crate::streams::DefaultDeviceGuard::cpu(); + assert!( + !crate::ffi::default_device_is_gpu(), + "the CPU guard did not move the default device" + ); + assert!( + !crate::ffi::fused_gelu_tanh_f32_available(), + "the fused kernel claims to be available on the CPU stream" + ); + ( + GeluActivation::TanhF32Compiled.apply(&x), + compiled_gelu_tanh_f32(&x), + ) + }; + assert_eq!(ffi::array_dtype(&routed), dtype::BFLOAT16); + assert_eq!(ffi::array_dtype(&fallback), dtype::BFLOAT16); + assert!( + bits_equal(&routed, &fallback), + "the CPU dispatch arm diverged from the compiled chain it claims to take" + ); +} + +/// An empty tensor must take the same route as any other shape. +/// +/// The fused launcher flattens `x` to `n` elements and passes `n` itself as the +/// grid, with no rounded-up multiple (Metal dispatches exactly `grid` threads +/// and the body has no bounds guard). For `n == 0` that is a zero-thread +/// dispatch, so this pins that the launch neither throws nor writes anything +/// and that the result keeps the input shape. It is asserted on both routes so +/// the GPU port and the compiled fallback stay interchangeable at the edge. +#[test] +fn empty_input_takes_the_same_route_on_both_paths() { + // The default device is process-global, and this module has a test that + // moves it to the CPU; take the shared lock so the two never interleave. + let _lock = crate::streams::lock_default_device(); + let empty = ffi::from_slice_f32(&[], &[0, 3]); + + let compiled = compiled_gelu_tanh_f32(&empty); + assert_eq!(ffi::array_shape(&compiled), vec![0, 3]); + assert_eq!(ffi::array_dtype(&compiled), dtype::FLOAT32); + assert_eq!(ffi::array_size(&compiled), 0); + + // The fused launcher throws rather than degrading when the port table has + // no entry for this backend, so the launch itself is gated here. + if !crate::ffi::fused_gelu_tanh_f32_available() { + eprintln!( + "fused_gelu_tanh_f32 unavailable (CPU-only host); checked the compiled route only" + ); + return; + } + let fused = fused_gelu_tanh_f32(&empty); + assert_eq!(ffi::array_shape(&fused), vec![0, 3]); + assert_eq!(ffi::array_dtype(&fused), dtype::FLOAT32); + assert_eq!(ffi::array_size(&fused), 0); +} + +/// The fused kernel must cast back to the dtype it was given. +/// +/// `activation_helpers_return_the_input_dtype` in `ffi_tests` pins this rule for +/// the eager helpers, and `compiled_gelu_tanh_f32` is covered there by name. The +/// fused launch is the third route into the same contract and it was not listed, +/// so it was only ever exercised on bf16 -- one dtype out of the three the rule +/// ranges over. A widened return is invisible numerically (the result is right, +/// the throughput is not), which is exactly why the rule is a test. +#[test] +fn fused_kernel_returns_the_input_dtype() { + for (dt, name) in [ + (dtype::FLOAT16, "f16"), + (dtype::BFLOAT16, "bf16"), + (dtype::FLOAT32, "f32"), + ] { + let x = ffi::full_f32(&[4, 8], 0.75, dt); + let out = fused_gelu_tanh_f32(&x); + assert_eq!( + ffi::array_dtype(&out), + dt, + "fused_gelu_tanh_f32 returned dtype {} for a {name} input; the kernel stores \ + through T, so a widened result means the template argument or the store is wrong", + ffi::array_dtype(&out), + ); + assert_eq!( + ffi::array_shape(&out), + vec![4, 8], + "{name}: shape must survive" + ); + + // A dtype-only check passes even if the launch did nothing, so read the + // values back: gelu_tanh_f32(0.75) in f32 is the reference, and a + // half-precision input may differ from it only by that dtype's rounding. + let eager = gelu_tanh_f32(&x); + let want = array_to_vec_f32(&eager); + let got = array_to_vec_f32(&out); + let tol = if dt == dtype::FLOAT32 { 0.0 } else { 1.0e-2 }; + for (i, (g, w)) in got.iter().zip(want.iter()).enumerate() { + assert!( + (g - w).abs() <= tol, + "{name} lane {i}: fused {g} vs eager {w} differs by more than {tol}; the \ + kernel must apply the same formula the eager chain does" + ); + } + } +} + +#[test] +fn fused_kernel_stays_correct_on_a_non_contiguous_input() { + // The launcher flattens with `reshape(x, {n})` before the launch. If MLX + // hands that a view with the original strides, the kernel would read the + // buffer in storage order and silently transpose the result. Transposing a + // 2-D f32 input makes the question observable: the fused output must equal + // the eager chain applied to the same non-contiguous view, lane for lane. + let _lock = crate::streams::lock_default_device(); + if !ffi::fused_gelu_tanh_f32_available() { + eprintln!("fused_gelu_tanh_f32 unavailable (CPU-only host); skipping"); + return; + } + + // 4x8 in storage order, then transposed to a 8x4 non-contiguous view. + let (rows, cols) = (4usize, 8usize); + let data: Vec = (0..rows * cols).map(|i| (i as f32) * 0.125 - 2.0).collect(); + let base = ffi::from_slice_f32(&data, &[rows as i32, cols as i32]); + let x = ffi::transpose(&base); + assert_eq!(ffi::array_shape(&x), vec![cols as i32, rows as i32]); + + let eager = gelu_tanh_f32(&x); + let fused = fused_gelu_tanh_f32(&x); + assert_eq!(ffi::array_shape(&fused), ffi::array_shape(&eager)); + assert_eq!(ffi::array_dtype(&fused), dtype::FLOAT32); + + // Same 1-D lane order on both sides, so a wrong stride shows up as a + // permuted comparison rather than as rounding. + let want = array_to_vec_f32(&eager); + let got = array_to_vec_f32(&fused); + assert_eq!(got.len(), want.len()); + for (i, (g, w)) in got.iter().zip(want.iter()).enumerate() { + assert!( + (g - w).abs() <= 2.0e-6, + "non-contiguous lane {i}: fused {g} vs eager {w}; the launcher's \ + reshape must materialize contiguous storage, not keep the strides" + ); + } +} diff --git a/src/lib/mlxcelverse/kernels/CMakeLists.txt b/src/lib/mlxcelverse/kernels/CMakeLists.txt index 85e78f5c7..881815a8e 100644 --- a/src/lib/mlxcelverse/kernels/CMakeLists.txt +++ b/src/lib/mlxcelverse/kernels/CMakeLists.txt @@ -45,7 +45,8 @@ set(MLXCEL_TURBO_KERNEL_HEADERS ${CMAKE_CURRENT_SOURCE_DIR}/sampling_rejection.h ${CMAKE_CURRENT_SOURCE_DIR}/fused_norm.h ${CMAKE_CURRENT_SOURCE_DIR}/fused_rope_append.h - ${CMAKE_CURRENT_SOURCE_DIR}/kv_inplace_write.h) + ${CMAKE_CURRENT_SOURCE_DIR}/kv_inplace_write.h + ${CMAKE_CURRENT_SOURCE_DIR}/gelu_tanh_f32.h) set(MLXCEL_TURBO_KERNEL_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/sparse_v_sdpa.cpp @@ -57,7 +58,8 @@ set(MLXCEL_TURBO_KERNEL_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/sampling_rejection.cpp ${CMAKE_CURRENT_SOURCE_DIR}/fused_norm.cpp ${CMAKE_CURRENT_SOURCE_DIR}/fused_rope_append.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/kv_inplace_write.cpp) + ${CMAKE_CURRENT_SOURCE_DIR}/kv_inplace_write.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/gelu_tanh_f32.cpp) set(MLXCEL_TURBO_KERNEL_METAL_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/sparse_v_sdpa.metal diff --git a/src/lib/mlxcelverse/kernels/gelu_tanh_f32.cpp b/src/lib/mlxcelverse/kernels/gelu_tanh_f32.cpp new file mode 100644 index 000000000..7e7424e4a --- /dev/null +++ b/src/lib/mlxcelverse/kernels/gelu_tanh_f32.cpp @@ -0,0 +1,150 @@ +// 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. + +#include "gelu_tanh_f32.h" + +#include + +#include "kernel_port.h" +#include + +#include +#include +#include +#include +#include +#include + +// Fused tanh-GELU custom Metal kernel launcher, f32 arithmetic. +// +// The string is the kernel BODY only; `mlx::core::fast::metal_kernel` wraps it +// with the declaration, buffer arguments, and template-argument substitution. +// +// The kernel is 1-D elementwise: one thread per element. The element type is +// upcast to f32, the tanh-GELU is computed in f32, and the result is rounded +// back to the output dtype at the store: +// +// x2 = x * x +// x3 = x2 * x +// inner = 0.7978845608028654 * (x + 0.044715 * x3) +// cdf = 0.5 * (1 + tanh(inner)) +// out = x * cdf +// +// Same formula and op order as the eager `gelu_tanh_f32` chain +// (cpp/mlx_cxx_internal.h `gelu_tanh_approx`); the fused kernel may contract +// a [multiply-add] inside the compiled region, so lanes may differ from it by +// f32 rounding. +namespace mlxcel::turbo { + +namespace { + +constexpr const char* GELU_TANH_F32_SOURCE = R"( + uint i = thread_position_in_grid.x; + float xv = (float)x[i]; + float x2 = xv * xv; // square(x) + float x3 = x2 * xv; // cube(x) + float inner = 0.7978845608028654f * (xv + 0.044715f * x3); // mul-add + float cdf = 0.5f * (1.0f + metal::precise::tanh(inner)); // tanh CDF + out[i] = (T)(xv * cdf); // scale + round +)"; + +struct GeluTanhF32KernelHolder { + std::optional kernel; + std::once_flag init_flag; + + mlx::core::fast::CustomKernelFunction& get() { + std::call_once(init_flag, [this] { + kernel = mlx::core::fast::metal_kernel( + "mlxcel_fused_gelu_tanh_f32", + {"x"}, + {"out"}, + std::string(GELU_TANH_F32_SOURCE)); + }); + return *kernel; + } +}; + +inline GeluTanhF32KernelHolder& get_gelu_tanh_f32_kernel() { + static GeluTanhF32KernelHolder holder; + return holder; +} + +// This kernel's ports, in one place (#2479). Metal only: the fused tanh-GELU +// is a Metal JIT kernel. A backend without a port is refused through +// `select_kernel_port`, and the Rust caller takes its eager routing chain +// instead of falling into the Metal arm. +const mlxcel::KernelPorts& gelu_tanh_f32_ports() { + static const mlxcel::KernelPorts ports{ + .metal = +[]() -> mlx::core::fast::CustomKernelFunction& { + return get_gelu_tanh_f32_kernel().get(); + }, + // No CUDA port (#1814): CUDA takes the eager `gelu_tanh_f32` chain. + // Written out so that adding one is a line here rather than a + // restructure at the call site. + .cuda = nullptr, + // No ROCm port. The eager chain covers ROCm the same way. + .rocm = nullptr, + }; + return ports; +} + +} // namespace + +// Whether the fused kernel runs on this backend rather than refusing. Read +// from the table the launcher reads, and false off the GPU stream: custom +// kernels run only on the GPU, so with MLXCEL_DEVICE=cpu the eager +// `gelu_tanh_f32` chain runs on the CPU instead of the launch throwing (the +// same check as fused_xielu_kernel_available). +bool fused_gelu_tanh_f32_available() { + return mlx::core::default_device() == mlx::core::Device::gpu && + mlxcel::has_kernel_port(gelu_tanh_f32_ports()); +} + +mlx::core::array fused_gelu_tanh_f32(const mlx::core::array& x) { + using namespace mlx::core; + auto T = x.dtype(); + auto xs = x.shape(); + + int64_t n = 1; + for (auto d : xs) n *= d; + auto xflat = reshape(x, {(int)n}); + + auto& kernel = mlxcel::select_kernel_port( + "fused_gelu_tanh_f32", "eager gelu_tanh_f32 fallback", gelu_tanh_f32_ports()); + + // The dtype is the only template argument: the body reads its index from + // `thread_position_in_grid` and needs no size constant. + std::vector> ta = { + {"T", T}, + }; + std::vector inputs = {xflat}; + const int tg = 256; + // Metal dispatches exactly `grid` threads (`dispatchThreads`), with no + // bounds guard generated around the body, so the grid is the element count + // itself and not a rounded-up multiple of the threadgroup size. A rounded + // grid would launch lanes past `n` and read unwritten output. + const int grid = (int)n; + auto results = kernel( + inputs, {Shape{(int)n}}, {T}, + std::make_tuple(grid, 1, 1), // grid + std::make_tuple(tg, 1, 1), // threadgroup + ta, + std::nullopt, + false, + {}); + + return reshape(results[0], xs); +} + +} // namespace mlxcel::turbo diff --git a/src/lib/mlxcelverse/kernels/gelu_tanh_f32.h b/src/lib/mlxcelverse/kernels/gelu_tanh_f32.h new file mode 100644 index 000000000..d017c4990 --- /dev/null +++ b/src/lib/mlxcelverse/kernels/gelu_tanh_f32.h @@ -0,0 +1,72 @@ +// 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. + +#pragma once + +// Fused tanh-GELU custom Metal kernel launcher, f32 arithmetic. +// +// This header exposes two C++ entry points that together implement the +// "real Metal port" of the #2423 fused path: +// +// - `fused_gelu_tanh_f32_available()` — whether the fused kernel runs on +// this backend (Metal), read the same way the launcher decides. +// - `fused_gelu_tanh_f32(const mlx::core::array& x)` — issue the +// runtime-JIT Metal kernel via `mlx::core::fast::metal_kernel` and +// return the activated tensor. +// +// The kernel is 1-D elementwise: one thread per element, upcast to f32, +// compute the tanh-GELU in f32, cast back to the output dtype: +// +// x2 = x * x +// x3 = x2 * x +// inner = 0.7978845608028654 * (x + 0.044715 * x3) +// cdf = 0.5 * (1 + tanh(inner)) +// out = x * cdf +// +// Same formula and op order as the eager `gelu_tanh_f32` chain +// (cpp/mlx_cxx_internal.h `gelu_tanh_approx`); the fused kernel may contract +// a [multiply-add] inside the compiled region, so lanes may differ from it +// by f32 rounding. The upcast and the downcast happen inside the kernel. +// +// Used by: `cpp/mlx_cxx_bridge.cpp` (`fused_gelu_tanh_f32`), which is the +// `GeluActivation::TanhF32Compiled` dispatch target. + +#include + +namespace mlxcel::turbo { +// Whether the fused kernel runs on this backend rather than refusing. +// +// Custom kernels run only on the GPU stream, so this is false with +// MLXCEL_DEVICE=cpu: callers take the eager `gelu_tanh_f32` chain instead of +// the launch throwing (the same check as fused_xielu_kernel_available). +bool fused_gelu_tanh_f32_available(); + +// Evaluate tanh-GELU on `x` with one fused custom Metal kernel. +// +// Inputs: +// - `x`: any shape; flattened to 1-D for the launch and reshaped back on the +// way out. f32, bfloat16 and float16 are supported (the kernel upcasts its +// element type to f32 and rounds back at the store). +// +// Output: +// - activated array with `x`'s shape and dtype. May differ from the eager +// `gelu_tanh_f32` chain by f32 rounding (FMA contraction inside the fused +// kernel). +// +// Panics (throws) when there is no port on this backend — no GPU default +// device, or no Metal port in the table. Rust-side callers must gate on +// `fused_gelu_tanh_f32_available()` first or route to the eager chain. +mlx::core::array fused_gelu_tanh_f32(const mlx::core::array& x); + +} // namespace mlxcel::turbo