Jais-3-MoVA-36B-A4B (MLX, bf16)
bf16 MLX conversion of IFM/Jais-3-MoVA-36B-A4B, converted from revision 50e83d4. 74.9 GB on disk. For Apple silicon.
Jais-3-MoVA-36B-A4B is a mixture-of-experts Jais-3 model with Mixture-of-Values (MoVA) attention: 36B total parameters, about 4B active per token, 512K context window.
Requirements
mlx-lm doesn't support the k2_horizon architecture (used by Jais-3) yet. Until support lands, this repo ships the MLX model code (k2_horizon.py), which mlx-lm loads through the model_file entry in config.json. So the released mlx-lm works as-is:
Pass --trust-remote-code (or trust_remote_code=True). The file runs on your machine, so read it first if you like; its comments describe how it differs from the PyTorch code.
How it was converted
No quantization. The original checkpoint is stored in fp32; this version rounds the weights to bf16 (round-to-nearest-even), the precision MLX models usually run in, and stacks the expert weights for gather_mm.
Conversion check
Checked against the PyTorch code (modeling_k2_horizon.py) in fp32 on the real weights of this model (rounded to bf16 on both sides, since the original checkpoint is fp32). The check used the model truncated to its first 5 layers (the 3 dense layers and the first 2 MoE + MoVA layers, plus the final norm and lm_head), with the same truncation on both sides, on English, code, Arabic and math prompts: the logits match to a relative error of 6e-6 or better, the next-token predictions agree at every position, and greedy decoding with the KV cache gives the same tokens (64/64 on every prompt). The remaining MoE + MoVA layers run the same code path.
Smoke-tested after conversion with released mlx-lm 0.32.0 on English and Arabic chat prompts: answers end normally (the chat template and stop tokens work).
Benchmarks (all Jais-3-MoVA-36B-A4B MLX variants)
Perplexity on 128 × 512 tokens (WikiText-2 test, and the first articles of the Arabic Wikipedia dump 20231101.ar; lower is better) and generation speed on an Apple M4 Max, all measured the same way:
The bf16 and 8-bit configurations could not be tested.
Perplexity is a coarse signal. Test the versions on your own workload before picking one.
Usage
mlx_lm.chat and mlx_lm.server take the same --trust-remote-code flag.
The model thinks before it answers, inside <ifm|think> … </ifm|think>. Leave room for that in max_tokens. The effort level is set with reasoning_effort in the chat template: "high" (default), "medium" or "low", e.g. apply_chat_template(messages, add_generation_prompt=True, reasoning_effort="low").
Notes:
- Server output: mlx-lm doesn't recognize the
<ifm|think> tags yet, so mlx_lm.server returns the thinking text inside content rather than in a separate reasoning field. The tool-call format isn't parsed yet either.
- Chat template change: the original template raises an error when an earlier assistant message has no thinking field, which is what OpenAI-style clients send. The template here renders empty thinking for those messages instead. Nothing else was changed.
License
Apache-2.0, inherited from the base model. Refer to the original model card for architecture, benchmarks and intended use.
Citation