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

docs/supported-models.md

@@ -51,6 +51,7 @@ Implemented model families include:
 - StarCoder2, StableLM, SmolLM3, Baichuan, MiniCPM, MiniCPM3, MiniMax,
   Ministral3, Nemotron, Nemotron-NAS, Step 3.5, MiMo
 - Mellum / Mellum 2 (`mellum`, JetBrains code model: sliding/full hybrid attention driven by `layer_types`, with QK-RMSNorm and a sparse softmax-routed MoE (`norm_topk_prob`) in every layer. Full-attention layers use YaRN-scaled RoPE plus a standard `KVCache`; sliding-attention layers use default RoPE plus a `RotatingKVCache` windowed to `sliding_window`. Supports tied and untied LM heads. Validated against `JetBrains/Mellum2-12B-A2.5B-Base`.)
+- Laguna XS 2.1 / XS.2 / S 2.1 (`laguna`, Poolside code models: a pre-norm decoder with QK-RMSNorm, a hybrid of `full_attention` layers (YaRN over the first `partial_rotary_factor` share of each head, explicit `attention_factor`, dense `KVCache`) and `sliding_attention` layers (plain RoPE at base 10000, `RotatingKVCache` windowed to `sliding_window` 512), per-layer query-head counts from `num_attention_heads_per_layer`, a per-head softplus attention output gate (`g_proj`, `gating: "per-head"`), an optional per-head attention sink on sliding layers, a dense SwiGLU MLP on layer 0, and elsewhere a sigmoid-routed (or `sqrtsoftplus`) top-k MoE whose expert selection adds `e_score_correction_bias` while the weights stay bias-free, renormalized and scaled by `moe_routed_scaling_factor`, plus an always-on shared expert. The published NVFP4 checkpoints keep the experts and the shared expert in the `compressed-tensors` `nvfp4-pack-quantized` layout (`weight_packed` / `weight_scale` / `weight_global_scale`); the loader reinterprets the packed E2M1 bytes as MLX native NVFP4 words, keeps every E4M3 block scale byte for byte, and applies `1 / weight_global_scale` per expert after the routed matmul (the `SwitchLinear` `global_scale` sidecar) rather than folding it into the 3-mantissa-bit block scales. Pre-stacked MLX conversions (`switch_mlp.*`, a fused `gate_up_proj` is split) and per-expert bf16 exports load through the same path. The chat template emits the tokenizer's own BOS spelling (`〈|EOS|〉`, id 2, which is also an EOS) first, so the tokenize sites suppress `add_special_tokens` for a prompt that starts with the tokenizer's BOS string (`MlxcelTokenizer::prompt_carries_bos`); generation stops on ids 2 and 24 (`</assistant>`, a plain added token), and `<think>` / `</think>` are the reasoning markers. The YaRN block's `attention_factor` is used when the config carries one and derived from `factor` otherwise, which is what `transformers`' `_compute_yarn_parameters` does and therefore what the checkpoint's own `modeling_laguna.py` inherits; mlx-lm's port drops the key and always derives, so `Laguna-XS.2-4bit`, which declares `attention_factor: 1.0` where the derived value would be 1.3466, rotates differently under the two runtimes. XS 2.1 declares the derived value and agrees either way. Not wired for tensor parallel. Validated against `poolside/Laguna-XS-2.1-NVFP4` (native compressed-tensors NVFP4, 234 planes transcoded at load) and `mlx-community/Laguna-XS.2-4bit` (pre-stacked affine, per-layer 8-bit router planes reconciled from the tensor shapes). Laguna S 2.1 was not available on the validation host.)
 - Inkling (`inkling_mm_model` and `inkling`): the text backbone combines hybrid sliding/global NoPE attention, learned banded relative-position bias and log-position temperature, four f32 short-convolution states per layer, dense-first SwiGLU, logsigmoid-softmax routed and shared experts, muP-scaled logits, and padded-vocabulary trimming. Image input is supported through both the CLI `--image` path and OpenAI-compatible server `image_url` content blocks when the checkpoint carries `vision_config` and `model.visual.*` weights. The HMLP tower maps row-major 40x40 RGB tiles duplicated across two temporal slots directly to text-width soft tokens; exact-width images intentionally receive a trailing all-padding tile column, and each rendered image placeholder expands to the number of tiles preprocessing actually produced before ordered feature scatter. Audio input is supported for checkpoints that also carry `audio_config` and `model.audio.*` weights through CLI `--audio` and server `input_audio`, alone or with still images. The bounded host path decodes, downmixes, and resamples WAV input to 16 kHz, produces reference-compatible dMel IDs, scatters image rows first and audio rows second, and keeps every multimodal request on classic prepared-embedding prefill. Video input is supported through CLI `--video` and server `video_url`: sampled frames default to 2 fps, odd sequences repeat their final frame, and at most 16 evenly spaced adjacent pairs are allocated request-wide without crossing clip boundaries. Each clip restarts its timestamp origin, and timestamps use actual probed source-frame time after sampling. Before decoding, the request is capped at eight clips, 32 unique selected frames, and 512 MiB of decoded pixel storage; only selected frames are decoded and retained. The second frame replaces temporal slot 1 only in the suffix of HMLP rows belonging to video tiles. Companion still images remain ahead of complete timestamped Inkling content parts inside the current user turn and retain duplicated temporal planes. Native chained MTP is available for B = 1 with the original `model.mtp.layers.*` tensors for standalone text targets and text-only requests on HMLP-wrapped checkpoints. Speculative decode uses the wrapper's `text` backbone, while image-, audio-, and video-bearing requests stay on the classic prepared-embedding prefill path before continuing through the shared text target. MTP shares the target embedding, final norm, and LM head while restoring and replaying exact KV plus four-convolution state after verification. The loader accepts original bf16/f32, native ModelOpt NVFP4, and pre-converted affine MLX text weights, plus dense or affine-quantized HMLP projections and dense or affine-quantized audio embeddings. Fused kernels, padded batching, and image-feature caching remain separate work. Validation covers f64-referenced dMel extraction and quantization, summed-channel audio embeddings, deterministic HMLP planning/folding, tiling/normalization, placeholder cardinality, image-then-audio scatter, current-user audio insertion with history, clip-local adjacent-frame selection, temporal suffix splicing, pre-decode resource caps, timestamp drift, raw-weight detection and normalization, the synthetic text graph, caches, sanitizer, chat markers, deterministic tiny-model greedy MTP parity, exact snapshot/replay, wrapper dispatch, and prepared-prefill preservation. The public 153.5 GB target, roughly 170.7 GB native NVFP4 checkpoint, and 4.5 GB MTP shard were not available on the validation host, so real image/audio/video answer quality, MTP throughput, and acceptance length remain unmeasured.)
 - TeleChat3 (`telechat3`, TeleAI: structurally a stock Llama decoder (RMSNorm, GQA, SwiGLU, untied `lm_head`), so what matters is the one thing that is not stock. `Telechat3-36B-Thinking` ships `rope_scaling.rope_type = "telechat3-yarn"` with `factor: 4.0`, `original_max_position_embeddings: 8192`, `beta_fast: 32` and `beta_slow: 1`. Upstream routes that into the same `YarnRoPE` it uses for `"yarn"` and `"deepseek_yarn"`: the vendor prefix names the checkpoint that produced the config, not a different algorithm, so the shared YaRN reader accepts all three spellings. This is why the family does not reuse the `llama3` decoder. `llama3::ModelArgs` reads its `rope_scaling` field since #1355, but only for the `default`, `linear` and `llama3` schemes, and it carries no `beta_fast` / `beta_slow`, so routing TeleChat3 through that path would warn once and then rotate at the unscaled base while every tensor kept the right shape; YaRN and default RoPE agree closely at small offsets and only diverge past `original_max_position_embeddings`, so no short-prompt output could expose it. TeleChat3 also keeps `rope_theta` at the top level of `config.json` while the shared reader looks for it inside the scaling block (defaulting to 500000, not this family's 1000000), so the base is injected into the block before the table is built; a table built at the wrong base is finite, correctly shaped and wrong. `attention_bias` puts bias terms on q/k/v/o and is plumbed because the family declares it, but the published 36B checkpoint sets it **false**, so it is not what distinguishes this family in practice. Like every other text family here it satisfies the maskless-prefill contract by routing a multi-token prefill with no caller mask through `causal_attention`. Not wired for tensor parallel, for the same reason it does not reuse the llama3 decoder: the TP runtime parses `config.json` straight into `llama3::ModelArgs`, which has no YaRN branch, so a sharded TeleChat3 would drop the scaling. Validation checkpoint: `mlx-community/Telechat3-36B-Thinking-4bit`.)
 - Apertus (`apertus`, Swiss AI: Llama-style dense transformer with an xIELU activation MLP (no gate), QK-norm, llama3 RoPE scaling, and untied embeddings)

src/bin/bench_decode.rs

@@ -166,7 +166,7 @@ fn load_cli_prompt(
 fn tokenize_prompt(tokenizer: &MlxcelTokenizer, prompt: &str) -> Result<Vec<i32>> {
     // Matches the CLI generate path: chat templates that render a BOS token
     // should not receive a duplicate special token from the tokenizer.
-    let add_special = !prompt.starts_with("<bos>") && !prompt.starts_with("<s>");
+    let add_special = !tokenizer.prompt_carries_bos(prompt);
     let ids = tokenizer
         .encode(prompt, add_special)
         .map_err(|err| anyhow::anyhow!("tokenization failed: {err}"))?;

src/bin/speculative_bench.rs

@@ -342,7 +342,7 @@ fn resolve_model_dir(name: &str) -> PathBuf {
 
 /// Encode a prompt with a tokenizer, mirroring `tests/tensor_parallel_real_models.rs`.
 fn encode_prompt(tokenizer: &MlxcelTokenizer, prompt: &str) -> Vec<i32> {
-    let add_special = !prompt.starts_with("<bos>") && !prompt.starts_with("<s>");
+    let add_special = !tokenizer.prompt_carries_bos(prompt);
     tokenizer
         .encode(prompt, add_special)
         .expect("tokenizer.encode must succeed on a valid utf-8 prompt")

src/cli/context_args.rs

@@ -131,8 +131,8 @@ impl ContextCompatArgs {
             return Err(
                 "--swa-full requests a full-size sliding-window-attention cache, which mlxcel \
                  does not have. Sliding-window families (Gemma 3/4, Exaone 4, gpt-oss, \
-                 RecurrentGemma, Step 3.5, Mellum, Ministral 3, AFMoE, DeepSeek V4) build their \
-                 own ring caches from the checkpoint's sliding_window inside each model's cache \
+                 RecurrentGemma, Step 3.5, Mellum, Laguna, Ministral 3, AFMoE, DeepSeek V4) build \
+                 their own ring caches from the checkpoint's sliding_window inside each model's cache \
                  constructor, and the capability the flag buys in llama-server (KV state \
                  save/restore and context shifting over SWA layers) is gated on those caches \
                  being scheduler-owned, not on their size. Drop the flag (or pass --swa-full \

src/commands/chat.rs

@@ -612,7 +612,7 @@ fn stream_turn<M: LanguageModel>(
     sampling_config: &SamplingConfig,
     show_reasoning: bool,
 ) -> Result<String> {
-    let add_special = !prompt.starts_with("<bos>") && !prompt.starts_with("<s>");
+    let add_special = !tokenizer.prompt_carries_bos(prompt);
     let prompt_tokens: Vec<i32> = tokenizer
         .encode(prompt, add_special)
         .map_err(|e| anyhow!("Tokenization failed: {e}"))?

src/commands/generate.rs

@@ -868,10 +868,11 @@ fn tokenize_prompt(
     tokenizer: &mlxcel::tokenizer::MlxcelTokenizer,
     prompt: &str,
 ) -> Result<Vec<i32>> {
-    // If the prompt already starts with a BOS token string (e.g. from a chat
+    // If the prompt already starts with the BOS token string (e.g. from a chat
     // template that embeds <bos>), skip add_special_tokens to avoid double-BOS.
-    // Matches mlx-lm generate.py behaviour.
-    let add_special = !prompt.starts_with("<bos>") && !prompt.starts_with("<s>");
+    // Matches mlx-lm generate.py behaviour; the rule lives on the tokenizer so
+    // the server paths tokenize identically.
+    let add_special = !tokenizer.prompt_carries_bos(prompt);
     let prompt_token_ids = tokenizer
         .encode(prompt, add_special)
         .map_err(|e| anyhow::anyhow!("Tokenization failed: {}", e))?;

src/distributed/tensor_parallel/inference.rs

@@ -336,6 +336,7 @@ fn fallback_architecture(model_type: ModelType) -> &'static str {
         ModelType::TeleChat3 => "telechat3",
         ModelType::StarCoder2 => "starcoder2",
         ModelType::Mellum => "mellum",
+        ModelType::Laguna => "laguna",
         ModelType::MiniCPM | ModelType::MiniCPMOVLM => "minicpm",
         ModelType::MiniCPM3 => "minicpm3",
         ModelType::StableLM => "stablelm",

src/execution/memory_estimate.rs

@@ -1046,6 +1046,7 @@ fn inspect_registry_id(model_type: ModelType) -> &'static str {
         ModelType::GptNeoX => "gpt_neox",
         ModelType::StarCoder2 => "starcoder2",
         ModelType::Mellum => "mellum",
+        ModelType::Laguna => "laguna",
         ModelType::Helium => "helium",
         ModelType::TeleChat3 => "telechat3",
         ModelType::MiniCPM => "minicpm",

src/lib/mlxcel-core/src/utils.rs

@@ -770,7 +770,7 @@ pub fn create_causal_mask_with_window(
 
 /// Create a sliding-window causal mask sized to the *full* key axis, without
 /// the `min(size + offset, window)` cap applied by [`create_causal_mask_with_window`].
-/// Used by: Gemma 3, Gemma 4 single-pass prefill longer than the sliding window, Cohere2/Gemma3n/Olmo3 dense prefill (#413)
+/// Used by: Gemma 3, Gemma 4 single-pass prefill longer than the sliding window, Cohere2/Gemma3n/Olmo3 dense prefill (#413), Laguna sliding layers
 ///
 /// # Arguments
 /// * `size` - Size of the query sequence
@@ -1014,7 +1014,7 @@ pub fn relu_squared(x: &MlxArray) -> UniquePtr<MlxArray> {
 ///// Numerically stable softplus activation: log(1 + exp(x)).
 /// Uses logaddexp(x, 0) internally to match Python's mx.logaddexp(x, 0).
 /// This avoids float16 overflow for values >= ~11.09 (exp(x) > float16 max).
-/// Used by: Mamba, Mamba2, Jamba, GatedDelta, RecurrentGemma
+/// Used by: Mamba, Mamba2, Jamba, GatedDelta, RecurrentGemma, Laguna
 #[inline]
 pub fn softplus(x: &MlxArray) -> UniquePtr<MlxArray> {
     ffi::softplus(x)

src/loaded_model.rs

@@ -195,6 +195,7 @@ pub enum LoadedModel {
     StarCoder2(models::StarCoder2Model),
     // Mixed full/sliding caches use a wrapper
     Mellum(models::MellumWrapper),
+    Laguna(models::LagunaWrapper),
     MiniCPM(models::MiniCPMModel),
     MiniCPM3(models::MiniCPM3Model),
     StableLM(models::StableLMModel),
@@ -363,6 +364,7 @@ macro_rules! delegate_language_model {
             LoadedModel::TeleChat3(inner) => LanguageModel::$method(inner, $($arg),*),
             LoadedModel::StarCoder2(inner) => LanguageModel::$method(inner, $($arg),*),
             LoadedModel::Mellum(inner) => LanguageModel::$method(inner, $($arg),*),
+            LoadedModel::Laguna(inner) => LanguageModel::$method(inner, $($arg),*),
             LoadedModel::MiniCPM(inner) => LanguageModel::$method(inner, $($arg),*),
             LoadedModel::MiniCPM3(inner) => LanguageModel::$method(inner, $($arg),*),
             LoadedModel::StableLM(inner) => LanguageModel::$method(inner, $($arg),*),

src/model_metadata.rs

@@ -231,6 +231,7 @@ macro_rules! for_each_model_registration {
             TeleChat3 => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::TeleChat3Model::load, args: models::telechat3::ModelArgs, weight_builder: models::TeleChat3Model::from_weights, wrap: LoadedModel::TeleChat3 } };
             StarCoder2 => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::StarCoder2Model::load, args: models::starcoder2::StarCoder2Config, weight_builder: models::StarCoder2Model::from_weights, wrap: LoadedModel::StarCoder2 } };
             Mellum => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::MellumModel::load, args: models::mellum::ModelArgs, weight_builder: models::MellumModel::from_weights, wrap: |m| LoadedModel::Mellum(models::MellumWrapper::new(m)) } };
+            Laguna => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::LagunaModel::load, args: models::laguna::ModelArgs, weight_builder: models::LagunaModel::from_weights, wrap: |m| LoadedModel::Laguna(models::LagunaWrapper::new(m)) } };
             MiniCPM => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::MiniCPMModel::load, args: models::minicpm::ModelArgs, weight_builder: models::MiniCPMModel::from_weights, wrap: LoadedModel::MiniCPM } };
             MiniCPM3 => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::MiniCPM3Model::load, args: models::minicpm3::ModelArgs, weight_builder: models::MiniCPM3Model::from_weights, wrap: LoadedModel::MiniCPM3 } };
             StableLM => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::StableLMModel::load, args: models::stablelm::ModelArgs, weight_builder: models::StableLMModel::from_weights, wrap: LoadedModel::StableLM } };

src/models/detection.rs

@@ -687,6 +687,7 @@ pub fn get_model_type(model_path: &Path) -> Result<ModelType> {
         "telechat3" => Ok(ModelType::TeleChat3),
         "starcoder2" => Ok(ModelType::StarCoder2),
         "mellum" => Ok(ModelType::Mellum),
+        "laguna" => Ok(ModelType::Laguna),
         "minicpm" => Ok(ModelType::MiniCPM),
         "minicpm3" => Ok(ModelType::MiniCPM3),
         "stablelm" => Ok(ModelType::StableLM),

src/models/detection_tests.rs

@@ -626,6 +626,32 @@ fn mellum_model_type_is_detected() {
     fs::remove_dir_all(model_dir).unwrap();
 }
 
+#[test]
+fn laguna_model_type_is_detected() {
+    let model_dir = temp_path("laguna_code");
+    fs::create_dir_all(&model_dir).unwrap();
+    fs::write(
+        model_dir.join("config.json"),
+        r#"{
+            "model_type": "laguna",
+            "architectures": ["LagunaForCausalLM"],
+            "hidden_size": 2048,
+            "head_dim": 128,
+            "num_hidden_layers": 40,
+            "num_attention_heads": 48,
+            "num_key_value_heads": 8,
+            "num_experts": 256,
+            "vocab_size": 100352
+        }"#,
+    )
+    .unwrap();
+
+    let detected = super::detection::get_model_type(&model_dir).unwrap();
+    assert_eq!(detected, ModelType::Laguna);
+
+    fs::remove_dir_all(model_dir).unwrap();
+}
+
 #[test]
 fn gemma4_detection_stays_on_text_route_without_vision_weights() {
     let model_dir = temp_path("gemma4_text_route");