Add Laguna (Poolside) architecture support

[JuanColilla]

Why

poolside/Laguna-S-2.1 (118B total / ~8B active MoE, 256 experts top-10 routed
plus one shared expert, OpenMDW-1.1) — and its siblings Laguna-XS-2.1
(33B-A3B) and Laguna-M.1 (225B-A23B) — have no mlx-lm architecture module
today. model_type: "laguna" in their config.json maps only to remote
custom_code classes (configuration_laguna.LagunaConfig /
modeling_laguna.LagunaForCausalLM) hosted in the HF repo, which mlx-lm
cannot run natively. This PR adds mlx_lm/models/laguna.py, so Laguna
checkpoints load and convert through the same path as every other
architecture — no trust_remote_code required, which matters more than
usual now that transformers gates custom_code execution behind
trust_remote_code=True (CVE-2026-5843, merged 2026-06-11).

Evidence of real demand: poolside has already uploaded an mlx-quantized
checkpoint (poolside/Laguna-S-2.1-NVFP4-mlx) whose weight-key layout
(mlp.switch_mlp.*, mlp.gate.*) already matches this port's target
naming — i.e. this is the naming convention the checkpoint publisher expects
mlx-lm support to use.

What

Laguna combines several features mlx-lm already supports individually,
just not previously combined in one model:

• Sigmoid-scored MoE router with an auxiliary-loss-free correction bias
  (same pattern as glm4_moe.py's MoEGate), 256 routed experts (top-10)
  plus one shared expert, via SwitchGLU.
• Alternating full/sliding-window attention with separate RoPE per layer
  type (YaRN for full-attention layers, default for sliding — same
  make_cache/alternating-mask pattern as gpt_oss.py).
• Per-layer attention head-count override (num_attention_heads_per_layer).
• Optional softplus output gating on attention (attn_out *= softplus(g_proj(x)),
  per-head or per-element) and optional learnable attention sinks on
  sliding layers (gpt_oss.py precedent for the sinks mechanism).
• QK-RMSNorm.

Also included: tests/test_models.py::test_laguna (exercises every
non-standard feature above through the existing shape/dtype/cache/batch>1
harness), and scripts/laguna_quant_predicate.py — a custom per-module
quantization predicate demonstrating how to produce the mixed-precision
variants some Laguna publishers may want (attention/router/embeddings at
one bit width, routed experts at a lower one — 8-bit/{4,6}-bit-experts, or
4-bit/2-bit-experts for a footprint that actually fits a 64 GB Mac, since
routed experts hold ~116B of Laguna S-2.1's ~117.6B stored parameters),
for recipes the standard mlx_lm.convert -q --q-bits N flag can't express
directly. The three uniform variants (4-bit/6-bit/8-bit everywhere) need no
custom predicate.

How tested

• tests/test_models.py::test_laguna (new, in this PR): a tiny random-init
  config that forces every layer through a genuinely distinct code path —
  dense MLP vs. MoE, full vs. sliding attention, YaRN vs. default RoPE,
  different per-layer head counts (8 vs. 12) — verified through the existing
  model_test_runner shape/dtype/cache/batch>1/deepcopy checks.
• Numerical parity against the reference transformers implementation.
  A local harness (not included in this PR — it depends on torch/
  transformers, which mlx-lm does not otherwise require) built a matched
  3-layer config in both the reference and this port from an identical
  random seed, ran this port's real Model.sanitize() on the reference's
  actual per-expert weight layout (not a hand-constructed shortcut around
  it), loaded with strict=True, and compared logits:
  max abs diff 5.1e-4, mean abs diff 1.3e-4 (float32, tolerance 1e-3).
  The router's auxiliary-loss-free correction bias was set to a nonzero
  random tensor before comparison — with it left at its zero default the
  first pass, an earlier review round found sanitize() was silently
  dropping this bias's real on-disk key (mlp.experts.e_score_correction_bias
  → mlp.gate.e_score_correction_bias) rather than remapping it, which a
  zero-valued bias could not have exposed; that gap is now closed and
  covered by a dedicated unit test, and the harness's discriminating power
  was confirmed by deliberately breaking the bias-add sign and observing the
  diff jump to 0.025 (fails, as it should).
• Structural shapes were additionally cross-checked against the real
  poolside/Laguna-S-2.1 checkpoint's safetensors headers (HTTP range
  request — no 219 GiB download), confirming the per-layer head-count
  override (48 vs. 72 heads) and per-expert unfused weight layout this port
  assumes are correct.

Known limitations / explicitly out of scope

