Super-squash branch 'main' using huggingface_hub

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,70 @@
+---
+license: apache-2.0
+base_model: poolside/Laguna-M.1
+base_model_relation: quantized
+library_name: mlx
+pipeline_tag: text-generation
+tags:
+  - mlx
+  - apple-silicon
+  - laguna
+  - laguna-m.1
+  - moe
+  - quantized
+  - conversational
+---
+
+# Laguna-M.1-MLX-Q3
+
+A community **MLX** (Apple Silicon) build of Poolside's
+[**Laguna M.1**](https://huggingface.co/poolside/Laguna-M.1) — a 225B-total /
+23B-active Mixture-of-Experts model — quantized to **3-bit** so it runs locally on
+a 128 GB Mac.
+
+Poolside ships official support for vLLM, SGLang, Transformers and TRT-LLM (all CUDA
+/ datacenter). This is the **Apple Silicon / edge** path.
+
+> Not affiliated with Poolside. Weights © Poolside, Inc., released under Apache 2.0.
+> Converted from the public release `poolside/Laguna-M.1-FP8` (weights verified
+> byte-identical to the public checkpoint; tokenizer, chat template and generation
+> config taken from the public release).
+
+## Use
+
+`laguna` isn't upstream in mlx-lm yet
+([PR #1415](https://github.com/ml-explore/mlx-lm/pull/1415) pending). Until it
+lands, install mlx-lm from the branch that ships the model file **and the
+tool-call parser fix** (needed for tool use in agentic clients), pinned to a commit:
+
+```bash
+pip install "git+https://github.com/eauchs/mlx-lm.git@5b0c0667f4a8c25ee9bb9ef729ab64822d1b246e"
+
+mlx_lm.generate --model ox-ox/Laguna-M.1-MLX-Q3 \
+  --prompt "Write a Python retry wrapper with exponential backoff." \
+  --temp 1.0 --top-k 20
+```
+
+Once the PR is merged, this is just `pip install -U mlx-lm` + the `mlx_lm.generate`
+line — no extra step.
+
+## Performance
+
+| Quant | Hardware | Throughput | Peak RAM |
+|-------|----------|-----------|----------|
+| Q3 | Apple M3 Max 128 GB | ~26.5 tok/s | ~100 GB |
+
+## Architecture
+
+70 layers (3 dense SwiGLU + 67 sparse MoE), 256 experts + 1 shared, top-k=16 sigmoid
+routing; global attention 64 Q / 8 KV heads, head_dim 128, QK-norm + softplus output
+gating; RoPE with YaRN. See the [official model card](https://huggingface.co/poolside/Laguna-M.1)
+for full details and benchmarks.
+
+## License
+
+Apache 2.0 (weights and conversion). © Poolside, Inc. for Laguna M.1.
+
+## Citation
+
+Conversion by Théophile Lafargue ([@eauchs](https://x.com/eauchs)). Model by
+Poolside — see their [technical report](https://huggingface.co/poolside/Laguna-M.1).

chat_template.jinja

@@ -0,0 +1,132 @@
+{#- Iteration on laguna_glm_thinking_v5/chat_template.jinja -#}
+{#- Adds a default system message (used when no system message is provided in `messages`). -#}
+{{- "〈|EOS|〉" -}}
+{%- set enable_thinking = enable_thinking | default(false) -%}
+{%- set render_assistant_messages_raw = render_assistant_messages_raw | default(false) -%}
+{%- set add_generation_prompt = add_generation_prompt | default(false) -%}
+
+{#- ───── header (system message) ───── -#}
+{%- set system_message = "You are a helpful, conversationally-fluent assistant made by Poolside. You are here to be helpful to users through natural language conversations." -%}
+{%- if messages and messages[0].role == "system" -%}
+  {%- set system_message = messages[0].content -%}
+{%- endif -%}
+
+{%- if (system_message and system_message.strip()) or tools -%}
+  {{- "<system>\n" -}}
+
+  {%- if system_message and system_message.strip() -%}
+    {{- "\n" -}}
+    {{- system_message.rstrip() -}}
+  {%- endif -%}
+
+  {%- if tools -%}
+    {{- "\n\n### Tools\n\n" -}}
+    {%- set ns = namespace(tool_string="You may call functions to assist with the user query.\n"
+          ~ "All available function signatures are listed below:\n"
+          ~ "<available_tools>\n") -%}
+    {%- for tool in tools -%}
+      {%- set ns.tool_string = ns.tool_string ~ (tool | tojson) ~ "\n" -%}
+    {%- endfor -%}
+    {%- if enable_thinking -%}
+      {%- set tool_string = ns.tool_string + "</available_tools>\n\n" ~
+            "Wrap your thinking in '<think>', '</think>' tags, followed by a function call. For each function call, return an unescaped XML-like object with function name and arguments within '<tool_call>' and '</tool_call>' tags, like here:\n" ~
+            "<think> your thoughts here </think>\n" ~
+            "<tool_call>function-name\n<arg_key>argument-key</arg_key>\n<arg_value>value-of-argument-key</arg_value>\n" ~
+            "</tool_call>" -%}
+    {%- else -%}
+      {%- set tool_string = ns.tool_string + "</available_tools>\n\n" ~
+            "For each function call, return an unescaped XML-like object " ~
+            "with function name and arguments within '<tool_call>' and '</tool_call>' tags, like here:\n" ~
+            "<tool_call>function-name\n<arg_key>argument-key</arg_key>\n<arg_value>value-of-argument-key</arg_value>\n" ~
+            "</tool_call>" -%}
+    {%- endif -%}
+    {{- tool_string -}}
+  {%- endif -%}
+
+  {{- "\n</system>\n" -}}
+{%- endif -%}
+
+{#- ───── main loop ───── -#}
+{%- for message in messages -%}
+  {%- set content = message.content if message.content is string else "" -%}
+  {%- if message.role == "user" -%}
+    {{- "<user>\n" + content + "\n</user>\n" -}}
+  {%- elif message.role == "assistant" -%}
+    {%- generation -%}
+      {{- "<assistant>\n" -}}
+      {%- if render_assistant_messages_raw -%}
+        {#- Raw mode: prepend the generation prompt token, then dump content verbatim. -#}
+        {#- The generation prompt is <think> when enable_thinking, </think> otherwise. -#}
+        {#- Only prepend if content doesn't already start with it. -#}
+        {%- if enable_thinking -%}
+          {%- if not content.startswith('<think>') -%}
+            {{- '<think>' -}}
+          {%- endif -%}
+        {%- else -%}
+          {%- if not content.startswith('</think>') -%}
+            {{- '</think>' -}}
+          {%- endif -%}
+        {%- endif -%}
+        {{- content -}}
+        {#- Append closing tag if content doesn't already end with it. -#}
+        {%- if not content.endswith('</assistant>\n') and not content.endswith('</assistant>') -%}
+          {{- '\n</assistant>' -}}
+        {%- endif -%}
+        {{- "\n" -}}
+      {%- else -%}
+        {#- Extract reasoning content from message.reasoning (vLLM field name) or message.reasoning_content, or from <think> tags -#}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning is string %}
+          {%- set reasoning_content = message.reasoning %}
+        {%- elif message.reasoning_content is string %}
+          {%- set reasoning_content = message.reasoning_content %}
+        {%- endif %}
+        {#- Always strip <think> tags from content if present to avoid duplication -#}
+        {%- if '</think>' in content %}
+          {%- if not reasoning_content %}
+            {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+          {%- endif %}
+          {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+        {%- endif %}
+        {#- Display reasoning content for all messages -#}
+        {%- if reasoning_content -%}
+          {{- '<think>\n' + reasoning_content.strip() + '\n</think>\n' -}}
+        {%- else -%}
+          {{- '</think>\n' -}}
+        {%- endif -%}
+        {#- Display main content -#}
+        {%- if content.strip() -%}
+          {{- content.strip() ~ "\n" -}}
+        {%- endif -%}
+        {%- if message.tool_calls -%}
+          {%- for tool_call in message.tool_calls -%}
+            {%- set function_data = tool_call.function -%}
+            {{- '<tool_call>' + function_data.name }}
+            {% set _args = function_data.arguments %}
+            {%- for k, v in _args.items() -%}
+              {{- "<arg_key>" ~ k ~ "</arg_key>\n" -}}
+              {{- "<arg_value>"}}{{ v | tojson(ensure_ascii=False) if v is not string else v }}{{ "</arg_value>\n" -}}
+            {%- endfor -%}
+            {{- "</tool_call>\n" -}}
+          {%- endfor -%}
+        {%- endif -%}
+        {{- "</assistant>\n" -}}
+      {%- endif -%}
+    {%- endgeneration -%}
+  {%- elif message.role == "tool" -%}
+    {{- "<tool_response>\n" + content + "\n</tool_response>\n" -}}
+  {%- elif message.role == "system" and loop.index0 != 0 -%}
+    {#- Render additional system messages (skip the first one which is handled separately in the header) -#}
+    {{- "<system>\n" + content + "\n</system>\n" -}}
+  {%- endif -%}
+{%- endfor -%}
+{#- ───── generation prompt ───── -#}
+{%- if add_generation_prompt -%}
+  {{- "<assistant>\n" -}}
+  {#- ───── Include reasoning mode directive ───── -#}
+  {%- if not enable_thinking %}
+    {{- '</think>' -}}
+  {%- else %}
+    {{- '<think>' -}}
+  {%- endif %}
+{%- endif -%}

config.json

@@ -0,0 +1,136 @@
+{
+    "architectures": [
+        "LagunaForCausalLM"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "auto_map": {
+        "AutoConfig": "configuration_laguna.LagunaConfig",
+        "AutoModelForCausalLM": "modeling_laguna.LagunaForCausalLM"
+    },
+    "bos_token_id": 2,
+    "eos_token_id": [
+        2,
+        24
+    ],
+    "gating": "per-element",
+    "head_dim": 128,
+    "hidden_size": 4096,
+    "intermediate_size": 16384,
+    "max_position_embeddings": 131072,
+    "mlp_layer_types": [
+        "dense",
+        "dense",
+        "dense",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse"
+    ],
+    "model_type": "laguna",
+    "moe_apply_router_weight_on_input": false,
+    "moe_intermediate_size": 1024,
+    "moe_routed_scaling_factor": 1.0,
+    "num_attention_heads": 64,
+    "num_experts": 256,
+    "num_experts_per_tok": 16,
+    "num_hidden_layers": 70,
+    "num_key_value_heads": 8,
+    "pad_token_id": 9,
+    "partial_rotary_factor": 1.0,
+    "quantization": {
+        "group_size": 64,
+        "bits": 3
+    },
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+        "full_attention": {
+            "rope_theta": 500000.0,
+            "rope_type": "yarn",
+            "factor": 32.0,
+            "original_max_position_embeddings": 4096,
+            "beta_slow": 1.0,
+            "beta_fast": 64.0,
+            "attention_factor": 1.0
+        }
+    },
+    "rope_scaling": {
+        "rope_type": "yarn",
+        "factor": 32.0,
+        "original_max_position_embeddings": 4096,
+        "beta_slow": 1.0,
+        "beta_fast": 64.0,
+        "attention_factor": 1.0
+    },
+    "rope_theta": 500000.0,
+    "router_aux_loss_coef": 0.0,
+    "shared_expert_intermediate_size": 1024,
+    "tie_word_embeddings": false,
+    "torch_dtype": "bfloat16",
+    "use_bidirectional_attention": false,
+    "use_cache": true,
+    "vocab_size": 100352
+}

configuration_laguna.py

@@ -0,0 +1,186 @@
+# ruff: noqa
+# Copyright 2025 Poolside and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Laguna configuration for transformers 4.56-4.x (used by vLLM).
+
+This uses rope_theta + rope_scaling (legacy format) instead of
+rope_parameters (v5 format).
+"""
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class LagunaConfig(PretrainedConfig):
+    r"""
+    Configuration class for Laguna model.
+
+    Laguna is Poolside's MoE architecture with:
+    - Attention output gating (softplus gate)
+    - Sigmoid routing instead of softmax
+    - No QKV bias
+    - Explicit head_dim parameter
+
+    Args:
+        head_dim (`int`, *optional*, defaults to 128):
+            Dimension of attention heads. Laguna uses explicit head_dim rather than
+            computing it from hidden_size // num_attention_heads.
+        qkv_bias (`bool`, *optional*, defaults to `False`):
+            Whether to add bias to QKV projections. Laguna uses no QKV bias.
+        attention_bias (`bool`, *optional*, defaults to `False`):
+            Whether to add bias to attention output projection. Laguna uses no attention bias.
+        gating (`bool`, *optional*, defaults to `True`):
+            Whether to use softplus output gating on attention. When True, a g_proj linear
+            layer is added and attn_output = attn_output * softplus(g_proj(x)).
+        vocab_size (`int`, *optional*, defaults to 100352):
+            Vocabulary size of the Laguna model.
+        hidden_size (`int`, *optional*, defaults to 2048):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 8192):
+            Dimension of the MLP representations for dense layers.
+        num_hidden_layers (`int`, *optional*, defaults to 48):
+            Number of hidden layers in the Transformer.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads.
+        num_key_value_heads (`int`, *optional*, defaults to 8):
+            Number of key-value heads for GQA.
+        max_position_embeddings (`int`, *optional*, defaults to 4096):
+            Maximum sequence length.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-6):
+            Epsilon for RMSNorm layers.
+        rope_theta (`float`, *optional*, defaults to 500000.0):
+            Base frequency for RoPE embeddings.
+        rope_scaling (`dict`, *optional*):
+            RoPE scaling configuration (e.g. YaRN, linear).
+        sliding_window (`int`, *optional*):
+            Sliding window attention size. Used by layers whose type in ``layer_types``
+            is ``"sliding_attention"``. When ``None``, all layers use full attention.
+        layer_types (`list[str]`, *optional*):
+            Per-layer attention type. Each element should be ``"sliding_attention"`` or
+            ``"global_attention"``. Length must equal ``num_hidden_layers``. When ``None``,
+            all layers default to global attention.
+        swa_attention_sink_enabled (`bool`, *optional*, defaults to `False`):
+            Whether to enable learnable attention sinks on sliding-window attention layers.
+        num_experts (`int`, *optional*, defaults to 256):
+            Number of routed experts.
+        num_experts_per_tok (`int`, *optional*, defaults to 16):
+            Number of experts selected per token (top-k).
+        moe_intermediate_size (`int`, *optional*, defaults to 1024):
+            Intermediate size of routed experts.
+        shared_expert_intermediate_size (`int`, *optional*, defaults to 1024):
+            Intermediate size of the shared expert.
+        norm_topk_prob (`bool`, *optional*, defaults to `True`):
+            Whether to normalize top-k routing probabilities.
+        decoder_sparse_step (`int`, *optional*, defaults to 1):
+            Frequency of MoE layers (1 = every layer is MoE after mlp_only_layers).
+        mlp_only_layers (`list[int]`, *optional*, defaults to `[0]`):
+            Layer indices that use dense MLP instead of MoE.
+        router_aux_loss_coef (`float`, *optional*, defaults to 0.001):
+            Auxiliary loss coefficient for load balancing.
+    """
+
+    model_type = "laguna"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.g_proj": "colwise",  # Laguna-specific gating projection
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    def __init__(
+        self,
+        vocab_size: int = 100352,
+        hidden_size: int = 2048,
+        intermediate_size: int = 8192,
+        num_hidden_layers: int = 48,
+        num_attention_heads: int = 32,
+        num_key_value_heads: int = 8,
+        head_dim: int = 128,
+        qkv_bias: bool = False,
+        attention_bias: bool = False,
+        gating: bool = True,
+        hidden_act: str = "silu",
+        max_position_embeddings: int = 4096,
+        initializer_range: float = 0.02,
+        rms_norm_eps: float = 1e-6,
+        use_cache: bool = True,
+        tie_word_embeddings: bool = False,
+        rope_theta: float = 500000.0,
+        rope_scaling: dict | None = None,
+        attention_dropout: float = 0.0,
+        sliding_window: int | None = None,
+        layer_types: list[str] | None = None,
+        swa_attention_sink_enabled: bool = False,
+        num_experts: int = 256,
+        num_experts_per_tok: int = 16,
+        moe_intermediate_size: int = 1024,
+        shared_expert_intermediate_size: int = 1024,
+        norm_topk_prob: bool = True,
+        decoder_sparse_step: int = 1,
+        mlp_only_layers: list[int] | None = None,
+        router_aux_loss_coef: float = 0.001,
+        output_router_logits: bool = False,
+        **kwargs,
+    ):
+        # Default mlp_only_layers: first layer is dense (moe_first_k_dense_replace=1)
+        if mlp_only_layers is None:
+            mlp_only_layers = [0]
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.qkv_bias = qkv_bias
+        self.attention_bias = attention_bias
+        self.gating = gating
+        self.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.attention_dropout = attention_dropout
+        # Sliding window attention arguments
+        self.sliding_window = sliding_window
+        self.layer_types = layer_types
+        self.swa_attention_sink_enabled = swa_attention_sink_enabled
+        # MoE arguments
+        self.num_experts = num_experts
+        self.num_experts_per_tok = num_experts_per_tok
+        self.moe_intermediate_size = moe_intermediate_size
+        self.shared_expert_intermediate_size = shared_expert_intermediate_size
+        self.norm_topk_prob = norm_topk_prob
+        self.decoder_sparse_step = decoder_sparse_step
+        self.mlp_only_layers = mlp_only_layers
+        self.router_aux_loss_coef = router_aux_loss_coef
+        self.output_router_logits = output_router_logits
+
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+
+
+__all__ = ["LagunaConfig"]

generation_config.json

@@ -0,0 +1,13 @@
+{
+  "bos_token_id": 2,
+  "do_sample": true,
+  "eos_token_id": [
+    2,
+    24
+  ],
+  "max_new_tokens": 4096,
+  "pad_token_id": 9,
+  "temperature": 1.0,
+  "top_p": 1.0,
+  "min_p": 0.0
+}

laguna.py

@@ -0,0 +1,357 @@
+# Copyright © 2025 Poolside AI / MLX port by Théophile Lafargue (@eauchs, x.com/eauchs)
+#
+# Laguna M.1 — MLX implementation
+# Architecture: MoE (256 experts, top-16, sigmoid routing)
+# Attention: GQA with QK-norm + softplus output gating (g_proj)
+# RoPE: YaRN (factor=32, original_max=4096, beta_fast=64, beta_slow=1)
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str = "laguna"
+    vocab_size: int = 100352
+    hidden_size: int = 4096
+    intermediate_size: int = 16384
+    num_hidden_layers: int = 70
+    num_attention_heads: int = 64
+    num_key_value_heads: int = 8
+    head_dim: int = 128
+    max_position_embeddings: int = 131072
+    rms_norm_eps: float = 1e-6
+    rope_theta: float = 500000.0
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    attention_bias: bool = False
+    attention_dropout: float = 0.0
+    tie_word_embeddings: bool = False
+    # MoE
+    num_experts: int = 256
+    num_experts_per_tok: int = 16
+    moe_intermediate_size: int = 1024
+    shared_expert_intermediate_size: int = 1024
+    norm_topk_prob: bool = True
+    moe_routed_scaling_factor: float = 1.0
+    # Layer types: "dense" or "sparse" per layer
+    mlp_layer_types: Optional[List[str]] = None
+    # Gating: "per-element" enables softplus output gating on attention
+    gating: Optional[str] = "per-element"
+
+
+class LagunaAttention(nn.Module):
+    """Laguna attention: GQA + QK-norm + softplus output gating."""
+
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        dim = args.hidden_size
+        self.n_heads = args.num_attention_heads
+        self.n_kv_heads = args.num_key_value_heads
+        self.head_dim = args.head_dim
+        self.scale = self.head_dim ** -0.5
+
+        self.q_proj = nn.Linear(dim, self.n_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.n_heads * self.head_dim, dim, bias=False)
+
+        # Laguna-specific: output gating projection
+        self.g_proj = nn.Linear(dim, self.n_heads * self.head_dim, bias=False)
+
+        # QK normalization (per-head RMSNorm before RoPE)
+        self.q_norm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
+        self.k_norm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
+
+        self.rope = initialize_rope(
+            dims=self.head_dim,
+            base=args.rope_theta,
+            traditional=False,
+            max_position_embeddings=args.max_position_embeddings,
+            scaling_config=args.rope_scaling,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries = self.q_proj(x)
+        keys = self.k_proj(x)
+        values = self.v_proj(x)
+
+        # Reshape to [B, n_heads, L, head_dim]
+        queries = self.q_norm(
+            queries.reshape(B, L, self.n_heads, -1)
+        ).transpose(0, 2, 1, 3)
+        keys = self.k_norm(
+            keys.reshape(B, L, self.n_kv_heads, -1)
+        ).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        # RoPE
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        # Attention
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+
+        # Laguna-specific: softplus gating BEFORE o_proj
+        gate = nn.softplus(self.g_proj(x))
+        output = output * gate
+
+        return self.o_proj(output)
+
+
+class LagunaMLP(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class LagunaRouter(nn.Module):
+    """Sigmoid-based top-k router (not softmax)."""
+
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.top_k = args.num_experts_per_tok
+        self.num_experts = args.num_experts
+        self.norm_topk_prob = args.norm_topk_prob
+        self.weight = mx.zeros((self.num_experts, args.hidden_size))
+        self.e_score_correction_bias = mx.zeros((self.num_experts,))
+
+    def __call__(self, x: mx.array):
+        # Sigmoid scoring in float32 for precision
+        router_logits = x @ self.weight.T
+        scores = mx.sigmoid(router_logits.astype(mx.float32))
+
+        # Apply correction bias for expert selection (not for final weights)
+        corrected_scores = scores + self.e_score_correction_bias
+
+        # Top-k selection on corrected scores
+        k = self.top_k
+        inds = mx.argpartition(-corrected_scores, kth=k - 1, axis=-1)[..., :k]
+        # Use original scores (without bias) for weighting
+        routing_weights = mx.take_along_axis(scores, inds, axis=-1)
+
+        if self.norm_topk_prob:
+            routing_weights = routing_weights / mx.sum(
+                routing_weights, axis=-1, keepdims=True
+            )
+
+        return inds, routing_weights.astype(x.dtype)
+
+
+class LagunaSparseMoeBlock(nn.Module):
+    """Laguna MoE: sigmoid router + SwitchGLU experts + shared expert."""
+
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.gate = LagunaRouter(args)
+        self.switch_mlp = SwitchGLU(
+            args.hidden_size,
+            args.moe_intermediate_size,
+            args.num_experts,
+        )
+        self.shared_expert = LagunaMLP(
+            args.hidden_size, args.shared_expert_intermediate_size
+        )
+
+    def __call__(self, x: mx.array) -> mx.array:
+        inds, scores = self.gate(x)
+
+        # Routed experts via SwitchGLU
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
+
+        # Add shared expert
+        y = y + self.shared_expert(x)
+
+        return y
+
+
+class LagunaDecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.self_attn = LagunaAttention(args)
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+
+        # Dense or sparse MLP based on mlp_layer_types
+        is_sparse = (
+            args.mlp_layer_types is not None
+            and layer_idx < len(args.mlp_layer_types)
+            and args.mlp_layer_types[layer_idx] == "sparse"
+        )
+        if is_sparse:
+            self.mlp = LagunaSparseMoeBlock(args)
+        else:
+            self.mlp = LagunaMLP(args.hidden_size, args.intermediate_size)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        return h + r
+
+
+class LagunaModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            LagunaDecoderLayer(args, i) for i in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+    ) -> mx.array:
+        h = self.embed_tokens(inputs)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        mask = create_attention_mask(h, cache[0])
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = LagunaModel(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+    ):
+        out = self.model(inputs, cache)
+        if self.args.tie_word_embeddings:
+            return self.model.embed_tokens.as_linear(out)
+        return self.lm_head(out)
+
+    def sanitize(self, weights):
+        # Dequantize FP8 compressed-tensors: uint8 + weight_scale → bfloat16
+        new_weights = {}
+        for k, v in weights.items():
+            if k.endswith("weight_scale"):
+                continue  # handled with corresponding weight
+            if k.endswith(".weight") and f"{k}_scale" in weights:
+                scale = weights[f"{k}_scale"]
+                new_weights[k] = self._dequant_fp8_block(v, scale)
+            else:
+                new_weights[k] = v
+        weights = new_weights
+
+        # Remap e_score_correction_bias:
+        # HF: model.layers.X.mlp.experts.e_score_correction_bias
+        # MLX: model.layers.X.mlp.gate.e_score_correction_bias
+        remapped = {}
+        for k, v in weights.items():
+            if "mlp.experts.e_score_correction_bias" in k:
+                new_key = k.replace(
+                    "mlp.experts.e_score_correction_bias",
+                    "mlp.gate.e_score_correction_bias",
+                )
+                remapped[new_key] = v
+            else:
+                remapped[k] = v
+        weights = remapped
+
+        # Stack per-expert weights into SwitchGLU format
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}"
+            for n in ["up_proj", "down_proj", "gate_proj"]:
+                key0 = f"{prefix}.mlp.experts.0.{n}.weight"
+                if key0 in weights:
+                    to_join = [
+                        weights.pop(f"{prefix}.mlp.experts.{e}.{n}.weight")
+                        for e in range(self.args.num_experts)
+                    ]
+                    weights[f"{prefix}.mlp.switch_mlp.{n}.weight"] = mx.stack(
+                        to_join
+                    )
+
+        # Remove unused keys
+        weights = {
+            k: v
+            for k, v in weights.items()
+            if "rotary_emb.inv_freq" not in k
+        }
+
+        return weights
+
+    @staticmethod
+    def _dequant_fp8_block(weight: mx.array, scale: mx.array) -> mx.array:
+        """Dequantize FP8 block-quantized weight (uint8 + scale → bfloat16).
+
+        Block size is 128x128 (inferred from scale shape vs weight shape).
+        """
+        m, n = weight.shape
+        sm, sn = scale.shape
+        bs_m = m // sm  # block size along rows
+        bs_n = n // sn  # block size along cols
+
+        # Reinterpret uint8 as float8 E4M3, then cast to bfloat16
+        weight_bf16 = mx.from_fp8(weight, mx.bfloat16)
+
+        # Apply block-wise scale
+        weight_bf16 = weight_bf16.reshape(sm, bs_m, sn, bs_n)
+        weight_bf16 = weight_bf16 * scale[:, None, :, None]
+        weight_bf16 = weight_bf16.reshape(m, n)
+
+        return weight_bf16
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            # Don't quantize router gate weights — keep full precision
+            if "mlp.gate.weight" in path:
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate

model-00000.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:289d02fc06a716236c5574faf8299a2b57c6c35e131dd8cdcc33b318cb89081f
+size 1644175670

model-00001.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6883ec5004f6a0f891892d9743eab15839dabb87219baa4cd2bf3a0278e1a88f
+size 135810626

model-00002.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:67fe2b2e953b88ac2e677444bb7d98bedde977c359f72bcb2b0f1a9158db3b77
+size 135810606

model-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e7172422290ce2d10224789d100547342746164213a2a12eb8855fe54cd6e29f
+size 135810598

model-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f59123e6a5be39b3e22767acde50728331381949ceaa309f408adc578ef8e5ee
+size 1464621066

model-00005.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b1249787a894391e3c8003106c7fa4412c9fef96bda364aa322c763baa5adacc
+size 1464621118

model-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7896e367cfcbf371fbade5332080136a67f3fbb7a48c20ed6ce175281bad3805
+size 1464621074

model-00007.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c9edd2ed48a70e3f5f33e8ecdb16cd6531a7930ff56b358ecaa641152f6bef9c
+size 1464621052

model-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7d99308fee141b7dec27bcfe9ccb8e0b464ad28e567f6eb3864963e9c50a092c
+size 1464621064

model-00009.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1de2f6b54f7d28bc23db7294870e8241e6ba1a5becddf87c3f87ccf4aa526591
+size 1464621064

model-00010.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dfcf2f8a7c8bee7a00c25b06c6a27264237dc348be585bf488f7ea14c2d743b4
+size 1464621072

model-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:43fddd5e9a40e7c55fd409edb407500dc15c8405dfc5205b9479247ffb6b00e0
+size 1464621115

model-00012.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1fea6040cdcfcc3c3f1aaac94fedf55c958e4ed21abda9af6cb8f631b43396d0
+size 1464621113

model-00013.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7495e941f0617de41038b69afc0f3b45debc8abd28d6fec0ae0c2a0244bc25ab
+size 1464621115

model-00014.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f9eaa0c634656292f68eda91232860071898f4bf3fa2458a13997fc00e74cfc3
+size 1464621147

model-00015.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aaac5d0d89dd4912a2656719ad408259f0514a7e127c60fe90b8cb88b0ca400f
+size 1464621119

model-00016.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5c20a269e34555331131e1d7314689f121cd39f8086d89dce73f080b285338ce
+size 1464621099

model-00017.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:777721efce68f71f4c9b92d8aaab9e4c42f9cda4bfa2ef232cd8e14aa688c94e
+size 1464621151

model-00018.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cf6944fd2bbe141b7acf70bf66764d251b32de1d936d72bc49a357197d57bb0f
+size 1464621127

model-00019.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7084ee3642b757eeb6ad376863c54c5df086c35c5ec88329a3affc11d4c1f86c
+size 1464621111

model-00020.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:364cc3f94ac63ee413ec9756f73358f8a3c065f2d411d555ad39eb92a4f1d9b6
+size 1464621131

model-00021.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b90c7a8cf19dc5b816e5c318480db5b17ea7a34361ae8cc82a8cafdb956edc3f
+size 1464621139

model-00022.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c1b86cfdc1b8d8959fbae7718a65f218bacf89839ad2e5073f4fa8bb0ee49e40
+size 1464621099

model-00023.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ce85a6ddc63df580dd08e3fba629140e8b7e3f644894ba7593de2b42cb768a59
+size 1464621101

model-00024.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:af66c4dfaeb860604f8eee4fea2d72aaa2a71e2a3fc54641dd42d489ff380e32
+size 1464621115

model-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:943998d5af4c439940888bc0b3ddccd24922f953d4de70833054aca993323cd6
+size 1464621123

model-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a5da51efb166cb877fbd0a0427726baa266dca616f67d6acbe996f1a874b4a68
+size 1464621149

model-00027.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a1e66e74fedb60fbef11c57620030b800a7faab81f77e15895e1ebb21115b4e1
+size 1464621125

model-00028.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:900189d65c3cbe5b0d21c881717c238fef028650e9d97c86ce6558f457575e90
+size 1464621129

model-00029.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ad3e5586b16a345c02887d661c072b38c18c03a64fe0affe93f93ea9e257fff6
+size 1464621105

model-00030.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e635491b95e8171efb622c08d1bf874f42a65c5c7e30907245f8b81e27d1ae43
+size 1464621121

model-00031.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:79c104139be06236652966fd043c89c27e33977ec7a3eb7dc6d9fb45271ad23e
+size 1464621097

model-00032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1b206c21e5fcda44424fcd98132e5f1e2209269d120c69ee126ebf18d101549b
+size 1464621121

model-00033.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7eb82761a82a1c91c2d8c96b49f428a314c504dc94c42c131c3cbefd4bd07113
+size 1464621117

model-00034.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2e379d82236fcead29d1c1bb57d7f2b9e32705b512577af812fa8636ea978e77
+size 1464621099

model-00035.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bdcdd537fd1f3ca4f6dcbbe152d0492ff6004ac50998518064ab5e521b7ec2be
+size 1464621109

model-00036.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5a675b978c670dd93f7b3fc7c6c028d5c8ae6b4c0917e787d5fd5812e587ceea
+size 1464621147

model-00037.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8878d9a2538e1957863b141d850cc5d490ea2848ab9a0a1ffbd53eac96e5bb56
+size 1464621105

model-00038.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e403b45c3e106108a15f0bf1bdadb35a7f1e1da34dc99daffb03dcbd8c3c8cdc
+size 1464621095

model-00039.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cf0fa1750fc082162fed306a88689f601e2c4785afcf91ecec92c51a41c25c6d
+size 1464621097

model-00040.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d61d869ae3c6c095af4397b186bf9b233b3a945a8f2b3b29a4e9e96943215cf8
+size 1464621141

model-00041.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:70c255f87050008a1e722201c12c20f707ff9d7b40fc33519d72ec4feb23eee5
+size 1464621093

model-00042.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:150ec5af9a70cf8d53c87ec8262c158c457191a19a4b070878c6c0f7ed6a2b9a
+size 1464621105