feat(models): add Arcee AFMoE (afmoe) text model support #845

Corrected against references/mlx-lm (commit 15b522f, 2026-07-11) on 2026-08-05.
The original body describes AFMoE as "structurally a qwen3_moe/deepseek MoE derivative with QK-norm". mlx_lm/models/afmoe.py and the arcee-ai/Trinity-Nano-Preview config carry six features the original body does not mention, several of which change the model class.

Checkpoint confirmed

arcee-ai/Trinity-* is the AFMoE family (model_type: afmoe). Validated targets, in ascending size:

• mlx-community/Trinity-Nano-Preview-4bit (3.2 GB, 56 layers, hidden 1024, 128 experts) — downloaded, recommended for the port
• mlx-community/Trinity-Mini-4bit (15.7 GB)

Note arcee-ai/AFM-4.5B is not this family; it is model_type: arcee.

What the original body missed

• Hybrid sliding/full attention. layer_types interleaves sliding_attention and full_attention (global_attn_every_n_layers: 4, sliding_window: 2048). This makes the port a hybrid-cache model, not a plain MoE decoder: sliding layers need a RotatingKVCache and full layers a standard KVCache, which in this tree means the ModelOwnedSequenceState<Cache> wrapper that src/models/gemma3.rs uses, not the simple LanguageModel cache path.
• Full-attention layers are NoPE. Attention.__init__ builds self.rope only when is_local_attention is true; global layers apply no positional encoding at all. Applying RoPE everywhere is the obvious mistake and is invisible on short prompts.
• Sigmoid output gate on attention. gate = mx.sigmoid(self.gate_proj(x)); output = output * gate before o_proj. This is an extra gate_proj weight per attention block, absent from qwen3_moe.
• Four norms per layer, not two: input_layernorm, post_attention_layernorm, pre_mlp_layernorm, post_mlp_layernorm. The block is h = x + post_attn_norm(attn(input_norm(x))) then out = h + post_mlp_norm(mlp(pre_mlp_norm(h))), a sandwich layout.
• muP embedding scale. When mup_enabled (true on Trinity), the embedding output is multiplied by sqrt(hidden_size) before the stack.
• Dense prefix. The first num_dense_layers (2 on Trinity) use a plain MLP; MoE starts at layer 2.

Routing detail, for the record: score_func: sigmoid, a expert_bias added for selection only while the returned weights are gathered from the unbiased scores, then route_norm renormalization and a route_scale multiply (2.826 on Trinity). Group routing exists but n_group: 1 on Trinity, so it never fires there.

Checkpoint weight layout

Experts ship pre-stacked as model.layers.{i}.mlp.experts.{gate,up,down}_proj, which is the layout SwitchGLU::from_weights already consumes directly. There is no per-expert split to write. The router is model.layers.{i}.mlp.router.gate, the selection bias is model.layers.{i}.mlp.expert_bias, and the shared expert is model.layers.{i}.mlp.shared_experts.*.

The Trinity config also carries per-tensor quantization overrides (model.layers.2.mlp.router.gate at 8 bits while the rest is 4-bit), so the loaders must infer per-tensor bit width rather than trusting the top-level bits.