feat(models): add Xiaomi MiMo v2 (mimo_v2_flash) text model support
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
> **Corrected against `references/mlx-lm` (commit `15b522f`, 2026-07-11) on 2026-08-05.**
> Labelled `modelsize:xlarge` + `status:blocked`: the only public checkpoint does not fit the development machine. See "Blocked on hardware" below.
## Blocked on hardware
`XiaomiMiMo/MiMo-V2-Flash` is the only checkpoint declaring `model_type: mimo_v2_flash`, and it is roughly 309B parameters: 48 layers, hidden 4096, **256 routed experts**, `moe_intermediate_size` 2048. `mlx-community/MiMo-V2-Flash-4bit` is **162 GB on disk**, against a 128 GB unified-memory dev box (107.5 GB wired limit). No 3-bit or 2-bit MLX conversion exists.
`XiaomiMiMo/MiMo-V2.5` and `MiMo-V2.5-Pro` are **not** substitutes: they declare `model_type: mimo_v2`, a different family. There is no smaller `mimo_v2_flash` checkpoint.
So the port can be written and unit-tested, but the project's real-checkpoint gate cannot be met on current hardware. Deferring until either a smaller checkpoint is published or a larger machine is available.
## Architecture, verified against upstream
The original body is broadly accurate. Corrections and additions:
- The interleave is driven by an explicit **`hybrid_layer_pattern`** list (`config.hybrid_layer_pattern[idx] == 1` selects sliding), not by a `sliding_window_pattern` modulus as in gemma3. The window is `sliding_window_size`.
- Dual RoPE base is confirmed: `rope_theta = args.swa_rope_theta` when `is_sliding_window`, else `args.rope_theta`. This is the primary divergence, as stated.
- Partial RoPE is `int(args.partial_rotary_factor * head_dim)`. Note `partial_rotary_factor` is typed `int` in upstream's `ModelArgs` but used as a multiplier, so parse defensively.
- Routing activation is **sigmoid** (`scores = mx.sigmoid(gates.astype(mx.float32))`), not softmax. The original body does not say which.
- `n_shared_experts` is `Optional`; when set, the shared MLP width is `moe_intermediate_size * n_shared_experts` and its output is added to the routed mixture.
- Layers are dense or MoE per `is_moe`, so there is a dense prefix as in the other MoE families here.
`src/models/mimo.rs` (MiMo v1) is a weak structural reference: v1 is a dense multi-token-prediction model with none of the MoE, sliding-window, or dual-rope machinery. `src/models/gemma3.rs` (rotating cache) and `src/models/cohere2_moe.rs` (shared expert) are the closer templates, as the original body notes.
---
Original issue body (retained for history)
## Summary
Add native mlxcel support for Xiaomi MiMo v2 (`model_type: mimo_v2_flash`), a MoE decoder with shared experts, partial RoPE, and interleaved sliding-window / full attention. The distinctive feature is a dual RoPE base: sliding-window layers use `swa_rope_theta` while full-attention layers use `rope_theta`. 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/mimo_v2_flash.py
## Public checkpoint
- XiaomiMiMo MiMo-V2 family (Hugging Face `XiaomiMiMo/*`).
- The availability of an `mlx-community` conversion is UNCONFIRMED at time of writing. Before validation, confirm an existing MLX SafeTensors conversion or produce one. Real-checkpoint validation is a hard gate; the issue is not done on synthetic weights alone.
## Architecture notes
- MoE decoder with shared experts (`n_shared_experts`), routed top-k experts, plus always-on shared expert MLP.
- PARTIAL RoPE via `partial_rotary_factor`: only the first `head_dim * partial_rotary_factor` slots are rotated.
- Interleaved attention: some layers use sliding-window attention, others use full (global) attention.
- Dual RoPE base: sliding-window layers rotate with `swa_rope_theta`, full-attention layers rotate with `rope_theta`. This is the primary divergence from existing ports.
- LayerNorm epsilon carried as `layernorm_epsilon`.
## Implementation plan
Cite and reuse existing mlxcel modules:
- Start from `src/models/mimo.rs` (MiMo v1) as the structural reference for the decoder skeleton and config.
- Sliding-window attention + rotating cache: reuse the pattern in `src/models/gemma3.rs` (`RotatingKVCache`, `create_sliding_window_prefill_mask`, and the `sliding_window` / `sliding_window_pattern` layer-selection logic).
- Partial RoPE: reuse the `partial_rotary_factor` handling in `src/models/gemma4.rs` (`default_partial_rotary_factor`, the `head_dim * partial_rotary_factor` rope-dims computation).
- MoE path: reuse `src/models/switch_layers.rs` for routed experts; for the always-on shared expert, follow the shared-expert pattern in `src/models/cohere2_moe.rs` (`moe_num_shared_experts` / `num_shared_experts`).
- The distinctive work is (1) selecting `swa_rope_theta` vs `rope_theta` per layer according to whether the layer is sliding-window or full, and (2) the sliding/full layer interleave wiring.
## Touchpoints & acceptance criteria
Follow `docs/adding-models.md`. Done only when it loads and generates from a real checkpoint.
- [ ] Config struct + serde parse for `mimo_v2_flash` (`n_shared_experts`, `partial_rotary_factor`, `rope_theta`, `swa_rope_theta`, `layernorm_epsilon`, sliding-window fields), with sensible `#[serde(default)]`s.
- [ ] `from_weights` constructs the model from a real `WeightMap`.
- [ ] `sanitize` handles any weight-key remaps / stacking the checkpoint needs.
- [ ] Detection arm added in `src/models/detection.rs` (`"mimo_v2_flash" => ...`).
- [ ] 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-rope dims, per-layer rope-theta selection, sliding-vs-full interleave).
- [ ] `docs/supported-models.md` updated with the new family.
- [ ] Real-checkpoint validation: `mlxcel list` shows the arch and `./target/release/mlxcel generate -m -p "..." -n ` produces coherent output.
## Effort
MEDIUM. The decoder, MoE, shared expert, partial RoPE, and sliding-window cache all exist as reusable building blocks; the new logic is the dual rope_theta selection and the sliding/full interleave.
Contributor guide
Research direction
Start with docs/adding-models.md, then compare src/models/mimo.rs with the closer templates in src/models/gemma3.rs and src/models/cohere2_moe.rs and review the upstream mlx-lm implementation. Add the model implementation, detection and registration entries, beside-unit tests, and supported-models documentation. Done requires real-checkpoint loading and coherent generation, but validation is blocked by the checkpoint's hardware requirements.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- ai, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100