[inureyes]

Summary

Ports four upstream numeric fixes in families mlxcel ships, as one PR per the issue.

• Gemma 3n (Precompute offset for Gemma3n ml-explore/mlx-lm#1886): the RoPE offset is read once per forward from the first owner cache and threaded through DecoderLayer into Gemma3nAttention::forward, so KV-shared layers no longer rotate queries at their owner's already-advanced offset.
• PhiMoE (Fix phimoe's sparsemixer ml-explore/mlx-lm#1939, fix(server): gate webui-only helpers so --no-default-features builds clean #1889, Fix PhiMoE routing and plain RoPE configs Blaizzy/mlx-vlm#2419): sparsemixer routing replaces softmax over the top-k logits; the longrope block goes through phi3::SuRope (new pub(crate) SuRope::from_scaling, rotary-prefix scale moved onto SuRopeTable), with rope_table_regime for the prompt cache. Null rope_scaling keeps plain RoPE; num_experts_per_tok > num_local_experts fails at load. The fused decode kernel (run_fused_moe_two_kernel) casts scores to f32 and multiplies; it does not renormalize (the normalizing code near mlx_cxx_kernels.cpp:1491 is the DeepSeek fused_moe_forward, a different entry point). A test pins fused == graph with scores summing to 1.5025.
• Gemma 4 audio (Fix Gemma 4 audio encoder attending one frame too far into the past Blaizzy/mlx-vlm#2385): causal_valid_mask builds the mask from dist = q + max_past - c; the old construction admitted dist == max_past.
• Phi-3-small (Fix incorrect masking axis in phi3small ml-explore/mlx-lm#1938): output_suppressed_token_ids returns the 91 padding ids below vocab_size; forward logits stay raw.
• logit_trace gains an audio arm (MLXCEL_TRACE_AUDIO, MLXCEL_TRACE_AUDIO_PROMPT); docs and CHANGELOG updated.

Tests

New tests in gemma3n_tests.rs (new), phimoe_tests.rs, encoder_tests.rs, phi3small_tests.rs. The Gemma 3n chunking test and the mask predicate test fail on main; the old-mask helper is asserted to differ at (12,13,0). Deviation from the issue: mlxcel's Gemma 3n MLP and AltUp input round to bf16, so the single 12-row pass differs from narrower passes by one bf16 ulp (2.7e-3 fixed, 2.7e-2 to 6.7e-2 reverted). Multi-pass splits are gated at 1e-3 (1e-6 fixed, at least 5.6e-3 reverted) and the single pass at 1e-2. Measured on CUDA only, with MLX_ENABLE_TF32=0.

Real-checkpoint traces (GB10, fixed vs reverted, --decided 2)

Checkpoint	Width	Context	Decided disagreement	PPL fixed / reverted
gemma3n-e2b-4bit	1	512	0/151	26.44 / 26.20
gemma3n-e2b-4bit	5	512	3/177 (1.7%)	28.48 / 27.36
gemma3n-e2b-4bit	32	512	14/334 (4.2%)	31.75 / 33.09
phi-3.5-moe-4bit	5	512	2/141 (1.4%)	7.91 / 9.62
phi-3.5-moe-4bit	1	512	1/48 (2.1%)	15.65 / 19.50
phi-3.5-moe-4bit	5	4200 (long table)	118/136 (86.8%)	7.70 / 3008.7
gemma-4-e2b-it-4bit + audio	prefill	100 audio tokens	0/78	1.657 / 1.657

Gemma 3n at width 1 shows no decided disagreement: shared layers write no KV, so the only error is a one-position query shift in 10 of 30 layers; the effect grows with width. The audio clip is a deterministic 4 s synthetic chirp/tone/noise WAV (tests/fixtures/audio/tone.wav is 0.25 s); the arms differ (max logit delta 0.875) but not at decided positions. A transformers float32 reference for Phi-3.5-MoE (42B) is not feasible on this host, and mlx-lm is not a valid short-regime oracle (long table only).

Other checks: mlxcel generate on Phi-3.5-MoE is fluent at 128 tokens. On phi-3-small-8k-instruct-aq4_64 (substitute for mlx-community Phi-3-small-8k-instruct-4bit), 256 sampled tokens through the suppression bias emitted no listed id.

Follow-ups (not in this PR)

• mlxcel-xla gemma3n_emit_ops.rs deliberately mirrors the old shared-layer query offset; MLX-vs-XLA Gemma 3n parity will now diverge. Not buildable here.
• phimoe::ModelArgs reads layer_norm_eps (default 1e-6); the checkpoint spells it rms_norm_eps: 1e-5.
• The feat: default the single-sequence CLI/bench prefill to chunked (2048) with a per-model opt-out policy #674 chunked-prefill opt-out on Gemma3nModel may be this bug; flipping it needs the parity battery.
• xla-diagnostics call sites in gemma3n.rs were updated by hand and not compiled.

Closes #2379

[inureyes]

Implementation Review Summary

Intent

Port four upstream numeric fixes (Gemma 3n KV-shared RoPE offset, PhiMoE sparsemixer and LongRoPE, Gemma 4 audio local mask, Phi-3-small padding-id suppression) in one PR (#2379).

Findings Addressed

• No CRITICAL or HIGH findings, so no fix commits were made. Head reviewed: 436af3bb.

Remaining Items

• The Gemma 4 audio mask change also reaches EmbeddingGemma 2, which builds the same crate::audio::AudioEncoder (src/models/embedding_gemma2.rs:232), so its audio embeddings move too. That matches the reference, but the CHANGELOG line and PR body name only Gemma 4. (MEDIUM, docs only, below the auto-fix threshold)
• null_or_absent_rope_scaling_keeps_plain_rope uses the PHI35_MOE_4BIT_CONFIG fixture for the "absent" case, but the real Phi-3.5-MoE-instruct-4bit ships a longrope block. The test is correct; the fixture name is misleading. (LOW)
• The sparsemixer unit tests run f32 gates only. bf16 matches upstream by construction (the -inf and the 2 * jitter_eps threshold are built in the gates dtype, the same as mlx-lm's Python-scalar promotion), and the phi-3.5-moe trace runs the bf16 path. (LOW)
• The new mask also drops dist == -max_future, matching upstream _build_causal_valid_mask. This has no effect on the shipped config (attention_context_right = 0). (Informational)
• The follow-ups listed in the PR body (the mlxcel-xla gemma3n emitter offset, PhiMoE layer_norm_eps vs rms_norm_eps, the feat: default the single-sequence CLI/bench prefill to chunked (2048) with a per-model opt-out policy #674 chunked-prefill opt-out, the uncompiled xla-diagnostics call sites) are known, out-of-scope items. By reading the code, the four xla-diagnostics sites have offset in scope through use super::* and pass it in the right argument position.

Verification

• All stated requirements implemented: each port matches the upstream source (mlx-lm sparsemixer, mlx-vlm _build_causal_valid_mask); validate() runs in from_weights, which load also uses; Phi-3-small ids are filtered to < vocab_size
• No placeholder/mock code remaining
• Integrated into project code flow: rope_table_regime and output_suppressed_token_ids are delegated by LoadedModel, CLI generate and the server scheduler use the suppression list, and all three Gemma 3n decoder loops pass the offset
• Project conventions followed: cargo fmt --check is clean
• Existing modules reused where applicable: PhiMoE goes through phi3::SuRope::from_scaling and SuRopeTable::scale_rotary_prefix, and phi3 passes the same arguments as before
• No unintended structural changes
• Tests pass: 42/42 on GB10 CUDA with MLX_ENABLE_TF32=0. With the Gemma 3n offset line reverted, kv_shared_layers_ignore_chunking fails (5.6e-3 against the 1e-3 multi-pass gate).

[inureyes]

Implementation Review Summary

Intent

Port four upstream numeric fixes (Gemma 3n KV-shared RoPE offset, PhiMoE sparsemixer + LongRoPE, Gemma 4 audio local mask, Phi-3-small dummy id suppression) and add a logit_trace audio arm (#2379).

Findings Addressed

None needed: no CRITICAL or HIGH findings. Reviewed at 87dcfdff.

Remaining Items

• sparsemixer with num_experts_per_tok == 0 reaches concatenate_many on an empty list; validate() only rejects > num_local_experts. Upstream has the same gap. (LOW)
• kv_shared_layers_ignore_chunking drives Gemma3nLanguageModel::forward only; the identical offset change in forward_with_inputs (VLM) and forward_diagnostics has no test of its own. (LOW)
• PhiMoeModel::rope_table_regime is not tested at model level; only SuRope::regime is. (LOW)

Checked against upstream

• Gemma 3n: offset read once from caches[0] before the layer loop, same as mlx-lm cache[0].offset; all three compiled forward loops thread it, owners use it for compute_kv.
• PhiMoE: sparsemixer matches mlx-lm phimoe.py line for line (one-hot where in place of put_along_axis). The fused kernel casts scores to f32 and sums over k without renormalizing; the new test discriminates (scores sum to 1.5025). RoPE call sequence (pre-scale, fast_rope_with_freqs(x, rope_dims, false, 1.0, offset, freqs), table_for(offset, l)) matches phi3. rope_table_regime reaches the prompt cache through the LoadedModel delegate.
• phi3: from_args now forwards the same five values to from_scaling, and scale_rotary_prefix is the old closure moved; behavior unchanged.
• Gemma 4 audio mask: matches mlx-vlm _build_causal_valid_mask.
• Phi-3-small: ids match _dummy_tokenizer_ids; CLI generate and the server scheduler consume output_suppressed_token_ids; forward logits stay raw.

Verification

• All stated requirements implemented
• No placeholder/mock code remaining
• Integrated into project code flow
• Project conventions followed
• Existing modules reused where applicable
• No unintended structural changes
• Tests pass (GB10 CUDA, MLX_ENABLE_TF32=0, filtered: models::phimoe 12, models::phi3:: 13, models::phi3small 12, audio::encoder 4, models::gemma3n::tests 1, rope_parameters_lift 5; CI clippy still pending at review time)

[inureyes]

Security and performance review (head 87dcfdf)

No CRITICAL or HIGH findings, so no fix commits were made.

Security. No attacker-controlled input reaches the new code. PhiMoE rejects num_experts_per_tok > num_local_experts in from_weights, and the directory loader goes through that path too. The Phi-3-small dummy ids are filtered to < vocab_size and read once when the scheduler is constructed. The logit_trace audio arm only reads clip and prompt paths the operator passes on the command line of an example binary.

Performance.

• MEDIUM (informational): sparsemixer builds about 12 small lazy ops per round, times top_k, per MoE layer per token, where the old argpartition + softmax routing built about 4. That is roughly 27 extra ops per layer for Phi-3.5-MoE (top_k 2). It matches mlx-lm's uncompiled sparsemixer op for op, so this is the cost of correct routing and not a regression against the reference. No PhiMoE checkpoint is on this host, so decode was not benchmarked. A tok/s A/B on a launch-bound host is worth doing when one is available.
• Fused single-token decode is still taken. The indices are uint32 [1, top_k] and the scores keep the gate dtype, the same as before. fused_expert_kernel_keeps_sparsemixer_scores ran the CUDA kernel on GB10 and passed.
• Gemma 3n rope_offset is one scalar read per forward, so per-layer cost is unchanged.
• LOW: router_jitter_noise <= 0 produces NaN scores, and num_experts_per_tok == 0 reaches an empty concatenate. Both come from the config, both are the same as upstream, and the second one also failed under the old routing.
• LOW: building expert_ids as uint32 would drop the per-round astype. Not worth a push alone.

Verification (GB10, CUDA, test-fast, TF32 off): models::phimoe 12/12, models::gemma3n::tests 1/1, audio::encoder::tests 4/4, models::phi3small::tests 12/12.