fix(models): close the GELU formula-rule gaps left after #2785 #2794

Problem / Background

PR #2785 (issue #2548) settled the GELU formula per family with the rule in docs/gelu-call-sites.md section 0b: the checkpoint config decides; an omitted key falls back to the Hugging Face config-class default for that model_type, evaluated as transformers ACT2FN; where mlx-lm / mlx-vlm departs from Hugging Face, mlxcel follows Hugging Face and the doc row records it. Its review found five places outside #2548's row list that still break the rule or are left over from it. This issue builds on #2785: implement it on main after #2785 merges (until then, read the code on branch fix/issue-2548-gelu-formula-decisions; every line number below is from that branch). Background: #2423 (GELU inventory), #2780 (single-kernel tanh GELU, separate).

Implement all five items in one PR. Run the local gate once at the end, not per item, and keep pushes to a minimum (ideally one): a Rust-touching push costs a full CI cycle.

Reference values used below, checked against transformers 5.22.1: GPTBigCodeConfig().activation_function == "gelu_pytorch_tanh", Idefics2VisionConfig().hidden_act == "gelu_pytorch_tanh"; ACT2FN maps gelu/gelu_python to erf, gelu_pytorch_tanh/gelu_new/gelu_fast to tanh, quick_gelu to x * sigmoid(1.702 x). The shared mapper is GeluActivation::from_hf_name (src/lib/mlxcel-core/src/utils.rs:1161-1168), with in_f32_compiled() (:1184-1189) turning Tanh into TanhF32Compiled.

1. GPT-BigCode runs erf; config and HF default say tanh (row T5)

Current: MLP::forward (src/models/gpt_bigcode.rs:581-591) hard-codes mlxcel_core::utils::gelu_approx (erf), with a comment justifying it by mlx-lm nn.gelu. ModelArgs (:87-127) does not read activation_function. The real santacoder fixture in src/models/gpt_bigcode_tests.rs:32 declares "activation_function": "gelu_pytorch_tanh", and HF builds the MLP as ACT2FN[config.activation_function]. Same situation StarCoder2 was in before #2785 (row T6). There is no XLA GPT-BigCode graph (src/lib/mlxcel-xla/assets/ has starcoder2 and granite only), so this is eager-only.

Fix: Add #[serde(default)] pub activation_function: Option<String> to ModelArgs. Resolve it once at load: absent or null means gelu_pytorch_tanh; otherwise GeluActivation::from_hf_name(name).map(GeluActivation::in_f32_compiled); None (a non-GELU name such as relu) is an Err from ModelArgs::validate (:222) with the message GPT-BigCode: unsupported activation_function "<name>"; expected a GELU variant (gelu, gelu_new, gelu_fast, gelu_pytorch_tanh, quick_gelu). Expose it as pub fn activation(&self) -> Result<GeluActivation, String> (the Nemotron-NAS pattern in src/models/nemotron_nas_activation.rs), have validate() call it, store the resolved value on MLP (filled in MLP::from_weights, :593-622), and call self.activation.apply(&h) in forward. Rewrite the comment to state the HF rule and that mlx-lm nn.gelu departs. Rejected: hard-coding TanhF32Compiled as src/models/starcoder2.rs does, because GPT-BigCode configs carry the key and rule 1 says the config decides.

Acceptance:

• parses_the_real_santacoder_config asserts args.activation() == Ok(GeluActivation::TanhF32Compiled).
• New test: a bare config (no key) also yields TanhF32Compiled; "gelu" yields Erf; "quick_gelu" yields Fast; "relu" makes validate() return Err containing unsupported activation_function. Fails on current code (field does not exist).
• New formula test: MLP::forward on a tiny fixture (the existing tiny_args() / hf_weights() helpers in gpt_bigcode_tests.rs) equals c_proj(TanhF32Compiled(c_fc(x))) within 1e-5 and differs from the erf result by more than 1e-4. Fails on current code.
• No GPT-BigCode checkpoint exists under models/ on the GB10 host (ls models | grep -iE 'bigcode|santacoder' is empty). The Metal benches use models/mlx/gpt_bigcode-santacoder; run the smoke (greedy mlxcel generate plus one mlxcel-server request) on a Mac host that has it, or state in the PR that it was not run and why.

