feat(models): add Phixtral (phixtral) text model support #844

Summary

Add native mlxcel support for Phixtral (model_type: phixtral), a Mixtral-style sparse MoE built on the Phi-2 backbone. Each block keeps the Phi-2 single-LayerNorm parallel-residual layout (attention and MLP share one input LayerNorm and run in parallel) and uses partial-rotary RoPEAttention, but replaces the dense MLP with top-k routing over num_local_experts experts. Porting brings the family to the same load-and-generate path as the other text models.

Upstream reference

• mlx-lm: https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/models/phixtral.py

Public checkpoint

• mlabonne/phixtral-4x2_8 (4 experts, top-2). Confirm an existing MLX SafeTensors conversion or produce one before validation. Real-checkpoint validation is a hard gate.

Architecture notes

• Mixtral-style sparse MoE on a Phi-2 backbone.
• RoPEAttention with partial rotary (rotary_dim): only the first rotary_dim slots are rotated.
• nn.LayerNorm (with bias), not RMSNorm.
• Phi-2 block shape: a single input LayerNorm per block feeding attention and MLP in PARALLEL (parallel residual), summed into the residual.
• Routing: an nn.Linear gate selects top-num_experts_per_tok (default 2) of num_local_experts experts per token.

Implementation plan

Cite and reuse existing mlxcel modules:

• Phi attention + LayerNorm + parallel-residual block: reuse src/models/phi3.rs / src/models/phi3small.rs for the Phi attention and LayerNorm handling. Note src/models/phimoe.rs already combines a Phi-family backbone (LayerNorm with bias, attention proj bias) with MoE routing and is the closest existing reference; start from it.
• Mixtral-style top-k MoE: reuse the routing pattern in src/models/mixtral.rs and the expert compute in src/models/switch_layers.rs.
• The distinctive work is combining Phi-2 partial rotary (rotary_dim) + the single-LayerNorm parallel-residual block with switch-based top-k MoE (gate as nn.Linear, softmax over top-k logits).

Touchpoints & acceptance criteria

Follow docs/adding-models.md. Done only when it loads and generates from a real checkpoint.

• Config struct + serde parse for phixtral (rotary_dim, num_local_experts, num_experts_per_tok, LayerNorm eps, Phi-2 dims), with sensible #[serde(default)]s.
• from_weights constructs the model from a real WeightMap, including gate + per-expert weights.
• sanitize handles the checkpoint's expert key layout / stacking.
• Detection arm added in src/models/detection.rs ("phixtral" => ...).
• Registration added via for_each_model_registration! in src/model_metadata.rs.
• TP / distributed arch-string wired if the family is served in that path.
• Unit tests in a _tests.rs file beside the implementation (config parse, partial-rotary dims, parallel-residual block, top-k gate selection).
• docs/supported-models.md updated with the new family.
• Real-checkpoint validation: mlxcel list shows the arch and ./target/release/mlxcel generate -m <path> -p "..." -n <N> produces coherent output.

Effort

MEDIUM. Phi attention/LayerNorm and Mixtral top-k routing both exist (phimoe.rs already fuses a Phi backbone with MoE); the new logic is Phi-2 partial rotary + parallel residual combined with the switch MoE gate.