[inureyes]

Summary

Gemma3n sparse layers (0-9, sparsity 0.95) rounded the bf16 gelu_topk three different ways: the eager compiled kernel promoted the whole tape to f32, and the XLA tape left three constants unrounded with a two-rounding cube. Both sides now move to mlx-lm's bf16 carrier, per the issue's settled Decision.

What changed

• The eager compiled_gelu_topk kernel (per-op and fused paths) mirrors mlx-lm's graph op for op: the f32 multiplier cast to the input dtype like std_multiplier.astype, the zero and nn.gelu_approx constants built in it like Python's weak scalars, the cube one power op. The shared gelu_tanh_approx helper stays as-is (its x*x*x cube serves the dense GeGLU paths).
• The XLA sparse_gelu_stages tape rounds the multiplier, root and cubic constants to bf16 at emit time and computes the cube as one f32 product, exact before a single rounding (36 to 34 converts).
• The pinned E2B probe's manual tape mirrors the same bf16 carrier and pins the compiled kernel against both the per-op tape and the FMA-site tape (MLX's CUDA compiler contracts the two mul+add sites into FMAs).
• New ffi_gelu_topk_tests module pins the kernel bit for bit against mlx-lm's carrier, with an f32-carrier reference kept so the test can tell carriers apart.
• docs/gelu-call-sites.md rows G15/G17 and the compiled_gelu_topk / gemma3n_mlp_forward rows describe the bf16 carrier and drop the known-gap note.

Test plan

• cargo test -p mlxcel-core --profile test-fast --features cuda -- --test-threads=1 gelu_topk: 3 passed. The bit-exact test reports FMA-site tape 0 mismatches, per-op 99/849, f32 carrier 212/883.
• mlxcel-xla sparse_gelu (34 converts), host-oracle and bf16_scalar tests pass; cargo fmt --check clean, clippy clean on the touched packages.
• GPU reference probe (CUDA GB10 under gpu-lock, MLX_ENABLE_TF32=0, mlx-lm's compiled gelu_topk as reference, widths 2048/8192): old f32 carrier 370/1247 mismatches, xxx cube 8/821, per-op nn.gelu_approx(z) 0/816, bf16 carrier with fused FMA sites 0/0. On CUDA the implementation targets the FMA-contracted compiled kernel; the XLA tape reproduces the per-op tape (MLX CPU semantics).
• Pinned E2B IREE probe on spark-101 (the issue's "spark-101 has no IREE install" is stale: the source+build tree at ~/.cache/mlxcel/iree-cuda-3.12.0rc20260721 works; IREE_CUDA_HOME + GEMMA3N_MODEL_DIR=models/gemma3n-e2b-4bit, --test-threads=1): bit_mismatches=0 on every layer0_sparse_* stage and fma_sites_bit_mismatches=0/24576 on layer0_sparse_activation_manual_vs_compiled (per-op 333/24576, the CUDA FMA residual). The probe overall still aborts at pre-existing, unrelated checks (below).
• Greedy smoke on gemma3n-e2b-4bit and gemma3n-e4b-bf16, through mlxcel generate and mlxcel-server /v1/chat/completions, temperature 0, 64 tokens, two prompts: token sequences identical to mlx-lm greedy on the same rendered prompt (prompt_tokens 14/16 match). The only text difference is a trailing space delta the mlxcel SentencePiece detokenizer drops before a newline; display only.

Pre-existing findings (not from this change)

• probe_real_projection's weight-carrier check fails on CUDA for every 4-bit weight: the Rust diagnostic pins MLX CPU dequant semantics (bf16(bf16(q*scale)+bias), 0 mismatches vs mx.dequantize on CPU) while the probe runs MLX on GPU; proven device-dependent via an oracle-venv GPU-vs-CPU check. This aborts pinned_e2b_real_projection_qmv_probe before the gelu stages; the stage results above were taken with that one assertion temporarily relaxed locally, and the edit is reverted.
• AltUp prediction/correction coefficient stages differ by f32 noise (~5.6e-6 / 1.5e-5) that cascades 1 bf16 ulp into the predictions; unreachable on main since the probe aborts earlier. Both warrant follow-up issues.

Closes #2795