2. Idefics2 (S6) and Granite 4 Vision (S7) force quick_gelu; HF runs tanh

Current: Both loaders pass use_fast_gelu = true to SigLipVisionModel::from_weights_with_quant_and_gelu (src/vision/encoders/siglip.rs:588-599), which makes select_mlp_activation (:68-83) return FastSigmoid regardless of hidden_act: src/loading/vlm_idefics2.rs:198-209 and src/loading/vlm_granite4_vision.rs:137-148. The comments cite mlx-vlm FastGELUMLP, which departs from HF.

• Idefics2: local models/idefics2-8b-4bit omits vision_config.hidden_act (its vision_config.model_type is idefics2), so the HF default gelu_pytorch_tanh applies. default_missing_hidden_act_by_model_type does not cover this model type, so the loader must resolve it explicitly.
• Granite 4 Vision: local models/granite-4.0-3b-vision-4bit declares vision_config.hidden_act = gelu_pytorch_tanh (vision_config.model_type = siglip_vision_model), and its remote modeling.py:251 builds the tower with AutoModel.from_config(config.vision_config), i.e. HF SiglipVisionModel with tanh.

Fix: In vlm_idefics2.rs, make vision_config mutable and call vision_config.default_missing_hidden_act_to_siglip(&raw_vision_config) (src/vision/config.rs:92) on the raw vision_config JSON, then pass use_fast_gelu = false. In vlm_granite4_vision.rs, call default_missing_hidden_act_by_model_type on the raw value (harmless when declared) and pass false. Rewrite both comments to cite the HF class and the mlx-vlm departure. Keep the use_fast_gelu parameter (LLM-jp-VL vlm_llmjp_vl.rs:300 and llmjp_vl_test_support.rs call the same entry point); after this change no production caller passes true, so note that in its doc comment. The Granite 4 Q-Former FFN (row V59, src/vision/connectors/granite4_vision.rs:273, erf) is not touched.

Acceptance:

• Loader-level config test for each family (a helper that returns the resolved VisionMlpActivation from a config JSON, extracted from the loader if needed): Idefics2 with the key absent resolves to a tanh variant; Granite 4 with gelu_pytorch_tanh resolves to a tanh variant; neither resolves to FastSigmoid. Fails on current code.
• Greedy mlxcel generate with one image on models/idefics2-8b-4bit and on models/granite-4.0-3b-vision-4bit produces a coherent description, plus one mlxcel-server /v1/chat/completions request with an image_url part on each. Record outputs in the PR.

3. Shared VisionConfig silently runs erf for any unknown hidden_act

Current: VisionHiddenActivation (src/vision/config.rs:38-60) has #[serde(other)] ExactGelu, so gelu_new, gelu_fast (tanh in ACT2FN) and non-GELU names such as relu or silu all parse to erf. The test at src/vision/config.rs:226-229 pins relu -> ExactGelu. The OpenXLA LlavaVisionConfig (src/lib/mlxcel-xla/src/emitter/vision_config.rs:424-441) rejects unknown names with unsupported vision hidden_act, so eager and XLA disagree. Users of the shared type: vlm_granite4_vision.rs:122, vlm_granite_vision.rs:103, vlm_idefics2.rs:162, vlm_llava.rs (through VLMConfig), vlm_llmjp_vl.rs:229, vlm_siglip.rs:191,292, vlm_smolvlm.rs:263,504, vlm_gemma.rs:164,270 (through VLMConfig), src/models/siglip_text.rs:89 (Option<VisionHiddenActivation>).

