feat(models): add the Laguna family with compressed-tensors NVFP4 experts

[inureyes]

Summary

Adds Poolside's Laguna family (model_type: "laguna", XS 2.1 / XS.2 / S 2.1): a pre-norm decoder with QK-RMSNorm, YaRN partial-rotary full-attention layers interleaved with window-512 sliding layers, per-layer query-head counts, a per-head softplus attention output gate, and a sigmoid or sqrt-softplus routed top-k MoE with a correction bias, a shared expert and moe_routed_scaling_factor. The published NVFP4 checkpoints keep every expert and the shared expert in the compressed-tensors nvfp4-pack-quantized layout, which nothing in src/ could read before.

What changed

• src/models/laguna{,_layers,_sanitize}.rs: config, per-layer RoPE resolution (per-layer-type entry overlaid on loose top-level scalars, explicit attention_factor preferred over the derived one, partial_rotary_factor fallback chain), the model shell and the sanitize. The transcode stacks the per-expert weight_packed bytes and reinterprets them as MLX native NVFP4 uint32 words on device, keeps the E4M3 block scales byte for byte, and carries 1 / weight_global_scale as a per-expert .global_scale sidecar. Pre-stacked MLX conversions (fused gate_up_proj split back), per-expert bf16 exports and tied heads load through the same idempotent sanitize.
• src/models/switch_layers.rs: SwitchLinear::Quantized gains global_scale: Option<[num_experts] f32>, applied to the selected experts' rows after gather_qmm on both layouts; planes carrying it stay off the fused affine decode kernel. Folding 1 / global into the 3-mantissa-bit E4M3 block scales was rejected because the quotient is generally not representable.
• Each attention layer builds its own prefill mask from the key count its cache returns, so sliding layers stay aligned with both the plain rotating cache and the buffered variant that speculative rollback (feat(speculative): add the Laguna DFlash drafter (sliding-window gated attention over captured target hidden states) #1351) will arm.
• MlxcelTokenizer::prompt_carries_bos: every tokenize site now suppresses add_special_tokens when the rendered prompt already opens with a BOS. Laguna's template emits 〈|EOS|〉 (id 2) first and the checkpoint's TemplateProcessing post-processor would otherwise prepend id 2 again. The two token-count routes follow the same rule, which also removes an off-by-one the whole Llama 3 lineage had.
• Config geometry is bounded at load (ModelArgs::validate, validate_router_geometry, and per-expert plane checks in the sanitize) because several of these values reach MLX as raw arguments. A config.json that declares sparse layers but omits num_experts_per_tok used to hand argpartition an out-of-range pivot, which MLX throws on and the cxx bridge does not catch, aborting the process after the whole checkpoint was resident; an under-declared num_experts let the router index past the stacked planes, which nothing downstream range-checks.
• Detection, registry, metadata, LoadedModel, memory estimate and TP tables carry the family; docs/supported-models.md documents it, and mlxcel arch lists Laguna XS / S (Poolside code, hybrid MoE).

Validation

Host: M5 Max, 128 GB, Metal. Everything below was re-run on the final commit.

• cargo test --profile test-fast --features metal,accelerate --lib models::: 1647 passed, 0 failed, 65 ignored; the models::laguna_tests subset is 21 tests. Narrow clippy and cargo fmt --all -- --check are clean, and CI is green. An earlier run on the GB10 (sm_121) CUDA host passed the same suites under --features cuda.
• models/laguna-xs-2.1-nvfp4: MLXCEL_FUSED_MOE=0 mlxcel generate -m models/laguna-xs-2.1-nvfp4 -p "Write a Python retry wrapper with exponential backoff." -n 64 --no-chat-template --temp 0 transcodes 234 compressed-tensors NVFP4 planes, loads in 3.2 s at 16.52 GB resident, and writes fluent Python at 53 to 79 tok/s. The default fused run is byte-identical, as expected: planes carrying .global_scale never enter the affine fused decode kernel.
• models/laguna-xs.2-4bit (pre-stacked affine path): same command, fluent Python at 121 to 128 tok/s unfused and 108 tok/s fused. The two greedy continuations diverge from about token 10, which is the pre-existing fused affine decode kernel rather than this port: models/qwen3-30b-a3b-4bit diverges at the same point under the same A/B on this host.
• Server on port 19347 with models/laguna-xs-2.1-nvfp4: /apply-template renders a prompt opening with 〈|EOS|〉, and /tokenize returns 51 ids with add_special=false ([2, 97, ...]) against 52 with add_special=true ([2, 2, ...]). A /v1/chat/completions request with enable_thinking=false reports prompt_tokens = 51 and ends finish_reason: "stop" after 672 of 700 allowed tokens with no </assistant> in the text, so id 24 fired the stop. Both count routes answer input_tokens: 51.

Not verified

• Token-exactness against an external reference is not established: the pinned mlx-lm checkout (0.31.3) carries no laguna.py, so there is nothing on either host to diff greedy ids against. The fused/unfused A/B above is the only comparison that was run.
• Laguna S 2.1 NVFP4 (about 100 GB) was not downloaded; its 48-layer config path is covered by unit tests only.
• cargo test --workspace --profile test-fast --features metal,accelerate was not run whole on this shared host; CI runs the workspace gate.
• The adapter route (loading/config_backed.rs) reaches LagunaModel::from_weights without the sanitize, so an NVFP4 checkpoint fails there with a missing-weight error. Laguna registers adapter: None and every other ConfigBacked family is wired the same way, so this is left as is.

Closes #1347

[inureyes]

Security and performance review

config.json is untrusted input on the mlxcel generate -m <org>/<repo> path, and the Laguna router and NVFP4 transcode took values from it without a bound. Fixed in d00262e.

HIGH: num_experts_per_tok reaches argpartition as an out-of-range pivot. It is a serde default, so a config that declares sparse layers through mlp_layer_types but omits the key yields 0, and router_select computed argpartition(kth = num_experts), one past the end of the score row. MLX throws on that, and the throw crosses the cxx bridge as an uncatchable abort on the first routed token, after the whole checkpoint is already resident. A value above the router width was quieter: kth went negative, slice took a negative start, and the block routed fewer than k experts with nothing reported.

HIGH: the router can index past the stacked expert planes. The router width comes from mlp.gate.proj while laguna_sanitize stacks exactly config.num_experts planes, so an under-declared num_experts leaves indices the planes do not cover, and neither gather_qmm nor the take behind the per-expert .global_scale range-checks a positive index. It also sent every leftover plane through the dense transcode as a linear nothing reads. afmoe::validate_stacked_experts guards the same hazard.

MEDIUM: an uncatchable abort and a panic on the load path. Per-expert weight_packed / weight_scale planes were stacked without checking that they agree on shape and dtype, so mlx::core::stack throws mid-load, and a mixed dtype would instead promote and be reinterpreted at the wrong width. Separately, head_dim < 2 panicked inside Ord::clamp while resolving the partial-rotary width.

Left alone: prompt_carries_bos recomputes bos_token_id() (two tokenizer encodes) per request at eight call sites. Real, but microseconds against a request, so a cache is not worth the added state.

Validation: cargo test --profile test-fast --features metal,accelerate --lib models:: (1647 passed), clippy and fmt clean.