Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 9 additions & 6 deletions docs/gelu-call-sites.md
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ The logit-trace gate cannot tell erf from tanh (see section 0), so these are con
| T2 | Phixtral | erf | **tanh (f32, compiled)** | `gelu_new`; mlx-lm `SwitchMLP` default `approx="precise"`. No local checkpoint |
| T3 | GPT-2 | erf | **tanh (f32, compiled)** | `activation_function = gelu_new` (`GPT2Config` default); mlx-lm `nn.gelu_approx`. No local checkpoint |
| T4 | GPT-NeoX / Pythia | erf | erf, kept | HF `hidden_act = gelu` (`GPTNeoXConfig` default); mlx-lm uses tanh (departs from HF) |
| T5 | GPT-BigCode | erf | erf, unchanged (out of #2548 scope) | HF says `gelu_pytorch_tanh`, mlx-lm erf: the same situation as T6, left for a follow-up |
| T5 | GPT-BigCode | erf | **from `activation_function`** (#2794): `gelu_pytorch_tanh` / `gelu_new` / `gelu_fast` and an absent key -> tanh (f32, compiled), `gelu` -> erf, `quick_gelu` -> fast, any other name is a load error | HF builds the MLP as `ACT2FN[config.activation_function]` and `GPTBigCodeConfig` defaults to `gelu_pytorch_tanh` (santacoder declares it); mlx-lm `nn.gelu` (erf) departs from HF. Eager only (no XLA GPT-BigCode graph); no local checkpoint, smoke not run |
| T6 | StarCoder2 (eager) | erf | **tanh (f32, compiled)** | config `hidden_act = gelu_pytorch_tanh` (`Starcoder2Config` default); the OpenXLA StarCoder2 graph, frozen after it ran token-exact against the HF fp32 oracle, was already tanh. mlx-lm `nn.gelu` departs from HF. The issue framed the XLA tanh as the bug; rule 1 makes eager the side that moves |
| T8 | Nemotron-NAS | `gelu_new` / `gelu_fast` -> SiLU, unknown -> SiLU | GELU names through `GeluActivation::from_hf_name` (tanh names -> tanh f32 compiled), `silu`, `relu`; any other name is a load error | wrong activation, not an erf/tanh choice |
| T10, T11 | ModernBERT MLP and classifier head | erf | erf, kept | HF `hidden_activation = gelu`, `classifier_activation = gelu`; mlx-vlm tanh departs |
Expand All @@ -79,14 +79,17 @@ The logit-trace gate cannot tell erf from tanh (see section 0), so these are con
| V33 | Qwen3-VL-MoE MLP / merger | MLP tanh, merger erf | kept | HF `Qwen3VLMoeVisionMLP` uses `ACT2FN[gelu_pytorch_tanh]`, the patch merger `nn.GELU()`; mlx-vlm has them swapped |
| V34 | Qwen3-Omni vision MLP | tanh | kept | `thinker_config.vision_config.hidden_act = gelu_pytorch_tanh`; mlx-vlm `nn.GELU()` departs |
| S3 | Aya Vision SigLIP | erf (key absent) | **tanh (f32, compiled)** | `SiglipVisionConfig` default `gelu_pytorch_tanh`; mlx-vlm `GELUMLP` |
| S6 | Idefics2 SigLIP | fast (forced by `use_fast_gelu = true`) | **tanh (f32, compiled)** (#2794) | HF `Idefics2VisionTransformer` uses `ACT2FN[hidden_act]` and `Idefics2VisionConfig` defaults to `gelu_pytorch_tanh`; `idefics2-8b-4bit` omits the key and its tower `model_type` is `idefics2`, so the loader resolves it with `default_missing_hidden_act_to_siglip`. mlx-vlm `FastGELUMLP` departs from HF |
| S7 | Granite 4 Vision SigLIP | fast (forced by `use_fast_gelu = true`) | **tanh (f32, compiled)** (#2794) | `granite-4.0-3b-vision-4bit` declares `gelu_pytorch_tanh`, and its remote `modeling.py` builds the tower with `AutoModel.from_config(vision_config)` (HF `SiglipVisionModel`); mlx-vlm `FastGELUMLP` departs from HF. The Q-Former FFN (V59, erf) is unchanged |
| S8 | Granite Vision 3.x SigLIP | erf (key absent) | **tanh (f32, compiled)** | `SiglipVisionConfig` default; mlx-vlm `FastGELUMLP` departs from HF |
| S9 | LLaVA 1.5 / LLaVA-NeXT CLIP | erf (key absent) | **fast**, eager and the XLA llava vision path | `CLIPVisionConfig` default `quick_gelu`; mlx-vlm `FastGELUMLP` |
| S10 | LLaVA-Interleave SigLIP | tanh (declared) | tanh, kept | `vision_config.hidden_act = gelu_pytorch_tanh`; mlx-vlm fast departs |
| S11 | Bunny | erf (hard-coded) | **tanh (f32, compiled)** | the tower is `google/siglip-so400m-patch14-384` (not CLIP), whose remote config declares `gelu_pytorch_tanh`; mlx-vlm `FastGELUMLP` departs |
| S12 | Phi-3-V / Phi-3.5-V CLIP | erf (hard-coded) | **fast** | remote code builds its CLIP with `hidden_act="quick_gelu"`; mlx-vlm fast |
| S14 | Llama-Nemotron-VL-Embed SigLIP | tanh (declared) | tanh, kept | `vision_config.hidden_act = gelu_pytorch_tanh`; mlx-vlm (granite_vision `FastGELUMLP`) departs |
| S-cfg | Shared `VisionConfig` / `VisionHiddenActivation` (every SigLIP / CLIP tower and the SigLIP text tower) | `gelu_new`, `gelu_fast` and non-GELU names (`relu`, `silu`) parsed as erf; a declared `gelu_pytorch_tanh` ran the eager ~11-pass f32 helper (`TanhF32`) | (#2794) declared names map through `GeluActivation::from_hf_name` (`gelu_new` / `gelu_fast` -> tanh, `gelu_python` -> erf); any other name is a load error, also in Bunny's sparse-config fallback; a declared tanh runs the compiled kernel (`TanhF32Compiled`), the same as an omitted key resolved to the SigLIP default | `ACT2FN` names; the OpenXLA `LlavaVisionConfig` already rejected unknown names and now accepts the same GELU set. The declared and omitted tanh paths were split only to pin the declared one bit-for-bit to the private copy #2423 replaced; the two agree up to f32 rounding |

Shared resolution: `VisionConfig::default_missing_hidden_act_by_model_type` resolves an omitted key by `vision_config.model_type` (SigLIP / SigLIP2 -> tanh, CLIP -> `quick_gelu`, anything else unchanged); the serde default stays erf for the other towers. A declared `quick_gelu` now parses as fast in both the eager `VisionConfig` and the OpenXLA llava vision config (it parsed as erf, or was rejected by XLA, before). The fused single-kernel tanh GELU raised in the #2548 thread is a separate follow-up; nothing here changes `compiled_gelu_tanh_f32`.
Shared resolution: `VisionConfig::default_missing_hidden_act_by_model_type` resolves an omitted key by `vision_config.model_type` (SigLIP / SigLIP2 -> tanh, CLIP -> `quick_gelu`, anything else unchanged); `default_missing_hidden_act_to_siglip` resolves it to tanh for loaders whose tower `model_type` is not a SigLIP one (Gemma 3, PaliGemma, Idefics2); the serde default stays erf for the other towers. A declared `quick_gelu` now parses as fast in both the eager `VisionConfig` and the OpenXLA llava vision config (it parsed as erf, or was rejected by XLA, before). Since #2794 every declared name goes through `GeluActivation::from_hf_name` and a non-GELU name is a load error in both, and every tanh resolution, declared or defaulted, runs the compiled f32 kernel. No production loader passes `use_fast_gelu = true` any more. The fused single-kernel tanh GELU raised in the #2548 thread is a separate follow-up; nothing here changes `compiled_gelu_tanh_f32`.

Sections 1 to 4 below are the #2423 inventory at its own revision; where a row's Notes cell ends with a **#2548** tag, the table above is the current state.

Expand All @@ -96,7 +99,7 @@ Sections 1 to 4 below are the #2423 inventory at its own revision; where a row's
|---|---|---|---|
| `mlxcel_core::gelu` (FFI) | `cpp/mlx_cxx_bridge.cpp:3286` | erf | f32 constants, result cast back to input dtype |
| `mlxcel_core::gelu_approx` (FFI) and `mlxcel_core::utils::gelu_approx` | `mlx_cxx_bridge.cpp:3303`; wrapper `src/lib/mlxcel-core/src/utils.rs:1049` | **erf** (despite the name) | Comment at :3304-3315 claims erf "matches Python nn.GELU(approx="precise")"; that is false, precise = tanh. Doc at `utils.rs:1046` and `lib.rs:1056` also say "Approximate GELU (erf-based)". |
| `compiled_gelu` | `mlx_cxx_bridge.cpp:1871-1893` | erf | Only Rust caller: StarCoder2 |
| `compiled_gelu` | `mlx_cxx_bridge.cpp:1871-1893` | erf | Only Rust caller: StarCoder2. **#2794:** removed (dead after StarCoder2 moved to tanh in #2548) |
| `compiled_gelu_approx` | `mlx_cxx_bridge.cpp:1895-1920`, FFI `lib.rs:775` | **erf** | No non-test Rust caller ("Used by: legacy/tests") |
| `compiled_geglu_activation` (`get_compiled_geglu`) | `mlx_cxx_bridge.cpp:1922-1951` | erf GeGLU | Not called from Rust directly, but `compiled_gelu_mlp_forward_fp16` uses the same compiled lambda (see below) |
| `compiled_geglu_approx_activation` | `mlx_cxx_bridge.cpp:1953-1975` | tanh GeGLU (`gelu_tanh_approx`) | Gemma family fallback / non-fused path |
Expand Down Expand Up @@ -155,7 +158,7 @@ Legend for formula columns: **erf**, **tanh**, **fast**. "cfg" means the formula
| T2 | `src/models/phixtral.rs:609`, `:616` | `SwitchMLP::forward` (sort / no-sort) | Phixtral | **erf** | tanh (`SwitchMLP` default `nn.GELU(approx="precise")`) | [mlx_lm/models/switch_layers.py:229](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/switch_layers.py#L229), [mlx_lm/models/phixtral.py:93](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/phixtral.py#L93) | **MISMATCH** | No local ckpt **#2548:** now tanh (f32, compiled). |
| T3 | `src/models/gpt2.rs:579` | `MLP::forward` | GPT-2 | **erf** | tanh | [mlx_lm/models/gpt2.py:80](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gpt2.py#L80) | **MISMATCH** | HF `activation_function=gelu_new` (tanh) too; module doc :28-33 acknowledges it. No local ckpt **#2548:** now tanh (f32, compiled). |
| T4 | `src/models/gpt_neox.rs:883` | `MLP::forward` | GPT-NeoX / Pythia | **erf** | tanh | [mlx_lm/models/gpt_neox.py:100](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gpt_neox.py#L100) | **MISMATCH** | Deliberate: HF Pythia `hidden_act=gelu` (erf); comment :866-882. No local ckpt **#2548:** erf kept (HF `gelu`). |
| T5 | `src/models/gpt_bigcode.rs:569` | `MLP::forward` | GPT-BigCode | erf | erf (`nn.gelu`) | [mlx_lm/models/gpt_bigcode.py:99](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gpt_bigcode.py#L99) | MATCH | The issue's known-correct case; HF says `gelu_pytorch_tanh` **#2548:** unchanged, follow-up. |
| T5 | `src/models/gpt_bigcode.rs:569` | `MLP::forward` | GPT-BigCode | erf | erf (`nn.gelu`) | [mlx_lm/models/gpt_bigcode.py:99](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gpt_bigcode.py#L99) | MATCH | The issue's known-correct case; HF says `gelu_pytorch_tanh` **#2548:** unchanged, follow-up. **#2794:** now from `activation_function` (tanh f32 compiled for santacoder). |
| T6 | `src/models/starcoder2.rs:216` | MLP (`compiled_gelu`) | StarCoder2 | erf | erf | [mlx_lm/models/starcoder2.py:82](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/starcoder2.py#L82) | MATCH | HF `hidden_act=gelu_pytorch_tanh`. Local: `starcoder2-3b-4bit`. (XLA emitter uses tanh here, see Appendix) **#2548:** eager now tanh (f32, compiled); XLA unchanged. |
| T7 | `src/models/phi3small.rs:475` -> `utils.rs:1060` `gegelu` | GeGELU | Phi-3-small | fast | fast (`a * mx.sigmoid(1.702 a)`) | [mlx_lm/models/phi3small.py:49](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/phi3small.py#L49) | MATCH | Local: `phi-3-small-8k-instruct-aq4_64` |
| T8 | `src/models/nemotron_nas.rs:364` | `NASMLP::forward` (cfg `act_fn`) | Nemotron-NAS | cfg: `gelu` -> erf; `gelu_new`/`gelu_fast` -> **silu** (falls to `_`) | `gelu` -> erf, `gelu_new`/`gelu_fast` -> tanh | [mlx_lm/models/nemotron-nas.py:95-97](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/nemotron-nas.py#L95-L97) ([mlx_vlm/models/nemotron_nas/language.py:27-32](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/nemotron_nas/language.py#L27-L32)) | MATCH for `gelu`; **MISMATCH** (silu) for `gelu_new`/`gelu_fast` | Wrong activation entirely for those strings, not an erf/tanh question. Local `nemotron-nas-30b-4bit` is actually `model_type=nemotron_h` **#2548:** GELU names via `from_hf_name`, unknown names rejected. |
Expand Down Expand Up @@ -256,8 +259,8 @@ Since #2423 the Gemma 3 VLM and PaliGemma loaders resolve a missing key with `Vi
| S3 | `src/loading/vlm_siglip.rs:219` | Aya Vision | **erf** (key absent) | tanh (`GELUMLP`) | [mlx_vlm/models/aya_vision/vision.py:10](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/aya_vision/vision.py#L10) | **MISMATCH** | `aya-vision-8b`: absent **#2548:** now tanh (f32, compiled). |
| S4 | `src/loading/vlm_smolvlm.rs:309` | SmolVLM / Idefics3 | tanh | tanh (`GELUMLP` via idefics3) | [mlx_vlm/models/idefics3/vision.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/idefics3/vision.py#L8) ([mlx_vlm/models/smolvlm/__init__.py:4](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/smolvlm/__init__.py#L4)) | MATCH | `smolvlm-instruct-bf16`, `idefics3-8b-llama3-4bit`: `gelu_pytorch_tanh` |
| S5 | `src/models/colidefics3.rs:162` | ColIdefics3 | cfg (tanh when declared) | tanh | [mlx_vlm/models/colidefics3/__init__.py:2](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/colidefics3/__init__.py#L2) -> idefics3 | MATCH if declared, else MISMATCH | no local ckpt |
| S6 | `src/loading/vlm_idefics2.rs:201` (`use_fast_gelu = true`) | Idefics2 | fast | fast (`FastGELUMLP`) | [mlx_vlm/models/idefics2/vision.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/idefics2/vision.py#L8) | MATCH | `idefics2-8b-4bit`: absent (forced fast) |
| S7 | `src/loading/vlm_granite4_vision.rs:141` (`use_fast_gelu = true`) | Granite 4 Vision | fast | fast | [mlx_vlm/models/granite4_vision/vision.py:77](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/granite4_vision/vision.py#L77) | MATCH | `granite-4.0-3b-vision-4bit`: `gelu_pytorch_tanh` (HF says tanh; both mlx-vlm and mlxcel use fast) |
| S6 | `src/loading/vlm_idefics2.rs:201` (`use_fast_gelu = true`) | Idefics2 | fast | fast (`FastGELUMLP`) | [mlx_vlm/models/idefics2/vision.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/idefics2/vision.py#L8) | MATCH | `idefics2-8b-4bit`: absent (forced fast) **#2794:** now tanh (f32, compiled), follows HF over mlx-vlm. |
| S7 | `src/loading/vlm_granite4_vision.rs:141` (`use_fast_gelu = true`) | Granite 4 Vision | fast | fast | [mlx_vlm/models/granite4_vision/vision.py:77](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/granite4_vision/vision.py#L77) | MATCH | `granite-4.0-3b-vision-4bit`: `gelu_pytorch_tanh` (HF says tanh; both mlx-vlm and mlxcel use fast) **#2794:** now tanh (f32, compiled), follows HF over mlx-vlm. |
| S8 | `src/loading/vlm_granite_vision.rs:118` | Granite Vision 3.x | **erf** (key absent) | fast (`FastGELUMLP`) | [mlx_vlm/models/granite_vision/vision.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/granite_vision/vision.py#L8) | **MISMATCH** (erf vs fast) | `granite-vision-3.2-2b-4bit`: absent **#2548:** now tanh (f32, compiled), follows HF over mlx-vlm. |
| S9 | `src/loading/vlm_llava.rs:698` (and host preprocessor `:444`) | LLaVA 1.5 / LLaVA-NeXT (CLIP) | **erf** (key absent; HF `CLIPVisionConfig().hidden_act == "quick_gelu"`, verified in transformers 5.20.1) | fast (`FastGELUMLP`) | [mlx_vlm/models/llava/vision.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/llava/vision.py#L8), [mlx_vlm/models/llava_next/vision.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/llava_next/vision.py#L8) | **MISMATCH** (erf vs fast) | `llava-1.5-7b-4bit`, `llava-next-mistral-7b-4bit`: absent **#2548:** now fast, eager and XLA. |
| S10 | same, SigLIP-towered `llava` | LLaVA-Interleave (Qwen, SigLIP) | tanh | fast | [mlx_vlm/models/llava/vision.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/llava/vision.py#L8) | **MISMATCH** (tanh vs fast) | `llava-interleave-qwen-0.5b-bf16`: `gelu_pytorch_tanh` **#2548:** tanh kept (declared). |
Expand Down
28 changes: 2 additions & 26 deletions src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1906,30 +1906,6 @@ std::unique_ptr<MlxArray> compiled_gpt_oss_swiglu_activation(
return std::make_unique<MlxArray>(std::move(result[0]));
}

// Compiled GELU: x * 0.5 * (1 + erf(x / sqrt(2)))
// Used by: FFI tests only (StarCoder2 moved to the tanh GELU in #2548)
namespace {
static std::function<std::vector<array>(const std::vector<array>&)> get_compiled_gelu() {
auto fn = [](const std::vector<array>& inputs) -> std::vector<array> {
const auto& x = inputs[0];
auto sqrt2 = array(std::sqrt(2.0f));
auto half = array(0.5f);
auto one = array(1.0f);
auto erf_val = mlx::core::erf(mlx::core::divide(x, sqrt2));
auto scale = mlx::core::multiply(half, mlx::core::add(one, erf_val));
auto result = mlx::core::multiply(x, scale);
return {mlx::core::astype(result, x.dtype())};
};
return compile_shapeless_audited("compiled_gelu", fn);
}
}

std::unique_ptr<MlxArray> compiled_gelu(const MlxArray& x) {
static auto compiled_fn = get_compiled_gelu();
auto result = compiled_fn({x.inner});
return std::make_unique<MlxArray>(std::move(result[0]));
}

// Compiled `gelu_approx`: despite the name this is the exact erf GELU,
// x * 0.5 * (1 + erf(x / sqrt(2))), i.e. MLX `nn.gelu` / `nn.GELU()`. It is
// NOT MLX `nn.gelu_approx` (the tanh form); for that use `gelu_tanh` or
Expand Down Expand Up @@ -3354,8 +3330,8 @@ std::unique_ptr<MlxArray> gelu(const MlxArray& a) {
// gelu(x) = x * 0.5 * (1 + erf(x / sqrt(2)))
//
// The f32 constants promote a half-precision input, so the result is cast
// back to the input dtype before it is returned. `compiled_gelu` has always
// done this; leaving it out here let the f32 activation reach the next
// back to the input dtype before it is returned, as the compiled GELU
// helpers do; leaving it out here let the f32 activation reach the next
// projection, where MLX promotes the f16 weight to f32 to match it. See
// `gelu_approx` for the measurement.
auto sqrt2 = array(std::sqrt(2.0f));
Expand Down
4 changes: 0 additions & 4 deletions src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h
Original file line number Diff line number Diff line change
Expand Up @@ -578,10 +578,6 @@ std::unique_ptr<MlxArray> qwen4_hyper_inject(const MlxArray& branch, const MlxAr
std::unique_ptr<MlxArray> compiled_scaled_rms_norm(const MlxArray& x, const MlxArray& scale, float eps);
std::unique_ptr<MlxArray> compiled_silu(const MlxArray& x);

// Compiled gelu: x * 0.5 * (1 + erf(x / sqrt(2))) — single fused kernel
// Used by: FFI tests only (StarCoder2 moved to the tanh GELU in #2548)
std::unique_ptr<MlxArray> compiled_gelu(const MlxArray& x);

// Compiled gelu_approx: despite the name, exact erf GELU
// (x * 0.5 * (1 + erf(x / sqrt(2)))), MLX `nn.gelu`, not MLX's tanh
// `nn.gelu_approx`; use `gelu_tanh` for that (#2423).
Expand Down
Loading
Loading