Fix: Replace the derived variant mapping with #[serde(try_from = "String")] (keep #[default] ExactGelu for an absent key through null_default; SiglipDefaultTanh is removed by item 4, and if item 3 lands first it stays unparsed until then). TryFrom<String> goes through GeluActivation::from_hf_name: Erf -> ExactGelu, Tanh -> GeluPytorchTanh, Fast -> QuickGelu, None -> Err("unsupported vision hidden_act \"<name>\"; expected gelu, gelu_python, gelu_pytorch_tanh, gelu_new, gelu_fast or quick_gelu"). Non-GELU names become a load error rather than gaining support: no SigLIP/CLIP checkpoint uses them (a scan of models/*/config.json finds silu only on Pixtral, Qwen2.5-VL and GLM-4V towers, which have their own config types). In parse_bunny_vision_config (src/loading/vlm_llava.rs:650-680) the parse is .ok(), which would swallow the new error and fall back to a hard-coded config; change its signature to return anyhow::Result<VisionConfig>, add ? at the caller (vlm_llava.rs:833), and propagate an error whose message contains unsupported vision hidden_act, while a config that is simply missing tower fields still takes the fallback (the existing tests parse_bunny_vision_config_uses_defaults_for_sparse_configs and parse_bunny_vision_config_preserves_full_configs in src/loading/vlm_llava_tests.rs:123,154 unwrap the new Result). Mirror the accepted set in XLA: add gelu_python (erf) and gelu_new / gelu_fast (tanh) to the match at vision_config.rs:424-441.

Acceptance:

• src/vision/config.rs test rewritten: gelu_new and gelu_fast parse to GeluPytorchTanh; gelu and gelu_python to ExactGelu; quick_gelu to QuickGelu; absent to ExactGelu; relu and silu fail to deserialize with unsupported vision hidden_act in the error. Fails on current code.
• siglip_text_tests.rs gains a case: a text tower declaring relu is a load error.
• A Bunny config with vision_config.hidden_act = "relu" makes parse_bunny_vision_config return Err; one with missing tower fields still returns the fallback config.
• XLA LlavaVisionConfig test: gelu_new parses to GeluPytorchTanh; relu is still rejected (the existing assertion at vision_config.rs:176-181).
• Every local checkpoint whose loader uses the shared VisionConfig still loads (covered by the smoke runs in items 2 and 4; add models/smolvlm-instruct-bf16).