• DFlash speculative-decoding draft checkpoints are not supported by this
  module, even though they share the same config.json model_type: "laguna" string. poolside/Laguna-S-2.1-DFlash (and its siblings) use a
  structurally different architecture (architectures: ["DFlashLagunaForCausalLM"]
  — dense, no MoE, all-sliding-window, with EAGLE-style auxiliary hidden
  states and block-causal masking for multi-token draft generation), and
  mlx-lm dispatches purely on model_type (mlx_lm/utils.py), not
  architectures. Converting a DFlash checkpoint with this module would not
  produce a working draft model. A separate architecture module would be
  needed for DFlash support; out of scope here (a dedicated
  dflash_laguna speculator is already proposed separately in Add dflash_laguna EAGLE-3 speculator for Laguna #1531).
• Attention sinks: the reference stores a per-layer sink parameter as
  self.sink (singular) when swa_attention_sink_enabled=True; this port
  uses self.sinks and sanitize() does not remap the name. None of the
  three published Laguna checkpoints (S-2.1, XS-2.1, M.1) set this flag, so
  it's inert today; a checkpoint that does would need the name reconciled
  first.
• attention_bias is read from config by this port, but the reference
  hardcodes no bias on Q/K/V/O regardless of config — inert for all three
  published checkpoints (all set it false), but worth aligning if a future
  variant sets it true.
• The parity harness compares a single full-prompt forward pass; it doesn't
  exercise incremental/cached decode (RotatingKVCache stepping) or router
  logit soft-capping (no published checkpoint sets a nonzero value there
  either).
• Loading a real Laguna checkpoint's tokenizer logs a transformers
  warning about an "incorrect regex pattern" and suggests
  fix_mistral_regex=True. This is a false positive of transformers's own
  detection heuristic (it fires for any config.json missing
  transformers_version, which Laguna's does, regardless of model_type)
  and not a real tokenization bug: round-tripping code samples with
  contractions, mixed case, and multi-newline runs through Laguna's real
  tokenizer produces byte-identical output with and without the flag, and
  the flag's own fix targets a different pipeline stage than the one
  Laguna's pre-tokenizer actually differs on. Safe to ignore.

[pierre427]

Some independent measured datapoints in case they're useful here — I ported Laguna-XS-2.1 (33B-A3B, 8-bit) and ran an optimization battery on an M5 Max (mlx 0.32.0-dev). Everything below is on XS-2.1; the router findings are architectural and should carry to the S-2.1 / M.1 siblings.

Windowed KV — confirms the sliding/full make_cache split in this PR is worth it, and the payoff grows fast with context. Decode, greedy, 200-tok gen:

ctx	full-KV	windowed	speedup	KV (full→win)
4k	53.1	58.2	1.10×	713 → 241 MB
16k	42.9	52.9	1.23×	2726 → 744 MB
32k	17.2	46.7	2.72×	5453 → 1426 MB

Prefill is ~2× on the sliding layers too. Prefill last-token logits are bitwise-identical to full-KV; generation ≤512 ctx is token-for-token identical; >512 differs only by the standard RotatingKVCache "≥ max_size" semantics (quality-neutral, not a bug).

Expert dynamic-k is a dead end on this router — worth not chasing. Profiling the sigmoid router over a real generation (39 MoE layers), cumulative routed mass is top-1 22% / top-4 64% / top-6 83% / top-7 92% — spread, not peaky. So trimming top_k barely moves decode: top-6 = 1.09× for KL 2.4e-2 vs top-8, top-5 = 1.09× for KL 6.7e-2. The always-on shared expert + sigmoid routing don't garden like the peakier softmax routers in e.g. gpt-oss.

6-bit body / 8-bit router is a strong quant target. Using a quant predicate like the laguna_quant_predicate.py here to pin the 39 routers to 8-bit and quantize the body to 6-bit, vs full 8-bit on XS-2.1: 24% less peak RAM (27.2 vs 35.6 GB), ~10% faster decode (109.8 vs 99.7 tok/s), 100% greedy agreement over 200 tok, KL 0.008. Effectively lossless.

Minor: the tokenizer threw a Mistral-style regex warning on load — likely benign, but worth an edge-case tokenization check.

Bench harnesses and exact configs available if useful.

[pierre427]

Adding clean S-2.1 data after a reboot, in case it helps decide the Laguna support direction here.

Local setup:

• model: poolside/Laguna-S-2.1 converted to MLX q4 affine, group size 64
• artifact index reports ~117.6B total params; local q4 artifact is ~62 GB
• router/gate projections kept at 8-bit in the local quant (mlp.gate.proj, group size 64)
• runtime: M5 Max / 128 GB, MLX 0.32.0.dev20260708+1ac2fbac, mlx-lm 0.31.3
• benchmark: clean post-reboot, no competing model process, direct model.make_cache() path, 128 generated tokens, no DFlash

Plain vs prompt-lookup decode on a repeated-code synthetic prompt:

ctx	plain first token	plain decode	PLD first token	PLD decode	PLD behavior
512	0.71s	68.15 tok/s	0.54s	56.93 tok/s	stayed PLD
4k	2.43s	66.11 tok/s	2.58s	57.25 tok/s	rate-gated off
16k	12.47s	57.75 tok/s	13.16s	50.28 tok/s	rate-gated off
32k	35.30s	49.40 tok/s	35.13s	39.33 tok/s	rate-gated off

Takeaways from this S-2.1 q4 run:

1. The mixed cache policy in this PR is the right default: sliding layers should use bounded rotating KV and full-attention layers should keep ordinary KV. Long-context decode stays usable at 32k with ~1.7 GB KV in this setup.
2. Plain decode should probably be the default for fresh generation on S-2.1 q4. PLD remains useful for copy-heavy edit/RAG workloads, but it should stay adaptive/rate-gated rather than blanket-on.
3. The conversion path should preserve higher precision for routing/control tensors. In this local q4, keeping mlp.gate.proj at 8-bit is cheap and seems like the right quality-preserving serving point.
4. DFlash should be treated as a separate follow-up once the S-2.1 DFlash checkpoint is wired; these numbers are target-only, no draft/speculator.

So my practical serving recipe for S-2.1 is: native Laguna arch + mixed sliding/full KV cache + q4 body/q8 routing, plain decode as baseline, prefix-cache for repeated prefixes, and PLD/DFlash as gated workload-specific accelerators rather than the architecture default.

[JuanColilla]

Adding clean S-2.1 data after a reboot, in case it helps decide the Laguna support direction here.

Local setup:

• model: poolside/Laguna-S-2.1 converted to MLX q4 affine, group size 64
• artifact index reports ~117.6B total params; local q4 artifact is ~62 GB
• router/gate projections kept at 8-bit in the local quant (mlp.gate.proj, group size 64)
• runtime: M5 Max / 128 GB, MLX 0.32.0.dev20260708+1ac2fbac, mlx-lm 0.31.3
• benchmark: clean post-reboot, no competing model process, direct model.make_cache() path, 128 generated tokens, no DFlash

Plain vs prompt-lookup decode on a repeated-code synthetic prompt:

ctx plain first token plain decode PLD first token PLD decode PLD behavior
512 0.71s 68.15 tok/s 0.54s 56.93 tok/s stayed PLD
4k 2.43s 66.11 tok/s 2.58s 57.25 tok/s rate-gated off
16k 12.47s 57.75 tok/s 13.16s 50.28 tok/s rate-gated off
32k 35.30s 49.40 tok/s 35.13s 39.33 tok/s rate-gated off
Takeaways from this S-2.1 q4 run:

1. The mixed cache policy in this PR is the right default: sliding layers should use bounded rotating KV and full-attention layers should keep ordinary KV. Long-context decode stays usable at 32k with ~1.7 GB KV in this setup.
2. Plain decode should probably be the default for fresh generation on S-2.1 q4. PLD remains useful for copy-heavy edit/RAG workloads, but it should stay adaptive/rate-gated rather than blanket-on.
3. The conversion path should preserve higher precision for routing/control tensors. In this local q4, keeping mlp.gate.proj at 8-bit is cheap and seems like the right quality-preserving serving point.
4. DFlash should be treated as a separate follow-up once the S-2.1 DFlash checkpoint is wired; these numbers are target-only, no draft/speculator.

So my practical serving recipe for S-2.1 is: native Laguna arch + mixed sliding/full KV cache + q4 body/q8 routing, plain decode as baseline, prefix-cache for repeated prefixes, and PLD/DFlash as gated workload-specific accelerators rather than the architecture default.

Thanks for taking the time to run a full battery on real hardware and share exact numbers, this is genuinely useful validation, especially the windowed-KV scaling curve and the router mass-distribution profiling. Replying to the two points that needed a closer look:

6-bit body / 8-bit router: dug into this, and there's actually no custom predicate needed for it at all. Router.weight in this PR (mlx_lm/models/laguna.py) is a raw array, not an nn.Linear, so it has no to_quantized. mlx_lm.utils.quantize_model's wrapper checks hasattr(module, "to_quantized") before ever calling a custom quant_predicate, so the router, and every RMSNorm in the model, is always left at full precision, regardless of what a predicate returns for those paths. Concretely: plain mlx_lm.convert(..., quantize=True, q_bits=6), with no custom predicate at all, already gives you "router protected / body at 6-bit" — and the router ends up at full float32/bf16 rather than 8-bit, which is probably why your KL numbers came out so clean (0.008, 100% greedy agreement).

Pushed a small follow-up commit that corrects laguna_quant_predicate.py's docstring/comment (it previously implied the router went through attention_bitly happens)and adds a test that exercises the real quantize_moderather thanonly unit-testing the predicate function in isolation, that gap is exactly how the original comment went unnoticed.
Tokenizer regex warning: confirmed it's a transformers- realtokenization issue. transformers' Mistral-regex heuristic firis missingtransformers_version, which Laguna's is, regardless of model_typeMistral-like"without actually checking its pre-tokenizer. Verified with the realround-tripping code samples (contractions, mixed case, multi-n-identicaloutput with and without fix_mistral_regex=True, andthat flag's o position inthe pre-tokenizer pipeline than the one Laguna's actually diff this to thePR description's "Known limitations" section so it doesn't conful checkpoint.

Appreciate dynamic-top-k data point too! Good to have that one closed off rather than rediscovered later.