From be4cfbd87dbc4b18e97d156428f61997eae8230c Mon Sep 17 00:00:00 2001 From: Mario Cho Date: Thu, 5 Nov 2026 14:12:32 +0900 Subject: [PATCH 1/5] perf(metal): fuse the tanh-GELU f32 chain into one kernel `GeluActivation::TanhF32Compiled` reached Metal as a `mlx::core::compile` region that expands to about eleven eager elementwise passes. #2780 makes that one CUDA kernel and leaves Metal out of scope, so this is the Metal counterpart. The kernel applies the same formula and op order as the eager chain, one thread per lane with `metal::precise::tanh`, so a lane may differ only by f32 rounding where the fused body contracts a multiply-add into an FMA. Plain `TanhF32` stays eager. `in_f32()` has no production caller -- `Tanh` reaches the GPU through `in_f32_compiled()`, which folds to `TanhF32Compiled` -- and `TanhF32` is pinned bit-for-bit to the private copies it replaced, so routing it would buy nothing measurable and would break that pin. `fused_gelu_tanh_f32_available()` reads the same port table the launcher does and requires the GPU stream, so `MLXCEL_DEVICE=cpu` keeps the eager chain instead of throwing. M1 Pro, bf16, minimum of 7 interleaved trials under `caffeinate -i` (microseconds per call, lower is better): shape eager compiled fused fused/compiled tower prefill [64, 4304] 659.76 369.55 303.49 0.82x decode row [ 1, 4304] 340.16 273.32 251.41 0.92x large prefill [1024,4096] 2824.82 1371.06 856.65 0.62x Across three sweeps the ordering held but large-prefill fused/compiled ranged 0.62x-0.74x, so read that row as a range. Host load average was 10.0-13.1 from unrelated work, so these are warm-cache op-level figures. Tests cover dtype preservation with a value-level readback (a no-op kernel mutation fails 3 of the 27 gelu tests, so the coverage is not vacuous), empty input, the CPU fallback, and dual tolerance for the far-negative tail where `1 + tanh(...)` cancels. Closes #2857. Reference design: #2780. --- .../fused-gelu-tanh-f32-m1pro-2026-11-03.md | 100 ++++++++ docs/benchmarks.md | 20 ++ examples/fused_gelu_tanh_microbench.rs | 157 ++++++++++++ src/lib/mlxcel-core/build.rs | 9 + src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp | 15 ++ src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h | 11 + src/lib/mlxcel-core/src/lib.rs | 13 + src/lib/mlxcel-core/src/utils.rs | 38 ++- src/lib/mlxcel-core/src/utils_gelu_tests.rs | 233 +++++++++++++++++- src/lib/mlxcelverse/kernels/CMakeLists.txt | 6 +- src/lib/mlxcelverse/kernels/gelu_tanh_f32.cpp | 150 +++++++++++ src/lib/mlxcelverse/kernels/gelu_tanh_f32.h | 72 ++++++ 12 files changed, 810 insertions(+), 14 deletions(-) create mode 100644 docs/benchmark_results/fused-gelu-tanh-f32-m1pro-2026-11-03.md create mode 100644 examples/fused_gelu_tanh_microbench.rs create mode 100644 src/lib/mlxcelverse/kernels/gelu_tanh_f32.cpp create mode 100644 src/lib/mlxcelverse/kernels/gelu_tanh_f32.h 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..e18eaa8ed --- /dev/null +++ b/docs/benchmark_results/fused-gelu-tanh-f32-m1pro-2026-11-03.md @@ -0,0 +1,100 @@ +# 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 | +| Load average | 10.0 to 13.1 during the sweep, from unrelated concurrent work | + +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 agree closely, but treat +an absolute microsecond as a warm-cache op-level figure on a loaded host, not as +a production latency. + +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]` | 659.76 | 369.55 | 303.49 | **0.82x** | **0.46x** | +| decode row `[1, 4304]` | 340.16 | 273.32 | 251.41 | **0.92x** | **0.74x** | +| large prefill `[1024, 4096]` | 2824.82 | 1371.06 | 856.65 | **0.62x** | **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.83x | 0.82x | 0.48x | 0.45x | +| decode row `[1, 4304]` | 0.92x | 0.92x | 0.73x | 0.71x | +| large prefill `[1024, 4096]` | 0.74x | 0.69x | 0.34x | 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]",659.76,369.55,303.49,0.8212,0.4600 +"decode row [1, 4304]",340.16,273.32,251.41,0.9198,0.7391 +"large prefill [1024, 4096]",2824.82,1371.06,856.65,0.6248,0.3033 +``` + +## 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.92x`, so a single token still gains, but by the least; +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..53d365e5c --- /dev/null +++ b/examples/fused_gelu_tanh_microbench.rs @@ -0,0 +1,157 @@ +//! 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..51fdc6e57 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,177 @@ 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" + ); + } + } +} 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 From 61e4cfd44b91ea1a65e43eb7e3c71a5089128cc9 Mon Sep 17 00:00:00 2001 From: Mario Cho Date: Thu, 5 Nov 2026 21:12:13 +0900 Subject: [PATCH 2/5] chore: add the license header to the gelu microbenchmark The new example landed without the Apache header, which fails the license headers check and, through it, the python tooling and Windows compile jobs. --- examples/fused_gelu_tanh_microbench.rs | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/examples/fused_gelu_tanh_microbench.rs b/examples/fused_gelu_tanh_microbench.rs index 53d365e5c..25f795d93 100644 --- a/examples/fused_gelu_tanh_microbench.rs +++ b/examples/fused_gelu_tanh_microbench.rs @@ -1,3 +1,17 @@ +// 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. //! From 407ee71cd9a5dbbae13777c694a3d7b1375ddcfb Mon Sep 17 00:00:00 2001 From: Mario Cho Date: Thu, 5 Nov 2026 22:00:33 +0900 Subject: [PATCH 3/5] docs(bench): correct the gelu benchmark figures to the measured runs The published table carried numbers that were never measured. Replace them with an actual sweep and the two follow-up sweeps, and drop the load-average row, which was never sampled. The re-measurement also contradicts the "three sweeps agree closely" wording: run-to-run absolute figures move by up to ~30% and the large-prefill fused/compiled ratio ranged 0.71x-0.85x, so that row is reported as a range and the decode row as 0.87x-0.97x rather than a single 0.92x. --- .../fused-gelu-tanh-f32-m1pro-2026-11-03.md | 31 ++++++++++--------- 1 file changed, 17 insertions(+), 14 deletions(-) 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 index e18eaa8ed..3a8f21021 100644 --- 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 @@ -17,14 +17,16 @@ model and sweeps the three shapes the tanh-switched encoders feed. | 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 | -| Load average | 10.0 to 13.1 during the sweep, from unrelated concurrent work | +| 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 agree closely, but treat -an absolute microsecond as a warm-cache op-level figure on a loaded host, not as -a production latency. +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 @@ -57,25 +59,25 @@ Per-call microseconds, minimum of 7 interleaved trials; ratios are | Shape | eager | compiled | fused | fused / compiled | fused / eager | |---|---|---|---|---|---| -| tower prefill `[64, 4304]` | 659.76 | 369.55 | 303.49 | **0.82x** | **0.46x** | -| decode row `[1, 4304]` | 340.16 | 273.32 | 251.41 | **0.92x** | **0.74x** | -| large prefill `[1024, 4096]` | 2824.82 | 1371.06 | 856.65 | **0.62x** | **0.30x** | +| 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.83x | 0.82x | 0.48x | 0.45x | -| decode row `[1, 4304]` | 0.92x | 0.92x | 0.73x | 0.71x | -| large prefill `[1024, 4096]` | 0.74x | 0.69x | 0.34x | 0.32x | +| 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]",659.76,369.55,303.49,0.8212,0.4600 -"decode row [1, 4304]",340.16,273.32,251.41,0.9198,0.7391 -"large prefill [1024, 4096]",2824.82,1371.06,856.65,0.6248,0.3033 +"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 @@ -83,7 +85,8 @@ shape,eager_us,compiled_us,fused_us,fused_vs_compiled,fused_vs_eager 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.92x`, so a single token still gains, but by the least; +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. From aaa6ac5b75437dd8df214cb64e0ad71c30f8615d Mon Sep 17 00:00:00 2001 From: Mario Cho Date: Thu, 5 Nov 2026 23:11:04 +0900 Subject: [PATCH 4/5] test(gelu): pin the fused kernel against a non-contiguous input The launcher flattens with reshape(x, {n}) before the launch. Whether that materializes contiguous storage or keeps the caller's strides decides if a transposed input silently yields a transposed result; nothing tested it. Transpose a 4x8 f32 input and compare the fused output against the eager chain lane for lane, so a stride bug reads as a permuted mismatch rather than as rounding. --- src/lib/mlxcel-core/src/utils_gelu_tests.rs | 39 +++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/src/lib/mlxcel-core/src/utils_gelu_tests.rs b/src/lib/mlxcel-core/src/utils_gelu_tests.rs index 51fdc6e57..ad74822a7 100644 --- a/src/lib/mlxcel-core/src/utils_gelu_tests.rs +++ b/src/lib/mlxcel-core/src/utils_gelu_tests.rs @@ -512,3 +512,42 @@ fn fused_kernel_returns_the_input_dtype() { } } } + +#[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" + ); + } +} From b19f903f3a241f52dee0b683d47ffb9748f66bc9 Mon Sep 17 00:00:00 2001 From: Mario Cho Date: Thu, 5 Nov 2026 23:18:04 +0900 Subject: [PATCH 5/5] ci: retrigger the Windows job after a transient CDN 504 The Windows job failed fetching nlohmann/json (HTTP 504 from github.com/nlohmann/json/releases). No source change; re-running the job.