fix(models): port four upstream numeric fixes (Gemma 3n KV-shared RoPE offset, PhiMoE sparsemixer and LongRoPE, Gemma 4 audio mask, Phi-3-small dummy ids) #2379

Problem / Background

Upstream fixed four numeric defects in families that mlxcel also ships, and none of the fixes have been ported. Each one silently produces different logits from the reference. They share the model-layer test surface (_tests.rs beside each model plus the logit trace), so they belong in one PR. Implement all four in one PR, run the local gate once at the end, and keep pushes to a minimum.

1. Gemma 3n: KV-shared layers rotate queries at the owner's advanced offset

Upstream: ml-explore/mlx-lm#1886 (commit 36de496), https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/models/gemma3n.py.

• Current: Gemma3nAttention::forward reads let cache_offset = cache.offset (src/models/gemma3n.rs:471). For a KV-shared layer, cache is the owner layer's cache (caches[cache_idx] via layer_idx_to_cache_idx, :1578-1600), which the owner has already advanced by L in this forward. Queries are then RoPE'd at offset + L (:496-503) against keys at <= offset. That is off by L on every shared layer, including off by one on every decode step.
• Fix: read the offset once before the layer loop, from the first non-shared layer's cache before any layer runs. Thread it as an offset: i32 parameter through the decoder layer's forward into Gemma3nAttention::forward, replacing the read at :471. Owner layers use the same value for compute_kv. With --max-kv-size trimming, keep using the monotonic offset, not live_len().
• Test: in a new src/models/gemma3n_tests.rs (wired with #[cfg(test)] #[path = "gemma3n_tests.rs"] mod tests;, the file does not exist yet), use a tiny config with num_kv_shared_layers = 2, mirroring upstream test_gemma3n_kv_shared_layers_ignore_chunking. The last-token logits from one prefill of 12 tokens must equal (within 1e-4) prefill in chunks [5, 4, 3] and token-by-token decode. This fails on main.
• Trace: logit trace on mlx-community/gemma-3n-E2B-it-4bit, fixed vs reverted arm. The reverted arm must show higher decided-position disagreement against the fixed arm at width 1.

2. PhiMoE: sparsemixer routing and LongRoPE

Upstream: ml-explore/mlx-lm#1939 (commit 5cfec4c) for sparsemixer, ml-explore/mlx-lm#1889 (commit d0aa560) for null rope_scaling, and Blaizzy/mlx-vlm#2419 for both. See https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/models/phimoe.py.

• Current routing: SparseMoeBlock::forward (src/models/phimoe.rs:157-171) takes top-k by argpartition and applies softmax over the k selected logits. HF PhiMoE uses sparsemixer, so expert weights differ on every token.
• Current RoPE: plain fast_rope at rope_theta (:248-249; module doc at :22 says "simplified from SuScaledRoPE"). ModelArgs (:34-75) has no rope_scaling, max_position_embeddings, original_max_position_embeddings or router_jitter_noise fields. mlx-community/Phi-3.5-MoE-instruct-4bit ships rope_scaling: {type: "longrope", short_factor[64], long_factor[64], short_mscale: 1.2432, long_mscale: 1.2432, original_max_position_embeddings: 4096}, max_position_embeddings: 131072, and router_jitter_noise: 0.01. Every position is therefore rotated with the wrong frequencies and without the mscale.
• Fix, routing: port upstream sparsemixer(gates, top_k, jitter_eps) as a free function in phimoe.rs. Run top_k iterations; in each, take top = max(masked) and ind = argmax(masked). Compute factor = max(|gates|, top), then band = where((top - gates) / factor > 2 * jitter_eps, -inf, masked), then score = take(softmax_precise(band), ind). Finally mask ind to -inf in masked, using a where(arange == ind, ...) one-hot if no put_along_axis exists in mlxcel-core. jitter_eps comes from router_jitter_noise (default 0.01). Reject num_experts_per_tok > num_local_experts at load. The fused decode path (:174-180, forward_fused_kernel) must receive the sparsemixer scores unchanged. Confirm it does not renormalize them.
• Fix, RoPE: reuse phi3::SuRope, which already implements transformers' short/long table selection by position (fix(phi3): select LongRoPE short or long factors by sequence position and read mscale overrides #1358). Do not copy mlx-lm's long-table-only SuScaledRoPE. SuRope::from_args (src/models/phi3.rs:232) is private and takes phi3::ModelArgs. Extract a pub(crate) constructor from (&phi3::RopeScaling, rope_dims, rope_theta, max_position_embeddings, original_max), add the four fields to phimoe::ModelArgs, and apply the table the way phi3 does (table_for at :302, fast_rope_with_freqs at :440-448, Q/K pre-scale). A null or absent rope_scaling keeps plain RoPE (Phi-mini-MoE). Also implement rope_table_regime for PhiMoE as phi3 does (phi3.rs:860), so the prompt cache does not reuse a prefix across the short/long boundary.
• Tests: in src/models/phimoe_tests.rs, port upstream's sparsemixer case. Gates [[3, 1, 0, -1], [3, 2.99, 0, -1]], top_k 2, eps 0.01 must give indices [[0, 1], [0, 1]] and scores [[1.0, 1.0], [0.5025, 1.0]] within 1e-4. Add a config test that a longrope block builds a SuRope with both tables and a null block builds none.
• Trace: logit trace on mlx-community/Phi-3.5-MoE-instruct-4bit at 512 tokens of context and again past 4096 (long table). The fixed arm must not be worse than the reverted arm against a transformers float32 reference, if one can be produced out of band. The reverted arm must show decided-position disagreement against the fixed arm.

3. Gemma 4 audio encoder: local mask reaches one frame too far back

Upstream: Blaizzy/mlx-vlm#2385 (merge 5fe604cf), https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/models/gemma4/audio.py (_build_causal_valid_mask).

• Current: build_causal_valid_mask (src/audio/encoder.rs:615-634) builds tril(ones(ctx, chunk)).T * tril(ones(chunk, ctx), k = max_past + max_future). That admits dist == max_past, one key further back than transformers' sliding_window_mask_function allows. max_past_horizon is attention_context_left - 1 (src/audio/config.rs:101). The Gemma 3n encoder was not changed upstream and is out of scope.
• Fix: build the mask from dist = q + max_past - c over q in [0, chunk) and c in [0, ctx), as (0 <= dist < max_past) | (dist < 0 && -dist < max_future), returned as bool [chunk, ctx].
• Test: in src/audio/encoder_tests.rs, compare against the predicate computed in plain Rust for the upstream parameter set (chunk, left, right) in (12,13,0) (12,1,0) (8,6,4) (4,2,2) (6,4,1) (5,1,5) (10,10,3). The old formula must fail at least the (12,13,0) case, which is the shipped checkpoint's config.
• Trace: run the logit trace over an audio-prompted Gemma 4 checkpoint (mlx-community/gemma-4-e2b-it-4bit) with a fixed audio clip and a teacher-forced transcript. Report fixed vs reverted decided-position disagreement.

4. Phi-3-small: dummy vocabulary ids are never suppressed

Upstream: ml-explore/mlx-lm#1938 (commit 011f2f6), https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/models/phi3small.py (_dummy_tokenizer_ids).

• Current: Phi3SmallModel (src/models/phi3small.rs:570-605) returns raw logits, and its LanguageModel impl (:671) does not override output_suppressed_token_ids. The reference masks ids 100256, 100258, 100259, 100260, 100264, 100265 and 100267..=100351 to -inf.
• Fix: override output_suppressed_token_ids (trait doc at src/lib/mlxcel-core/src/generate.rs:245-271) to return those ids filtered to < vocab_size, so synthetic configs with small vocabularies do not index out of range. This is the integration path. CLI generate (src/commands/generate.rs:1731) and the server scheduler (src/server/batch/scheduler/config.rs:93) already merge these ids into the TokenBiasMap. Do not mask inside forward, which would also change logprob output that other families leave raw.
• Test: in the phi3small tests, a config with vocab_size = 100352 returns exactly the 91 ids. A config with vocab_size = 1000 returns none.
• No logit trace: the forward logits do not change.

Scope

In scope: src/models/gemma3n.rs, src/models/phimoe.rs, src/models/phi3.rs (constructor extraction only), src/audio/encoder.rs, src/models/phi3small.rs, their _tests.rs files, and the PhiMoE entry in docs/supported-models.md (routing and RoPE description).

Out of scope: the Gemma 3n audio encoder mask, and PhiMoE tensor or pipeline parallel.

Acceptance Criteria

• Item 1: the chunking-invariance test passes and fails on main, and the Gemma 3n trace is reported. (PR fix(models): port upstream numeric fixes for four model families #2457)
• Item 2: the sparsemixer and RoPE-config tests pass, the PhiMoE traces (short and long context) are reported, and mlxcel generate -m <Phi-3.5-MoE-instruct-4bit> is fluent at 128 tokens. (PR fix(models): port upstream numeric fixes for four model families #2457)
• Item 3: the mask predicate test passes for all seven parameter sets and fails on main for (12,13,0), and the audio trace is reported. (PR fix(models): port upstream numeric fixes for four model families #2457)
• Item 4: the suppression tests pass, and mlx-community/Phi-3-small-8k-instruct-4bit generates without emitting any listed id over a 256-token run. (PR fix(models): port upstream numeric fixes for four model families #2457)
• Each trace follows docs/benchmarks.md#judging-a-change-that-moves-the-numbers: gate on decided-position disagreement, trace at the width the code runs at, and include the reverted arm. (PR fix(models): port upstream numeric fixes for four model families #2457)

Verification

cargo test --workspace --profile test-fast --features metal,accelerate
cargo clippy --workspace --all-targets -- -D warnings
cargo fmt --check
cargo build --release --features metal,accelerate --example logit_trace
./target/release/examples/logit_trace models/<checkpoint> docs/architecture.md 5 60 8 512 > fixed.tsv
python3 scripts/compare_logit_traces.py reverted.tsv fixed.tsv --decided 2