4. Declared gelu_pytorch_tanh on the shared SigLIP MLP runs the eager ~11-pass kernel (LOW, perf, from #2423)

Current: select_mlp_activation maps GeluPytorchTanh to VisionMlpActivation::PytorchTanh (siglip.rs:76), which gelu() turns into GeluActivation::TanhF32 (:103), the eager f32 helper of about eleven elementwise passes. The omitted-key path (SiglipDefaultTanh -> PytorchTanhCompiled -> TanhF32Compiled, :77, :104) already uses the compiled kernel. The split exists only because the declared path was pinned bit-for-bit to the private copy #2423 replaced (src/vision/config.rs:44-48). This affects every SigLIP tower that declares the key (SmolVLM / Idefics3, LLaVA-Interleave, LLM-jp-VL, Llama-Nemotron-VL-Embed, Granite 4 after item 2) and the SigLIP text tower (siglip_text.rs:124-126).

Fix: Map GeluPytorchTanh to PytorchTanhCompiled, then delete VisionMlpActivation::PytorchTanh and fold SiglipDefaultTanh into GeluPytorchTanh (remove the variant and make default_missing_hidden_act_to_siglip / default_missing_hidden_act_by_model_type set GeluPytorchTanh; update vlm_llava.rs:680 and the doc comments at config.rs:44-48 and siglip.rs:55-60). GeluActivation::TanhF32 stays: Molmo2 (molmo2.rs:94) and Qwen3-VL (qwen3_vl.rs:772) still use it, and its bit-exact test in utils_gelu_tests.rs:90-95 is unaffected.

Acceptance:

• siglip.rs selection test (:881-899) asserts GeluPytorchTanh -> PytorchTanhCompiled; siglip_text_tests.rs:287,320 updated likewise. Fails on current code.
• config.rs tests missing_hidden_act_resolves_to_the_siglip_default_only_when_absent (:233-251) and missing_hidden_act_resolves_to_the_tower_class_default (:261 onward) expect GeluPytorchTanh where they now expect SiglipDefaultTanh.
• pytorch_tanh_gelu_matches_hugging_face_f32_golden (siglip.rs:903-921) runs on PytorchTanhCompiled and keeps its 2e-6 tolerance.
• encoder_block_shared_with_vision_is_unchanged (siglip_block_tests.rs:201-215, 1e-6 golden, fixture declares GeluPytorchTanh at :126): if the compiled kernel's FMA contraction moves it past 1e-6, widen the bound to the measured drift rounded up to the next power of ten (state the measured value in a comment); do not recapture the golden.
• Greedy mlxcel generate with an image on models/smolvlm-instruct-bf16 gives the same text before and after, or a difference explained by f32 rounding; one mlxcel-server image request on the same model.

5. C++ compiled_gelu is dead outside FFI tests (LOW, cleanup)

Current: After #2785 its only Rust caller (StarCoder2) is gone. Definition src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:1908-1931, header mlx_cxx_bridge.h:583, cxx decl src/lib/mlxcel-core/src/lib.rs:792-794, tests src/lib/mlxcel-core/src/ffi_tests.rs:1956 (test_compiled_gelu), :1987 (test_compiled_gelu_preserves_bf16_and_f16_dtype), :2158 (test_compiled_gelu_matches_gelu), and the compile-audit lists at :2720 and :2927.

Fix: Remove the function, its declarations, the three tests and both audit-list entries. Reword the comment in gelu() at mlx_cxx_bridge.cpp:3357 that cites compiled_gelu. Out of scope: compiled_gelu_approx is also tests-only; leave it.

Acceptance:

• git grep -nw compiled_gelu -- src returns nothing.
• The compile-audit test in ffi_tests.rs still passes with the shortened expected list.

Docs

Update docs/gelu-call-sites.md: section 0b rows T5, S6, S7 (Before / After / Reason, with the mlx-vlm / mlx-lm departure), a new row for the shared VisionConfig unknown-name rule and the declared-tanh compiled switch, the "Shared resolution" paragraph, the inventory rows T5 (:158), S6 (:259), S7 (:260), and the compiled_gelu row in section 1. Update the GeluActivation "Used by" doc comment in utils.rs to add GPT-BigCode.

Scope

In scope: the files named above. Out of scope: the single-kernel tanh GELU (#2780); compiled_gelu_approx; removing the use_fast_gelu parameter; any XLA GPT-BigCode graph.

Verification (GB10, once at the end)

cargo fmt --all -- --check
cargo clippy --profile test-fast --features cuda --all-targets -- -D warnings
MLX_ENABLE_TF32=0 gpu-lock run --tag gelu-followup -- cargo test --profile test-fast --features cuda --lib -- --test-threads=1 gpt_bigcode siglip vision::config idefics2 granite4 llava
MLX_ENABLE_TF32=0 gpu-lock run --tag gelu-followup -- cargo test --profile test-fast --features cuda -p mlxcel-core -- --test-threads=1 gelu
cargo test -p mlxcel-xla vision_config

Build without gpu-lock (sccache holds the lock otherwise); wrap only test binaries, mlxcel generate and mlxcel-server runs in gpu-lock run. Smoke runs: greedy (--temp 0) mlxcel generate --image and one mlxcel-server /v1/chat/completions image request each on models/idefics2-8b-4bit, models/granite-4.0-3b-vision-4bit and models/smolvlm-instruct-bf16; GPT-BigCode as described in item 1. A pass is coherent output and no load error.