diff --git a/TECHNICAL_REPORTS/2785-gelu-formula-decisions-20261031.en.md b/TECHNICAL_REPORTS/2785-gelu-formula-decisions-20261031.en.md new file mode 100644 index 000000000..454df309e --- /dev/null +++ b/TECHNICAL_REPORTS/2785-gelu-formula-decisions-20261031.en.md @@ -0,0 +1,61 @@ +# PR #2785: Settle the GELU formula for the families #2423 left open + +**Date**: 2026-10-31 (generated before merge; the PR was open at the time) +**Status**: In review +**Languages**: Rust, C++ +**Risk Level**: Medium (about 30 families change their GELU formula; each change follows the checkpoint config, and every family with a local checkpoint was smoke-tested on the CLI and the server) + +## Executive Summary + +#2423 inventoried every GELU call site and switched a first group of families to the formula their checkpoints declare, leaving the rest as open decisions in issue #2548. This PR closes them with one written rule: the checkpoint config decides, and where the key is omitted, the Hugging Face config-class default for that `model_type` decides, as `ACT2FN` evaluates it. Where the mlx-lm or mlx-vlm port departs from Hugging Face, mlxcel follows Hugging Face and records the divergence. `docs/gelu-call-sites.md` section 0b holds one row per decision with its reason. + +## 1. Problem Statement + +`mlxcel_core::gelu_approx` computes the exact erf GELU despite its name, while MLX `nn.gelu_approx` is the tanh form. Many ports copied the name and so ran erf where the checkpoint declares `gelu_new` or `gelu_pytorch_tanh`, or where a CLIP tower's omitted `hidden_act` means `quick_gelu`. The erf/tanh gap is below 1e-3 per element and the logit-trace gate cannot see it, so these were config decisions rather than measured regressions. Two sites were worse than a formula choice: Nemotron-NAS sent `gelu_new`, `gelu_fast` and any unknown name to SiLU, and Gemma 4's tensor-parallel per-layer gate computed erf while the single-host path computed tanh. + +## 2. Technical Review + +- **Formula per family.** Each changed row was checked against the local checkpoint's `config.json` or the remote modeling code: StarCoder2 and Gemma v1 declare `gelu_pytorch_tanh`, Phi-2 `gelu_new`, Molmo's remote code builds its CLIP ViT with `image_mlp_activations="quick_gelu"`, DeepSeek-VL2's `siglip_vit.py` uses `nn.GELU(approximate="tanh")`, and Phi-3.5-V names `openai/clip-vit-large-patch14-336`. +- **Shared kernel blast radius.** `compiled_gelu_mlp_forward_fp16` moved from the erf to the tanh GeGLU. Its only callers are Gemma v1, the Gemma 4 dense MLP and the Gemma 4 assistant drafter, all `gelu_pytorch_tanh`. +- **Shared vision config.** An omitted `vision_config.hidden_act` resolves by tower `model_type` (SigLIP to tanh, CLIP to `quick_gelu`) in the loaders that need it; the serde default stays erf for the other towers. A declared `quick_gelu` now parses as fast; no local checkpoint declared it explicitly, so no existing tower changed through that path. +- **Gemma3n eager vs OpenXLA tape.** The kernel, the fused MLP sparse branch, the XLA `sparse_gelu_stages` tape and the host oracle moved to tanh together. One structural difference predates this PR: the eager kernel multiplies a bf16 `stddev` by an f32 `array(std_multiplier)`, so MLX promotes everything from the cutoff onward to f32 and rounds once at the end, while the tape rounds every intermediate through bf16 (as mlx-lm's Python weak scalars would). The erf tape had the same shape. It needs a same-input comparison on an IREE host. + +## 3. Technical Decisions + +### 3.1 StarCoder2: move eager, not the XLA graph + +The issue treated the OpenXLA StarCoder2 tanh graph as the bug because eager and mlx-lm use erf. The checkpoint declares `gelu_pytorch_tanh`, which is also the `Starcoder2Config` default, and the XLA goldens were frozen after the graph ran token-exact against an HF fp32 oracle built with that activation (`spike/openxla/dense_arch_check.py`, PR #561). Under the decision rule eager is the side that moves, so the frozen goldens stay untouched. + +### 3.2 Tanh variant per site + +Families switched from erf use `GeluActivation::TanhF32Compiled`, the #2423 choice that adds no prefill time. Gemma sites use the existing input-dtype tanh GeGLU kernels so the bf16 and quantized Gemma paths compute the same thing. The private copies in the Nemotron codec (erf) and the VoiceChat TTS backbone (tanh with `power`) fold onto `GeluActivation`; both folds change the last bit of some half-precision elements and are covered by tolerance tests. + +### 3.3 Nemotron-NAS rejects unknown activations + +`hidden_act` is parsed once at load into `NasActivation`. GELU names go through `GeluActivation::from_hf_name`; any name that is not SiLU, ReLU or a GELU variant is a load error instead of a silent SiLU. + +## 4. Change Summary + +| Item | Value | +|---|---| +| Files changed | 50 | +| Lines | +948 / -305 | +| New tests | FFI tests for the tanh `gelu_topk` and the tanh dense GeGLU, tower-default resolution (eager and XLA), MiniCPM-o activation parsing, Nemotron-NAS activation parsing, `quick_gelu` emitter op, tolerance tests for the two folds | + +| Category | Summary | +|---|---| +| Gemma | v1 bf16, Gemma 4 dense bf16 and drafter, Gemma 4 TP gate, Gemma3n sparse layers (eager and XLA), RecurrentGemma move to tanh | +| Text LMs | Phi, Phixtral, GPT-2, StarCoder2 (eager), Moondream2/3 dense MLP move to tanh | +| Vision | Nine SigLIP-style towers and the Moondream vision MLP and projector move to tanh; Qwen2-VL, Molmo, LLaVA 1.5 / NeXT and Phi-3.5-V CLIP move to `quick_gelu`; MiniCPM-o reads its config | +| Kept with reason | GPT-NeoX, ModernBERT, Llama 4 vision, Kokoro PLBert, Molmo2, Qwen3-VL-MoE, Qwen3-Omni, LLaVA-Interleave, Llama-Nemotron-VL-Embed | + +## 5. Follow-up Actions + +- GPT-BigCode (T5) declares `gelu_pytorch_tanh` and still runs erf; it is now the one row that knowingly breaks the rule this PR writes down. +- Idefics2 (S6) and Granite 4 Vision (S7) force `quick_gelu` (matching mlx-vlm) while Hugging Face evaluates tanh; the inventory marks them MATCH against mlx-vlm, so they were outside #2548. +- Run the `xla-diagnostics` Gemma3n graphs on an IREE host and settle the eager-vs-tape carrier difference above. +- The fused single-kernel tanh GELU from the issue thread is tracked separately in #2780. + +## Appendix: Validation + +Narrow CUDA test filters on GB10 (`--test-threads=1`) pass for every touched module, the mlxcel-xla emitter tests pass with the StarCoder2 goldens byte-exact, and clippy and fmt are clean. The full workspace suite on this branch has one failure, `tests/step3p7_parity.rs` `text_decode_delegates_sequence_local_native_batching`, which fails identically on the base commit. Every changed family with a local checkpoint ran one greedy `mlxcel generate` and one `mlxcel-server` request with coherent output; families without a checkpoint (Gemma v1 bf16, RecurrentGemma, Phixtral, GPT-2, Phi-4-MM, Nemotron-NAS with a GELU activation, VoiceChat TTS, the Gemma 4 TP gate) rely on the synthetic tests named in the PR. diff --git a/TECHNICAL_REPORTS/2785-gelu-formula-decisions-20261031.ko.md b/TECHNICAL_REPORTS/2785-gelu-formula-decisions-20261031.ko.md new file mode 100644 index 000000000..ccc2868ef --- /dev/null +++ b/TECHNICAL_REPORTS/2785-gelu-formula-decisions-20261031.ko.md @@ -0,0 +1,61 @@ +# PR #2785: #2423이 미결로 남긴 계열의 GELU 공식 확정 + +**작성일**: 2026-10-31 (머지 전 작성, 당시 PR은 열린 상태) +**상태**: 리뷰 중 +**언어**: Rust, C++ +**위험도**: 중간 (약 30개 계열의 GELU 공식이 바뀐다. 각 변경은 체크포인트 config를 따르며, 로컬 체크포인트가 있는 계열은 모두 CLI와 서버에서 스모크 테스트를 거쳤다) + +## 요약 + +#2423은 모든 GELU 호출 지점을 목록화하고 첫 번째 그룹을 체크포인트가 선언한 공식으로 바꾼 뒤, 나머지를 이슈 #2548의 미결 항목으로 남겼다. 이 PR은 하나의 명시된 규칙으로 이를 마무리한다. 체크포인트 config가 결정하고, 키가 생략된 경우에는 해당 `model_type`의 Hugging Face config 클래스 기본값이 `ACT2FN` 기준으로 결정한다. mlx-lm 또는 mlx-vlm 포트가 Hugging Face와 다르면 mlxcel은 Hugging Face를 따르고 그 차이를 기록한다. `docs/gelu-call-sites.md` 0b 절에 결정마다 한 행과 이유가 있다. + +## 1. 문제 정의 + +`mlxcel_core::gelu_approx`는 이름과 달리 정확한 erf GELU를 계산하고, MLX `nn.gelu_approx`는 tanh 형태다. 많은 포트가 이름만 옮겨 와서, 체크포인트가 `gelu_new`나 `gelu_pytorch_tanh`를 선언하거나 CLIP 타워의 생략된 `hidden_act`가 `quick_gelu`를 뜻하는 곳에서도 erf를 실행했다. erf와 tanh의 차이는 원소당 1e-3 미만이고 logit-trace 게이트로는 구분되지 않으므로, 측정된 회귀가 아니라 config 판단의 문제였다. 두 지점은 공식 선택보다 심각했다. Nemotron-NAS는 `gelu_new`, `gelu_fast`와 모르는 이름을 모두 SiLU로 보냈고, Gemma 4의 tensor-parallel per-layer gate는 단일 호스트 경로가 tanh인데도 erf를 계산했다. + +## 2. 기술적 검토 사항 + +- **계열별 공식.** 바뀐 행마다 로컬 체크포인트의 `config.json`이나 원격 모델링 코드와 대조했다. StarCoder2와 Gemma v1은 `gelu_pytorch_tanh`, Phi-2는 `gelu_new`를 선언하고, Molmo 원격 코드는 CLIP ViT를 `image_mlp_activations="quick_gelu"`로 만들며, DeepSeek-VL2의 `siglip_vit.py`는 `nn.GELU(approximate="tanh")`를 쓰고, Phi-3.5-V는 `openai/clip-vit-large-patch14-336`를 지정한다. +- **공유 커널의 영향 범위.** `compiled_gelu_mlp_forward_fp16`이 erf GeGLU에서 tanh GeGLU로 바뀌었다. 호출자는 Gemma v1, Gemma 4 dense MLP, Gemma 4 assistant drafter뿐이며 모두 `gelu_pytorch_tanh`다. +- **공유 vision config.** 생략된 `vision_config.hidden_act`는 필요한 로더에서 타워 `model_type`으로 해석된다 (SigLIP은 tanh, CLIP은 `quick_gelu`). 다른 타워를 위해 serde 기본값은 erf로 유지된다. 선언된 `quick_gelu`는 이제 fast로 파싱되며, 이를 명시적으로 선언한 로컬 체크포인트는 없어서 이 경로로 바뀐 기존 타워는 없다. +- **Gemma3n eager와 OpenXLA tape.** 커널, fused MLP의 sparse 분기, XLA `sparse_gelu_stages` tape, host oracle이 함께 tanh로 바뀌었다. 이 PR 이전부터 있던 구조적 차이가 하나 있다. eager 커널은 bf16 `stddev`에 f32 `array(std_multiplier)`를 곱하므로 MLX가 cutoff 이후 전부를 f32로 승격하고 마지막에 한 번만 반올림한다. 반면 tape은 (mlx-lm의 Python weak scalar처럼) 모든 중간값을 bf16으로 반올림한다. erf tape도 같은 구조였다. IREE 호스트에서 같은 입력으로 비교해야 한다. + +## 3. 기술적 선택과 그 이유 + +### 3.1 StarCoder2: XLA 그래프가 아니라 eager를 옮긴다 + +이슈는 eager와 mlx-lm이 erf를 쓰므로 OpenXLA StarCoder2의 tanh 그래프를 버그로 보았다. 그러나 체크포인트는 `gelu_pytorch_tanh`를 선언하고 이는 `Starcoder2Config` 기본값이기도 하며, XLA golden은 그 activation으로 만든 HF fp32 oracle과 token-exact로 일치한 뒤 고정되었다 (`spike/openxla/dense_arch_check.py`, PR #561). 결정 규칙에 따라 옮겨야 할 쪽은 eager이며, 고정된 golden은 그대로 둔다. + +### 3.2 지점별 tanh 변형 + +erf에서 바뀐 계열은 prefill 시간을 늘리지 않는 #2423의 선택인 `GeluActivation::TanhF32Compiled`를 쓴다. Gemma 지점은 기존 입력 dtype tanh GeGLU 커널을 써서 bf16 경로와 양자화 경로가 같은 계산을 하게 한다. Nemotron codec (erf)과 VoiceChat TTS backbone (`power`를 쓰는 tanh)의 private 복사본은 `GeluActivation`으로 합쳐졌다. 두 경우 모두 일부 half-precision 원소의 마지막 비트가 달라질 수 있으며 tolerance 테스트로 범위를 묶었다. + +### 3.3 Nemotron-NAS는 모르는 activation을 거부한다 + +`hidden_act`는 로드 시 한 번 `NasActivation`으로 파싱된다. GELU 이름은 `GeluActivation::from_hf_name`을 거치고, SiLU, ReLU, GELU 변형이 아닌 이름은 조용히 SiLU를 쓰는 대신 로드 오류가 된다. + +## 4. 변경 요약 + +| 항목 | 값 | +|---|---| +| 변경 파일 | 50 | +| 라인 | +948 / -305 | +| 새 테스트 | tanh `gelu_topk`와 tanh dense GeGLU의 FFI 테스트, 타워 기본값 해석 (eager와 XLA), MiniCPM-o activation 파싱, Nemotron-NAS activation 파싱, `quick_gelu` emitter op, 두 통합에 대한 tolerance 테스트 | + +| 카테고리 | 요약 | +|---|---| +| Gemma | v1 bf16, Gemma 4 dense bf16과 drafter, Gemma 4 TP gate, Gemma3n sparse layer (eager와 XLA), RecurrentGemma가 tanh로 이동 | +| 텍스트 LM | Phi, Phixtral, GPT-2, StarCoder2 (eager), Moondream2/3 dense MLP가 tanh로 이동 | +| Vision | SigLIP 계열 타워 9개와 Moondream vision MLP 및 projector가 tanh로, Qwen2-VL, Molmo, LLaVA 1.5 / NeXT, Phi-3.5-V CLIP이 `quick_gelu`로 이동, MiniCPM-o는 config를 읽음 | +| 이유와 함께 유지 | GPT-NeoX, ModernBERT, Llama 4 vision, Kokoro PLBert, Molmo2, Qwen3-VL-MoE, Qwen3-Omni, LLaVA-Interleave, Llama-Nemotron-VL-Embed | + +## 5. 후속 조치 + +- GPT-BigCode (T5)는 `gelu_pytorch_tanh`를 선언하지만 여전히 erf로 실행된다. 이 PR이 문서화한 규칙을 알고도 어기는 유일한 행이다. +- Idefics2 (S6)와 Granite 4 Vision (S7)은 (mlx-vlm과 같이) `quick_gelu`를 강제하지만 Hugging Face는 tanh로 계산한다. 인벤토리가 mlx-vlm 기준 MATCH로 표시해 #2548 범위 밖이었다. +- IREE 호스트에서 `xla-diagnostics` Gemma3n 그래프를 실행해 위의 eager와 tape 간 carrier 차이를 정리한다. +- 이슈 스레드의 fused single-kernel tanh GELU는 #2780에서 별도로 추적한다. + +## 부록: 검증 + +GB10에서 바뀐 모듈마다 좁은 CUDA 테스트 필터 (`--test-threads=1`)가 통과하고, mlxcel-xla emitter 테스트는 StarCoder2 golden이 byte-exact인 상태로 통과하며, clippy와 fmt가 깨끗하다. 이 브랜치의 전체 workspace 테스트는 실패가 하나 (`tests/step3p7_parity.rs` `text_decode_delegates_sequence_local_native_batching`) 있으며 base 커밋에서도 똑같이 실패한다. 로컬 체크포인트가 있는 바뀐 계열은 모두 greedy `mlxcel generate` 한 번과 `mlxcel-server` 요청 한 번에서 일관된 출력을 냈다. 체크포인트가 없는 계열 (Gemma v1 bf16, RecurrentGemma, Phixtral, GPT-2, Phi-4-MM, GELU activation의 Nemotron-NAS, VoiceChat TTS, Gemma 4 TP gate)은 PR에 적힌 합성 테스트에 의존한다. diff --git a/docs/gelu-call-sites.md b/docs/gelu-call-sites.md index d4b963bd8..0ae27aba5 100644 --- a/docs/gelu-call-sites.md +++ b/docs/gelu-call-sites.md @@ -21,14 +21,74 @@ New code chooses its formula through `mlxcel_core::utils::GeluActivation` (`Erf` | MiniCPM-V 4.6 encoder MLP | V5 (4.6 only) | `hidden_act = gelu_pytorch_tanh` (tanh) | tanh | erf | **tanh (f32, compiled), from `hidden_act`** | | MiniCPM-V 4.6 window merger | V3 | `gelu_pytorch_tanh`, hard-coded | tanh | erf | **tanh (f32, compiled)** | | MiniCPM-V 4.6 downsample merger | V4 | `nn.GELU()` (erf) | tanh | erf | erf, kept (mlx-vlm departs from HF) | -| MiniCPM-o encoder MLP | V5 (MiniCPM-o) | not checked (no local checkpoint) | tanh | erf | erf, kept (follow-up) | +| MiniCPM-o encoder MLP | V5 (MiniCPM-o) | not checked (no local checkpoint) | tanh | erf | erf, kept (follow-up; #2548 reads `vision_config.hidden_act`, tanh) | | InternViT MLP | V1 | `hidden_act = gelu` (erf) | tanh | erf | erf, kept (mlx-vlm departs from HF) | | Gemma 3 VLM SigLIP | S1 | key absent, `SiglipVisionConfig` default `gelu_pytorch_tanh` | tanh | erf | **tanh (f32, compiled)** | | PaliGemma SigLIP | S2 | key absent, same default | tanh | erf | **tanh (f32, compiled)** | | Gemma 4 vision GeGLU (also DiffusionGemma, EmbeddingGemma2) | V17 | `hidden_activation = gelu_pytorch_tanh` (or absent, same default) | tanh | erf | **tanh (f32, compiled), from `hidden_activation`** | -| Gemma3n sparse layers | G15, G17 | `gelu_pytorch_tanh` | tanh | erf | erf, kept (follow-up: kernel, fused path and XLA oracle move together) | +| Gemma3n sparse layers | G15, G17 | `gelu_pytorch_tanh` | tanh | erf | erf, kept (follow-up: kernel, fused path and XLA oracle move together; done in #2548, tanh) | -Every other MISMATCH row below is unchanged and listed in the follow-up issue. The private tanh copies in section 2e were folded into `GeluActivation` keeping their precision (bit-equal, see `utils_gelu_tests.rs`). +Every other MISMATCH row below was left to the follow-up issue #2548, whose decisions are in section 0b. The private tanh copies in section 2e were folded into `GeluActivation` keeping their precision (bit-equal, see `utils_gelu_tests.rs`). + +## 0b. Final decisions for the families #2423 left open (#2548) + +Decision rule, applied to every row below: + +1. The checkpoint decides: its `config.json` (`hidden_act`, `hidden_activation`, `vision_config.hidden_act`, or the family's own key) and, where the key is missing, the default of the Hugging Face config class for that `model_type` (Hugging Face drops a key equal to the class default), evaluated as `transformers.activations.ACT2FN` does. Remote-code checkpoints are read from their own modeling files. +2. Cross-check against the mlx-lm / mlx-vlm port. Where the MLX port departs from Hugging Face, mlxcel follows Hugging Face and the row records the divergence. +3. Each row carries a one-line reason. + +The logit-trace gate cannot tell erf from tanh (see section 0), so these are config decisions, smoke-tested with one `mlxcel generate` and one `mlxcel-server` request per family with a local checkpoint (results in the #2548 PR). Variants: "tanh (f32, compiled)" is `GeluActivation::TanhF32Compiled`, the #2423 choice for switched families; "tanh (input dtype)" is the C++ `gelu_tanh_approx` helper (`GeluActivation::Tanh`, or the fused Gemma GeGLU kernels), used where the sibling Gemma kernels and the XLA bf16 tape already compute it that way; "fast" is `GeluActivation::Fast` (`quick_gelu`). + +| Row | Family | Before | After | Reason | +|---|---|---|---|---| +| G1, G3 | Gemma v1 quantized / mixed | tanh (input dtype) | tanh, kept | config `gelu_pytorch_tanh`; HF `GemmaConfig` rewrites even `gelu` to `gelu_pytorch_tanh`. mlx-lm `gemma.py` uses erf (departs from HF) | +| G2, G3b (PaliGemma v1 bf16) | Gemma v1 bf16 | erf | **tanh (input dtype)**: `compiled_gelu_mlp_forward_fp16` now runs the tanh GeGLU kernel | same as G1; the quantized and bf16 paths now agree | +| G7 | Gemma 4 dense MLP, bf16 weights | erf | **tanh (input dtype)**, same kernel as G2 | `hidden_activation = gelu_pytorch_tanh`; matches the quantized path (G6) | +| G11 | Gemma 4 per-layer gate, tensor-parallel runtime | erf | **tanh (input dtype)**: `compiled_geglu_approx_activation`, the single-host fallback kernel | `gelu_pytorch_tanh`; TP and single-host now agree | +| G12 | Gemma 4 assistant drafter, bf16 | erf | **tanh (input dtype)**, same kernel as G2 | reuses the Gemma 4 decoder layer (`gelu_pytorch_tanh`) | +| G15, G17 | Gemma3n sparse layers | erf | **tanh (f32 intermediates for bf16 input, single rounding)** in `compiled_gelu_topk` (fused and per-op paths); the XLA `sparse_gelu_stages` tape and the `gemma3n_math::sparse_gelu` host oracle moved with it | `hidden_activation = gelu_pytorch_tanh`; mlx-lm `gelu_topk` is `nn.gelu_approx` | +| G18 | Gemma3n `gelu_topk` dead early return | erf | **tanh (input dtype)** | unreachable; kept consistent with the dense layers | +| G20 | RecurrentGemma recurrent y-branch | erf | **tanh (f32, compiled)** | `RecurrentGemmaConfig` default `gelu_pytorch_tanh`; mlx-lm `nn.gelu_approx`. No local checkpoint | +| G21 | RecurrentGemma MLP | erf | **tanh (input dtype)**: the Gemma tanh GeGLU kernel | same as G20 | +| G24 | Nemotron VoiceChat TTS backbone and char MLP | tanh (private, `power` cube) | **tanh (input dtype)**, folded onto `GeluActivation::Tanh` | removes the `power` copy; bounded against `mlx.nn.gelu_approx` by a tolerance test, TTS smoke not run (no checkpoint) | +| T1 | Phi-1 / 1.5 / 2 | erf | **tanh (f32, compiled)** | config `hidden_act = gelu_new` (`PhiConfig` default); mlx-lm `nn.gelu_approx` | +| 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 | +| 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 | +| T13 | Kokoro PLBert | tanh (input dtype) | tanh, kept | ALBERT `AlbertConfig` default `gelu_new`; the checkpoint's `plbert` block omits the key. mlx-audio `nn.GELU()` departs | +| T15 | Nemotron VoiceChat codec ConvNeXt | erf (private weak-scalar copy) | erf, folded onto `GeluActivation::Erf` | mlx-audio `nn.gelu`; tolerance test bounds the fold | +| T17 | Moondream3 LM dense MLP | erf | **tanh (f32, compiled)** | mlx-vlm `TanhGELUMLP`; Moondream remote code `F.gelu(approximate="tanh")` | +| T19 | Moondream2 LM | erf | **tanh (f32, compiled)** | remote code `gelu_approx` = `F.gelu(approximate="tanh")`; mlx-vlm `TanhGELUMLP` | +| V5 (MiniCPM-o) | MiniCPM-o encoder MLP | erf | **from `vision_config.hidden_act`** (`gelu_pytorch_tanh` -> tanh f32 compiled; missing -> SigLIP default tanh; non-GELU -> load error) | `minicpm-o-4.5-4bit` declares `gelu_pytorch_tanh` | +| V6 | Phi-4-MM / Phi4-SigLIP vision MLP | erf | **tanh (f32, compiled)** | `Phi4MultimodalVisionConfig` default `gelu_pytorch_tanh`; mlx-vlm `GELUMLP`. No local checkpoint | +| V11 | DeepSeek-VL2 SigLIP MLP | erf | **tanh (f32, compiled)** | DeepSeek-VL2 `siglip_vit.py` builds `nn.GELU(approximate="tanh")`; mlx-vlm `approx="precise"` | +| V13 | Llama 4 vision | erf | erf, kept | `vision_config.hidden_act = gelu` (`Llama4VisionConfig` default); mlx-vlm hard-codes fast (departs from HF) | +| V14 | Youtu-VL vision MLP | erf | **tanh (f32, compiled)** | `vision_config.hidden_act = gelu_pytorch_tanh`; the XLA path was already tanh, eager and XLA now agree | +| V16 | Jina-VLM ViT MLP | erf | **tanh (f32, compiled)** | `vision_config.block_config.ffn_config.activation_type = gelu_pytorch_tanh` | +| V19 | Qwen2-VL ViT MLP | erf (XLA: tanh) | **fast**, eager and XLA (`emitter/qwen2_vl.rs`) | `Qwen2VLVisionConfig` default `quick_gelu` (key omitted); mlx-vlm `approx="fast"` | +| V22 | Molmo ViT MLP | erf (XLA: tanh) | **fast**, eager and XLA (`emitter/molmo_vision.rs`) | remote code `image_mlp_activations="quick_gelu"`; mlx-vlm fast | +| V23 | Molmo2 ViT | tanh (f32) | tanh, kept | `vit_config.hidden_act = gelu_pytorch_tanh`, HF `Molmo2VisionMLP` uses `ACT2FN`; mlx-vlm maps it to fast (departs) | +| V25 | LFM2-VL / LFM2.5-VL SigLIP2 MLP | erf | **tanh (f32, compiled)** | `vision_config.hidden_act = gelu_pytorch_tanh` | +| V27, V28 | Moondream2 / 3 vision MLP and projector | erf | **tanh (f32, compiled)** | remote code `F.gelu(approximate="tanh")` for both; mlx-vlm tanh | +| V29 | PaddleOCR-VL vision MLP | erf | **tanh (f32, compiled)** | `vision_config.hidden_act = gelu_pytorch_tanh` (`PaddleOCRVisionConfig` default) | +| 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` | +| 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 | + +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`. + +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. ## 1. Core helper semantics (what each entry point actually computes) @@ -61,55 +121,55 @@ Legend for formula columns: **erf**, **tanh**, **fast**. "cfg" means the formula | # | Site | Function | Family | mlxcel | Upstream | Upstream source | Verdict | Notes / local ckpt | |---|---|---|---|---|---|---|---|---| -| G1 | `src/models/gemma.rs:258` | `GemmaMLP::forward`, quantized branch (`compiled_gelu_approx_mlp_forward`) | Gemma v1 | tanh | erf (`nn.gelu`) | [mlx_lm/models/gemma.py:98](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma.py#L98) (mlx-vlm [mlx_vlm/models/gemma/language.py:84](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/gemma/language.py#L84) also erf) | **MISMATCH** | Upstream ignores HF `hidden_activation=gelu_pytorch_tanh`; mlxcel follows HF here. Local: `gemma-2b-4bit` (quantized, so this path runs) | -| G2 | `src/models/gemma.rs:277` | same, non-quantized branch (`compiled_gelu_mlp_fp16` -> erf) | Gemma v1 | erf | erf | [mlx_lm/models/gemma.py:98](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma.py#L98) | MATCH | mlxcel is internally inconsistent: quantized = tanh, bf16 = erf | -| G3 | `src/models/gemma.rs:289` | same, fallback (`compiled_geglu_approx_activation`) | Gemma v1 | tanh | erf | [mlx_lm/models/gemma.py:98](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma.py#L98) | MISMATCH | Reached only for mixed quantized/non-quantized projections | -| G3b | `src/models/gemma.rs` and `src/models/gemma2.rs` reached through `src/loading/vlm_siglip.rs:148-167` (`build_paligemma_text_model`) | PaliGemma text tower (`text_config.model_type` gemma -> `GemmaModel`, gemma2 -> `Gemma2Model`) | PaliGemma / PaliGemma2 | gemma: quantized tanh, bf16 erf (G1/G2); gemma2: tanh | cfg: `gelu` -> erf, `gelu_pytorch_tanh`/`gelu_new`/`gelu_fast`/`tanh` -> tanh (default `gelu_pytorch_tanh`) | [mlx_vlm/models/paligemma/language.py:16-25](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/paligemma/language.py#L16-L25), :152 | PaliGemma2: MATCH. PaliGemma v1: quantized MATCH, bf16 **MISMATCH** | The mlx-vlm PaliGemma reference honors the config, unlike [mlx_lm/models/gemma.py:98](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma.py#L98) (G1). Local: `paligemma2-3b-6bit` (`text_config.hidden_act = gelu_pytorch_tanh`, gemma2 tower) | +| G1 | `src/models/gemma.rs:258` | `GemmaMLP::forward`, quantized branch (`compiled_gelu_approx_mlp_forward`) | Gemma v1 | tanh | erf (`nn.gelu`) | [mlx_lm/models/gemma.py:98](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma.py#L98) (mlx-vlm [mlx_vlm/models/gemma/language.py:84](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/gemma/language.py#L84) also erf) | **MISMATCH** | Upstream ignores HF `hidden_activation=gelu_pytorch_tanh`; mlxcel follows HF here. Local: `gemma-2b-4bit` (quantized, so this path runs) **#2548:** kept tanh. | +| G2 | `src/models/gemma.rs:277` | same, non-quantized branch (`compiled_gelu_mlp_fp16` -> erf) | Gemma v1 | erf | erf | [mlx_lm/models/gemma.py:98](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma.py#L98) | MATCH | mlxcel is internally inconsistent: quantized = tanh, bf16 = erf **#2548:** now tanh (input dtype). | +| G3 | `src/models/gemma.rs:289` | same, fallback (`compiled_geglu_approx_activation`) | Gemma v1 | tanh | erf | [mlx_lm/models/gemma.py:98](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma.py#L98) | MISMATCH | Reached only for mixed quantized/non-quantized projections **#2548:** kept tanh. | +| G3b | `src/models/gemma.rs` and `src/models/gemma2.rs` reached through `src/loading/vlm_siglip.rs:148-167` (`build_paligemma_text_model`) | PaliGemma text tower (`text_config.model_type` gemma -> `GemmaModel`, gemma2 -> `Gemma2Model`) | PaliGemma / PaliGemma2 | gemma: quantized tanh, bf16 erf (G1/G2); gemma2: tanh | cfg: `gelu` -> erf, `gelu_pytorch_tanh`/`gelu_new`/`gelu_fast`/`tanh` -> tanh (default `gelu_pytorch_tanh`) | [mlx_vlm/models/paligemma/language.py:16-25](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/paligemma/language.py#L16-L25), :152 | PaliGemma2: MATCH. PaliGemma v1: quantized MATCH, bf16 **MISMATCH** | The mlx-vlm PaliGemma reference honors the config, unlike [mlx_lm/models/gemma.py:98](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma.py#L98) (G1). Local: `paligemma2-3b-6bit` (`text_config.hidden_act = gelu_pytorch_tanh`, gemma2 tower) **#2548:** PaliGemma v1 bf16 now tanh. | | G4 | `src/models/gemma2.rs:281`, `:302` | `MLP::forward` quantized / fallback | Gemma2 | tanh | tanh (`nn.gelu_approx`) | [mlx_lm/models/gemma2.py:119](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma2.py#L119) | MATCH | No fp16 erf path in gemma2.rs. Local: `gemma-2-9b-8bit`, `gemma2-2b-4bit` | | G5 | `src/models/gemma3.rs:461`, `:482` | `MLP::forward` quantized / fallback | Gemma3 text (and Gemma3 VLM text tower) | tanh | tanh | [mlx_lm/models/gemma3_text.py:122](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma3_text.py#L122) | MATCH | Local: `gemma3-1b-4bit`, `gemma-3-4b-it-4bit`, `gemma3-4b-4bit` | | G6 | `src/models/gemma4.rs:930`, `:950` | `MLP::forward` quantized (`..._global_scale` / plain) | Gemma4 dense MLP | tanh | tanh | [mlx_lm/models/gemma4_text.py:96](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma4_text.py#L96) (mlx-vlm [mlx_vlm/models/gemma4/language.py:24](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/gemma4/language.py#L24)) | MATCH | All local gemma-4 quantized/NVFP4 checkpoints | -| G7 | `src/models/gemma4.rs:974` | `MLP::forward` non-quantized (`compiled_gelu_mlp_fp16` -> erf) | Gemma4 dense MLP, bf16/f16 weights | **erf** | tanh | [mlx_lm/models/gemma4_text.py:96](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma4_text.py#L96) | **MISMATCH** | Not in the issue. Any bf16 Gemma4 target, or any quantized checkpoint that leaves an MLP unquantized, runs erf. No bf16 Gemma4 target locally | +| G7 | `src/models/gemma4.rs:974` | `MLP::forward` non-quantized (`compiled_gelu_mlp_fp16` -> erf) | Gemma4 dense MLP, bf16/f16 weights | **erf** | tanh | [mlx_lm/models/gemma4_text.py:96](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma4_text.py#L96) | **MISMATCH** | Not in the issue. Any bf16 Gemma4 target, or any quantized checkpoint that leaves an MLP unquantized, runs erf. No bf16 Gemma4 target locally **#2548:** now tanh (input dtype). | | G8 | `src/models/gemma4.rs:981` | fallback (`compiled_geglu_approx_activation`) | Gemma4 | tanh | tanh | [mlx_lm/models/gemma4_text.py:96](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma4_text.py#L96) | MATCH | | | G9 | `src/models/gemma4.rs:611`, `:632`, `:656`, `:718` | SwitchGeGLU experts (`compiled_geglu_approx_activation`, `compiled_switch_qgeglu_forward`, `fused_moe_geglu_kernel`) | Gemma4 26B-A4B MoE | tanh | tanh (GeGLU class) | [mlx_lm/models/gemma4_text.py:96](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma4_text.py#L96), :149-165 | MATCH | Local: `gemma-4-26b-a4b-it-4bit`, `-qat-4bit` | | G10 | `src/models/gemma4.rs:2948`, `:2968` | per-layer-input gate (`compiled_per_layer_input_gate` / `compiled_geglu_approx_activation`) | Gemma4 E2B/E4B | tanh | tanh | [mlx_lm/models/gemma4_text.py:382](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma4_text.py#L382) ([mlx_vlm/models/gemma4/language.py:359](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/gemma4/language.py#L359)) | MATCH | | -| G11 | `src/distributed/tensor_parallel/llama_runtime.rs:1110` | Gemma4 per-layer gate inside TP `forward_impl` | Gemma4 E2B/E4B under tensor-parallel runtime | **erf** (`gelu_approx`) | tanh | [mlx_lm/models/gemma4_text.py:382](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma4_text.py#L382) | **MISMATCH** | Single-host path (G10) is tanh; the TP copy is erf. Local: `gemma-4-e2b-it-*`, `gemma-4-e4b-it-*` | -| G12 | `src/lib/mlxcel-core/src/drafter/gemma4_assistant/layer.rs:104` | drafter MLP, non-quantized (`compiled_gelu_mlp_fp16` -> erf) | Gemma4 assistant drafter | **erf** | tanh (reuses `gemma4.language.DecoderLayer`) | [mlx_vlm/speculative/drafters/gemma4_assistant/gemma4_assistant.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/speculative/drafters/gemma4_assistant/gemma4_assistant.py#L8) -> [mlx_vlm/models/gemma4/language.py:24](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/gemma4/language.py#L24) (https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/speculative/drafters/gemma4_assistant/gemma4_assistant.py) | **MISMATCH** | Local: `gemma-4-31b-it-assistant-bf16` (bf16, so this path runs) | +| G11 | `src/distributed/tensor_parallel/llama_runtime.rs:1110` | Gemma4 per-layer gate inside TP `forward_impl` | Gemma4 E2B/E4B under tensor-parallel runtime | **erf** (`gelu_approx`) | tanh | [mlx_lm/models/gemma4_text.py:382](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma4_text.py#L382) | **MISMATCH** | Single-host path (G10) is tanh; the TP copy is erf. Local: `gemma-4-e2b-it-*`, `gemma-4-e4b-it-*` **#2548:** now tanh, same kernel as single-host. | +| G12 | `src/lib/mlxcel-core/src/drafter/gemma4_assistant/layer.rs:104` | drafter MLP, non-quantized (`compiled_gelu_mlp_fp16` -> erf) | Gemma4 assistant drafter | **erf** | tanh (reuses `gemma4.language.DecoderLayer`) | [mlx_vlm/speculative/drafters/gemma4_assistant/gemma4_assistant.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/speculative/drafters/gemma4_assistant/gemma4_assistant.py#L8) -> [mlx_vlm/models/gemma4/language.py:24](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/gemma4/language.py#L24) (https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/speculative/drafters/gemma4_assistant/gemma4_assistant.py) | **MISMATCH** | Local: `gemma-4-31b-it-assistant-bf16` (bf16, so this path runs) **#2548:** now tanh (input dtype). | | G13 | `src/lib/mlxcel-core/src/drafter/gemma4_assistant/layer.rs:114` | drafter MLP, quantized fallback | Gemma4 assistant drafter | tanh | tanh | as G12 | MATCH | Local: `gemma-4-12b-it-assistant-4bit` | | G14 | `src/models/gemma3n.rs:816` -> `gemma3n_mlp_forward` (bridge :3024) | fused bf16 MLP, dense layers (`activation_sparsity == 0`) | Gemma3n | tanh | tanh | [mlx_lm/models/gemma3n.py:195](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma3n.py#L195) ([mlx_vlm/models/gemma3n/language.py:212](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/gemma3n/language.py#L212)) | MATCH | Gated by `use_fused_decode_path()` (non-NA hardware) | -| G15 | `src/models/gemma3n.rs:816` -> `gemma3n_mlp_forward` sparse branch | fused bf16 MLP, sparse layers | Gemma3n | **erf** (`compiled_gelu_topk`) | tanh (`gelu_topk` -> `nn.gelu_approx`) | [mlx_lm/models/gemma3n.py:161-165](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma3n.py#L161-L165), :193 ([mlx_vlm/models/gemma3n/language.py:176-180](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/gemma3n/language.py#L176-L180)) | **MISMATCH** | Layers 0-9 carry `activation_sparsity_pattern = 0.95` in the local checkpoints | +| G15 | `src/models/gemma3n.rs:816` -> `gemma3n_mlp_forward` sparse branch | fused bf16 MLP, sparse layers | Gemma3n | **erf** (`compiled_gelu_topk`) | tanh (`gelu_topk` -> `nn.gelu_approx`) | [mlx_lm/models/gemma3n.py:161-165](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma3n.py#L161-L165), :193 ([mlx_vlm/models/gemma3n/language.py:176-180](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/gemma3n/language.py#L176-L180)) | **MISMATCH** | Layers 0-9 carry `activation_sparsity_pattern = 0.95` in the local checkpoints **#2548:** now tanh, eager and XLA. Rounding: for bf16 inputs the eager kernel promotes the cutoff, shift and tanh to f32 (bf16 stddev times an f32 `std_multiplier` array) and rounds once, while the XLA tape rounds each intermediate through bf16 (as the erf tape did). Known pre-existing eager/XLA gap; needs a same-input comparison on an IREE host. | | G16 | `src/models/gemma3n.rs:840` | per-op MLP, dense layers | Gemma3n | tanh | tanh | [mlx_lm/models/gemma3n.py:195](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma3n.py#L195) | MATCH | | -| G17 | `src/models/gemma3n.rs:842` -> `gelu_topk` :785 (`compiled_gelu_topk`) | per-op MLP, sparse layers | Gemma3n | **erf** | tanh | [mlx_lm/models/gemma3n.py:161-165](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma3n.py#L161-L165) | **MISMATCH** | Local: `gemma3n-e2b-4bit`, `gemma3n-e4b-4bit`, `gemma3n-e4b-bf16` | -| G18 | `src/models/gemma3n.rs:780` | `MLP::gelu_topk` early return (`gelu_approx`) | Gemma3n | erf | tanh | [mlx_lm/models/gemma3n.py:195](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma3n.py#L195) | MISMATCH (dead) | Unreachable in production: `forward` calls `gelu_topk` only when sparsity > 0; the two test callers (:3326, :4712) assert sparsity > 0 | +| G17 | `src/models/gemma3n.rs:842` -> `gelu_topk` :785 (`compiled_gelu_topk`) | per-op MLP, sparse layers | Gemma3n | **erf** | tanh | [mlx_lm/models/gemma3n.py:161-165](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma3n.py#L161-L165) | **MISMATCH** | Local: `gemma3n-e2b-4bit`, `gemma3n-e4b-4bit`, `gemma3n-e4b-bf16` **#2548:** now tanh, eager and XLA. Rounding: for bf16 inputs the eager kernel promotes the cutoff, shift and tanh to f32 (bf16 stddev times an f32 `std_multiplier` array) and rounds once, while the XLA tape rounds each intermediate through bf16 (as the erf tape did). Known pre-existing eager/XLA gap; needs a same-input comparison on an IREE host. | +| G18 | `src/models/gemma3n.rs:780` | `MLP::gelu_topk` early return (`gelu_approx`) | Gemma3n | erf | tanh | [mlx_lm/models/gemma3n.py:195](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma3n.py#L195) | MISMATCH (dead) | Unreachable in production: `forward` calls `gelu_topk` only when sparsity > 0; the two test callers (:3326, :4712) assert sparsity > 0 **#2548:** now tanh. | | G19 | `src/models/gemma3n.rs:1058`, `:1185` | AltUp per-layer-input gate | Gemma3n | tanh | tanh (`first_prediction = nn.gelu_approx(...)`) | [mlx_lm/models/gemma3n.py:359](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/gemma3n.py#L359) | MATCH | (`:3316`, `:3865`, `:4711` are `cfg(all(test, feature = "xla-diagnostics"))`, erf inline in :3316/:4711, excluded) | -| G20 | `src/models/recurrent_gemma.rs:345` | `RecurrentBlock` y-branch | RecurrentGemma | **erf** | tanh | [mlx_lm/models/recurrent_gemma.py:205](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/recurrent_gemma.py#L205) | **MISMATCH** | HF `hidden_activation=gelu_pytorch_tanh`. No local ckpt | -| G21 | `src/models/recurrent_gemma.rs:552` | `MLPBlock::forward` | RecurrentGemma | **erf** | tanh | [mlx_lm/models/recurrent_gemma.py:284](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/recurrent_gemma.py#L284) | **MISMATCH** | No local ckpt | +| G20 | `src/models/recurrent_gemma.rs:345` | `RecurrentBlock` y-branch | RecurrentGemma | **erf** | tanh | [mlx_lm/models/recurrent_gemma.py:205](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/recurrent_gemma.py#L205) | **MISMATCH** | HF `hidden_activation=gelu_pytorch_tanh`. No local ckpt **#2548:** now tanh (f32, compiled). | +| G21 | `src/models/recurrent_gemma.rs:552` | `MLPBlock::forward` | RecurrentGemma | **erf** | tanh | [mlx_lm/models/recurrent_gemma.py:284](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/recurrent_gemma.py#L284) | **MISMATCH** | No local ckpt **#2548:** now tanh GeGLU kernel. | | G22 | `src/models/embedding_gemma2_layers.rs:145`, `:192` | GeGLU FFN and PLE gate | EmbeddingGemma2 | tanh | tanh | [mlx_vlm/models/embedding_gemma2/language.py:92](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/embedding_gemma2/language.py#L92), :126 | MATCH | | | G23 | `src/models/diffusion_gemma/mod.rs:506` | GeGLU | DiffusionGemma | tanh | tanh | [mlx_vlm/models/diffusion_gemma/language.py:24](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/diffusion_gemma/language.py#L24) | MATCH | Local: `diffusiongemma-26b-a4b-it-4bit` | -| G24 | `src/models/gemma3_backbone.rs:59`, used at `:83` and `src/models/nemotron_voicechat/tts/norm_mlp.rs:148` | private `gelu_approx` (tanh with `power(x, 3)` in input dtype) | Nemotron VoiceChat TTS (Gemma3 backbone + char encoder MLP) | tanh | tanh | [mlx_vlm/models/nemotron_voicechat/tts.py:49](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/nemotron_voicechat/tts.py#L49) (mlx-audio [mlx_audio/sts/models/nemotron_voicechat/tts.py:35](https://github.com/Blaizzy/mlx-audio/blob/v0.5.9/mlx_audio/sts/models/nemotron_voicechat/tts.py#L35)) | MATCH | Uses `power`, the op the `gelu_approx` comment blames for bf16 NaN; it is in production here | +| G24 | `src/models/gemma3_backbone.rs:59`, used at `:83` and `src/models/nemotron_voicechat/tts/norm_mlp.rs:148` | private `gelu_approx` (tanh with `power(x, 3)` in input dtype) | Nemotron VoiceChat TTS (Gemma3 backbone + char encoder MLP) | tanh | tanh | [mlx_vlm/models/nemotron_voicechat/tts.py:49](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/nemotron_voicechat/tts.py#L49) (mlx-audio [mlx_audio/sts/models/nemotron_voicechat/tts.py:35](https://github.com/Blaizzy/mlx-audio/blob/v0.5.9/mlx_audio/sts/models/nemotron_voicechat/tts.py#L35)) | MATCH | Uses `power`, the op the `gelu_approx` comment blames for bf16 NaN; it is in production here **#2548:** folded onto `GeluActivation::Tanh`. | ### 2b. Other text LMs, encoders, and audio | # | Site | Function | Family | mlxcel | Upstream | Upstream source | Verdict | Notes / local ckpt | |---|---|---|---|---|---|---|---|---| -| T1 | `src/models/phi.rs:268` | `MLP::forward` | Phi (Phi-1/1.5/2) | **erf** | tanh | [mlx_lm/models/phi.py:116](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/phi.py#L116) | **MISMATCH** | HF `hidden_act=gelu_new` (tanh) too. Local: `phi-2-4bit` | -| 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 | -| 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 | -| 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 | -| 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` | -| 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) | +| T1 | `src/models/phi.rs:268` | `MLP::forward` | Phi (Phi-1/1.5/2) | **erf** | tanh | [mlx_lm/models/phi.py:116](https://github.com/ml-explore/mlx-lm/blob/c39139d70f0798834fc721c28a1aa8e1374ea33c/mlx_lm/models/phi.py#L116) | **MISMATCH** | HF `hidden_act=gelu_new` (tanh) too. Local: `phi-2-4bit` **#2548:** now tanh (f32, compiled). | +| 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. | +| 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` | +| 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. | | T9 | `src/models/bert.rs:290-291` (`Activation::apply`), tanh copy `bert.rs:303` | BERT FFN (cfg `hidden_act`); also `deberta_v2.rs:235` | BERT, XLM-R, DeBERTa-v2 | cfg: `gelu`/`gelu_python` -> erf, `gelu_new`/`gelu_fast`/`gelu_pytorch_tanh` -> tanh | erf always (`nn.GELU()`) | [mlx_vlm/models/bert/bert.py:86](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/bert/bert.py#L86), [mlx_vlm/models/xlm_roberta/xlm_roberta.py:84](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/xlm_roberta/xlm_roberta.py#L84), [mlx_vlm/models/gliner2_5/deberta.py:161](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/gliner2_5/deberta.py#L161) | MATCH for `gelu`; MISMATCH (HF-faithful) for tanh strings | Local `rampart-mlx-*` (`hidden_act=gelu`) and `gliner2.5-small-v1` (`gelu`) both hit MATCH | -| T10 | `src/models/modernbert.rs:358` | `geglu` (MLP) | ModernBERT | **erf** | tanh (`nn.GELU(approx="precise")`) | [mlx_vlm/models/modernbert/modernbert.py:27](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/modernbert/modernbert.py#L27) | **MISMATCH** | HF ModernBERT `hidden_activation=gelu` (erf); mlxcel follows HF. No local ckpt | -| T11 | `src/models/modernbert_heads.rs:254` | classifier head | ModernBERT reranker | **erf** | tanh | [mlx_vlm/models/modernbert/modernbert.py:169](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/modernbert/modernbert.py#L169) | **MISMATCH** | HF `classifier_activation=gelu` (erf) | +| T10 | `src/models/modernbert.rs:358` | `geglu` (MLP) | ModernBERT | **erf** | tanh (`nn.GELU(approx="precise")`) | [mlx_vlm/models/modernbert/modernbert.py:27](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/modernbert/modernbert.py#L27) | **MISMATCH** | HF ModernBERT `hidden_activation=gelu` (erf); mlxcel follows HF. No local ckpt **#2548:** erf kept (HF `gelu`). | +| T11 | `src/models/modernbert_heads.rs:254` | classifier head | ModernBERT reranker | **erf** | tanh | [mlx_vlm/models/modernbert/modernbert.py:169](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/modernbert/modernbert.py#L169) | **MISMATCH** | HF `classifier_activation=gelu` (erf) **#2548:** erf kept (HF `gelu`). | | T12 | `src/token_classification/gliner2_5_scorer.rs:190` | FiLM scorer | GLiNER2.5 | erf | erf | [mlx_vlm/models/gliner2_5/boundary.py:312](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/gliner2_5/boundary.py#L312) | MATCH | Local: `gliner2.5-small-v1` | -| T13 | `src/models/kokoro/bert.rs:165` -> `src/models/kokoro/ops.rs:223` `gelu_new` | PLBert (ALBERT) FFN | Kokoro TTS | tanh | erf (`nn.GELU()`) | [mlx_audio/tts/models/kokoro/modules.py:534](https://github.com/Blaizzy/mlx-audio/blob/v0.5.9/mlx_audio/tts/models/kokoro/modules.py#L534), :567 | **MISMATCH** | mlxcel uses the ALBERT `gelu_new` the weights were trained with; local config's `plbert` block has no `hidden_act`. Local: `kokoro-82m` | +| T13 | `src/models/kokoro/bert.rs:165` -> `src/models/kokoro/ops.rs:223` `gelu_new` | PLBert (ALBERT) FFN | Kokoro TTS | tanh | erf (`nn.GELU()`) | [mlx_audio/tts/models/kokoro/modules.py:534](https://github.com/Blaizzy/mlx-audio/blob/v0.5.9/mlx_audio/tts/models/kokoro/modules.py#L534), :567 | **MISMATCH** | mlxcel uses the ALBERT `gelu_new` the weights were trained with; local config's `plbert` block has no `hidden_act`. Local: `kokoro-82m` **#2548:** tanh kept (ALBERT `gelu_new`). | | T14 | `src/models/whisper/encoder.rs:106`, `:110`; `src/models/whisper/layers.rs:238` | conv stem + MLP | Whisper | erf | erf | [mlx_audio/stt/models/whisper/whisper.py:415](https://github.com/Blaizzy/mlx-audio/blob/v0.5.9/mlx_audio/stt/models/whisper/whisper.py#L415), :439-440 | MATCH | Local: `whisper-base` | -| T15 | `src/audio/nemotron_codec/layers.rs:233` (`gelu_exact`, private erf), used :344 | ConvNeXt block | Nemotron VoiceChat codec | erf | erf | [mlx_audio/codec/models/nemotron_voicechat/codec.py:140](https://github.com/Blaizzy/mlx-audio/blob/v0.5.9/mlx_audio/codec/models/nemotron_voicechat/codec.py#L140) | MATCH | | +| T15 | `src/audio/nemotron_codec/layers.rs:233` (`gelu_exact`, private erf), used :344 | ConvNeXt block | Nemotron VoiceChat codec | erf | erf | [mlx_audio/codec/models/nemotron_voicechat/codec.py:140](https://github.com/Blaizzy/mlx-audio/blob/v0.5.9/mlx_audio/codec/models/nemotron_voicechat/codec.py#L140) | MATCH | **#2548:** folded onto `GeluActivation::Erf`. | | T16 | `src/models/moondream3.rs:258` | `tau_values` | Moondream3 LM | erf | erf | [mlx_vlm/models/moondream3/language.py:23](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/moondream3/language.py#L23) | MATCH | | -| T17 | `src/models/moondream3.rs:311` | `DenseMlp::forward` | Moondream3 LM | **erf** | tanh (`TanhGELUMLP`) | [mlx_vlm/models/moondream3/language.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/moondream3/language.py#L8) | **MISMATCH** | No local ckpt | +| T17 | `src/models/moondream3.rs:311` | `DenseMlp::forward` | Moondream3 LM | **erf** | tanh (`TanhGELUMLP`) | [mlx_vlm/models/moondream3/language.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/moondream3/language.py#L8) | **MISMATCH** | Local: `moondream3-preview-mlx-4bit`, smoke-tested after the change (correct answer on the orange image). **#2548:** now tanh (f32, compiled). | | T18 | `src/models/moondream3.rs:385` | `SparseMoeMlp` GeGLU | Moondream3 LM | erf | erf | [mlx_vlm/models/moondream3/language.py:205](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/moondream3/language.py#L205) | MATCH | | -| T19 | `src/models/moondream2.rs:270` | `Moondream2Mlp::forward` | Moondream2 LM | **erf** | tanh (`TanhGELUMLP`) | [mlx_vlm/models/moondream2/language.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/moondream2/language.py#L8) | **MISMATCH** | Local comment says `nn.GELU(approx="tanh")`. Local: `moondream2` | +| T19 | `src/models/moondream2.rs:270` | `Moondream2Mlp::forward` | Moondream2 LM | **erf** | tanh (`TanhGELUMLP`) | [mlx_vlm/models/moondream2/language.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/moondream2/language.py#L8) | **MISMATCH** | Local comment says `nn.GELU(approx="tanh")`. Local: `moondream2` **#2548:** now tanh (f32, compiled). | | T20 | `src/models/florence2/layers.rs:289`, `:360` | encoder / decoder FFN | Florence-2 LM | erf | erf | [mlx_vlm/models/florence2/language.py:137](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/florence2/language.py#L137), :163 | MATCH | | | T21 | `src/models/nemotron_parse/encoder.rs:208`, `decoder.rs:103` | MLP | Nemotron-Parse | erf | erf | [mlx_vlm/models/nemotron_parse/vision.py:47](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/nemotron_parse/vision.py#L47), `language.py` :121 | MATCH | | | T22 | `src/audio/qwen3_omni_moe/mod.rs:227`, `:395-397`, `:451` | audio encoder FFN, conv stem, projector | Qwen3-Omni audio | erf | erf | [mlx_vlm/models/qwen3_omni_moe/audio.py:94](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/qwen3_omni_moe/audio.py#L94), :187, :257-259 | MATCH | Local: `qwen3-omni-30b-a3b-instruct-4bit` | @@ -125,36 +185,36 @@ Legend for formula columns: **erf**, **tanh**, **fast**. "cfg" means the formula | V2 | `src/vision/internvl.rs:108` | `mlp1` connector | InternVL | erf | erf | [mlx_vlm/models/internvl_chat/internvl_chat.py:30](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/internvl_chat/internvl_chat.py#L30) | MATCH | (HF-layout `internvl` projector upstream is `approx="precise"`, [mlx_vlm/models/internvl/internvl.py:21](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/internvl/internvl.py#L21); that layout is #2393, not in tree) | | V3 | `src/vision/encoders/minicpmv4_6.rs:343` | ViT merger (window) | MiniCPM-V 4.6 | **erf** | tanh | [mlx_vlm/models/minicpmv4_6/minicpmv4_6.py:114](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/minicpmv4_6/minicpmv4_6.py#L114) | **changed in #2423** to tanh (f32, compiled kernel); HF `MiniCPMV4_6ViTWindowAttentionMerger` hard-codes `gelu_pytorch_tanh` | Issue item. Local: `minicpm-v-4.6-bf16`, `minicpm-v-4.6-mxfp4` (`vision_config.hidden_act=gelu_pytorch_tanh`) | | V4 | `src/vision/encoders/minicpmv4_6.rs:393` | `MergerBlock` | MiniCPM-V 4.6 | **erf** | tanh | [mlx_vlm/models/minicpmv4_6/minicpmv4_6.py:154](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/minicpmv4_6/minicpmv4_6.py#L154) | MISMATCH vs mlx-vlm, matches HF (`MiniCPMV4_6DownsampleMLP` uses `nn.GELU()`, erf); **kept erf in #2423** | Issue item. Same ckpts. The 4.6 encoder layers reuse `MiniCPMOVisionModel`, so their MLP is V5 | -| V5 | `src/vision/encoders/minicpmo.rs:237` | `MiniCPMOMlp` | MiniCPM-o, and MiniCPM-V 4.6 encoder layers (`minicpmv4_6.rs` wraps `MiniCPMOVisionModel`) | **erf** | tanh ([mlx_vlm/models/minicpmo/vision.py:67](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/minicpmo/vision.py#L67) `approx="precise"`; 4.6: `_gelu_from_name("gelu_pytorch_tanh")` -> `approx="tanh"`, [mlx_vlm/models/minicpmv4_6/vision.py:11-14](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/minicpmv4_6/vision.py#L11-L14), :73) | [mlx_vlm/models/minicpmo/vision.py](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/minicpmo/vision.py), [mlx_vlm/models/minicpmv4_6/vision.py](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/minicpmv4_6/vision.py) | **changed in #2423** for MiniCPM-V 4.6 (tanh f32 from `hidden_act`, compiled kernel); MiniCPM-o still erf (MISMATCH, follow-up) | Issue item. Local: `minicpm-v-4.6-bf16`, `minicpm-v-4.6-mxfp4` (no MiniCPM-o ckpt) | -| V6 | `src/vision/encoders/phi4_siglip.rs:160` | Phi4 SigLIP2 MLP (hard-coded, ignores config) | Phi-4-MM vision and Phi4-SigLIP | **erf** | tanh (`GELUMLP`) | [mlx_vlm/models/phi4mm/vision.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/phi4mm/vision.py#L8), [mlx_vlm/models/phi4_siglip/vision.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/phi4_siglip/vision.py#L8) | **MISMATCH** | Upstream config default `hidden_act=gelu_pytorch_tanh` ([mlx_vlm/models/phi4_siglip/config.py:53](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/phi4_siglip/config.py#L53)). No local ckpt | +| V5 | `src/vision/encoders/minicpmo.rs:237` | `MiniCPMOMlp` | MiniCPM-o, and MiniCPM-V 4.6 encoder layers (`minicpmv4_6.rs` wraps `MiniCPMOVisionModel`) | **erf** | tanh ([mlx_vlm/models/minicpmo/vision.py:67](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/minicpmo/vision.py#L67) `approx="precise"`; 4.6: `_gelu_from_name("gelu_pytorch_tanh")` -> `approx="tanh"`, [mlx_vlm/models/minicpmv4_6/vision.py:11-14](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/minicpmv4_6/vision.py#L11-L14), :73) | [mlx_vlm/models/minicpmo/vision.py](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/minicpmo/vision.py), [mlx_vlm/models/minicpmv4_6/vision.py](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/minicpmv4_6/vision.py) | **changed in #2423** for MiniCPM-V 4.6 (tanh f32 from `hidden_act`, compiled kernel); MiniCPM-o still erf (MISMATCH, follow-up) | Issue item. Local: `minicpm-v-4.6-bf16`, `minicpm-v-4.6-mxfp4` (no MiniCPM-o ckpt) **#2548:** MiniCPM-o now reads `vision_config.hidden_act` (tanh). | +| V6 | `src/vision/encoders/phi4_siglip.rs:160` | Phi4 SigLIP2 MLP (hard-coded, ignores config) | Phi-4-MM vision and Phi4-SigLIP | **erf** | tanh (`GELUMLP`) | [mlx_vlm/models/phi4mm/vision.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/phi4mm/vision.py#L8), [mlx_vlm/models/phi4_siglip/vision.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/phi4_siglip/vision.py#L8) | **MISMATCH** | Upstream config default `hidden_act=gelu_pytorch_tanh` ([mlx_vlm/models/phi4_siglip/config.py:53](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/phi4_siglip/config.py#L53)). No local ckpt **#2548:** now tanh (f32, compiled). | | V7 | `src/vision/phi4_siglip_vl.rs:103` | projector | Phi4-SigLIP | erf | erf | [mlx_vlm/models/phi4_siglip/phi4_siglip.py:22](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/phi4_siglip/phi4_siglip.py#L22) | MATCH | | | V8 | `src/vision/phi4mm_vl.rs:446` | HD-transform projector | Phi-4-MM | erf | erf | [mlx_vlm/models/phi4mm/phi4mm.py:23](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/phi4mm/phi4mm.py#L23) | MATCH | | | V9 | `src/multimodal/phi4mm_xla_audio.rs:462` | HD-transform projector (XLA host path) | Phi-4-MM | erf | erf | [mlx_vlm/models/phi4mm/phi4mm.py:23](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/phi4mm/phi4mm.py#L23) | MATCH | | | V10 | `src/vision/phi3_vl.rs:192` | `img_projection` | Phi-3-V | erf | erf | [mlx_vlm/models/phi3_v/vision.py:188](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/phi3_v/vision.py#L188) | MATCH | | -| V11 | `src/vision/encoders/deepseek_vl2.rs:140` | SigLIP `Block` MLP | DeepSeek-VL2 | **erf** | tanh (`approx="precise"`) | [mlx_vlm/models/deepseek_vl_v2/vision.py:149](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/deepseek_vl_v2/vision.py#L149) | **MISMATCH** | Inline comment says "GELU tanh approximation". Local: `deepseek-vl2-small-4bit` (no vision `hidden_act` in config) | +| V11 | `src/vision/encoders/deepseek_vl2.rs:140` | SigLIP `Block` MLP | DeepSeek-VL2 | **erf** | tanh (`approx="precise"`) | [mlx_vlm/models/deepseek_vl_v2/vision.py:149](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/deepseek_vl_v2/vision.py#L149) | **MISMATCH** | Inline comment says "GELU tanh approximation". Local: `deepseek-vl2-small-4bit` (no vision `hidden_act` in config) **#2548:** now tanh (f32, compiled). | | V12 | `src/vision/connectors/deepseek_vl2.rs:93` | projector | DeepSeek-VL2 | erf | erf | [mlx_vlm/models/deepseek_vl_v2/deepseek_vl_v2.py:32](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/deepseek_vl_v2/deepseek_vl_v2.py#L32), :50, :57 | MATCH | | -| V13 | `src/vision/encoders/llama4.rs:440`, `:443` | `Llama4VisionMLP` (encoder + projector) | Llama 4 vision | **erf** | fast (`approx="fast"`) | [mlx_vlm/models/llama4/vision.py:155](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/llama4/vision.py#L155) | **MISMATCH** (erf vs fast) | Local `llama-4-scout-17b-4bit` has `vision_config.hidden_act = "gelu"` (erf): mlxcel matches HF, mlx-vlm hard-codes fast | -| V14 | `src/vision/encoders/youtu_vl_layers.rs:213` | Siglip2 MLP | Youtu-VL | **erf** | tanh (`nn.gelu_approx`) | [mlx_vlm/models/youtu_vl/vision.py:131](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/youtu_vl/vision.py#L131) | **MISMATCH** | Local: `youtu-vl-4b-instruct` (`gelu_pytorch_tanh`) | +| V13 | `src/vision/encoders/llama4.rs:440`, `:443` | `Llama4VisionMLP` (encoder + projector) | Llama 4 vision | **erf** | fast (`approx="fast"`) | [mlx_vlm/models/llama4/vision.py:155](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/llama4/vision.py#L155) | **MISMATCH** (erf vs fast) | Local `llama-4-scout-17b-4bit` has `vision_config.hidden_act = "gelu"` (erf): mlxcel matches HF, mlx-vlm hard-codes fast **#2548:** erf kept (config `gelu`). | +| V14 | `src/vision/encoders/youtu_vl_layers.rs:213` | Siglip2 MLP | Youtu-VL | **erf** | tanh (`nn.gelu_approx`) | [mlx_vlm/models/youtu_vl/vision.py:131](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/youtu_vl/vision.py#L131) | **MISMATCH** | Local: `youtu-vl-4b-instruct` (`gelu_pytorch_tanh`) **#2548:** now tanh (f32, compiled). | | V15 | `src/vision/encoders/youtu_vl_merger.rs:88` | patch merger | Youtu-VL | erf | erf | [mlx_vlm/models/youtu_vl/vision.py:372](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/youtu_vl/vision.py#L372) | MATCH | | -| V16 | `src/vision/encoders/jina_vlm.rs:166` | ViT MLP | Jina-VLM | **erf** | cfg: `gelu_pytorch_tanh` -> tanh, else erf | [mlx_vlm/models/jina_vlm/vision.py:55-58](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/jina_vlm/vision.py#L55-L58) | **MISMATCH** | Struct comment says "tanh-GELU". Local `jina-vlm-mlx`: `ffn_config.activation_type = gelu_pytorch_tanh` | +| V16 | `src/vision/encoders/jina_vlm.rs:166` | ViT MLP | Jina-VLM | **erf** | cfg: `gelu_pytorch_tanh` -> tanh, else erf | [mlx_vlm/models/jina_vlm/vision.py:55-58](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/jina_vlm/vision.py#L55-L58) | **MISMATCH** | Struct comment says "tanh-GELU". Local `jina-vlm-mlx`: `ffn_config.activation_type = gelu_pytorch_tanh` **#2548:** now tanh (f32, compiled). | | V17 | `src/vision/encoders/gemma4.rs:352` | vision GeGLU MLP | Gemma4 vision | **erf** | tanh | [mlx_vlm/models/gemma4/vision.py:243](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/gemma4/vision.py#L243) | **changed in #2423** to tanh (f32, compiled kernel) from `hidden_activation`; also reaches DiffusionGemma and EmbeddingGemma2, whose mlx-vlm ports reuse `gemma4.vision` | Local: every `gemma-4-*` VLM ckpt (`vision_config.hidden_activation = gelu_pytorch_tanh`) | | V18 | `src/vision/encoders/gemma3n.rs:92` | `RMSNormAct2d` | Gemma3n vision (MobileNet) | erf | erf | [mlx_vlm/models/gemma3n/vision.py:147](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/gemma3n/vision.py#L147) | MATCH | | -| V19 | `src/vision/encoders/qwen2_vl.rs:416` | ViT `MLP` | Qwen2-VL | **erf** | fast (`approx="fast"`) | [mlx_vlm/models/qwen2_vl/vision.py:161](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/qwen2_vl/vision.py#L161) | **MISMATCH** (erf vs fast) | HF Qwen2-VL vision default is `quick_gelu` (fast) per the transformers `Qwen2VLVisionConfig` class default (verified: transformers 5.20.1 `Qwen2VLVisionConfig().hidden_act == "quick_gelu"`; the local configs omit the key). Local: `qwen2-vl-2b`, `qwen2-vl-2b-4bit` | +| V19 | `src/vision/encoders/qwen2_vl.rs:416` | ViT `MLP` | Qwen2-VL | **erf** | fast (`approx="fast"`) | [mlx_vlm/models/qwen2_vl/vision.py:161](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/qwen2_vl/vision.py#L161) | **MISMATCH** (erf vs fast) | HF Qwen2-VL vision default is `quick_gelu` (fast) per the transformers `Qwen2VLVisionConfig` class default (verified: transformers 5.20.1 `Qwen2VLVisionConfig().hidden_act == "quick_gelu"`; the local configs omit the key). Local: `qwen2-vl-2b`, `qwen2-vl-2b-4bit` **#2548:** now fast, eager and XLA. | | V20 | `src/vision/encoders/qwen2_vl.rs:490` | `PatchMerger` | Qwen2-VL | erf | erf | [mlx_vlm/models/qwen2_vl/vision.py:106](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/qwen2_vl/vision.py#L106) | MATCH | | | V21 | `src/vision/encoders/qwen2_5_vl.rs:736` | `PatchMerger` | Qwen2.5-VL | erf | erf | [mlx_vlm/models/qwen2_5_vl/vision.py:107](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/qwen2_5_vl/vision.py#L107) | MATCH | Vision MLP is SwiGLU | -| V22 | `src/vision/encoders/molmo.rs:111` | `MolmoViTMLP` | Molmo | **erf** | fast | [mlx_vlm/models/molmo/vision.py:41](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/molmo/vision.py#L41) | **MISMATCH** (erf vs fast) | Struct comment says "fast GELU". Local: `molmo-7b` (no act key) | -| V23 | `src/vision/encoders/molmo2.rs:119` (private tanh `:90`) | `ViTMLP` | Molmo2 | tanh (f32) | fast (`_gelu_from_name` maps `gelu_pytorch_tanh` -> `approx="fast"`) | [mlx_vlm/models/molmo2/vision.py:11-14](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/molmo2/vision.py#L11-L14) | **MISMATCH** (tanh vs fast) | mlxcel follows HF (`vit_config.hidden_act = gelu_pytorch_tanh`); mlx-vlm maps it to fast. Local: `molmo2-4b` | +| V22 | `src/vision/encoders/molmo.rs:111` | `MolmoViTMLP` | Molmo | **erf** | fast | [mlx_vlm/models/molmo/vision.py:41](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/molmo/vision.py#L41) | **MISMATCH** (erf vs fast) | Struct comment says "fast GELU". Local: `molmo-7b` (no act key) **#2548:** now fast, eager and XLA. | +| V23 | `src/vision/encoders/molmo2.rs:119` (private tanh `:90`) | `ViTMLP` | Molmo2 | tanh (f32) | fast (`_gelu_from_name` maps `gelu_pytorch_tanh` -> `approx="fast"`) | [mlx_vlm/models/molmo2/vision.py:11-14](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/molmo2/vision.py#L11-L14) | **MISMATCH** (tanh vs fast) | mlxcel follows HF (`vit_config.hidden_act = gelu_pytorch_tanh`); mlx-vlm maps it to fast. Local: `molmo2-4b` **#2548:** tanh kept (HF). | | V24 | same `molmo2.rs` ViT, via `Molmo2VisionTransformer` | ViT MLP | MolmoPoint | tanh | tanh (`approx="tanh"` for `gelu_pytorch_tanh`) | [mlx_vlm/models/molmo_point/vision.py:14-17](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/molmo_point/vision.py#L14-L17) | MATCH | No local ckpt | -| V25 | `src/vision/encoders/lfm2_vl.rs:181` | SigLIP2 NaFlex MLP | LFM2-VL / LFM2.5-VL | **erf** | tanh (`GELUMLP`) | [mlx_vlm/models/lfm2_vl/vision.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/lfm2_vl/vision.py#L8) | **MISMATCH** | Local: `lfm2-vl-1.6b-4bit`, `lfm2-vl-450m-4bit`, `lfm2.5-vl-1.6b-4bit` | +| V25 | `src/vision/encoders/lfm2_vl.rs:181` | SigLIP2 NaFlex MLP | LFM2-VL / LFM2.5-VL | **erf** | tanh (`GELUMLP`) | [mlx_vlm/models/lfm2_vl/vision.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/lfm2_vl/vision.py#L8) | **MISMATCH** | Local: `lfm2-vl-1.6b-4bit`, `lfm2-vl-450m-4bit`, `lfm2.5-vl-1.6b-4bit` **#2548:** now tanh (f32, compiled). | | V26 | `src/vision/connectors/lfm2_vl.rs:128` | projector | LFM2-VL | erf | erf | [mlx_vlm/models/lfm2_vl/lfm2_vl.py:37](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/lfm2_vl/lfm2_vl.py#L37) | MATCH | | -| V27 | `src/vision/encoders/moondream3.rs:209` | vision `MLP` | Moondream2 VL and Moondream3 VL (shared encoder) | **erf** | tanh (`TanhGELUMLP`) | [mlx_vlm/models/moondream2/vision.py:6](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/moondream2/vision.py#L6), [mlx_vlm/models/moondream3/vision.py:6](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/moondream3/vision.py#L6) | **MISMATCH** | Local: `moondream2` | -| V28 | `src/vision/encoders/moondream3.rs:360` | `VisionProjection` | Moondream2 VL and Moondream3 VL | **erf** | tanh (`approx="tanh"`) | [mlx_vlm/models/moondream2/vision.py:90](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/moondream2/vision.py#L90), [mlx_vlm/models/moondream3/vision.py:90](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/moondream3/vision.py#L90) | **MISMATCH** | Local: `moondream2` | -| V29 | `src/vision/encoders/paddleocr_vl.rs:591` | vision `MLP` (via `mlxcel_core::gelu`) | PaddleOCR-VL | **erf** | tanh (`approx="precise"`) | [mlx_vlm/models/paddleocr_vl/vision.py:277](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/paddleocr_vl/vision.py#L277) | **MISMATCH** | Not in the issue (uses `gelu`, not `gelu_approx`). Local: `paddleocr-vl-bfloat16` (`vision_config.hidden_act = gelu_pytorch_tanh`) | +| V27 | `src/vision/encoders/moondream3.rs:209` | vision `MLP` | Moondream2 VL and Moondream3 VL (shared encoder) | **erf** | tanh (`TanhGELUMLP`) | [mlx_vlm/models/moondream2/vision.py:6](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/moondream2/vision.py#L6), [mlx_vlm/models/moondream3/vision.py:6](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/moondream3/vision.py#L6) | **MISMATCH** | Local: `moondream2` **#2548:** now tanh (f32, compiled). | +| V28 | `src/vision/encoders/moondream3.rs:360` | `VisionProjection` | Moondream2 VL and Moondream3 VL | **erf** | tanh (`approx="tanh"`) | [mlx_vlm/models/moondream2/vision.py:90](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/moondream2/vision.py#L90), [mlx_vlm/models/moondream3/vision.py:90](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/moondream3/vision.py#L90) | **MISMATCH** | Local: `moondream2` **#2548:** now tanh (f32, compiled). | +| V29 | `src/vision/encoders/paddleocr_vl.rs:591` | vision `MLP` (via `mlxcel_core::gelu`) | PaddleOCR-VL | **erf** | tanh (`approx="precise"`) | [mlx_vlm/models/paddleocr_vl/vision.py:277](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/paddleocr_vl/vision.py#L277) | **MISMATCH** | Not in the issue (uses `gelu`, not `gelu_approx`). Local: `paddleocr-vl-bfloat16` (`vision_config.hidden_act = gelu_pytorch_tanh`) **#2548:** now tanh (f32, compiled). | | V30 | `src/vision/connectors/paddleocr_vl.rs:86` | projector | PaddleOCR-VL | erf | erf | [mlx_vlm/models/paddleocr_vl/vision.py:151](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/paddleocr_vl/vision.py#L151) | MATCH | | | V31 | `src/vision/encoders/qwen3_vl.rs:811` (private tanh `:770`) | `VisionMLP` | Qwen3-VL, Qwen3.5 / 3.6 (VL), Cohere Compass | tanh (f32) | tanh | [mlx_vlm/models/qwen3_vl/vision.py:163](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/qwen3_vl/vision.py#L163); [mlx_vlm/models/qwen3_5/vision.py:1](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/qwen3_5/vision.py#L1) (imports qwen3_vl); [mlx_vlm/models/cohere_compass/vision.py:145](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/cohere_compass/vision.py#L145) | MATCH | Local: `qwen3-vl-*`, `qwen3.5-*`, `qwen3.6-*`, `north-micro-vision-instruct-4bit` | | V32 | `src/vision/encoders/qwen3_vl.rs:1024` | `PatchMerger` | same | erf | erf | [mlx_vlm/models/qwen3_vl/vision.py:109](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/qwen3_vl/vision.py#L109); [mlx_vlm/models/cohere_compass/vision.py:87](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/cohere_compass/vision.py#L87) | MATCH | | -| V33 | `qwen3_vl.rs:811` / `:1024` for `qwen3_vl_moe` | MLP / merger | Qwen3-VL-MoE | MLP tanh, merger erf | MLP **erf**, merger **tanh** (swapped vs qwen3_vl) | [mlx_vlm/models/qwen3_vl_moe/vision.py:164](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/qwen3_vl_moe/vision.py#L164) (MLP `nn.GELU()`), :109 (merger `approx="tanh"`) | **MISMATCH** (both) | mlx-vlm looks swapped relative to HF (`vision_config.hidden_act = gelu_pytorch_tanh` for the MLP). Local: `qwen3-vl-30b-a3b-4bit` | -| V34 | `qwen3_vl.rs:811` for `qwen3_omni_moe` | vision MLP | Qwen3-Omni vision | tanh | erf (`nn.GELU()`) | [mlx_vlm/models/qwen3_omni_moe/vision.py:163](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/qwen3_omni_moe/vision.py#L163) | **MISMATCH** | mlxcel follows HF (`thinker_config.vision_config.hidden_act = gelu_pytorch_tanh`). Merger :1024 erf vs `vision.py` :106 erf MATCH. Local: `qwen3-omni-30b-a3b-instruct-4bit` | +| V33 | `qwen3_vl.rs:811` / `:1024` for `qwen3_vl_moe` | MLP / merger | Qwen3-VL-MoE | MLP tanh, merger erf | MLP **erf**, merger **tanh** (swapped vs qwen3_vl) | [mlx_vlm/models/qwen3_vl_moe/vision.py:164](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/qwen3_vl_moe/vision.py#L164) (MLP `nn.GELU()`), :109 (merger `approx="tanh"`) | **MISMATCH** (both) | mlx-vlm looks swapped relative to HF (`vision_config.hidden_act = gelu_pytorch_tanh` for the MLP). Local: `qwen3-vl-30b-a3b-4bit` **#2548:** kept (matches HF). | +| V34 | `qwen3_vl.rs:811` for `qwen3_omni_moe` | vision MLP | Qwen3-Omni vision | tanh | erf (`nn.GELU()`) | [mlx_vlm/models/qwen3_omni_moe/vision.py:163](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/qwen3_omni_moe/vision.py#L163) | **MISMATCH** | mlxcel follows HF (`thinker_config.vision_config.hidden_act = gelu_pytorch_tanh`). Merger :1024 erf vs `vision.py` :106 erf MATCH. Local: `qwen3-omni-30b-a3b-instruct-4bit` **#2548:** tanh kept (HF). | | V35 | `src/vision/encoders/kimi_vl.rs:144` (`MoonViTMlpActivation::Gelu`, default) | MoonViT MLP | Kimi-VL, Kimi K2.5 | erf | erf | [mlx_vlm/models/kimi_vl/vision.py:228](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/kimi_vl/vision.py#L228), [mlx_vlm/models/kimi_k25/vision.py:245](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/kimi_k25/vision.py#L245) | MATCH | Local: `kimi-vl-a3b-thinking-4bit` | | V36 | `kimi_vl.rs:145` -> private `gelu_tanh` `:158` (set by `src/loading/vlm_locateanything.rs:248`) | MoonViT MLP | LocateAnything | tanh | tanh (`approx="precise"`) | [mlx_vlm/models/locateanything/vision.py:174](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/locateanything/vision.py#L174) | MATCH | Local: `locateanything-3b-4bit` | | V37 | `src/vision/encoders/moonvit3d.rs:251` -> `kimi_vl::gelu_tanh` | MoonViT3D MLP | Kimi K3 | tanh | tanh (`nn.gelu_approx`) | [mlx_vlm/models/kimi_k3/vision.py:118](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/kimi_k3/vision.py#L118) | MATCH | | @@ -193,18 +253,18 @@ Since #2423 the Gemma 3 VLM and PaliGemma loaders resolve a missing key with `Vi |---|---|---|---|---|---|---|---| | S1 | `src/loading/vlm_gemma.rs:177` | Gemma3 VLM | **erf** (key absent) | tanh (`GELUMLP`) | [mlx_vlm/models/gemma3/vision.py:9](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/gemma3/vision.py#L9) | **changed in #2423**: missing key resolves to `gelu_pytorch_tanh` (tanh f32, compiled kernel) | `gemma-3-4b-it-4bit`, `gemma3-4b-4bit`: absent | | S2 | `src/loading/vlm_siglip.rs:301` | PaliGemma / PaliGemma2 | **erf** (key absent) | tanh (`GELUMLP`) | [mlx_vlm/models/paligemma/vision.py:9](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/paligemma/vision.py#L9) | **changed in #2423**: missing key resolves to `gelu_pytorch_tanh` (tanh f32, compiled kernel) | `paligemma2-3b-6bit`: absent | -| 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 | +| 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) | -| 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 | -| 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 | -| 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` | -| S11 | `src/loading/vlm_llava.rs:816` (`parse_bunny_vision_config` hard-codes `ExactGelu`, `:653`) | Bunny | **erf** | fast (`FastGELUMLP`) | [mlx_vlm/models/llava_bunny/vision.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/llava_bunny/vision.py#L8) | **MISMATCH** (erf vs fast) | `bunny-llama3-8b-4bit`: no `vision_config` | -| S12 | `src/loading/vlm_special.rs:1753` (`phi3_vision_config()` hard-codes `ExactGelu`, `:56`) | Phi-3-V / Phi-3.5-V (CLIP) | **erf** | fast (`FastGELUMLP`) | [mlx_vlm/models/phi3_v/vision.py:9](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/phi3_v/vision.py#L9) | **MISMATCH** (erf vs fast) | `phi-3.5-vision-4bit` | +| 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). | +| S11 | `src/loading/vlm_llava.rs:816` (`parse_bunny_vision_config` hard-codes `ExactGelu`, `:653`) | Bunny | **erf** | fast (`FastGELUMLP`) | [mlx_vlm/models/llava_bunny/vision.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/llava_bunny/vision.py#L8) | **MISMATCH** (erf vs fast) | `bunny-llama3-8b-4bit`: no `vision_config` **#2548:** now tanh (SigLIP tower). | +| S12 | `src/loading/vlm_special.rs:1753` (`phi3_vision_config()` hard-codes `ExactGelu`, `:56`) | Phi-3-V / Phi-3.5-V (CLIP) | **erf** | fast (`FastGELUMLP`) | [mlx_vlm/models/phi3_v/vision.py:9](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/phi3_v/vision.py#L9) | **MISMATCH** (erf vs fast) | `phi-3.5-vision-4bit` **#2548:** now fast. | | S13 | `src/loading/vlm_llmjp_vl.rs:300` (`use_fast_gelu = false`) | LLM-jp-VL / JAGLE | tanh | tanh (`GELUMLP`) | [mlx_vlm/models/llmjpvl/vision.py:9](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/llmjpvl/vision.py#L9) | MATCH | `llm-jp-4-vl-9b-beta`, `jagle-vl-2.2b-jagle-finevision`: `gelu_pytorch_tanh` | -| S14 | `src/models/llama_nemotron_vl_embedding.rs:206` | Llama-Nemotron-VL-Embed | cfg | fast (uses `granite_vision.VisionModel`) | [mlx_vlm/models/llama_nemotron_vl_embedding/llama_nemotron_vl_embedding.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/llama_nemotron_vl_embedding/llama_nemotron_vl_embedding.py#L8) | MISMATCH unless cfg is fast (never: mlxcel has no cfg path to fast) | no local ckpt | +| S14 | `src/models/llama_nemotron_vl_embedding.rs:206` | Llama-Nemotron-VL-Embed | cfg | fast (uses `granite_vision.VisionModel`) | [mlx_vlm/models/llama_nemotron_vl_embedding/llama_nemotron_vl_embedding.py:8](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/llama_nemotron_vl_embedding/llama_nemotron_vl_embedding.py#L8) | MISMATCH unless cfg is fast (never: mlxcel has no cfg path to fast) | no local ckpt **#2548:** tanh kept (declared). | | S15 | `src/models/siglip_text.rs:124` (missing key -> tanh) | SigLIP text tower | tanh by default | tanh | [mlx_vlm/models/siglip/siglip.py:36](https://github.com/Blaizzy/mlx-vlm/blob/v0.7.8/mlx_vlm/models/siglip/siglip.py#L36) | MATCH | no local ckpt | ### 2d-bis. Config guards (where mlxcel reads the activation string only to accept or reject it) @@ -229,8 +289,8 @@ Since #2423 the Gemma 3 VLM and PaliGemma loaders resolve a missing key with `Vi | `src/vision/encoders/siglip.rs:92` `gelu_pytorch_tanh` | tanh | **f32** then cast back | SigLIP/CLIP when `hidden_act = gelu_pytorch_tanh` (S4, S10, S13, S15) | `GeluActivation::TanhF32` (bit-equal) | | `src/vision/encoders/kimi_vl.rs:158` `gelu_tanh` | tanh | input dtype, `x*x*x` | LocateAnything (V36), Kimi K3 MoonViT3D (V37) | `GeluActivation::Tanh` (bit-equal) | | `src/vision/encoders/qwen3_vl.rs:770` `gelu_pytorch_tanh` | tanh | **f32** then cast back | Qwen3-VL family (V31, V33, V34) | `GeluActivation::TanhF32` (bit-equal) | -| `src/models/gemma3_backbone.rs:59` `gelu_approx` | tanh | input dtype, **`power(x, 3)`** | Nemotron VoiceChat TTS (G24) | not folded: folding would replace `power` with `x*x*x` and change numerics (follow-up) | -| `src/audio/nemotron_codec/layers.rs:233` `gelu_exact` | erf | input dtype weak scalars | Nemotron VoiceChat codec (T15) | not folded (erf, weak-scalar evaluation; follow-up) | +| `src/models/gemma3_backbone.rs:59` `gelu_approx` | tanh | input dtype, **`power(x, 3)`** | Nemotron VoiceChat TTS (G24) | not folded in #2423; **folded in #2548** onto `GeluActivation::Tanh` (tolerance test, numerics change in the last bit) | +| `src/audio/nemotron_codec/layers.rs:233` `gelu_exact` | erf | input dtype weak scalars | Nemotron VoiceChat codec (T15) | not folded in #2423; **folded in #2548** onto `GeluActivation::Erf` (tolerance test) | The three f32-evaluated copies already sidestep the bf16 concern the `gelu_approx` comment cites; `gemma3_backbone.rs` runs the `power` form the comment says produced NaN, in production. @@ -238,15 +298,18 @@ The three f32-evaluated copies already sidestep the bf16 concern the `gelu_appro These build StableHLO directly and are unaffected by `gelu_approx` semantics; listed so a later fold-in does not miss them. -- `src/lib/mlxcel-xla/src/emitter/model.rs:1116` `gelu_tanh`: tanh, used for Gemma2/3 and OLMo GeGLU (`:1279`) **and the StarCoder2 dense MLP** (`:1242`). The eager MLX StarCoder2 path is erf (T6) and upstream is erf, so the XLA StarCoder2 path diverges from both. -- `src/lib/mlxcel-xla/src/emitter/youtu_vl.rs:75` tanh (vision MLP, `:187`) and `:79` exact (merger, `:251`): XLA Youtu-VL matches upstream; the eager path (V14) does not. -- `src/lib/mlxcel-xla/src/emitter/gemma3n_math.rs:212` tanh (dense) and `:230` erf (sparse `sparse_gelu`, `:245`): the host oracle mirrors the eager erf choice for sparse layers (G15/G17), not upstream. +- `src/lib/mlxcel-xla/src/emitter/model.rs:1116` `gelu_tanh`: tanh, used for Gemma2/3 and OLMo GeGLU (`:1279`) **and the StarCoder2 dense MLP** (`:1242`). Before #2548 the eager StarCoder2 path was erf (T6), like mlx-lm; #2548 moved eager to tanh, the HF-declared activation this graph was validated against, so the two now agree. +- `src/lib/mlxcel-xla/src/emitter/youtu_vl.rs:75` tanh (vision MLP, `:187`) and `:79` exact (merger, `:251`): XLA Youtu-VL matches upstream; since #2548 the eager path (V14) does too. +- `src/lib/mlxcel-xla/src/emitter/gemma3n_math.rs:212` tanh (dense) and `:230` erf (sparse `sparse_gelu`, `:245`): the host oracle mirrored the eager erf choice for sparse layers (G15/G17). Since #2548 `sparse_gelu` and the emitted `sparse_gelu_stages` tape (`gemma3n_emit_ops.rs`) are tanh, with the eager kernel. +- `src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs` (Qwen2-VL ViT MLP) and `emitter/molmo_vision.rs` (Molmo ViT MLP) emitted tanh while eager was erf and HF is `quick_gelu`; since #2548 both emit `numeric_ops::quick_gelu`, like the eager encoders (V19, V22). - `src/lib/mlxcel-xla/src/emitter/gemma3n_emit_ops.rs:408` `gelu`: tanh. - `src/lib/mlxcel-xla/src/emitter/phi4_audio.rs:436` `gelu`: erf. -- `src/lib/mlxcel-xla/src/emitter/vision_config.rs:397-399`: `gelu_pytorch_tanh` -> tanh, missing/`gelu` -> erf, anything else is rejected (so XLA refuses `quick_gelu` rather than substituting). +- `src/lib/mlxcel-xla/src/emitter/vision_config.rs:397-399`: `gelu_pytorch_tanh` -> tanh, missing/`gelu` -> erf, anything else is rejected (so XLA refused `quick_gelu` rather than substituting). Since #2548 `quick_gelu` is accepted and an omitted key resolves by `vision_config.model_type` (SigLIP -> tanh, CLIP -> `quick_gelu`), as the eager loaders do. ## 4. Summary +This summary is the #2423 inventory state; section 0b holds the final decisions. + ### MISMATCH families (candidates for erf-vs-tanh measurement) Primary set, mlxcel computes **erf** where upstream computes **tanh** (the issue's question): @@ -286,5 +349,3 @@ Secondary set, not erf-vs-tanh but still a reference mismatch: - wrong activation: Nemotron-NAS maps `gelu_new`/`gelu_fast` to silu (T8); no local checkpoint declares those. Measurement caveat: for InternVL, Llama 4 vision, GPT-NeoX, ModernBERT, Molmo2, Qwen3-Omni vision, Qwen3-VL-MoE, Kokoro and Gemma v1 the mlx-vlm / mlx-lm / mlx-audio reference itself departs from the checkpoint's HF `hidden_act`. A logit-trace gate against those references will show mlxcel's HF-faithful choice as divergence; decide per family whether the reference or HF is the target before reading the result. -stash@{0}: WIP on fix/xla-feature-build: a57ca3dd fix(server): handle PromptCacheWarmup in the OpenXLA worker -4658ae83 fix(engine): reset shared model caches when the CLI opens a sequence (#2542) diff --git a/src/audio/nemotron_codec/codec_tests.rs b/src/audio/nemotron_codec/codec_tests.rs index 080a68757..85d65bfb3 100644 --- a/src/audio/nemotron_codec/codec_tests.rs +++ b/src/audio/nemotron_codec/codec_tests.rs @@ -321,3 +321,47 @@ fn cache_update_semantics() { let wide = mlxcel_core::zeros(&[1, 1, 2], mlxcel_core::dtype::FLOAT32); assert!(cache.update(&wide, CacheKey::Block(1), 2, false).is_err()); } + +/// #2548 folded the private weak-scalar erf GELU onto `GeluActivation::Erf`. +/// f32 (the dtype the codec runs in) stays within f32 rounding of the former +/// expression; half-precision inputs within a couple of rounding steps. +#[test] +fn folded_codec_gelu_stays_within_rounding_of_the_weak_scalar_form() { + use super::layers::{CODEC_GELU, scalar_like}; + let weak_scalar_gelu = |x: &MlxArray| { + let scaled = mlxcel_core::divide(x, &scalar_like(std::f32::consts::SQRT_2, x)); + let one_plus = mlxcel_core::add(&mlxcel_core::erf(&scaled), &scalar_like(1.0, x)); + let prod = mlxcel_core::multiply(x, &one_plus); + mlxcel_core::divide(&prod, &scalar_like(2.0, x)) + }; + let grid: Vec = (0..=2048) + .map(|i| -8.0 + 16.0 * i as f32 / 2048.0) + .collect(); + let xs = mlxcel_core::from_slice_f32(&grid, &[grid.len() as i32]); + for (name, dt, abs, rel) in [ + ("f32", mlxcel_core::dtype::FLOAT32, 2e-6f32, 2e-6f32), + ( + "bf16", + mlxcel_core::dtype::BFLOAT16, + 2e-2, + 2.0 * 2f32.powi(-8), + ), + ( + "f16", + mlxcel_core::dtype::FLOAT16, + 4e-3, + 2.0 * 2f32.powi(-11), + ), + ] { + let x = mlxcel_core::astype(&xs, dt); + let got = CODEC_GELU.apply(&x); + assert_eq!(mlxcel_core::array_dtype(&got), dt, "{name}: dtype"); + let (got, want) = (to_vec(&got), to_vec(&weak_scalar_gelu(&x))); + for ((xi, g), w) in grid.iter().zip(&got).zip(&want) { + assert!( + (g - w).abs() <= abs + rel * w.abs(), + "{name}: gelu({xi}) = {g}, former form {w}" + ); + } + } +} diff --git a/src/audio/nemotron_codec/layers.rs b/src/audio/nemotron_codec/layers.rs index 295711b3e..ca216bd38 100644 --- a/src/audio/nemotron_codec/layers.rs +++ b/src/audio/nemotron_codec/layers.rs @@ -27,6 +27,7 @@ //! the operand dtype ([`scalar_like`]). That keeps the bf16 decoder numerics //! identical to the reference instead of silently promoting to f32. +use mlxcel_core::utils::GeluActivation; use mlxcel_core::weights::WeightMap; use mlxcel_core::{MlxArray, UniquePtr}; @@ -229,14 +230,12 @@ impl ConvTranspose1d { } } -/// `nn.gelu` (exact erf form) with weak scalars: `x * (1 + erf(x / sqrt2)) / 2`. -pub(crate) fn gelu_exact(x: &MlxArray) -> UniquePtr { - let scaled = mlxcel_core::divide(x, &scalar_like(std::f32::consts::SQRT_2, x)); - let erf = mlxcel_core::erf(&scaled); - let one_plus = mlxcel_core::add(&erf, &scalar_like(1.0, x)); - let prod = mlxcel_core::multiply(x, &one_plus); - mlxcel_core::divide(&prod, &scalar_like(2.0, x)) -} +/// The ConvNeXt block's GELU: mlx-audio's `nn.gelu`, the exact erf form, +/// through the shared [`GeluActivation::Erf`]. Until #2548 this was a private +/// copy evaluated with weak scalars in the input dtype; the shared helper +/// evaluates in f32 and casts back, so a half-precision input can differ in +/// the last bit (bounded in `codec_tests.rs`). +pub(crate) const CODEC_GELU: GeluActivation = GeluActivation::Erf; /// LayerNorm over the channel (last) axis of `[B, T, C]`, eps 1e-6. pub(crate) struct ChannelLayerNorm { @@ -341,7 +340,7 @@ impl ConvNeXtBlock1d { let hidden = self.dwconv.forward(&hidden)?; let hidden = self.norm.forward(&hidden); let hidden = self.pwconv1.forward(&hidden)?; - let hidden = gelu_exact(&hidden); + let hidden = CODEC_GELU.apply(&hidden); let hidden = self.pwconv2.forward(&hidden)?; Ok(mlxcel_core::add(x, &hidden)) } diff --git a/src/distributed/tensor_parallel/llama_runtime.rs b/src/distributed/tensor_parallel/llama_runtime.rs index 9cabf58b8..cc10278cc 100644 --- a/src/distributed/tensor_parallel/llama_runtime.rs +++ b/src/distributed/tensor_parallel/llama_runtime.rs @@ -1505,9 +1505,11 @@ impl TensorParallelGemma4Model { seq_len, text_config.hidden_size_per_layer_input as i32, ); + // Same tanh GeGLU kernel as the single-host fallback in + // `gemma4.rs` (the per-layer gate is `gelu_pytorch_tanh`); + // this TP copy computed erf before #2548. let gate = gate_proj.forward(&h); - let gate = mlxcel_core::gelu_approx(&gate); - let gate = mlxcel_core::multiply(&gate, &layer_input); + let gate = mlxcel_core::compiled_geglu_approx_activation(&gate, &layer_input); let gate = proj.forward(&gate); let gate = post_norm.forward(&gate); h = mlxcel_core::add(&residual, &gate); diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp index 98b1020fe..c4a1dced9 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp @@ -1907,7 +1907,7 @@ std::unique_ptr compiled_gpt_oss_swiglu_activation( } // Compiled GELU: x * 0.5 * (1 + erf(x / sqrt(2))) -// Used by: StarCoder2 +// Used by: FFI tests only (StarCoder2 moved to the tanh GELU in #2548) namespace { static std::function(const std::vector&)> get_compiled_gelu() { auto fn = [](const std::vector& inputs) -> std::vector { @@ -2039,11 +2039,12 @@ std::unique_ptr compiled_geglu_approx_activation( } // Compiled gelu_topk: sparse GELU with dynamic threshold, -// gelu(max(0, x - (mean + std * multiplier))) with the EXACT erf GELU. -// mlx-lm's `gelu_topk` applies `nn.gelu_approx` (tanh) and Hugging Face applies -// `gelu_pytorch_tanh`, so this kernel's activation departs from both; the -// switch needs this kernel, the fused `gemma3n_mlp_forward` sparse branch and -// the XLA Gemma3n host oracle changed together (#2423 follow-up). +// gelu_tanh(max(0, x - (mean + std * multiplier))). +// The activation is the tanh GELU in the input dtype (`gelu_tanh_approx`), what +// mlx-lm's `gelu_topk` (`nn.gelu_approx`) and Hugging Face's +// `gelu_pytorch_tanh` compute and what the dense Gemma3n GeGLU uses. It was the +// erf GELU before #2548; the XLA Gemma3n emitter (`sparse_gelu_stages`) and its +// host oracle (`gemma3n_math::sparse_gelu`) mirror this formula. // Used by: Gemma3n MLP layers with activation_sparsity > 0 namespace { static std::function(const std::vector&)> get_compiled_gelu_topk() { @@ -2064,13 +2065,8 @@ namespace { auto shifted = mlx::core::subtract(x, cutoff); auto zeroed = mlx::core::maximum(shifted, array(0.0f)); - // exact erf GELU (see the note above the namespace) - 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(zeroed, sqrt2)); - auto scale = mlx::core::multiply(half, mlx::core::add(one, erf_val)); - auto result = mlx::core::multiply(zeroed, scale); + // tanh GELU in the input dtype (see the note above the namespace) + auto result = gelu_tanh_approx(zeroed); return {mlx::core::astype(result, x.dtype())}; }; return compile_shapeless_audited("compiled_gelu_topk", fn); @@ -3041,15 +3037,15 @@ std::unique_ptr compiled_swiglu_mlp_forward_fp16( } // GELU MLP forward for non-quantized (FP16/BF16) weights: -// down_proj(gelu(gate_proj(x)) * up_proj(x)) +// down_proj(gelu_tanh(gate_proj(x)) * up_proj(x)) // // Matmul operations run outside the compile boundary because // mlx::core::compile with matmul+transpose can produce incorrect results // when the compiled graph is reused across layers with different weights. -// The GeGLU activation (gelu(gate) * up) uses the compiled geglu kernel -// which correctly fuses element-wise ops only. +// The GeGLU activation (gelu_tanh(gate) * up) uses the compiled tanh geglu +// kernel which correctly fuses element-wise ops only. // -// Used by: Gemma, Gemma4 and other GELU-gated FP models +// Used by: Gemma v1, Gemma4 dense MLP and the Gemma4 assistant drafter std::unique_ptr compiled_gelu_mlp_forward_fp16( const MlxArray& x, const MlxArray& gate_weight, @@ -3069,8 +3065,11 @@ std::unique_ptr compiled_gelu_mlp_forward_fp16( auto up = mlx::core::matmul(x.inner, up_t); if (up_bias) up = mlx::core::add(up, up_bias->inner); - // GeGLU activation: gelu(gate) * up — compiled element-wise fusion - static auto compiled_act = get_compiled_geglu(); + // GeGLU activation: gelu_tanh(gate) * up, compiled element-wise fusion. + // Every caller is a Gemma MLP whose checkpoint declares + // `gelu_pytorch_tanh`, so this uses the same tanh GeGLU as the quantized + // `compiled_gelu_approx_mlp_forward` (erf before #2548). + static auto compiled_act = get_compiled_geglu_approx(); auto activated = compiled_act({gate, up}); // down_proj(activated) = activated @ down_w.T [+ bias] diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h index cc661885e..4307bc122 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h @@ -579,7 +579,7 @@ std::unique_ptr compiled_scaled_rms_norm(const MlxArray& x, const MlxA std::unique_ptr compiled_silu(const MlxArray& x); // Compiled gelu: x * 0.5 * (1 + erf(x / sqrt(2))) — single fused kernel -// Used by: StarCoder2 and other precise GELU-based models +// Used by: FFI tests only (StarCoder2 moved to the tanh GELU in #2548) std::unique_ptr compiled_gelu(const MlxArray& x); // Compiled gelu_approx: despite the name, exact erf GELU @@ -594,7 +594,7 @@ std::unique_ptr compiled_gelu_approx(const MlxArray& x); std::unique_ptr compiled_gelu_tanh_f32(const MlxArray& x); // Compiled gelu_topk: sparse GELU with dynamic threshold — single fused kernel -// erf_gelu(max(0, x - (mean + std * multiplier))); upstream uses tanh (#2423) +// gelu_tanh(max(0, x - (mean + std * multiplier))) in the input dtype (#2548) // Used by: Gemma3n MLP layers with activation_sparsity > 0 std::unique_ptr compiled_gelu_topk( const MlxArray& x, @@ -801,9 +801,9 @@ std::unique_ptr compiled_swiglu_mlp_forward_fp16( ); // Compiled GELU MLP forward for non-quantized (FP16/BF16) weights: -// down_proj(gelu(gate_proj(x)) * up_proj(x)) -// Fuses gate_proj + gelu + up_proj + multiply + down_proj into compiled graph. -// Used by: Gemma, Gemma4 and other GELU-gated FP models +// down_proj(gelu_tanh(gate_proj(x)) * up_proj(x)) +// Runs the matmuls eagerly and the tanh GeGLU as one compiled kernel (#2548). +// Used by: Gemma v1, Gemma4 dense MLP and the Gemma4 assistant drafter std::unique_ptr compiled_gelu_mlp_forward_fp16( const MlxArray& x, const MlxArray& gate_weight, diff --git a/src/lib/mlxcel-core/src/ffi_tests.rs b/src/lib/mlxcel-core/src/ffi_tests.rs index edeab78f2..909611fe1 100644 --- a/src/lib/mlxcel-core/src/ffi_tests.rs +++ b/src/lib/mlxcel-core/src/ffi_tests.rs @@ -2019,6 +2019,119 @@ fn test_compiled_gelu_topk_preserves_bf16_and_f16_dtype() { } } +/// Max absolute difference between two arrays, both read back as f32. +fn max_abs_diff_2548(a: &MlxArray, b: &MlxArray) -> f32 { + let a = crate::utils::array_to_vec_f32(&astype(a, dtype::FLOAT32)); + let b = crate::utils::array_to_vec_f32(&astype(b, dtype::FLOAT32)); + assert_eq!(a.len(), b.len()); + a.iter() + .zip(&b) + .map(|(x, y)| (x - y).abs()) + .fold(0.0f32, f32::max) +} + +/// The cutoff-shifted input `max(x - (mean + std * multiplier), 0)` the +/// `gelu_topk` kernel activates, computed eagerly in f32. +fn gelu_topk_shifted_f32_2548(x: &MlxArray, std_multiplier: f32) -> UniquePtr { + let x = astype(x, dtype::FLOAT32); + let mean = mean_axis(&x, -1, true); + let var = mean_axis(&square(&subtract(&x, &mean)), -1, true); + let mult = full_f32(&[1], std_multiplier, dtype::FLOAT32); + let cutoff = add(&mean, &multiply(&sqrt(&var), &mult)); + maximum(&subtract(&x, &cutoff), &full_f32(&[1], 0.0, dtype::FLOAT32)) +} + +#[test] +fn test_compiled_gelu_topk_is_the_tanh_gelu_of_the_shifted_input() { + // #2548: mlx-lm `gelu_topk` and Hugging Face Gemma3n (`gelu_pytorch_tanh`) + // apply the tanh GELU after the cutoff; the kernel used erf before. + let values: Vec = (0..64).map(|i| (i as f32 - 20.0) * 0.173).collect(); + let x = from_slice_f32(&values, &[2, 32]); + let multiplier = -0.5; + let out = compiled_gelu_topk(&x, multiplier); + eval(&out); + + let shifted = gelu_topk_shifted_f32_2548(&x, multiplier); + let tanh_ref = gelu_tanh(&shifted); + let erf_ref = gelu(&shifted); + eval(&tanh_ref); + eval(&erf_ref); + + let to_tanh = max_abs_diff_2548(&out, &tanh_ref); + let to_erf = max_abs_diff_2548(&out, &erf_ref); + assert!(to_tanh <= 2e-6, "f32 topk vs tanh GELU: {to_tanh}"); + assert!( + to_erf > 1e-4, + "f32 topk must not be the erf GELU (max diff to erf {to_erf})" + ); + + // bf16: the kernel's intermediates follow MLX promotion against the f32 + // multiplier scalar. Report both carriers; the result must sit within one + // bf16 rounding step of the f32 tanh reference either way. + let x_bf16 = astype(&x, dtype::BFLOAT16); + let out_bf16 = compiled_gelu_topk(&x_bf16, multiplier); + eval(&out_bf16); + let shifted_bf16 = gelu_topk_shifted_f32_2548(&x_bf16, multiplier); + let f32_carrier = astype(&gelu_tanh(&shifted_bf16), dtype::BFLOAT16); + let bf16_carrier = gelu_tanh(&astype(&shifted_bf16, dtype::BFLOAT16)); + eval(&f32_carrier); + eval(&bf16_carrier); + let to_f32_carrier = max_abs_diff_2548(&out_bf16, &f32_carrier); + let to_bf16_carrier = max_abs_diff_2548(&out_bf16, &bf16_carrier); + eprintln!( + "compiled_gelu_topk bf16: max diff to f32-carrier tanh {to_f32_carrier}, \ + to bf16-carrier tanh {to_bf16_carrier}" + ); + assert!( + to_f32_carrier.min(to_bf16_carrier) <= 0.0625, + "bf16 topk is not the tanh GELU: {to_f32_carrier} / {to_bf16_carrier}" + ); +} + +#[test] +fn test_compiled_gelu_mlp_forward_fp16_uses_the_tanh_geglu() { + // #2548: the dense Gemma MLP kernel is the tanh GeGLU, like the quantized + // `compiled_gelu_approx_mlp_forward`; it used erf before. + let x_vals: Vec = (0..8).map(|i| (i as f32 - 3.5) * 0.4).collect(); + let w_vals: Vec = (0..48) + .map(|i| ((i * 7 % 13) as f32 - 6.0) * 0.25) + .collect(); + let x = from_slice_f32(&x_vals, &[1, 8]); + let gate_w = from_slice_f32(&w_vals[..48], &[6, 8]); + let up_w = from_slice_f32(&w_vals.iter().rev().copied().collect::>(), &[6, 8]); + let down_vals: Vec = (0..48).map(|i| ((i * 5 % 11) as f32 - 5.0) * 0.1).collect(); + let down_w = from_slice_f32(&down_vals, &[8, 6]); + + let out = unsafe { + compiled_gelu_mlp_forward_fp16( + &x, + &gate_w, + &up_w, + &down_w, + std::ptr::null(), + std::ptr::null(), + std::ptr::null(), + ) + }; + eval(&out); + + let gate = matmul(&x, &transpose(&gate_w)); + let up = matmul(&x, &transpose(&up_w)); + let down_t = transpose(&down_w); + let tanh_ref = matmul(&multiply(&gelu_tanh(&gate), &up), &down_t); + let erf_ref = matmul(&multiply(&gelu(&gate), &up), &down_t); + eval(&tanh_ref); + eval(&erf_ref); + + let to_tanh = max_abs_diff_2548(&out, &tanh_ref); + let to_erf = max_abs_diff_2548(&out, &erf_ref); + assert!(to_tanh <= 1e-5, "dense GeGLU vs tanh reference: {to_tanh}"); + assert!( + to_erf > 1e-4, + "dense GeGLU must not be the erf GeGLU (max diff to erf {to_erf})" + ); +} + #[test] fn test_gelu_approx_bf16_negative_values() { // Verify gelu_approx does not produce NaN for negative bf16 inputs. diff --git a/src/lib/mlxcel-core/src/layers.rs b/src/lib/mlxcel-core/src/layers.rs index 125b663fe..b358961b9 100644 --- a/src/lib/mlxcel-core/src/layers.rs +++ b/src/lib/mlxcel-core/src/layers.rs @@ -6220,12 +6220,14 @@ pub fn metal4_attention( /// GELU MLP forward for non-quantized (FP16/BF16) UnifiedLinear layers. /// /// When all three projections are non-quantized, calls the C++ path which -/// runs matmul operations directly and fuses only the GELU activation via -/// compiled element-wise kernels. +/// runs matmul operations directly and fuses only the tanh GeGLU activation +/// (`gelu_tanh(gate) * up`, the formula Gemma checkpoints declare as +/// `gelu_pytorch_tanh`) via compiled element-wise kernels. It computed the erf +/// GeGLU before #2548. /// /// Returns `None` if any projection is quantized (caller should use the quantized path). /// -/// Used by: Gemma, Gemma4 and other GELU-gated FP models +/// Used by: Gemma v1, Gemma4 dense MLP and the Gemma4 assistant drafter pub fn compiled_gelu_mlp_fp16( x: &MlxArray, gate_proj: &UnifiedLinear, diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index 59f693c1e..ab385de33 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/src/lib/mlxcel-core/src/lib.rs @@ -790,7 +790,7 @@ mod ffi { ) -> UniquePtr; /// Compiled gelu: x * 0.5 * (1 + erf(x / sqrt(2))) — single fused kernel - /// Used by: StarCoder2 and other precise GELU-based models + /// Used by: FFI tests only (StarCoder2 moved to the tanh GELU in #2548) fn compiled_gelu(x: &MlxArray) -> UniquePtr; /// Compiled `gelu_approx`: despite the name, the exact erf GELU @@ -814,8 +814,8 @@ mod ffi { fn compiled_geglu_approx_activation(gate: &MlxArray, x: &MlxArray) -> UniquePtr; /// Compiled gelu_topk: sparse GELU with dynamic threshold — single fused kernel - /// erf_gelu(max(0, x - (mean + std * multiplier))). Upstream applies the - /// tanh GELU here; see the C++ note (#2423). + /// gelu_tanh(max(0, x - (mean + std * multiplier))) in the input dtype, + /// matching mlx-lm's `gelu_topk` and Hugging Face (#2548). /// Used by: Gemma3n MLP layers with activation_sparsity > 0 fn compiled_gelu_topk(x: &MlxArray, std_multiplier: f32) -> UniquePtr; @@ -956,9 +956,10 @@ mod ffi { down_bias: *const MlxArray, ) -> UniquePtr; - /// Compiled GELU MLP forward for non-quantized (FP16/BF16) weights - /// Fuses gate_proj + gelu + up_proj + multiply + down_proj into compiled graph - /// Used by: Gemma, Gemma4 and other GELU-gated FP models + /// Compiled GELU MLP forward for non-quantized (FP16/BF16) weights: + /// `down_proj(gelu_tanh(gate_proj(x)) * up_proj(x))`, the tanh GeGLU + /// compiled, the matmuls eager (#2548; erf before). + /// Used by: Gemma v1, Gemma4 dense MLP and the Gemma4 assistant drafter unsafe fn compiled_gelu_mlp_forward_fp16( x: &MlxArray, gate_weight: &MlxArray, diff --git a/src/lib/mlxcel-core/src/utils.rs b/src/lib/mlxcel-core/src/utils.rs index 1fcf07f0c..9540c0efc 100644 --- a/src/lib/mlxcel-core/src/utils.rs +++ b/src/lib/mlxcel-core/src/utils.rs @@ -1122,9 +1122,13 @@ pub fn compiled_gelu_tanh_f32(x: &MlxArray) -> UniquePtr { /// MLP (LocateAnything, Kimi K3), the Molmo2 ViT, the Qwen3-VL vision blocks, /// the InternViT MLP, the MiniCPM-o / MiniCPM-V 4.6 encoder MLP and the 4.6 /// mergers, and the Gemma 4 vision GeGLU (also reached by DiffusionGemma and -/// EmbeddingGemma2) (#2423). Not used by the roughly 80 `gelu_approx` call -/// sites, which compute erf. Regenerate with -/// `grep -rln 'GeluActivation' src --include='*.rs'`. +/// EmbeddingGemma2) (#2423); the text MLPs of Phi, Phixtral, GPT-2, +/// StarCoder2, Moondream2/3 and RecurrentGemma, the Nemotron-NAS MLP, the +/// Nemotron VoiceChat TTS backbone and codec, the Gemma3n dead `gelu_topk` +/// branch, and the Phi-4 SigLIP, DeepSeek-VL2, Youtu-VL, Jina-VLM, LFM2-VL, +/// Moondream, PaddleOCR-VL, Qwen2-VL and Molmo vision MLPs (#2548). Not used +/// by the remaining `gelu_approx` call sites, which compute erf. Regenerate +/// with `grep -rln 'GeluActivation' src --include='*.rs'`. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] #[non_exhaustive] pub enum GeluActivation { diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_emit_ops.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_emit_ops.rs index a001fc440..ba5a9e94f 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_emit_ops.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_emit_ops.rs @@ -442,14 +442,16 @@ pub(super) struct SparseGeluStages { pub cutoff: Val, pub shifted_raw: Val, pub shifted: Val, - pub erf: Val, + pub tanh: Val, pub activated: Val, } /// Match the materialized BF16 carrier tape used by MLX's sparse GELU. /// /// Reductions accumulate in f32, but every observable elementwise result is -/// stored through BF16 before it feeds the next operation. +/// stored through BF16 before it feeds the next operation. The activation is +/// the tanh GELU, the formula of the eager `compiled_gelu_topk` kernel, mlx-lm's +/// `gelu_topk` and Hugging Face `gelu_pytorch_tanh` (erf before #2548). pub(super) fn sparse_gelu_stages( b: &mut Builder, x: &Val, @@ -488,18 +490,30 @@ pub(super) fn sparse_gelu_stages( let positive = b.compare("GT", &shifted_raw, &zeros, "FLOAT"); let shifted = b.select(&positive, &shifted_raw, &zeros); let shifted = round_bf16(b, &shifted); - let sqrt2 = b.const_f32(std::f32::consts::SQRT_2); - let sqrt2 = b.broadcast(&sqrt2, &[], vec![rows, width]); - let scaled = b.divide(&shifted, &sqrt2); - let scaled = round_bf16(b, &scaled); - let erf = b.erf(&scaled); - let erf = round_bf16(b, &erf); + // `gelu_tanh_approx`: x * (0.5 * (1 + tanh(sqrt(2/pi) * (x + 0.044715 x^3)))). + let shape = vec![rows, width]; + let root = b.const_f32((2.0f32 / std::f32::consts::PI).sqrt()); + let root = b.broadcast(&root, &[], shape.clone()); + let cubic = b.const_f32(0.044_715); + let cubic = b.broadcast(&cubic, &[], shape.clone()); + let squared_x = b.multiply(&shifted, &shifted); + let squared_x = round_bf16(b, &squared_x); + let cubed = b.multiply(&squared_x, &shifted); + let cubed = round_bf16(b, &cubed); + let cubic_term = b.multiply(&cubic, &cubed); + let cubic_term = round_bf16(b, &cubic_term); + let inner = b.add(&shifted, &cubic_term); + let inner = round_bf16(b, &inner); + let inner = b.multiply(&root, &inner); + let inner = round_bf16(b, &inner); + let tanh = b.gemma3n_tanh(&inner); + let tanh = round_bf16(b, &tanh); let half = b.const_f32(0.5); - let half = b.broadcast(&half, &[], vec![rows, width]); - let one = b.broadcast(one, &[], vec![rows, width]); - let one_erf = b.add(&one, &erf); - let one_erf = round_bf16(b, &one_erf); - let scale = b.multiply(&half, &one_erf); + let half = b.broadcast(&half, &[], shape.clone()); + let one = b.broadcast(one, &[], shape); + let one_tanh = b.add(&one, &tanh); + let one_tanh = round_bf16(b, &one_tanh); + let scale = b.multiply(&half, &one_tanh); let scale = round_bf16(b, &scale); let activated = b.multiply(&shifted, &scale); let activated = round_bf16(b, &activated); @@ -510,7 +524,7 @@ pub(super) fn sparse_gelu_stages( cutoff, shifted_raw, shifted, - erf, + tanh, activated, } } @@ -1120,10 +1134,13 @@ mod tests { let body = b.body(); assert_eq!( body.matches("stablehlo.convert").count(), - 28, - "14 dynamic MLX intermediates must each round through bf16" + 36, + "18 dynamic MLX intermediates must each round through bf16" ); assert_eq!(body.matches("stablehlo.sqrt").count(), 1); + // The activation is the tanh GELU (#2548), not erf. + assert_eq!(body.matches("stablehlo.tanh").count(), 1); + assert!(!body.contains("chlo.erf")); assert!(!body.contains("stablehlo.rsqrt")); } diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_math.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_math.rs index 85ed9f57a..1a5190740 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_math.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_math.rs @@ -215,6 +215,7 @@ pub(crate) fn gelu_approx(value: f32) -> f32 { 0.5 * value * (1.0 + inner.tanh()) } +#[cfg(test)] fn erf_approx(value: f32) -> f32 { // Abramowitz-Stegun 7.1.26 is sufficient for this independent f32 oracle. let sign = value.signum(); @@ -227,6 +228,7 @@ fn erf_approx(value: f32) -> f32 { sign * (1.0 - polynomial * (-x * x).exp()) } +#[cfg(test)] fn gelu_erf(value: f32) -> f32 { 0.5 * value * (1.0 + erf_approx(value / std::f32::consts::SQRT_2)) } @@ -242,7 +244,8 @@ pub(crate) fn sparse_gelu(values: &[f32], std_multiplier: f32) -> Vec { .iter() .map(|&value| { let shifted = (value - cutoff).max(0.0); - gelu_erf(shifted) + // tanh GELU, like the eager `compiled_gelu_topk` (#2548). + gelu_approx(shifted) }) .collect() } @@ -448,7 +451,9 @@ mod tests { let sparse = sparse_gelu(&[0.0, 1.0, 2.0], 0.0); assert_eq!(sparse[0], 0.0); assert_eq!(sparse[1], 0.0); - assert!((sparse[2] - gelu_erf(1.0)).abs() < 1e-6); + assert!((sparse[2] - gelu_approx(1.0)).abs() < 1e-6); + // tanh, not erf: the two differ by more than f32 noise at 1.0. + assert!((sparse[2] - gelu_erf(1.0)).abs() > 1e-5); } #[test] diff --git a/src/lib/mlxcel-xla/src/emitter/gemma3n_qmv.rs b/src/lib/mlxcel-xla/src/emitter/gemma3n_qmv.rs index 607c35d31..a0f599c8f 100644 --- a/src/lib/mlxcel-xla/src/emitter/gemma3n_qmv.rs +++ b/src/lib/mlxcel-xla/src/emitter/gemma3n_qmv.rs @@ -1058,7 +1058,7 @@ pub fn run_post_attention_diagnostic_probe( sparse.cutoff, sparse.shifted_raw, sparse.shifted, - sparse.erf, + sparse.tanh, sparse.activated, product, down, diff --git a/src/lib/mlxcel-xla/src/emitter/molmo_vision.rs b/src/lib/mlxcel-xla/src/emitter/molmo_vision.rs index 5eafe09af..c509edae9 100644 --- a/src/lib/mlxcel-xla/src/emitter/molmo_vision.rs +++ b/src/lib/mlxcel-xla/src/emitter/molmo_vision.rs @@ -18,7 +18,7 @@ use super::builder::{Builder, Ty, Val}; use super::molmo_vision_config::{ MolmoVisionConfig, MolmoVisionWeightDType, MolmoVisionWeightSpec, }; -use super::molmo_vision_ops::{scalar_broadcast, silu, softmax_last, tanh_gelu}; +use super::molmo_vision_ops::{scalar_broadcast, silu, softmax_last}; struct Args { values: Vec, @@ -180,7 +180,10 @@ fn encoder_layer( let residual = builder.add(hidden, &attention); let norm = layer_norm(builder, &residual, &args.take(), &args.take(), epsilon); let expanded = quant_linear(builder, args, config, &norm, true); - let activated = tanh_gelu(builder, &expanded); + // Molmo's remote code builds its CLIP ViT with + // `image_mlp_activations="quick_gelu"`, as the eager encoder evaluates it + // (#2548; tanh before). + let activated = super::numeric_ops::quick_gelu(builder, &expanded); let contracted = quant_linear(builder, args, config, &activated, true); builder.add(&residual, &contracted) } diff --git a/src/lib/mlxcel-xla/src/emitter/molmo_vision_ops.rs b/src/lib/mlxcel-xla/src/emitter/molmo_vision_ops.rs index 5b45f3020..75c4579f0 100644 --- a/src/lib/mlxcel-xla/src/emitter/molmo_vision_ops.rs +++ b/src/lib/mlxcel-xla/src/emitter/molmo_vision_ops.rs @@ -19,23 +19,6 @@ pub(super) fn scalar_broadcast(builder: &mut Builder, value: f32, shape: Vec Val { - let shape = value.ty.shape.clone(); - let half = scalar_broadcast(builder, 0.5, shape.clone()); - let one = scalar_broadcast(builder, 1.0, shape.clone()); - let coefficient = scalar_broadcast(builder, 0.044_715, shape.clone()); - let scale = scalar_broadcast(builder, 0.797_884_6, shape); - let squared = builder.multiply(value, value); - let cubed = builder.multiply(&squared, value); - let nonlinear = builder.multiply(&coefficient, &cubed); - let inner = builder.add(value, &nonlinear); - let scaled = builder.multiply(&scale, &inner); - let tanh = builder.tanh(&scaled); - let cdf = builder.add(&one, &tanh); - let half_value = builder.multiply(value, &half); - builder.multiply(&half_value, &cdf) -} - pub(super) fn softmax_last(builder: &mut Builder, scores: &Val) -> Val { let last = scores.ty.shape.len() - 1; let negative_infinity = builder.const_f32(f32::NEG_INFINITY); diff --git a/src/lib/mlxcel-xla/src/emitter/molmo_vision_tests.rs b/src/lib/mlxcel-xla/src/emitter/molmo_vision_tests.rs index c4dacde43..6cb4b164e 100644 --- a/src/lib/mlxcel-xla/src/emitter/molmo_vision_tests.rs +++ b/src/lib/mlxcel-xla/src/emitter/molmo_vision_tests.rs @@ -34,6 +34,9 @@ fn pinned_graph_contains_vit_pool_projector_and_static_output() { assert!(mlir.contains("stablehlo.dot_general")); assert!(mlir.contains("stablehlo.reduce")); assert!(mlir.contains("stablehlo.exponential")); + // ViT MLP: quick_gelu (#2548), not the tanh GELU. + assert!(mlir.contains(&crate::emitter::builder::f32_hex(1.702))); + assert!(!mlir.contains("stablehlo.tanh")); assert!(mlir.contains("stablehlo.compare")); assert!(mlir.contains("tensor<13x144x3584xf32>")); assert!(mlir.contains("molmo.image_masks")); diff --git a/src/lib/mlxcel-xla/src/emitter/numeric_ops.rs b/src/lib/mlxcel-xla/src/emitter/numeric_ops.rs index 9f062681b..c96e3ef05 100644 --- a/src/lib/mlxcel-xla/src/emitter/numeric_ops.rs +++ b/src/lib/mlxcel-xla/src/emitter/numeric_ops.rs @@ -203,6 +203,25 @@ pub(crate) fn exact_gelu(builder: &mut Builder, value: &Val) -> Val { builder.multiply(&half_value, &cdf) } +/// Hugging Face `quick_gelu` / MLX `nn.gelu_fast_approx`: +/// `x * sigmoid(1.702 x)`, with the sigmoid spelled `1 / (1 + exp(-z))`. +/// +/// Used by: the CLIP-towered LLaVA vision MLP, the Qwen2-VL ViT MLP and the +/// Molmo ViT MLP (#2548), matching the eager `GeluActivation::Fast`. +pub(crate) fn quick_gelu(builder: &mut Builder, value: &Val) -> Val { + let shape = value.ty.shape.clone(); + let one = builder.const_f32(1.0); + let one = builder.broadcast(&one, &[], shape.clone()); + let coefficient = builder.const_f32(1.702); + let coefficient = builder.broadcast(&coefficient, &[], shape); + let scaled = builder.multiply(&coefficient, value); + let negative = builder.negate(&scaled); + let exponential = builder.exponential(&negative); + let denominator = builder.add(&one, &exponential); + let sigmoid = builder.divide(&one, &denominator); + builder.multiply(value, &sigmoid) +} + /// PyTorch tanh-approximation GELU decomposition. pub(crate) fn tanh_gelu(builder: &mut Builder, value: &Val) -> Val { let shape = value.ty.shape.clone(); @@ -757,6 +776,18 @@ mod tests { assert!(!gelu_graph.contains("chlo.erf")); } + #[test] + fn quick_gelu_is_sigmoid_of_1_702_x() { + let input = Builder::arg(0, Ty::f32(vec![2, 3])); + let mut builder = Builder::new(); + let _ = quick_gelu(&mut builder, &input); + let body = builder.body(); + assert!(body.contains("stablehlo.exponential")); + assert!(body.contains(&crate::emitter::builder::f32_hex(1.702))); + assert!(!body.contains("chlo.erf")); + assert!(!body.contains("stablehlo.tanh")); + } + #[test] fn exact_gelu_uses_erf_not_tanh() { let input = Builder::arg(0, Ty::f32(vec![2, 3])); diff --git a/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs b/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs index 28eb04360..d37264c27 100644 --- a/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs +++ b/src/lib/mlxcel-xla/src/emitter/qwen2_vl.rs @@ -25,7 +25,7 @@ use std::path::Path; use serde_json::Value; use super::builder::{Builder, Precision, Ty, Val}; -use super::numeric_ops::{exact_gelu, layer_norm_2d as layer_norm, stable_softmax, tanh_gelu}; +use super::numeric_ops::{exact_gelu, layer_norm_2d as layer_norm, quick_gelu, stable_softmax}; /// Qualified flattened-patch capacities for the pinned Qwen2-VL image path. /// @@ -876,7 +876,9 @@ fn encoder_layer( config.layer_norm_eps, ); let fc1 = linear_2d(builder, &norm2, &args.take(), &args.take()); - let activated = tanh_gelu(builder, &fc1); + // `Qwen2VLVisionConfig().hidden_act == "quick_gelu"` (the shipped configs + // omit it), as the eager encoder evaluates it (#2548; tanh before). + let activated = quick_gelu(builder, &fc1); let fc2 = linear_2d(builder, &activated, &args.take(), &args.take()); builder.add(&residual, &fc2) } @@ -1077,6 +1079,9 @@ mod tests { assert!(mlir.contains("-> tensor<4x1536xf32>")); assert_eq!(mlir.matches("stablehlo.dot_general").count(), 15); assert_eq!(mlir.matches("chlo.erf").count(), 1); + // ViT MLP: quick_gelu (#2548), not the tanh GELU. + assert!(!mlir.contains("stablehlo.tanh")); + assert!(mlir.contains(&crate::emitter::builder::f32_hex(1.702))); } #[test] diff --git a/src/lib/mlxcel-xla/src/emitter/vision.rs b/src/lib/mlxcel-xla/src/emitter/vision.rs index 07a6b0898..ba6f4d7a9 100644 --- a/src/lib/mlxcel-xla/src/emitter/vision.rs +++ b/src/lib/mlxcel-xla/src/emitter/vision.rs @@ -22,7 +22,9 @@ //! linear weights remain `[out, in]`. use super::builder::{Builder, Ty, Val}; -use super::numeric_ops::{exact_gelu, layer_norm_2d as layer_norm, stable_softmax, tanh_gelu}; +use super::numeric_ops::{ + exact_gelu, layer_norm_2d as layer_norm, quick_gelu, stable_softmax, tanh_gelu, +}; use super::vision_config::{ LlavaVisionConfig, VisionActivation, VisionProjector, VisionWeightSpec, }; @@ -152,6 +154,7 @@ fn activate(builder: &mut Builder, value: &Val, activation: VisionActivation) -> match activation { VisionActivation::ExactGelu => exact_gelu(builder, value), VisionActivation::GeluPytorchTanh => tanh_gelu(builder, value), + VisionActivation::QuickGelu => quick_gelu(builder, value), } } diff --git a/src/lib/mlxcel-xla/src/emitter/vision_config.rs b/src/lib/mlxcel-xla/src/emitter/vision_config.rs index c69b64a9c..f8a35239f 100644 --- a/src/lib/mlxcel-xla/src/emitter/vision_config.rs +++ b/src/lib/mlxcel-xla/src/emitter/vision_config.rs @@ -22,6 +22,8 @@ use serde_json::Value; pub(crate) enum VisionActivation { ExactGelu, GeluPytorchTanh, + /// `quick_gelu`, declared or the `CLIPVisionConfig` default (#2548). + QuickGelu, } impl VisionActivation { @@ -29,6 +31,7 @@ impl VisionActivation { match self { Self::ExactGelu => "ExactGelu", Self::GeluPytorchTanh => "GeluPytorchTanh", + Self::QuickGelu => "QuickGelu", } } } @@ -170,7 +173,7 @@ mod tests { .contains("smaller than") ); root["vision_config"]["image_size"] = 28.into(); - root["vision_config"]["hidden_act"] = "quick_gelu".into(); + root["vision_config"]["hidden_act"] = "relu".into(); assert!( LlavaVisionConfig::from_json_str(&root.to_string()) .unwrap_err() @@ -185,6 +188,28 @@ mod tests { ); } + #[test] + fn omitted_hidden_act_resolves_to_the_tower_class_default() { + // #2548: HF drops a key equal to the class default, so a CLIP tower + // without `hidden_act` is `quick_gelu` and a SigLIP one is tanh. + let mut root: Value = serde_json::from_str(&config_json()).unwrap(); + root["vision_config"] + .as_object_mut() + .unwrap() + .remove("hidden_act"); + let siglip = LlavaVisionConfig::from_json_str(&root.to_string()).unwrap(); + assert_eq!(siglip.activation, VisionActivation::GeluPytorchTanh); + root["vision_config"]["model_type"] = "clip_vision_model".into(); + let clip = LlavaVisionConfig::from_json_str(&root.to_string()).unwrap(); + assert_eq!(clip.activation, VisionActivation::QuickGelu); + root["vision_config"]["hidden_act"] = "quick_gelu".into(); + let declared = LlavaVisionConfig::from_json_str(&root.to_string()).unwrap(); + assert_eq!(declared.activation, VisionActivation::QuickGelu); + root["vision_config"]["hidden_act"] = "gelu".into(); + let erf = LlavaVisionConfig::from_json_str(&root.to_string()).unwrap(); + assert_eq!(erf.activation, VisionActivation::ExactGelu); + } + #[test] fn pinned_siglip_contract_has_exact_schema() { let path = Path::new("/tmp/mlxcel-llava-hf-1090956d"); @@ -398,10 +423,16 @@ impl LlavaVisionConfig { } let activation = match vision.get("hidden_act").and_then(Value::as_str) { Some("gelu_pytorch_tanh") => VisionActivation::GeluPytorchTanh, - // Gemma 3 configs omit the key; Hugging Face's SigLIP default is - // the tanh GELU, and the eager loader resolves it the same way - // (`VisionConfig::default_missing_hidden_act_to_siglip`, #2423). - None if model_type == "gemma3" => VisionActivation::GeluPytorchTanh, + Some("quick_gelu") => VisionActivation::QuickGelu, + // An omitted key is the Hugging Face class default of the tower, + // resolved as the eager loaders do + // (`VisionConfig::default_missing_hidden_act_by_model_type`): + // SigLIP -> tanh (Gemma 3, #2423; SigLIP-towered LLaVA), CLIP -> + // `quick_gelu` (LLaVA 1.5 / LLaVA-NeXT, #2548). + None if model_type == "gemma3" || vision_model_type == "siglip_vision_model" => { + VisionActivation::GeluPytorchTanh + } + None if vision_model_type == "clip_vision_model" => VisionActivation::QuickGelu, None | Some("gelu") => VisionActivation::ExactGelu, other => { return Err(format!( diff --git a/src/loading/vlm_granite_vision.rs b/src/loading/vlm_granite_vision.rs index d16ee68ec..761c99b0e 100644 --- a/src/loading/vlm_granite_vision.rs +++ b/src/loading/vlm_granite_vision.rs @@ -96,13 +96,18 @@ pub(crate) fn load_granite_vision_vlm(model_path: &Path) -> Result let text_args: models::granite::ModelArgs = serde_json::from_value(text_config_val) .map_err(|e| anyhow::anyhow!("Failed to parse Granite text config: {}", e))?; - let vision_config: vision::config::VisionConfig = serde_json::from_value( - full_config - .get("vision_config") - .cloned() - .ok_or_else(|| anyhow::anyhow!("Missing vision_config in config.json"))?, - ) - .map_err(|e| anyhow::anyhow!("Failed to parse Granite vision config: {}", e))?; + let raw_vision_config = full_config + .get("vision_config") + .cloned() + .ok_or_else(|| anyhow::anyhow!("Missing vision_config in config.json"))?; + let mut vision_config: vision::config::VisionConfig = + serde_json::from_value(raw_vision_config.clone()) + .map_err(|e| anyhow::anyhow!("Failed to parse Granite vision config: {}", e))?; + // Granite Vision 3.x's SigLIP tower omits `hidden_act`; Hugging Face + // (`LlavaNext` + `SiglipVisionConfig`) evaluates `gelu_pytorch_tanh`. + // mlx-vlm's `FastGELUMLP` departs from it; this follows HF (#2548; erf + // before). + vision_config.default_missing_hidden_act_by_model_type(&raw_vision_config); let group_size = text_args.group_size(); let bits = text_args.bits(); diff --git a/src/loading/vlm_llava.rs b/src/loading/vlm_llava.rs index acda962cc..cdb7421cc 100644 --- a/src/loading/vlm_llava.rs +++ b/src/loading/vlm_llava.rs @@ -392,9 +392,10 @@ pub(crate) fn load_llava_host_preprocessor( let (config_str, full_config) = read_sanitized_vlm_config(model_path) .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; - let vlm_config: VLMConfig = parse_vlm_config(&config_str, "VLM config") + let mut vlm_config: VLMConfig = parse_vlm_config(&config_str, "VLM config") .map_err(|error| HostPreprocessorError::InvalidConfig(error.to_string()))?; require_llava_host_family(&vlm_config)?; + resolve_llava_vision_hidden_act(&mut vlm_config, &full_config); let text_model_type = vlm_config .text_config @@ -633,11 +634,28 @@ fn sanitize_bunny_weight_keys(weights: &mut WeightMap) { } } +/// Resolve an omitted `vision_config.hidden_act` to the tower's Hugging Face +/// class default: `quick_gelu` for the CLIP towers of LLaVA 1.5 / LLaVA-NeXT, +/// `gelu_pytorch_tanh` for a SigLIP tower (#2548; both were erf before). The +/// IREE preprocessor resolves the same rule in the OpenXLA vision config. +fn resolve_llava_vision_hidden_act(vlm_config: &mut VLMConfig, full_config: &Value) { + if let Some(raw) = full_config.get("vision_config") { + vlm_config + .vision_config + .default_missing_hidden_act_by_model_type(raw); + } +} + fn parse_bunny_vision_config( full_config: &Value, vision_config_val: &Value, ) -> vision::config::VisionConfig { - serde_json::from_value(vision_config_val.clone()).unwrap_or_else(|_| { + let parsed: Option = + serde_json::from_value(vision_config_val.clone()).ok(); + if let Some(mut config) = parsed { + config.default_missing_hidden_act_by_model_type(vision_config_val); + config + } else { let mm_hidden = full_config .get("mm_hidden_size") .and_then(|v| v.as_u64()) @@ -656,9 +674,12 @@ fn parse_bunny_vision_config( image_size: 384, num_channels: 3, layer_norm_eps: 1e-6, - hidden_act: vision::config::VisionHiddenActivation::ExactGelu, + // Bunny's `mm_vision_tower` is `google/siglip-so400m-patch14-384`, + // whose remote config declares `gelu_pytorch_tanh` (#2548; erf + // before). mlx-vlm's `FastGELUMLP` departs from it. + hidden_act: vision::config::VisionHiddenActivation::SiglipDefaultTanh, } - }) + } } /// Load a LLaVA VLM model (CLIP/SigLIP vision + Llama/Qwen2 text + MLP projector). @@ -668,7 +689,8 @@ pub(crate) fn load_llava_vlm(model_path: &Path) -> Result { use vision::encoders::siglip::SigLipVisionModel; let (config_str, full_config) = read_sanitized_vlm_config(model_path)?; - let vlm_config: VLMConfig = parse_vlm_config(&config_str, "VLM config")?; + let mut vlm_config: VLMConfig = parse_vlm_config(&config_str, "VLM config")?; + resolve_llava_vision_hidden_act(&mut vlm_config, &full_config); let mut weights = strip_language_model_prefix(load_vlm_weights_common(model_path, None)?); diff --git a/src/loading/vlm_siglip.rs b/src/loading/vlm_siglip.rs index 2d69d08c0..d456fc1b2 100644 --- a/src/loading/vlm_siglip.rs +++ b/src/loading/vlm_siglip.rs @@ -184,13 +184,16 @@ pub(crate) fn load_aya_vision_vlm(model_path: &Path) -> Result { sanitize_aya_weights(&mut weights); models::sanitize_tied_embeddings(&mut weights, &full_config); - let vision_config: vision::config::VisionConfig = serde_json::from_value( - full_config - .get("vision_config") - .cloned() - .ok_or_else(|| anyhow::anyhow!("Missing vision_config"))?, - ) - .map_err(|e| anyhow::anyhow!("Failed to parse vision_config: {}", e))?; + let raw_vision_config = full_config + .get("vision_config") + .cloned() + .ok_or_else(|| anyhow::anyhow!("Missing vision_config"))?; + let mut vision_config: vision::config::VisionConfig = + serde_json::from_value(raw_vision_config.clone()) + .map_err(|e| anyhow::anyhow!("Failed to parse vision_config: {}", e))?; + // Aya Vision's SigLIP tower omits `hidden_act`; the `SiglipVisionConfig` + // default is `gelu_pytorch_tanh` (#2548; erf before). + vision_config.default_missing_hidden_act_by_model_type(&raw_vision_config); let mut text_config = full_config .get("text_config") diff --git a/src/loading/vlm_special.rs b/src/loading/vlm_special.rs index f6eb4b108..27144fd2d 100644 --- a/src/loading/vlm_special.rs +++ b/src/loading/vlm_special.rs @@ -53,7 +53,10 @@ fn phi3_vision_config() -> vision::config::VisionConfig { num_channels: 3, patch_size: 14, layer_norm_eps: 1e-5, - hidden_act: vision::config::VisionHiddenActivation::ExactGelu, + // The Phi-3-V / Phi-3.5-V remote code builds its CLIP ViT-L/14-336 with + // `hidden_act="quick_gelu"` (`CLIP_VIT_LARGE_PATCH14_336_CONFIG`); erf + // before #2548. + hidden_act: vision::config::VisionHiddenActivation::QuickGelu, } } @@ -404,8 +407,15 @@ pub(crate) fn load_minicpmo_vlm(model_path: &Path) -> Result { let mut text_config: models::qwen3::ModelArgs = serde_json::from_value(full_config.clone()) .map_err(|e| anyhow::anyhow!("Failed to parse MiniCPM-o text config: {}", e))?; text_config.set_checkpoint_label(model_path); - let vision_config: MiniCPMOVisionConfig = + let mut vision_config: MiniCPMOVisionConfig = parse_required_vlm_subconfig(&full_config, "vision_config", "MiniCPM-o vision config")?; + // `vision_config.hidden_act` decides the encoder MLP GELU (#2548). + vision_config.mlp_activation = vision::encoders::minicpmo::mlp_activation_from_vision_config( + full_config + .get("vision_config") + .unwrap_or(&serde_json::Value::Null), + ) + .map_err(|e| anyhow::anyhow!(e))?; let processor_config = minicpmo_processor_config_value(model_path, &full_config) .unwrap_or_else(|| full_config.clone()); diff --git a/src/models/gemma3_backbone.rs b/src/models/gemma3_backbone.rs index 345a31195..02db33469 100644 --- a/src/models/gemma3_backbone.rs +++ b/src/models/gemma3_backbone.rs @@ -31,10 +31,11 @@ //! The layers reuse [`TransformerBlock`]'s weights and attention unchanged, //! so the per-layer RoPE base/scale rule, fused QKV projection, and SDPA //! dispatch are the same as the Gemma 3 text model's. The MLP activation is -//! the one exception: it runs [`gelu_approx`], an op-for-op port of -//! `mlx.nn.gelu_approx`, instead of the text model's fused GeGLU kernel, -//! because the two round differently and a host that samples from the -//! backbone state (the EAR-TTS RVQ `argmin`) sees the difference. One +//! the one exception: it runs [`BACKBONE_GELU`] (`mlx.nn.gelu_approx`, the +//! input-dtype tanh GELU) and multiplies by `up` eagerly, instead of the text +//! model's fused GeGLU kernel, because the fused kernel rounds once and a host +//! that samples from the backbone state (the EAR-TTS RVQ `argmin`) sees the +//! difference. One //! consequence of the reuse worth knowing: the global //! layers resolve their RoPE base through `rope_overrides`, so an operator //! `--rope-freq-base` override installed for a process also reaches a @@ -43,27 +44,19 @@ use crate::models::gemma3::{Cache, CacheInterface, ModelArgs, TransformerBlock}; use mlxcel_core::dtype; use mlxcel_core::layers::{GemmaRMSNorm, KVCache, RotatingKVCache}; +use mlxcel_core::utils::GeluActivation; use mlxcel_core::weights::WeightMap; use mlxcel_core::{MlxArray, UniquePtr}; -/// `mlx.nn.gelu_approx`, op for op: -/// `0.5 * x * (1 + tanh(sqrt(2 / pi) * (x + 0.044715 * x ** 3)))`, with every -/// literal in `x`'s dtype (Python's weak scalars) and the cube taken with -/// `power`. +/// The backbone MLP's GELU: `mlx.nn.gelu_approx`, the tanh form in `x`'s +/// dtype, through the shared [`GeluActivation::Tanh`]. /// -/// The crate's fused GeGLU (`compiled_geglu_approx_activation`) cubes with -/// two multiplies, groups the halving differently, and for bf16 rounds only -/// once; it disagrees with the reference on about half of all bf16 elements -/// and on the last bits of `f32` ones. -pub fn gelu_approx(x: &MlxArray) -> UniquePtr { - let d = mlxcel_core::array_dtype(x); - let lit = |v: f64| mlxcel_core::full_f32(&[], v as f32, d); - let cube = mlxcel_core::power(x, &lit(3.0)); - let inner = mlxcel_core::add(x, &mlxcel_core::multiply(&lit(0.044715), &cube)); - let inner = mlxcel_core::multiply(&lit((2.0 / std::f64::consts::PI).sqrt()), &inner); - let cdf = mlxcel_core::add(&lit(1.0), &mlxcel_core::tanh(&inner)); - mlxcel_core::multiply(&mlxcel_core::multiply(&lit(0.5), x), &cdf) -} +/// Until #2548 this was a private op-for-op port that cubed with `power`, the +/// op an earlier bf16 NaN was traced to. The shared helper cubes with +/// `x * x * x` and groups the halving as `x * (0.5 * (1 + tanh))`, so it can +/// differ from upstream's expression in the last bit of an element (bounded +/// by `gemma3_backbone_tests.rs`). +pub(crate) const BACKBONE_GELU: GeluActivation = GeluActivation::Tanh; /// One Gemma 3 decoder layer with the reference MLP activation. fn layer_forward( @@ -79,7 +72,7 @@ fn layer_forward( let normed = layer.pre_feedforward_layernorm.forward(&h); let mlp = &layer.mlp; let gated = mlxcel_core::multiply( - &gelu_approx(&mlp.gate_proj.forward(&normed)), + &BACKBONE_GELU.apply(&mlp.gate_proj.forward(&normed)), &mlp.up_proj.forward(&normed), ); let ff = mlp.down_proj.forward(&gated); diff --git a/src/models/gemma3_backbone_tests.rs b/src/models/gemma3_backbone_tests.rs index c7c70c2a4..905989727 100644 --- a/src/models/gemma3_backbone_tests.rs +++ b/src/models/gemma3_backbone_tests.rs @@ -172,8 +172,8 @@ fn wrong_cache_count_is_rejected() { /// `mx.power`; and every Python float becomes a weak scalar, which the MLX /// bindings build as `array(float(v), x.dtype)` (an `f32` rounding first, then /// a cast to `x`'s dtype). This is written from that line, not from -/// [`gelu_approx`], so the comparison below checks the port against the -/// upstream definition on whatever backend runs the test. +/// [`BACKBONE_GELU`], so the comparison below checks the shared helper against +/// the upstream definition on whatever backend runs the test. /// /// Upstream wraps the function in `@partial(mx.compile, shapeless=True)`. /// Both the port and this reference are the uncompiled expression: the @@ -197,49 +197,48 @@ fn gelu_approx_f64(x: f64) -> f64 { 0.5 * x * (1.0 + ((2.0 / std::f64::consts::PI).sqrt() * (x + 0.044715 * x.powi(3))).tanh()) } -fn bits(a: &MlxArray) -> Vec { - to_vec(a).iter().map(|v| v.to_bits()).collect() -} - #[test] -fn gelu_approx_matches_mlx_nn_bit_for_bit() { - // Bit-for-bit against the upstream expression, evaluated at run time on - // the same device with the same ops. Hardcoding values instead ties the - // test to one backend's `tanh`/`power` kernels: Metal and ROCm disagree by - // one f32 ulp at x = 4.1 (4.099957 vs 4.0999565), and both are correct - // for their backend. A dense grid over [-8, 8] makes a changed op - // (`x * x * x` for `power`, a literal rounded in another dtype, a single - // rounding at the end, the fused GeGLU kernel) show up as a differing ulp - // somewhere even when a handful of points happen to round the same way: - // on ROCm, `x * x * x` differs from `power` at only 2 of these 4097 f32 - // inputs. +fn backbone_gelu_stays_within_rounding_of_mlx_nn() { + // #2548 folded the private op-for-op port onto the shared + // `GeluActivation::Tanh` (`x * x * x` instead of `power`, a different + // grouping of the halving), so this is a tolerance test, not the former + // bit-for-bit one: on a dense grid over [-8, 8] every element must sit + // within a couple of the dtype's rounding steps of the upstream + // expression evaluated on the same device. let grid: Vec = (0..=4096) .map(|i| -8.0 + 16.0 * i as f32 / 4096.0) .collect(); let xs = mlxcel_core::from_slice_f32(&grid, &[grid.len() as i32]); - for (name, dt) in [ - ("f32", dtype::FLOAT32), - ("bf16", dtype::BFLOAT16), - ("f16", dtype::FLOAT16), + // (name, dtype, absolute bound, relative bound) + for (name, dt, abs, rel) in [ + ("f32", dtype::FLOAT32, 2e-6, 2e-6), + ("bf16", dtype::BFLOAT16, 2e-2, 2.0 * 2f64.powi(-8)), + ("f16", dtype::FLOAT16, 4e-3, 2.0 * 2f64.powi(-11)), ] { let x = mlxcel_core::astype(&xs, dt); - let got = gelu_approx(&x); + let got = BACKBONE_GELU.apply(&x); assert_eq!(mlxcel_core::array_dtype(&got), dt, "{name}: dtype"); let want = mlx_nn_gelu_approx_reference(&x); - let (got, want) = (bits(&got), bits(&want)); - let mismatches: Vec<_> = grid + let (got, want) = (to_vec(&got), to_vec(&want)); + let mut worst = 0.0f64; + let outside: Vec<_> = grid .iter() .zip(got.iter().zip(&want)) - .filter(|(_, (g, w))| g != w) - .map(|(x, (g, w))| (*x, f32::from_bits(*g), f32::from_bits(*w))) + .filter(|(_, (g, w))| { + let err = (f64::from(**g) - f64::from(**w)).abs(); + worst = worst.max(err); + err > abs + rel * f64::from(**w).abs() + }) + .map(|(x, (g, w))| (*x, *g, *w)) .collect(); assert!( - mismatches.is_empty(), - "{name}: {} of {} elements differ from mlx.nn.gelu_approx (x, got, want): {:?}", - mismatches.len(), + outside.is_empty(), + "{name}: {} of {} elements outside the bound (worst {worst:e}) (x, got, want): {:?}", + outside.len(), grid.len(), - &mismatches[..mismatches.len().min(8)] + &outside[..outside.len().min(8)] ); + eprintln!("{name}: worst |folded - mlx.nn| = {worst:e}"); } // Backend-independent sanity values, so the check above cannot pass by @@ -253,7 +252,7 @@ fn gelu_approx_matches_mlx_nn_bit_for_bit() { ("f32", dtype::FLOAT32, 1e-6, 1e-6), ("bf16", dtype::BFLOAT16, 4e-3, 1e-2), ] { - let got = to_vec(&gelu_approx(&mlxcel_core::astype(&xs, dt))); + let got = to_vec(&BACKBONE_GELU.apply(&mlxcel_core::astype(&xs, dt))); for (&xi, &g) in x.iter().zip(&got) { let want = gelu_approx_f64(f64::from(xi)); let err = (f64::from(g) - want).abs(); diff --git a/src/models/gemma3n.rs b/src/models/gemma3n.rs index e7c4498df..abbdc9d7e 100644 --- a/src/models/gemma3n.rs +++ b/src/models/gemma3n.rs @@ -780,13 +780,15 @@ impl MLP { /// Uses a compiled fused kernel matching Python's @mx.compile gelu_topk. fn gelu_topk(&self, x: &MlxArray) -> UniquePtr { if self.activation_sparsity <= 0.0 { - return mlxcel_core::gelu_approx(x); + // Unreachable from `forward` (it calls this only for sparse + // layers); kept on the dense layers' tanh GELU. + return mlxcel_core::utils::GeluActivation::Tanh.apply(x); } // Single compiled kernel: mean/std/cutoff/max/gelu fused together, - // mirroring mlx-lm's @partial(mx.compile, shapeless=True) gelu_topk - // except for the activation: this kernel applies erf GELU where - // mlx-lm applies `nn.gelu_approx` (tanh); tracked as a #2423 follow-up. + // mirroring mlx-lm's @partial(mx.compile, shapeless=True) gelu_topk, + // including its tanh GELU (`nn.gelu_approx`; HF `gelu_pytorch_tanh`). + // It applied the erf GELU before #2548. mlxcel_core::compiled_gelu_topk(x, self.std_multiplier) } @@ -3332,13 +3334,29 @@ mod xla_canonical_diagnostics { let sparse_shifted_raw = mlxcel_core::subtract(&gate, &sparse_cutoff); let sparse_zero = mlxcel_core::full_f32(&[1], 0.0, mlxcel_core::dtype::FLOAT32); let sparse_shifted = mlxcel_core::maximum(&sparse_shifted_raw, &sparse_zero); - let sparse_sqrt2 = - mlxcel_core::full_f32(&[1], std::f32::consts::SQRT_2, mlxcel_core::dtype::FLOAT32); - let sparse_erf = mlxcel_core::erf(&mlxcel_core::divide(&sparse_shifted, &sparse_sqrt2)); + // tanh GELU stage (#2548): tanh(sqrt(2/pi) * (x + 0.044715 x^3)). + let sparse_root = mlxcel_core::full_f32( + &[1], + (2.0f32 / std::f32::consts::PI).sqrt(), + mlxcel_core::dtype::FLOAT32, + ); + let sparse_cubic = mlxcel_core::full_f32(&[1], 0.044_715, mlxcel_core::dtype::FLOAT32); + let sparse_cube = mlxcel_core::multiply( + &mlxcel_core::multiply(&sparse_shifted, &sparse_shifted), + &sparse_shifted, + ); + let sparse_inner = mlxcel_core::multiply( + &sparse_root, + &mlxcel_core::add( + &sparse_shifted, + &mlxcel_core::multiply(&sparse_cubic, &sparse_cube), + ), + ); + let sparse_tanh = mlxcel_core::tanh(&sparse_inner); let sparse_half = mlxcel_core::full_f32(&[1], 0.5, mlxcel_core::dtype::FLOAT32); let sparse_one = mlxcel_core::full_f32(&[1], 1.0, mlxcel_core::dtype::FLOAT32); let sparse_scale = - mlxcel_core::multiply(&sparse_half, &mlxcel_core::add(&sparse_one, &sparse_erf)); + mlxcel_core::multiply(&sparse_half, &mlxcel_core::add(&sparse_one, &sparse_tanh)); let sparse_activation_manual = mlxcel_core::astype( &mlxcel_core::multiply(&sparse_shifted, &sparse_scale), mlxcel_core::array_dtype(&gate), @@ -3488,7 +3506,7 @@ mod xla_canonical_diagnostics { array_to_f32(&sparse_shifted_raw), ), ("layer0_sparse_shifted", array_to_f32(&sparse_shifted)), - ("layer0_sparse_erf", array_to_f32(&sparse_erf)), + ("layer0_sparse_tanh", array_to_f32(&sparse_tanh)), ("layer0_sparse_activation", array_to_f32(&sparse_activation)), ("layer0_sparse_product", array_to_f32(&sparse_product)), ("layer0_mlp_down", array_to_f32(&mlp_down)), @@ -4726,9 +4744,25 @@ mod xla_canonical_diagnostics { let sparse_shifted_raw = mlxcel_core::subtract(&gate, &sparse_cutoff); let sparse_zero = mlxcel_core::full_f32(&[1], 0.0, mlxcel_core::dtype::FLOAT32); let sparse_shifted = mlxcel_core::maximum(&sparse_shifted_raw, &sparse_zero); - let sparse_sqrt2 = - mlxcel_core::full_f32(&[1], std::f32::consts::SQRT_2, mlxcel_core::dtype::FLOAT32); - let sparse_erf = mlxcel_core::erf(&mlxcel_core::divide(&sparse_shifted, &sparse_sqrt2)); + // tanh GELU stage (#2548): tanh(sqrt(2/pi) * (x + 0.044715 x^3)). + let sparse_root = mlxcel_core::full_f32( + &[1], + (2.0f32 / std::f32::consts::PI).sqrt(), + mlxcel_core::dtype::FLOAT32, + ); + let sparse_cubic = mlxcel_core::full_f32(&[1], 0.044_715, mlxcel_core::dtype::FLOAT32); + let sparse_cube = mlxcel_core::multiply( + &mlxcel_core::multiply(&sparse_shifted, &sparse_shifted), + &sparse_shifted, + ); + let sparse_inner = mlxcel_core::multiply( + &sparse_root, + &mlxcel_core::add( + &sparse_shifted, + &mlxcel_core::multiply(&sparse_cubic, &sparse_cube), + ), + ); + let sparse_tanh = mlxcel_core::tanh(&sparse_inner); let sparse_activation = layer0.mlp.gelu_topk(&gate); let sparse_product = mlxcel_core::multiply(&sparse_activation, &up); let mlp_down = layer0.mlp.down_proj.forward(&sparse_product); @@ -4769,7 +4803,7 @@ mod xla_canonical_diagnostics { array_to_f32(&sparse_shifted_raw), ), ("decode_sparse_shifted", array_to_f32(&sparse_shifted)), - ("decode_sparse_erf", array_to_f32(&sparse_erf)), + ("decode_sparse_tanh", array_to_f32(&sparse_tanh)), ("decode_sparse_activation", array_to_f32(&sparse_activation)), ("decode_sparse_product", array_to_f32(&sparse_product)), ("decode_mlp_down", array_to_f32(&mlp_down)), diff --git a/src/models/gpt2.rs b/src/models/gpt2.rs index 6bb5b248a..e4d0cec01 100644 --- a/src/models/gpt2.rs +++ b/src/models/gpt2.rs @@ -25,12 +25,10 @@ //! - `LayerNorm` with bias (`layer_norm_epsilon`), not RMSNorm. //! - One fused `c_attn` projection that is split three ways on the last axis //! into Q, K and V. Multi-head attention: `n_kv_heads == n_head`. -//! - Simple MLP `c_proj(gelu_approx(c_fc(x)))` with an intermediate width of -//! `4 * n_embd`. `mlxcel_core::utils::gelu_approx` is the erf-based exact -//! GELU in this tree despite its name (see its own doc comment), not the -//! tanh approximation that upstream's `nn.gelu_approx` evaluates for the -//! checkpoint's `gelu_new` activation; the erf/tanh difference is under -//! 1e-3 and produces identical generated tokens. +//! - Simple MLP `c_proj(gelu(c_fc(x)))` with an intermediate width of +//! `4 * n_embd`. The GELU is the tanh form the checkpoint's +//! `activation_function = gelu_new` and upstream's `nn.gelu_approx` compute, +//! evaluated in f32 as one compiled kernel (it was erf before #2548). //! - Tied output head: logits come from `wte.as_linear`, there is no separate //! `lm_head` tensor in the checkpoint. //! @@ -586,7 +584,7 @@ impl SplitAttention for Attention { } } -// MLP (erf-based exact GELU via `gelu_approx`, no gate/up pattern). +// MLP (tanh GELU for `gelu_new`, no gate/up pattern). pub struct MLP { pub c_fc: UnifiedLinear, @@ -596,7 +594,7 @@ pub struct MLP { impl MLP { pub fn forward(&self, x: &MlxArray) -> UniquePtr { let h = self.c_fc.forward(x); - let h = mlxcel_core::utils::gelu_approx(&h); + let h = mlxcel_core::utils::GeluActivation::TanhF32Compiled.apply(&h); self.c_proj.forward(&h) } diff --git a/src/models/moondream2.rs b/src/models/moondream2.rs index 3609c1525..ca2359069 100644 --- a/src/models/moondream2.rs +++ b/src/models/moondream2.rs @@ -289,8 +289,10 @@ impl Moondream2Mlp { fn forward(&self, x: &MlxArray) -> UniquePtr { let hidden = self.fc1.forward(x); - // nn.GELU(approx="tanh") - let hidden = mlxcel_core::gelu_approx(&hidden); + // nn.GELU(approx="tanh"), as the remote code's `gelu_approx` + // (`F.gelu(approximate="tanh")`); evaluated in f32 as one compiled + // kernel (erf before #2548). + let hidden = mlxcel_core::utils::GeluActivation::TanhF32Compiled.apply(&hidden); self.fc2.forward(&hidden) } } diff --git a/src/models/moondream3.rs b/src/models/moondream3.rs index cd36de023..89163a1ae 100644 --- a/src/models/moondream3.rs +++ b/src/models/moondream3.rs @@ -383,7 +383,9 @@ impl DenseMlp { fn forward(&self, x: &MlxArray) -> UniquePtr { let hidden = self.fc1.forward(x); - let hidden = mlxcel_core::gelu_approx(&hidden); + // Tanh GELU (mlx-vlm `TanhGELUMLP`, the Moondream remote code's + // `F.gelu(approximate="tanh")`), compiled f32 (erf before #2548). + let hidden = mlxcel_core::utils::GeluActivation::TanhF32Compiled.apply(&hidden); self.fc2.forward(&hidden) } } diff --git a/src/models/nemotron_nas.rs b/src/models/nemotron_nas.rs index b84660bd9..71a984578 100644 --- a/src/models/nemotron_nas.rs +++ b/src/models/nemotron_nas.rs @@ -31,8 +31,7 @@ use mlxcel_core::layers::{KVCache, RMSNorm, UnifiedEmbedding, UnifiedLinear}; use mlxcel_core::utils::create_causal_mask; use mlxcel_core::weights::WeightMap; use mlxcel_core::{ - MlxArray, UniquePtr, add, array_shape, copy, fast_rope, gelu, multiply, relu, reshape, silu, - transpose_axes, + MlxArray, UniquePtr, add, array_shape, copy, fast_rope, multiply, reshape, transpose_axes, }; use serde::Deserialize; use std::cell::RefCell; @@ -352,7 +351,7 @@ struct NASMLP { gate_proj: UnifiedLinear, up_proj: UnifiedLinear, down_proj: UnifiedLinear, - act_fn: String, + act_fn: NasActivation, } impl NASMLP { @@ -360,13 +359,7 @@ impl NASMLP { let gate = self.gate_proj.forward(x); let up = self.up_proj.forward(x); - // Apply activation function - let activated = match self.act_fn.as_str() { - "silu" => silu(&gate), - "relu" => relu(&gate), - "gelu" => gelu(&gate), - _ => silu(&gate), - }; + let activated = self.act_fn.apply(&gate); let combined = multiply(&activated, &up); self.down_proj.forward(&combined) @@ -529,6 +522,9 @@ impl NemotronNASModel { let embeddings = UnifiedEmbedding::from_weights(&weights, "model.embed_tokens", group_size, bits)?; + // Resolve `hidden_act` once; an unknown name is a load error (#2548). + let act_fn = NasActivation::from_hf_name(&config.hidden_act)?; + // Build layers let mut layers = Vec::with_capacity(config.num_hidden_layers); for (i, block_config) in config.block_configs.iter().enumerate() { @@ -666,7 +662,7 @@ impl NemotronNASModel { gate_proj, up_proj, down_proj, - act_fn: config.hidden_act.clone(), + act_fn, }) }; @@ -866,6 +862,10 @@ impl LanguageModel for NemotronNASModel { } } +#[path = "nemotron_nas_activation.rs"] +mod activation; +use activation::NasActivation; + #[cfg(test)] #[path = "nemotron_nas_engine_tests.rs"] pub(crate) mod engine_tests; diff --git a/src/models/nemotron_nas_activation.rs b/src/models/nemotron_nas_activation.rs new file mode 100644 index 000000000..ba2eecf2a --- /dev/null +++ b/src/models/nemotron_nas_activation.rs @@ -0,0 +1,100 @@ +// 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. + +//! The gated-MLP activation a Nemotron-NAS checkpoint declares in `hidden_act`. +//! +//! Parsed once at load. Before #2548 the MLP matched the string per call and +//! sent every name it did not know, including the tanh GELU names `gelu_new` +//! and `gelu_fast`, to SiLU. GELU names now go through +//! [`GeluActivation::from_hf_name`] and any other unknown name is a load error. + +use mlxcel_core::utils::GeluActivation; +use mlxcel_core::{MlxArray, UniquePtr, gelu, relu, silu}; + +/// Used by: `NASMLP` in `nemotron_nas.rs`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum NasActivation { + Silu, + Relu, + Gelu(GeluActivation), +} + +impl NasActivation { + /// Map a Hugging Face `hidden_act` onto the activation `ACT2FN` computes, + /// or reject it. + pub(super) fn from_hf_name(name: &str) -> Result { + match name.trim().to_ascii_lowercase().as_str() { + "silu" | "swish" => Ok(Self::Silu), + "relu" => Ok(Self::Relu), + other => GeluActivation::from_hf_name(other) + // The tanh names run as the compiled f32 kernel the other + // families switched to tanh in #2423 / #2548 use. + .map(|act| Self::Gelu(act.in_f32_compiled())) + .ok_or_else(|| { + format!( + "Nemotron-NAS: unsupported hidden_act {name:?}; expected silu, relu, \ + or a GELU variant (gelu, gelu_new, gelu_fast, gelu_pytorch_tanh, \ + quick_gelu)" + ) + }), + } + } + + pub(super) fn apply(self, x: &MlxArray) -> UniquePtr { + match self { + Self::Silu => silu(x), + Self::Relu => relu(x), + // Keep the erf entry point this MLP always used for `gelu`. + Self::Gelu(GeluActivation::Erf) => gelu(x), + Self::Gelu(act) => act.apply(x), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn maps_gelu_names_through_the_shared_selector() { + assert_eq!( + NasActivation::from_hf_name("gelu"), + Ok(NasActivation::Gelu(GeluActivation::Erf)) + ); + for name in ["gelu_new", "gelu_fast", "gelu_pytorch_tanh"] { + assert_eq!( + NasActivation::from_hf_name(name), + Ok(NasActivation::Gelu(GeluActivation::TanhF32Compiled)), + "{name} is the tanh GELU, not SiLU" + ); + } + assert_eq!( + NasActivation::from_hf_name("quick_gelu"), + Ok(NasActivation::Gelu(GeluActivation::Fast)) + ); + } + + #[test] + fn keeps_silu_and_relu() { + assert_eq!(NasActivation::from_hf_name("silu"), Ok(NasActivation::Silu)); + assert_eq!(NasActivation::from_hf_name("relu"), Ok(NasActivation::Relu)); + } + + #[test] + fn rejects_unknown_names_instead_of_using_silu() { + let err = NasActivation::from_hf_name("tanh").unwrap_err(); + assert!(err.contains("unsupported hidden_act"), "{err}"); + assert!(NasActivation::from_hf_name("").is_err()); + } +} diff --git a/src/models/nemotron_voicechat/tts/norm_mlp.rs b/src/models/nemotron_voicechat/tts/norm_mlp.rs index d877dc46e..583e5db8e 100644 --- a/src/models/nemotron_voicechat/tts/norm_mlp.rs +++ b/src/models/nemotron_voicechat/tts/norm_mlp.rs @@ -29,7 +29,7 @@ use mlxcel_core::layers::UnifiedLinear; use mlxcel_core::weights::WeightMap; use mlxcel_core::{MlxArray, UniquePtr}; -use crate::models::gemma3_backbone::gelu_approx; +use crate::models::gemma3_backbone::BACKBONE_GELU; /// Copy `key` out of `weights`, or report it missing. pub(crate) fn weight(weights: &WeightMap, key: &str) -> Result, String> { @@ -119,7 +119,7 @@ impl OffsetRmsNorm { } /// The reference `MLP`: `down_proj(gelu_approx(gate_proj(x)) * up_proj(x))`, -/// with the op-for-op [`gelu_approx`] rather than the crate's fused GeGLU. +/// with the backbone's [`BACKBONE_GELU`] rather than the crate's fused GeGLU. pub struct GeluMlp { gate_proj: UnifiedLinear, up_proj: UnifiedLinear, @@ -145,7 +145,7 @@ impl GeluMlp { pub fn forward(&self, x: &MlxArray) -> UniquePtr { let gated = mlxcel_core::multiply( - &gelu_approx(&self.gate_proj.forward(x)), + &BACKBONE_GELU.apply(&self.gate_proj.forward(x)), &self.up_proj.forward(x), ); self.down_proj.forward(&gated) diff --git a/src/models/phi.rs b/src/models/phi.rs index 2efb43eb9..8293f2408 100644 --- a/src/models/phi.rs +++ b/src/models/phi.rs @@ -302,9 +302,11 @@ pub struct MLP { impl MLP { pub fn forward(&self, x: &MlxArray) -> UniquePtr { - // Simple MLP: fc2(gelu(fc1(x))) + // Simple MLP: fc2(gelu(fc1(x))). Phi-1/1.5/2 declare + // `hidden_act = gelu_new` (tanh; mlx-lm `nn.gelu_approx`), evaluated in + // f32 as one compiled kernel (erf before #2548). let h = self.fc1.forward(x); - let h = mlxcel_core::utils::gelu_approx(&h); + let h = mlxcel_core::utils::GeluActivation::TanhF32Compiled.apply(&h); self.fc2.forward(&h) } diff --git a/src/models/phixtral.rs b/src/models/phixtral.rs index 94e14c7a4..45c0ce000 100644 --- a/src/models/phixtral.rs +++ b/src/models/phixtral.rs @@ -604,6 +604,12 @@ impl SplitAttention for Attention { // Experts. +/// The expert GELU: tanh, as mlx-lm's `SwitchMLP` default +/// (`nn.GELU(approx="precise")`) and the checkpoints' `gelu_new` compute it, +/// evaluated in f32 as one compiled kernel (erf before #2548). +const PHIXTRAL_GELU: mlxcel_core::utils::GeluActivation = + mlxcel_core::utils::GeluActivation::TanhF32Compiled; + /// Upstream's `SwitchMLP`: `fc2(gelu(fc1(x)))` over the selected experts. /// /// Two differences from [`crate::models::switch_layers::SwitchGLU`] make it a @@ -649,14 +655,14 @@ impl PhixtralSwitchMlp { let (sorted_x, sorted_idx, inv_order) = gather_sort(&x_exp, indices); let h = self.fc1.forward(&sorted_x, &sorted_idx, true); let h = add_expert_bias(&h, &self.fc1_bias, &sorted_idx); - let h = mlxcel_core::utils::gelu_approx(&h); + let h = PHIXTRAL_GELU.apply(&h); let out = self.fc2.forward(&h, &sorted_idx, true); let out = add_expert_bias(&out, &self.fc2_bias, &sorted_idx); scatter_unsort(&out, &inv_order, &indices_shape) } else { let h = self.fc1.forward(&x_exp, indices, false); let h = add_expert_bias(&h, &self.fc1_bias, indices); - let h = mlxcel_core::utils::gelu_approx(&h); + let h = PHIXTRAL_GELU.apply(&h); let out = self.fc2.forward(&h, indices, false); let out = add_expert_bias(&out, &self.fc2_bias, indices); mlxcel_core::squeeze_axis(&out, -2) diff --git a/src/models/recurrent_gemma.rs b/src/models/recurrent_gemma.rs index cddb5d5af..8d77e490f 100644 --- a/src/models/recurrent_gemma.rs +++ b/src/models/recurrent_gemma.rs @@ -341,9 +341,12 @@ struct RecurrentBlock { impl RecurrentBlock { fn forward(&self, x: &MlxArray, cache: Option<&mut RGLRUCache>) -> UniquePtr { - // Y branch: linear + GELU + // Y branch: linear + GELU. RecurrentGemma declares + // `hidden_activation = gelu_pytorch_tanh` (the `RecurrentGemmaConfig` + // default; mlx-lm `nn.gelu_approx`), so this is the tanh GELU, as one + // compiled f32 kernel (erf before #2548). let y = self.linear_y.forward(x); - let y = mlxcel_core::gelu_approx(&y); + let y = mlxcel_core::utils::GeluActivation::TanhF32Compiled.apply(&y); // X branch: linear -> conv -> RGLRU let h = self.linear_x.forward(x); @@ -614,8 +617,9 @@ impl MLPBlock { let gate = self.gate_proj.forward(x); let up = self.up_proj.forward(x); - // GELU(gate) * up - let activated = mlxcel_core::multiply(&mlxcel_core::gelu_approx(&gate), &up); + // gelu_tanh(gate) * up: the Gemma tanh GeGLU kernel, matching mlx-lm's + // `nn.gelu_approx` and HF `gelu_pytorch_tanh` (erf before #2548). + let activated = mlxcel_core::compiled_geglu_approx_activation(&gate, &up); self.down_proj.forward(&activated) } diff --git a/src/models/starcoder2.rs b/src/models/starcoder2.rs index e6dc9a7fb..3c06ceb5e 100644 --- a/src/models/starcoder2.rs +++ b/src/models/starcoder2.rs @@ -235,8 +235,12 @@ impl StarCoder2MLP { pub fn forward(&self, x: &MlxArray) -> UniquePtr { let x = self.c_fc.forward(x); - // Use compiled GELU for kernel fusion (matches Python's gelu activation) - let x = mlxcel_core::compiled_gelu(&x); + // StarCoder2 checkpoints declare `hidden_act = gelu_pytorch_tanh` (also + // the `Starcoder2Config` default), so this is the tanh GELU, evaluated + // in f32 as one compiled kernel. mlx-lm uses erf (`nn.gelu`) and this + // path did too before #2548; the OpenXLA StarCoder2 graph, proven + // token-exact against the Hugging Face oracle, was already tanh. + let x = mlxcel_core::utils::GeluActivation::TanhF32Compiled.apply(&x); self.c_proj.forward(&x) } } diff --git a/src/vision/config.rs b/src/vision/config.rs index ceb30e763..7ed401a7e 100644 --- a/src/vision/config.rs +++ b/src/vision/config.rs @@ -31,9 +31,10 @@ where /// /// Existing and unknown values retain mlxcel's exact-erf GELU behavior. The /// explicit `gelu_pytorch_tanh` variant follows the Hugging Face polynomial. -/// A missing key also parses as `ExactGelu`; the loaders that have been -/// measured against their checkpoints resolve that case to the Hugging Face -/// class default with [`VisionConfig::default_missing_hidden_act_to_siglip`]. +/// A missing key also parses as `ExactGelu`; every loader of a SigLIP or CLIP +/// tower resolves that case to the Hugging Face class default with +/// [`VisionConfig::default_missing_hidden_act_by_model_type`] (or, for Gemma 3 +/// and PaliGemma, [`VisionConfig::default_missing_hidden_act_to_siglip`]). #[derive(Debug, Clone, Copy, Default, Deserialize, PartialEq, Eq)] pub enum VisionHiddenActivation { #[serde(rename = "gelu_pytorch_tanh")] @@ -46,6 +47,13 @@ pub enum VisionHiddenActivation { /// bit-for-bit to the private copy it replaced (#2423). #[serde(skip)] SiglipDefaultTanh, + /// `quick_gelu` (`x * sigmoid(1.702 x)`), declared or resolved from an + /// omitted key on a CLIP tower (the Hugging Face `CLIPVisionConfig` + /// default) by [`VisionConfig::default_missing_hidden_act_by_model_type`]. + /// Also set directly by loaders whose CLIP tower config is built in code + /// (Phi-3-V). Parsed as `ExactGelu` before #2548. + #[serde(rename = "quick_gelu")] + QuickGelu, #[default] #[serde(other)] ExactGelu, @@ -80,14 +88,42 @@ impl VisionConfig { /// serde default keeps exact erf for every other loader; this is opt-in per /// loader because the CLIP towers' class default is `quick_gelu`, not tanh. /// - /// Used by: the Gemma 3 VLM and PaliGemma loaders (#2423). Aya Vision and - /// Granite Vision 3.x omit the key too and are not switched yet. + /// Used by: the Gemma 3 VLM and PaliGemma loaders (#2423). pub(crate) fn default_missing_hidden_act_to_siglip(&mut self, raw: &serde_json::Value) { - let declared = raw.get("hidden_act").is_some_and(|value| !value.is_null()); - if !declared { + if !hidden_act_declared(raw) { self.hidden_act = VisionHiddenActivation::SiglipDefaultTanh; } } + + /// Resolve an omitted `hidden_act` to the Hugging Face class default of + /// the tower `model_type` names: `gelu_pytorch_tanh` for + /// `siglip_vision_model` / `siglip2_vision_model` (`SiglipVisionConfig`), + /// `quick_gelu` for `clip_vision_model` (`CLIPVisionConfig`). Any other + /// model type, and a declared key, are left as parsed. + /// + /// The serde default stays exact erf because the shared struct also + /// parses towers whose class default is neither (#2548). + /// + /// Used by: the LLaVA (and its host preprocessor), Bunny, Aya Vision and + /// Granite Vision 3.x loaders (#2548). + pub(crate) fn default_missing_hidden_act_by_model_type(&mut self, raw: &serde_json::Value) { + if hidden_act_declared(raw) { + return; + } + match self.model_type.as_str() { + "siglip_vision_model" | "siglip2_vision_model" => { + self.hidden_act = VisionHiddenActivation::SiglipDefaultTanh; + } + "clip_vision_model" => { + self.hidden_act = VisionHiddenActivation::QuickGelu; + } + _ => {} + } + } +} + +fn hidden_act_declared(raw: &serde_json::Value) -> bool { + raw.get("hidden_act").is_some_and(|value| !value.is_null()) } fn default_image_size() -> usize { @@ -184,12 +220,10 @@ mod tests { fn vision_hidden_act_selects_only_the_explicit_pytorch_tanh_variant() { let tanh: VisionConfig = serde_json::from_value(config(Some("gelu_pytorch_tanh"))).unwrap(); assert_eq!(tanh.hidden_act, VisionHiddenActivation::GeluPytorchTanh); + let quick: VisionConfig = serde_json::from_value(config(Some("quick_gelu"))).unwrap(); + assert_eq!(quick.hidden_act, VisionHiddenActivation::QuickGelu); - for value in [ - config(None), - config(Some("gelu")), - config(Some("quick_gelu")), - ] { + for value in [config(None), config(Some("gelu")), config(Some("relu"))] { let parsed: VisionConfig = serde_json::from_value(value).unwrap(); assert_eq!(parsed.hidden_act, VisionHiddenActivation::ExactGelu); } @@ -216,4 +250,48 @@ mod tests { parsed.default_missing_hidden_act_to_siglip(&null_act); assert_eq!(parsed.hidden_act, VisionHiddenActivation::SiglipDefaultTanh); } + + fn with_model_type(model_type: &str, hidden_act: Option<&str>) -> serde_json::Value { + let mut value = config(hidden_act); + value["model_type"] = serde_json::Value::String(model_type.to_string()); + value + } + + #[test] + fn missing_hidden_act_resolves_to_the_tower_class_default() { + let cases = [ + // HF drops a key equal to the class default: SigLIP -> tanh. + ( + "siglip_vision_model", + None, + VisionHiddenActivation::SiglipDefaultTanh, + ), + ( + "siglip2_vision_model", + None, + VisionHiddenActivation::SiglipDefaultTanh, + ), + // CLIPVisionConfig().hidden_act == "quick_gelu". + ("clip_vision_model", None, VisionHiddenActivation::QuickGelu), + // A declared key always wins. + ( + "clip_vision_model", + Some("gelu"), + VisionHiddenActivation::ExactGelu, + ), + ( + "siglip_vision_model", + Some("gelu_pytorch_tanh"), + VisionHiddenActivation::GeluPytorchTanh, + ), + // Unknown tower types keep the serde default. + ("vision", None, VisionHiddenActivation::ExactGelu), + ]; + for (model_type, act, expected) in cases { + let raw = with_model_type(model_type, act); + let mut parsed: VisionConfig = serde_json::from_value(raw.clone()).unwrap(); + parsed.default_missing_hidden_act_by_model_type(&raw); + assert_eq!(parsed.hidden_act, expected, "{raw}"); + } + } } diff --git a/src/vision/encoders/deepseek_vl2.rs b/src/vision/encoders/deepseek_vl2.rs index b4560848a..adb42daa0 100644 --- a/src/vision/encoders/deepseek_vl2.rs +++ b/src/vision/encoders/deepseek_vl2.rs @@ -139,9 +139,10 @@ impl Block { let y = self.attn.forward(&self.norm1.forward(x)); let x = mlxcel_core::add(x, &y); let y = self.fc1.forward(&self.norm2.forward(&x)); - // Exact erf GELU (`gelu_approx` is erf despite its name). mlx-vlm uses - // `nn.GELU(approx="precise")`, the tanh form; #2423 follow-up. - let y = mlxcel_core::gelu_approx(&y); + // Tanh GELU: DeepSeek-VL2's SigLIP ViT builds its MLP with + // `nn.GELU(approximate="tanh")` (`siglip_vit.py`), and mlx-vlm uses + // `approx="precise"`; f32, one compiled kernel (erf before #2548). + let y = mlxcel_core::utils::GeluActivation::TanhF32Compiled.apply(&y); let y = self.fc2.forward(&y); mlxcel_core::add(&x, &y) } diff --git a/src/vision/encoders/jina_vlm.rs b/src/vision/encoders/jina_vlm.rs index f885e4ad5..a86f9f376 100644 --- a/src/vision/encoders/jina_vlm.rs +++ b/src/vision/encoders/jina_vlm.rs @@ -165,7 +165,10 @@ struct JinaVlmViTMLP { impl JinaVlmViTMLP { fn forward(&self, x: &MlxArray) -> UniquePtr { let h = self.up.forward(x); - let h = mlxcel_core::gelu_approx(&h); + // `vision_config.block_config.ffn_config.activation_type = + // gelu_pytorch_tanh` (mlx-vlm maps it to tanh); f32, one compiled + // kernel (erf before #2548). + let h = mlxcel_core::utils::GeluActivation::TanhF32Compiled.apply(&h); self.down.forward(&h) } diff --git a/src/vision/encoders/lfm2_vl.rs b/src/vision/encoders/lfm2_vl.rs index eba84ae69..a053d1a47 100644 --- a/src/vision/encoders/lfm2_vl.rs +++ b/src/vision/encoders/lfm2_vl.rs @@ -176,9 +176,10 @@ impl Mlp { } fn forward(&self, x: &MlxArray) -> UniquePtr { - // hidden_act = "gelu_pytorch_tanh" (tanh-approximate GELU). + // hidden_act = "gelu_pytorch_tanh" (tanh-approximate GELU), in f32 as + // one compiled kernel (erf before #2548). let x = self.fc1.forward(x); - let x = mlxcel_core::gelu_approx(&x); + let x = mlxcel_core::utils::GeluActivation::TanhF32Compiled.apply(&x); self.fc2.forward(&x) } } diff --git a/src/vision/encoders/minicpmo.rs b/src/vision/encoders/minicpmo.rs index a4c515a5a..19ca21664 100644 --- a/src/vision/encoders/minicpmo.rs +++ b/src/vision/encoders/minicpmo.rs @@ -40,13 +40,35 @@ pub struct MiniCPMOVisionConfig { pub patch_size: usize, #[serde(default = "default_layer_norm_eps")] pub layer_norm_eps: f32, - /// GELU the encoder MLP applies. Not read from MiniCPM-o's config: it stays - /// exact erf there (no local MiniCPM-o checkpoint has been measured, #2423). - /// MiniCPM-V 4.6 sets it from its `vision_config.hidden_act`. + /// GELU the encoder MLP applies. The MiniCPM-o loader sets it from + /// `vision_config.hidden_act` with [`mlp_activation_from_vision_config`] + /// (#2548; erf before); MiniCPM-V 4.6 sets it from its own + /// `vision_config.hidden_act` (#2423). #[serde(skip)] pub mlp_activation: GeluActivation, } +/// The encoder MLP GELU a MiniCPM-o `vision_config` declares. +/// +/// The tower is a SigLIP (`siglip_vision_model`), so an omitted `hidden_act` +/// is the `SiglipVisionConfig` default `gelu_pytorch_tanh`. The tanh form runs +/// in f32 as one compiled kernel, like MiniCPM-V 4.6. A name that is not a +/// GELU variant is an error rather than a silent erf. +/// +/// Used by: the MiniCPM-o loader (`loading/vlm_special.rs`). +pub(crate) fn mlp_activation_from_vision_config( + raw: &serde_json::Value, +) -> Result { + match raw.get("hidden_act").and_then(serde_json::Value::as_str) { + None => Ok(GeluActivation::TanhF32Compiled), + Some(name) => GeluActivation::from_hf_name(name) + .map(GeluActivation::in_f32_compiled) + .ok_or_else(|| { + format!("MiniCPM-o vision_config.hidden_act {name:?} is not a GELU variant") + }), + } +} + fn default_num_channels() -> usize { 3 } diff --git a/src/vision/encoders/minicpmo_tests.rs b/src/vision/encoders/minicpmo_tests.rs index 80ad2ad02..4c43e45f4 100644 --- a/src/vision/encoders/minicpmo_tests.rs +++ b/src/vision/encoders/minicpmo_tests.rs @@ -29,3 +29,27 @@ fn get_2d_sincos_pos_embed_returns_expected_length() { let embed = get_2d_sincos_pos_embed(2, 3, 8); assert_eq!(embed.len(), 2 * 3 * 8); } + +#[test] +fn mlp_activation_follows_the_vision_config() { + use super::mlp_activation_from_vision_config; + use mlxcel_core::utils::GeluActivation; + + let tanh = serde_json::json!({"hidden_act": "gelu_pytorch_tanh"}); + assert_eq!( + mlp_activation_from_vision_config(&tanh), + Ok(GeluActivation::TanhF32Compiled) + ); + // Omitted: the SigLIP class default, tanh (#2548). + assert_eq!( + mlp_activation_from_vision_config(&serde_json::json!({})), + Ok(GeluActivation::TanhF32Compiled) + ); + let erf = serde_json::json!({"hidden_act": "gelu"}); + assert_eq!( + mlp_activation_from_vision_config(&erf), + Ok(GeluActivation::Erf) + ); + let silu = serde_json::json!({"hidden_act": "silu"}); + assert!(mlp_activation_from_vision_config(&silu).is_err()); +} diff --git a/src/vision/encoders/molmo.rs b/src/vision/encoders/molmo.rs index 5a4c6ffbc..c3a3eaeaf 100644 --- a/src/vision/encoders/molmo.rs +++ b/src/vision/encoders/molmo.rs @@ -101,7 +101,8 @@ impl MolmoVisionConfig { } } -// ViT MLP (quantized w1/w2 + fast GELU). +// ViT MLP (quantized w1/w2 + fast GELU). The remote code builds its CLIP ViT +// with `image_mlp_activations="quick_gelu"`; mlx-vlm uses `approx="fast"`. struct MolmoViTMLP { w1: UnifiedLinear, w2: UnifiedLinear, @@ -110,7 +111,8 @@ struct MolmoViTMLP { impl MolmoViTMLP { fn forward(&self, x: &MlxArray) -> UniquePtr { let h = self.w1.forward(x); - let h = mlxcel_core::gelu_approx(&h); + // quick_gelu (erf before #2548). + let h = mlxcel_core::utils::GeluActivation::Fast.apply(&h); self.w2.forward(&h) } diff --git a/src/vision/encoders/moondream3.rs b/src/vision/encoders/moondream3.rs index 108658b27..3f973f0cf 100644 --- a/src/vision/encoders/moondream3.rs +++ b/src/vision/encoders/moondream3.rs @@ -208,7 +208,10 @@ impl Moondream3VisionMlp { fn forward(&self, x: &MlxArray) -> UniquePtr { let hidden = self.fc1.forward(x); - let hidden = mlxcel_core::gelu_approx(&hidden); + // Tanh GELU (mlx-vlm `TanhGELUMLP`, remote code `gelu_approx` = + // `F.gelu(approximate="tanh")`); f32, one compiled kernel (erf before + // #2548). + let hidden = mlxcel_core::utils::GeluActivation::TanhF32Compiled.apply(&hidden); self.fc2.forward(&hidden) } } @@ -400,7 +403,8 @@ impl Moondream3VisionModel { let merged = mlxcel_core::reshape(&merged, &[1, enc_n * enc_n, dim * 2]); let projected = self.proj_fc1.forward(&merged); - let projected = mlxcel_core::gelu_approx(&projected); + // Projection MLP: the same tanh GELU (#2548; erf before). + let projected = mlxcel_core::utils::GeluActivation::TanhF32Compiled.apply(&projected); self.proj_fc2.forward(&projected) } diff --git a/src/vision/encoders/paddleocr_vl.rs b/src/vision/encoders/paddleocr_vl.rs index 51ca3054f..a1c627dc0 100644 --- a/src/vision/encoders/paddleocr_vl.rs +++ b/src/vision/encoders/paddleocr_vl.rs @@ -613,7 +613,10 @@ impl MLP { fn forward(&self, x: &MlxArray) -> UniquePtr { let h = self.fc1.forward(x); - let h = mlxcel_core::gelu(&h); + // `vision_config.hidden_act = gelu_pytorch_tanh` (also the HF + // `PaddleOCRVisionConfig` default; mlx-vlm `approx="precise"`): tanh, + // in f32 as one compiled kernel (erf before #2548). + let h = mlxcel_core::utils::GeluActivation::TanhF32Compiled.apply(&h); self.fc2.forward(&h) } } diff --git a/src/vision/encoders/phi4_siglip.rs b/src/vision/encoders/phi4_siglip.rs index 5c1aae8e2..28f4750fc 100644 --- a/src/vision/encoders/phi4_siglip.rs +++ b/src/vision/encoders/phi4_siglip.rs @@ -159,7 +159,10 @@ impl Phi4SigLipMlp { fn forward(&self, x: &MlxArray) -> UniquePtr { let x = self.fc1.forward(x); - let x = mlxcel_core::gelu_approx(&x); + // SigLIP2 tower: `hidden_act = gelu_pytorch_tanh` (the HF + // `Phi4MultimodalVisionConfig` and mlx-vlm config default), tanh in f32 + // as one compiled kernel (erf before #2548). + let x = mlxcel_core::utils::GeluActivation::TanhF32Compiled.apply(&x); self.fc2.forward(&x) } } diff --git a/src/vision/encoders/qwen2_vl.rs b/src/vision/encoders/qwen2_vl.rs index 35d09dd01..75761880c 100644 --- a/src/vision/encoders/qwen2_vl.rs +++ b/src/vision/encoders/qwen2_vl.rs @@ -413,7 +413,10 @@ impl VisionMLP { fn forward(&self, x: &MlxArray) -> UniquePtr { let h = self.fc1.forward(x); - let h = mlxcel_core::gelu_approx(&h); + // `quick_gelu`: the shipped configs omit `vision_config.hidden_act` + // and `Qwen2VLVisionConfig().hidden_act == "quick_gelu"`; mlx-vlm uses + // `approx="fast"` too (erf before #2548). + let h = mlxcel_core::utils::GeluActivation::Fast.apply(&h); self.fc2.forward(&h) } } diff --git a/src/vision/encoders/siglip.rs b/src/vision/encoders/siglip.rs index a0b924eb6..8b121b7ec 100644 --- a/src/vision/encoders/siglip.rs +++ b/src/vision/encoders/siglip.rs @@ -54,7 +54,8 @@ pub(crate) enum VisionMlpActivation { Exact, PytorchTanh, /// The tanh GELU of a SigLIP tower whose config omits `hidden_act` (Gemma 3, - /// PaliGemma), evaluated by the compiled f32 kernel (#2423). + /// PaliGemma, Aya Vision, Granite Vision 3.x, Bunny), evaluated by the + /// compiled f32 kernel (#2423, #2548). PytorchTanhCompiled, FastSigmoid, } @@ -74,6 +75,7 @@ pub(crate) fn select_mlp_activation( match hidden_act { VisionHiddenActivation::GeluPytorchTanh => VisionMlpActivation::PytorchTanh, VisionHiddenActivation::SiglipDefaultTanh => VisionMlpActivation::PytorchTanhCompiled, + VisionHiddenActivation::QuickGelu => VisionMlpActivation::FastSigmoid, VisionHiddenActivation::ExactGelu => VisionMlpActivation::Exact, } } @@ -891,6 +893,11 @@ mod tests { select_mlp_activation(VisionHiddenActivation::SiglipDefaultTanh, false), VisionMlpActivation::PytorchTanhCompiled ); + // CLIP's `quick_gelu` (#2548), declared or the class default. + assert_eq!( + select_mlp_activation(VisionHiddenActivation::QuickGelu, false), + VisionMlpActivation::FastSigmoid + ); } #[test] diff --git a/src/vision/encoders/youtu_vl_layers.rs b/src/vision/encoders/youtu_vl_layers.rs index 68a653a87..1addb299d 100644 --- a/src/vision/encoders/youtu_vl_layers.rs +++ b/src/vision/encoders/youtu_vl_layers.rs @@ -209,8 +209,10 @@ impl VisionMLP { pub(super) fn forward(&self, x: &MlxArray) -> UniquePtr { let x = self.fc1.forward(x); - // gelu_pytorch_tanh — matches `nn.gelu_approx` in upstream MLX usage. - let x = mlxcel_core::utils::gelu_approx(&x); + // `vision_config.hidden_act = gelu_pytorch_tanh`, mlx-vlm + // `nn.gelu_approx`, and the OpenXLA Youtu-VL graph: tanh, in f32 as one + // compiled kernel (erf before #2548). + let x = mlxcel_core::utils::GeluActivation::TanhF32Compiled.apply(&x); self.fc2.forward(&x) } }