feat(models): add the Laguna family (Poolside Laguna XS/S, hybrid sliding/full MoE with NVFP4 experts) #1347

Summary

mlxcel cannot load Poolside's Laguna coding models (model_type: "laguna", architectures: ["LagunaForCausalLM"]): Laguna XS 2.1 / XS.2 (40 layers, 256 experts, about 33B total) and Laguna S 2.1 (48 layers, about 117B total). The family is a pre-norm decoder with QK-RMSNorm, hybrid full_attention (YaRN, partial rotary) / sliding_attention (plain RoPE, window 512) layers, per-layer query-head counts, a per-head softplus attention output gate, and a sigmoid-routed top-k MoE with a correction bias, a shared expert and a moe_routed_scaling_factor. The published NVFP4 checkpoints store experts in the compressed-tensors nvfp4-pack-quantized layout (per-expert weight_packed / weight_scale / weight_global_scale). This change adds src/models/laguna.rs, the detection arm, the NVFP4 sanitize path, the sqrtsoftplus router variant, and the Laguna tokenization contract (the chat template owns the BOS token).

Current behavior

• src/models/detection.rs:150 dispatches on model_type; there is no "laguna" arm, so mlxcel generate -m models/Laguna-XS-2.1-NVFP4 fails with an unsupported model type.
• The nearest existing module is src/models/mellum.rs (JetBrains Mellum 2): layer_types full/sliding hybrid, per-layer-type rope_parameters, YaRN via compute_yarn_rope at src/models/mellum.rs:190, RotatingKVCache on sliding layers, SwitchGLU MoE. Mellum has no attention output gate, a single head count for every layer, a softmax router without correction bias or shared expert, and no partial_rotary_factor.
• src/models/sanitize.rs:564 (repack_nvfp4_weights_to_quantized) handles only the ModelOpt layout keyed by .weight_scale_2; the compressed-tensors layout (weight_packed / weight_scale / weight_global_scale) has no reader anywhere in src/.
• src/server/model_provider.rs:1418 suppresses add_special_tokens only when the rendered prompt starts with the literal <bos> or <s>:

let add_special = !plain_prompt.starts_with("<bos>") && !plain_prompt.starts_with("<s>");

The Laguna chat template emits 〈|EOS|〉 (id 2, which is both BOS and EOS) as its first token, and the checkpoint's tokenizer.json post_processor (TemplateProcessing, single: [〈|EOS|〉, A]) prepends id 2 again when add_special_tokens is true, so every chat request would start with a doubled id 2.

Expected behavior

Config (top level of config.json; types and defaults):

rope_parameters (XS 2.1):

"full_attention":    {"rope_theta": 500000.0, "rope_type": "yarn", "factor": 32.0,
                      "original_max_position_embeddings": 8192, "beta_slow": 1.0, "beta_fast": 64.0,
                      "attention_factor": 1.3465735902799727, "partial_rotary_factor": 0.5},
"sliding_attention": {"rope_type": "default", "rope_theta": 10000.0, "partial_rotary_factor": 1.0}

S 2.1 uses factor: 128, beta_fast: 32, attention_factor: 1.4852; XS.2 uses factor: 64, original_max_position_embeddings: 4096, and also carries original_max_position_embeddings at the top level of rope_parameters (a non-dict entry that applies to every layer type as a default). When a layer type has no entry, fall back to the top-level rope_theta (default 500000) with rope_type: default.

Per-layer forward (x: [B, L, H], n_h = num_attention_heads_per_layer[i], n_kv = 8, d = 128):

