Repository navigation
Fix phimoe's sparsemixer - #1939
Merged
Merged
Conversation
zcbenz
approved these changes
Oct 2, 2026
This was referenced Oct 16, 2026
inureyes
added a commit
to lablup/mlxcel
that referenced
this pull request
Oct 19, 2026
Gemma 3n: KV-shared layers read the RoPE offset from their owner's cache after the owner had already appended the current pass, so their queries were rotated L positions ahead of their keys (one position on every decode step). The offset is now read once per forward from the first owner cache and threaded through the decoder layer into the attention (ml-explore/mlx-lm#1886). PhiMoE: routing is HF sparsemixer instead of a softmax over the top-k logits (ml-explore/mlx-lm#1939), and the longrope block is applied through phi3::SuRope with short/long table selection and the configured mscale; a null or absent rope_scaling keeps plain RoPE. SuRope gains a pub(crate) from_scaling constructor and the rotary prefix scale moves onto SuRopeTable so both families share it. PhiMoE reports rope_table_regime and rejects num_experts_per_tok > num_local_experts at load. Gemma 4 audio: the local attention mask is built from the query-minus-key distance and no longer admits dist == max_past, one key further back than the reference (Blaizzy/mlx-vlm#2385). Phi-3-small: the vocabulary's 91 padding ids are returned from output_suppressed_token_ids so CLI and server generation never emit them; forward logits stay raw (ml-explore/mlx-lm#1938). logit_trace gains an opt-in audio arm (MLXCEL_TRACE_AUDIO) that traces a Gemma 4 transcript behind an audio clip. Refs #2379
inureyes
added a commit
to lablup/mlxcel
that referenced
this pull request
Oct 19, 2026
## Summary Ports four upstream numeric fixes in families mlxcel ships, as one PR per the issue. - **Gemma 3n** (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** (ml-explore/mlx-lm#1939, #1889, 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** (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** (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 #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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
While fixing #1884 , I decided to do a parity check and it did not match. Turns out we should've been using the HF's
sparsemixerinstead of a plain softmax. Now matches the HF reference on GPU with MLX_ENABLE_TF32=0 to