h_in = RMSNorm(x; input_layernorm, eps)
q = q_proj(h_in).reshape(B, L, n_h, d);  k = k_proj(h_in).reshape(B, L, n_kv, d);  v = v_proj(h_in).reshape(B, L, n_kv, d)
q = RMSNorm_d(q; q_norm.weight[128], eps);  k = RMSNorm_d(k; k_norm.weight[128], eps)   # per head, before RoPE
q, k -> [B, heads, L, d]
full layers:    rotate the first r = int(d * partial_rotary_factor) = 64 dims with YaRN frequencies (rotate-half, not interleaved);
                multiply q[..., :r] and k[..., :r] by mscale = attention_factor before rotation
                (mscale = 0.1 * ln(factor) + 1 when mscale=1, mscale_all_dim=0; equals the config's attention_factor)
sliding layers: rotate all 128 dims, base 10000, no scaling
attn = SDPA(q, k, v, scale = d^-0.5, mask)          # GQA: n_h / n_kv query heads per kv head
                full layers: causal mask, KVCache
                sliding layers: causal mask with window 512 (key allowed iff 0 <= q_pos - k_pos <= 512 during prefill), RotatingKVCache(max_size = 512)
                if sink is present (sliding layers only): attention sinks per head, as in gpt_oss.rs
o = attn.transpose -> [B, L, n_h * d]
gate = softplus(g_proj(h_in).astype(f32)).astype(o.dtype)            # [B, L, n_h] (per-head)
o = (o.reshape(B, L, n_h, d) * gate[..., None]).reshape(B, L, n_h * d)
x = x + o_proj(o)
m_in = RMSNorm(x; post_attention_layernorm)
x = x + MLP(m_in)

Dense MLP (layer 0): down_proj(silu(gate_proj(m_in)) * up_proj(m_in)), intermediate 8192.

MoE block (layers with mlp_layer_types[i] == "sparse"):

logits = (m_in @ gate.proj.weight.T).astype(f32)                      # [N, E], E = 256
if softcap > 0: logits = tanh(logits / softcap) * softcap
scores = sigmoid(logits)                     if score_func == "sigmoid"
       = sqrt(softplus(logits))              if score_func == "sqrtsoftplus"     # softplus(z) = ln(1 + e^z)
       = softmax(logits, axis=-1)            if moe_router_use_sigmoid == false (legacy)
sel = scores + e_score_correction_bias[E]                               # selection only
idx = argpartition(-sel, k-1)[..., :k]                                  # k = num_experts_per_tok
w = take_along_axis(scores, idx)                                        # weights come from scores WITHOUT the bias
if norm_topk_prob: w = w / sum(w, -1, keepdims)
y = sum_k w_k * SwiGLU_expert_{idx_k}(m_in)                             # expert: down(silu(gate(x)) * up(x)), intermediate 512
y = y * moe_routed_scaling_factor                                       # 2.5
out = y + shared_expert(m_in)                                           # SwiGLU MLP, intermediate 512, unquantized key names mlp.shared_expert.*

Final: logits = lm_head(RMSNorm(x; model.norm)) (lm_head.weight [100352, hidden], untied).

Caches: KVCache on full layers, RotatingKVCache::new(512) on sliding layers (mellum.rs pattern). Full layer 0 is the first entry, so the mask builder can key on layer_types.iter().position(|t| t == "full_attention").

Implementation plan

1. src/models/laguna.rs (new), pub struct ModelArgs via serde with the table above. rope_parameters is serde_json::Value so the per-layer-type dicts and the loose top-level scalars can both be read: layer_rope(layer_type) returns (base, rotated_dims, Option<YarnRope>): reuse crate::models::mellum::compute_yarn_rope(rotated_dims, &entry) (make it pub(crate) if needed) with rotated_dims = (head_dim as f32 * partial) as usize, partial from the entry, then top-level partial_rotary_factor, then 1.0. Apply the YarnRope.mscale to q[..., :rotated_dims] and k[..., :rotated_dims] before mlxcel_core::fast_rope with dims = rotated_dims and freqs = yarn.freqs; the remaining head_dim - rotated_dims features pass through untouched (MLX fast::rope with dims < head_dim already does this).
2. Attention::from_weights(weights, prefix, args, layer_idx): UnifiedLinear for q_proj, k_proj, v_proj, o_proj, g_proj (bias-free), RMSNorm for q_norm / k_norm ([head_dim]). Store n_heads per layer from num_attention_heads_per_layer. Gate: softplus via mlxcel_core::softplus on an f32 cast, multiply per head (reshape [B, L, n_h, d]), per-element when gating is a string other than "per-head". Optional sink (self_attn.sink [n_h]) on sliding layers when swa_attention_sink_enabled; wire it the way gpt_oss.rs passes sinks into SDPA.
3. LagunaSparseMoeBlock: gate: UnifiedLinear (mlp.gate.proj.weight, bf16, never quantized), e_score_correction_bias: UniquePtr<MlxArray> (kept f32), switch_mlp: SwitchGLU (SwitchLinear::from_weights_with_mode with the checkpoint's mode), shared_expert: MLP, routed_scaling_factor, score_func enum {Sigmoid, SqrtSoftplus, Softmax}. Router math as in Expected behavior; reuse crate::models::switch_layers::moe_weighted_sum. Keep MLXCEL_FUSED_MOE=0 semantics: the fused decode kernel is affine-only (switch_layers.rs:802), so NVFP4 experts take the forward path automatically.
4. Sanitize (Laguna::sanitize_weights), run before any from_weights, in this order:
   1. Drop language_model. prefixes if present, drop every key containing rotary_emb.inv_freq, and every key ending in .self_attn.k_scale, .self_attn.v_scale, .input_global_scale.
   2. Router: model.layers.N.mlp.gate.weight -> model.layers.N.mlp.gate.proj.weight; model.layers.N.mlp.experts.e_score_correction_bias -> model.layers.N.mlp.gate.e_score_correction_bias.
   3. compressed-tensors NVFP4 experts (present when quantization_config.format == "nvfp4-pack-quantized" or config_groups.group_0.format says so; record quantization = {group_size: 16, bits: 4, mode: "nvfp4"} for the rest of the loader). For each layer N and projection p in {gate_proj, up_proj, down_proj} and expert e in 0..num_experts, the checkpoint has:
      • model.layers.N.mlp.experts.{e}.{p}.weight_packed: U8 [out, in/2], two E2M1 codes per byte, low nibble first. Reinterpret the bytes as little-endian uint32 [out, in/8] (8 codes per word, code j of the word in bits 4j..4j+4); this is exactly MLX's nvfp4 weight layout, so no nibble shuffling.
      • .weight_scale: F8_E4M3 [out, in/16] (the safetensors dtype is F8_E4M3; keep the raw bytes, do not go through the load-time f16 promotion, or re-encode with f32_to_f8_e4m3 from src/models/sanitize.rs:127 which is lossless for values that came from E4M3).
      • .weight_global_scale: F32 [1] per expert and projection; the dequantized weight is code * E4M3(weight_scale) / weight_global_scale.
        Stack over experts into model.layers.N.mlp.switch_mlp.{p}.weight (uint32 [E, out, in/8]), .scales (uint8 [E, out, in/16]), and a new per-expert sidecar .global_scale (f32 [E], value 1 / weight_global_scale[e]). SwitchLinear gains an optional global_scale: Option<UniquePtr<MlxArray>> that multiplies the gathered output by take(global_scale, indices)[..., None] after gather_qmm (the same role QuantizedWeight::apply_global_scale plays for a dense linear). This keeps every code and block scale bit-exact. Folding 1/global into the E4M3 block scales (re-encoding scale / global to E4M3) is not acceptable here: the quotient is generally not representable in 3 mantissa bits and shifts expert outputs by up to 6 percent.
      • Shared expert: model.layers.N.mlp.shared_expert.{p}.weight_packed/.weight_scale/.weight_global_scale -> UnifiedLinear nvfp4 weight/scales plus QuantizedWeight::apply_global_scale(1 / global).
      • Unquantized tensors (attention projections, g_proj, q_norm/k_norm, layer norms, model.layers.0.mlp.*, embed_tokens, lm_head, mlp.gate.proj) stay bf16 and load through the plain UnifiedLinear::from_weights.
   4. Pre-stacked MLX checkpoints (poolside/Laguna-XS-2.1-NVFP4-mlx, mlx-community/Laguna-XS.2-4bit) already carry model.layers.N.mlp.switch_mlp.{gate_proj|up_proj|down_proj|gate_up_proj}.{weight,scales,biases}; when gate_up_proj is present split it back into gate_proj (rows 0..I) and up_proj (rows I..2I) along axis 1 so one SwitchGLU layout serves both. Also accept per-expert bf16 experts.{e}.{gate_proj,up_proj,down_proj}.weight (the bf16 originals) by stacking.
   5. Pop lm_head.weight when tie_word_embeddings.
      The sanitize must be idempotent (run twice, same map), which the existing sanitize tests assert for every family.
5. Detection: src/models/detection.rs add "laguna" => Ok(ModelType::Laguna); add ModelType::Laguna to src/models/mod.rs (enum, the (description, category) table near mod.rs:910, and the architecture list). Loader: src/loading/mod.rs arm calling LagunaModel::load. LanguageModel impl follows mellum.rs (make_caches returns the mixed KVCache / RotatingKVCache vector through the model-owned sequence state pattern in src/models/model_owned.rs if the trait's Vec<KVCache> cannot carry rotating caches directly; check how mellum.rs does it and copy that exactly).
6. Tokenization contract:
   • src/server/model_provider.rs:1418 and src/commands/generate.rs:810: extend the BOS-prefix suppression to the tokenizer's own BOS string: read bos_token from tokenizer_config.json (here 〈|EOS|〉) and treat a prompt that starts with it like one that starts with <bos>. Implement once in MlxcelTokenizer (fn prompt_carries_bos(&self, prompt: &str) -> bool) and call it from both sites so the scheduler and the dispatch thread keep producing identical ids (issue perf(server): incremental detokenization and per-token streaming overhead reduction #633 invariant).
   • Stop tokens: eos_token_id = [2, 24]; 24 is </assistant>, an ordinary (non-special) added token, so the stop check must run on ids, not on skip_special_tokens decoding.
   • Reasoning: the template emits <think> when enable_thinking is true and </think> otherwise; <think>...</think> is the reasoning span for stream_filter.rs (plain-text markers, no special ids).
   • Do not patch the pre-tokenizer regex. The checkpoint's tokenizer.json pre-tokenizer (Split with [^\r\n\p{L}\p{N}]?\p{L}+|..., then ByteLevel) is what the provider's own Hub loading path and the TensorRT-LLM / vLLM deployments use; the tokenizers crate loads it unchanged.
7. docs/supported-models.md: add Laguna (XS 2.1, XS.2, S 2.1; NVFP4 and MLX 4-bit).

Validation

(a) Unit tests, src/models/laguna_tests.rs:

• layer_rope_full_uses_yarn_partial_and_sliding_uses_default: for the XS 2.1 rope_parameters, full layers yield rotated_dims == 64, mscale within 1e-6 of 1.3465735902799727, sliding layers yield rotated_dims == 128, base 10000, no mscale.
• router_sigmoid_selects_with_bias_but_weights_without_bias: hand-built logits where the bias flips the top-k choice; assert indices follow scores + bias and weights are the bias-free scores renormalized and scaled by 2.5.
• router_sqrtsoftplus_matches_formula: score_func = "sqrtsoftplus" reproduces sqrt(ln(1 + exp(z))) within 1e-6 and stays finite for z = -100.
• per_head_gate_applies_softplus_per_head: a synthetic 2-head layer where g_proj gives logits [0, 10]; output head 0 is scaled by ln 2, head 1 by about 10.
• sanitize_compressed_tensors_nvfp4_stacks_experts_and_keeps_global_scale: build 2 experts with weight_packed U8 [4, 8], weight_scale [4, 1], weight_global_scale; assert switch_mlp.gate_proj.weight is uint32 [2, 4, 2] with the expected little-endian words, .scales bytes are verbatim, .global_scale == [1/g0, 1/g1], input_global_scale and k_scale/v_scale are gone, and running sanitize twice is a no-op.
• sanitize_splits_prestacked_gate_up: a gate_up_proj.weight [E, 2I, in/8] input yields gate_proj / up_proj halves.
• detection_laguna: model_type: "laguna" maps to ModelType::Laguna (add to src/models/detection_tests.rs).
• prompt_carries_bos_recognizes_tokenizer_bos_string in src/tokenizer tests: a prompt starting with 〈|EOS|〉 suppresses add_special.
• Causality: src/models/laguna_tests.rs::prefill_is_causal_on_both_layer_types with a random tiny config and a 96-token prompt: logits at position t must not change when tokens after t change (maskless forward contract; a 6-token prompt cannot expose a bidirectional prefill).

(b) Real checkpoint: poolside/Laguna-XS-2.1-NVFP4 (21.5 GB, native nvfp4, modelsize:medium) downloaded to models/Laguna-XS-2.1-NVFP4, plus mlx-community/Laguna-XS.2-4bit (affine 4-bit) for the pre-stacked path.

./target/release/mlxcel generate -m models/Laguna-XS-2.1-NVFP4 -p "Write a Python retry wrapper with exponential backoff." -n 64
MLXCEL_FUSED_MOE=0 ./target/release/mlxcel generate -m models/Laguna-XS.2-4bit -p "Write a Python retry wrapper with exponential backoff." -n 64
./target/release/mlxcel-server -m models/Laguna-XS-2.1-NVFP4   # then one /v1/chat/completions with enable_thinking=false

Acceptance: greedy output is fluent Python and token-exact for the first 64 tokens against a reference run of the same checkpoint on the same machine with the RAW (untemplated) prompt, with MLXCEL_FUSED_MOE=0; the chat request renders exactly one id 2 at the start of the prompt and stops on </assistant> (id 24). Laguna S 2.1 NVFP4 (poolside/Laguna-S-2.1-NVFP4, about 100 GB, modelsize:large) loads and produces fluent output on a 128 GB or larger machine (finite logits, no NaN; token-exactness not required for S).

Acceptance criteria

• mlxcel list shows Laguna and mlxcel generate runs Laguna XS 2.1 NVFP4 with the command above
• The NVFP4 sanitize keeps every E2M1 code and E4M3 block scale byte-exact and applies the per-expert global scale as a post-matmul multiplier
• Router supports sigmoid and sqrtsoftplus scoring with correction-bias selection, top-k renormalization and moe_routed_scaling_factor
• Sliding layers use RotatingKVCache, full layers use YaRN with partial_rotary_factor and attention_factor
• Chat requests carry a single BOS (id 2) and stop on ids 2 and 24
• docs/supported-models.md updated
• detection table in src/models/detection.rs updated with a test
• cargo test --workspace --profile test-fast --features metal,accelerate passes
• cargo clippy --workspace --all-targets -- -D warnings and cargo fmt --all -- --check pass

Out of scope

• The DFlash speculative drafter (feat(speculative): add the Laguna DFlash drafter (sliding-window gated attention over captured target hidden states) #1351).
• moe_apply_router_weight_on_input = true (no published checkpoint sets it; reject at load).
• Laguna M.1 (poolside/Laguna-M.1, not inspected here).

Dependencies

• Blocked by fix(server): make the chat-template tojson filter match Python json.dumps (kwargs and separators) #1309 (tool-call rendering; text generation without tools does not depend on it).