Laguna S 2.1 initial release

Co-authored-by: varunrandery <varunrandery@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,36 @@
+*.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
+model.safetensors.index.json filter=lfs diff=lfs merge=lfs -text

LICENSE.md

@@ -0,0 +1,49 @@
+OpenMDW License Agreement, version 1.1 (OpenMDW-1.1)
+
+By exercising rights granted to you under this agreement, you accept and agree
+to its terms.
+
+As used in this agreement, "Model Materials" means the materials provided to
+you under this agreement, consisting of: (1) one or more machine learning
+models (including architecture and parameters); and (2) all related artifacts
+(including associated data, documentation and software) that are provided to
+you hereunder.
+
+Subject to your compliance with this agreement, permission is hereby granted,
+free of charge, to deal in the Model Materials without restriction, including
+under all copyright, patent, database, and trade secret rights included or
+embodied therein.
+
+If you distribute any portion of the Model Materials, you shall retain in your
+distribution (1) a copy of this agreement, and (2) all copyright notices and
+other notices of origin included in the Model Materials that are applicable to
+your distribution.
+
+If you file, maintain, or voluntarily participate in a lawsuit against any
+person or entity asserting that the Model Materials directly or indirectly
+infringe any patent or copyright, then all rights and grants made to you
+hereunder are terminated, unless that lawsuit was in response to a
+corresponding lawsuit first brought against you.
+
+This agreement does not impose any restrictions or obligations with respect to
+any use, modification, or sharing of any outputs generated by using the Model
+Materials.
+
+THE MODEL MATERIALS ARE PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE, TITLE, NONINFRINGEMENT, ACCURACY, OR THE
+ABSENCE OF LATENT OR OTHER DEFECTS OR ERRORS, WHETHER OR NOT DISCOVERABLE, ALL
+TO THE GREATEST EXTENT PERMISSIBLE UNDER APPLICABLE LAW.
+
+YOU ARE SOLELY RESPONSIBLE FOR (1) CLEARING RIGHTS OF OTHER PERSONS THAT MAY
+APPLY TO THE MODEL MATERIALS OR ANY USE THEREOF, INCLUDING WITHOUT LIMITATION
+ANY PERSON'S COPYRIGHTS OR OTHER RIGHTS INCLUDED OR EMBODIED IN THE MODEL
+MATERIALS; (2) OBTAINING ANY NECESSARY CONSENTS, PERMISSIONS OR OTHER RIGHTS
+REQUIRED FOR ANY USE OF THE MODEL MATERIALS; OR (3) PERFORMING ANY DUE
+DILIGENCE OR UNDERTAKING ANY OTHER INVESTIGATIONS INTO THE MODEL MATERIALS OR
+ANYTHING INCORPORATED OR EMBODIED THEREIN.
+
+IN NO EVENT SHALL THE PROVIDERS OF THE MODEL MATERIALS BE LIABLE FOR ANY CLAIM,
+DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
+OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE MODEL MATERIALS, THE
+USE THEREOF OR OTHER DEALINGS THEREIN.

README.md

@@ -0,0 +1,266 @@
+---
+library_name: vllm
+inference: false
+extra_gated_description: >-
+  To learn more about how we process your personal data, please read our <a
+  href="https://poolside.ai/legal/privacy">Privacy Policy</a>.
+tags:
+- laguna-s-2.1
+license: openmdw-1.1
+pipeline_tag: text-generation
+base_model:
+- poolside/Laguna-S-2.1
+---
+
+<p align="center">
+  <img alt="poolside-banner" src="https://poolside.ai/assets/laguna/laguna-s-2-1-banner.svg" width="800px">
+</p>
+
+<p align="center">
+  <a href="https://openrouter.ai/poolside/laguna-s-2.1"><strong>Use on OpenRouter</strong></a> ·
+  <a href="https://vercel.com/ai-gateway/models/laguna-s-2.1"><strong>Use on Vercel AI Gateway</strong></a> ·
+  <a href="https://poolside.ai/blog/introducing-laguna-s-2-1"><strong>Release blog post</strong></a>
+</p>
+
+<br>
+
+# Laguna S 2.1-NVFP4
+Laguna S 2.1-NVFP4 is a 117.6B total parameter Mixture-of-Experts model with 8.5B activated parameters per token designed for agentic coding and long-horizon work on a local machine. It uses Sliding Window Attention with per-head gating in 36 out of 48 layers for fast inference and low KV cache requirements.
+
+
+## Highlights
+- **Mixed SWA and global attention layout**: Laguna S 2.1 uses softplus gating with per-layer rotary scales, enabling mixed SWA (Sliding Window Attention) and global attention layers in a 3:1 ratio (across 48 total layers)
+- **KV cache in FP8**: KV cache quantized to FP8, reducing memory per token
+- **Native reasoning support**: Interleaved thinking between tool calls with support for enabling and disabling thinking per-request
+- **Local-ready**: At 117.6B total parameters and 8.5B activated, the NVFP4 weights are roughly 71 GB. Available on [Ollama](https://ollama.com/laguna-s-2.1) and [llama.cpp](https://github.com/ggml-org/llama.cpp/pull/25165) (BF16 and Q4\_K\_M only)
+- **OpenMDW-1.1 license**: Use and modify the model and associated materials freely for commercial and non-commercial purposes ([learn more about OpenMDW](https://openmdw.ai/))
+
+---
+
+## Model overview
+
+- Training: pre-training, post-training and reinforcement learning stages
+- Number of parameters: 117.6B total with 8.5B activated per token
+- Optimizer: Muon
+- Layers: 48 layers (12 layers with global attention, 36 layers with sliding window attention)
+- Experts: 256 experts with 1 shared expert
+- Sliding Window: 512 tokens
+- Modality: text-to-text
+- Context window: 262,144 tokens
+- Reasoning support: interleaved thinking with preserved thinking
+
+
+## Benchmark results
+
+<p align="center">
+  <img alt="benchmarks" src="https://poolside.ai/assets/laguna/laguna-s-2-1-chart.svg" width="800px">
+</p>
+
+| Model | Size | Terminal-Bench 2.1 | SWE-bench Multilingual | SWE-Bench Pro (Public Dataset) | DeepSWE | SWE Atlas (Codebase QnA) | Toolathlon Verified |
+|---|---|---|---|---|---|---|---|
+| **Laguna S 2.1** | 118B-A8B | **70.2%** | **78.5%** | **59.4%** | **40.4%** | **46.2%** | **49.7%** |
+| Tencent Hy3 | 295B-A21B | 71.7% | 75.8% | 57.9% | - | - | - |
+| Inkling | 975B-A41B | 63.8% | - | 54.3% | - | - | 45.5%* |
+| Nemotron 3 Ultra | 550B-A55B | 56.4% | 67.7% | - | - | - | 34.3%* |
+| DeepSeek-V4-Pro Max | 1.6T-A49B | 64.0%* | 76.2% | 55.4% | 9.0%* | 27.2%* | 55.9%* |
+| Kimi K3 | 2800B-A50B | 88.3% | - | - | 69% | - | - |
+| Qwen 3.7 Max | - | 74.5%* | 78.3% | 60.6% | - | - | - |
+| Muse Spark 1.1 | - | 80% | - | 61.5% | 53.3% | 42.2%* | 75.6% |
+| Claude Fable 5 | - | 88% | - | 80.3% | 70% | - | - |
+
+Benchmarks as of 21 July 2026. Laguna S 2.1 in **bold**; a dash (-) marks a benchmark a model was not evaluated on. Scores marked * are as reported by third parties: Terminal-Bench 2.1 and DeepSWE via Artificial Analysis, SWE Atlas via Scale AI's official leaderboard, and Toolathlon Verified via its official leaderboard. Full evaluation trajectories: [trajectories.poolside.ai](https://trajectories.poolside.ai).
+
+## Usage
+
+### Context length
+
+This checkpoint ships configured for a 262,144-token (256K) context window. This is the configuration we recommend for best output quality.
+
+The weights are native 1M checkpoints: training included a long-context extension stage up to 1,048,576 tokens, and quantization was calibrated at the 1M configuration. If you need more than 256K of context, restore the 1M configuration by editing `config.json`:
+
+```json
+"rope_parameters": {
+  "full_attention": {
+    "factor": 128.0,
+    "attention_factor": 1.4852030263919618
+  }
+},
+"max_position_embeddings": 1048576
+```
+
+You may experience quality degradation with the 1M configuration. If you use it, we recommend sampling with `temperature <= 0.7` and `top_p <= 0.95`.
+
+### Recommended sampling
+
+For the best balance of quality and reliability we recommend sampling with `temperature 0.7` and `top_p 0.95`.
+
+
+### Local deployment
+
+Laguna S 2.1-NVFP4 is supported in vLLM, SGLang and Transformers, and TRT-LLM thanks to the support of the team at NVIDIA. Use Laguna-S 2.1 with Ollama (with MLX support) or Llama.cpp (BF16 and Q4\_K\_M only) for the best results on your local machine.
+
+#### vLLM
+
+> Serving requires vLLM 0.25.0 or later. See the main [Laguna S 2.1 model card](https://huggingface.co/poolside/Laguna-S-2.1) for the full recipe.
+
+The full vLLM recipe is on the main [Laguna S 2.1 model card](https://huggingface.co/poolside/Laguna-S-2.1) and on the [vLLM recipes page](https://recipes.vllm.ai/poolside/Laguna-S-2.1). Quantization is detected automatically from `quantization_config` in this checkpoint, so the same command works with `poolside/Laguna-S-2.1-NVFP4` substituted for the model ID. No extra flags required.
+
+> [!NOTE]
+> **Optional: speculative decoding with DFlash.** Pair with the quantization-matched draft model [poolside/Laguna-S-2.1-DFlash-NVFP4](https://huggingface.co/poolside/Laguna-S-2.1-DFlash-NVFP4) by adding `--speculative-config '{"model":"poolside/Laguna-S-2.1-DFlash-NVFP4","num_speculative_tokens":15,"method":"dflash"}'` to the serve command.
+
+#### DGX Spark & Mac Studio
+
+Ollama is the quickest way to run Laguna S 2.1 on a single 128 GB box, whether
+that is an NVIDIA DGX Spark (GB10) or an Apple Silicon Mac Studio:
+
+```bash
+ollama run laguna-s-2.1
+```
+
+It pulls the Q4\_K\_M GGUF (about 75 GB) and runs on the GB10 GPU on the Spark
+or on Metal on the Mac, with no build step. We measured about 12.6 tokens/s on
+the Spark and 17.6 on Apple Silicon.
+
+- You need about 128 GB of unified memory; the Q4\_K\_M weights are around
+  75 GB on their own.
+- The first load reads 75 GB off disk and can run past Ollama's 5-minute
+  default timeout. If it does, set `OLLAMA_LOAD_TIMEOUT=20m`.
+- Reasoning is on by default. Pass `"think": false` in the request to turn it off.
+- The higher-fidelity tags need a bigger box: `laguna-s-2.1:q8_0` is about
+  128 GB (so a 192 GB Mac Studio) and `laguna-s-2.1:f16` is about 235 GB (256 GB
+  or more). On the Spark's 128 GB, stick to Q4\_K\_M. Apple Silicon can also
+  run it through MLX.
+- Ollama is built on llama.cpp. To build the runtime yourself, poolside's
+  [`laguna` branch of llama.cpp](https://github.com/poolsideai/llama.cpp/tree/laguna)
+  serves the same GGUFs.
+
+**Maximum performance on the DGX Spark (native NVFP4 + DFlash)**
+
+For the fastest setup on the Spark, serve this checkpoint with vLLM instead. That
+gives you the native NVFP4 kernels and DFlash speculative decoding.
+
+*One-time setup:*
+
+```bash
+# Python headers: Triton JIT needs them and DGX OS ships without them.
+# Without this, the first server start dies in triton/runtime/build.py.
+sudo apt install -y python3.12-dev
+
+curl -LsSf https://astral.sh/uv/install.sh | sh
+uv venv ~/venvs/vllm025 -p 3.12
+
+# vLLM 0.25.1 with CUDA-13 torch (aarch64 wheels are on PyPI)
+uv pip install -p ~/venvs/vllm025 vllm==0.25.1 --torch-backend=cu130
+
+# FlashInfer nightly trio: without flashinfer-python the NVFP4 path is not
+# native; the jit-cache wheel avoids most first-start JIT compilation.
+uv pip install -p ~/venvs/vllm025 \
+  "flashinfer-python==0.6.15.dev20260712" \
+  "flashinfer-cubin==0.6.15.dev20260712" \
+  "flashinfer-jit-cache==0.6.15.dev20260712" \
+  --extra-index-url https://flashinfer.ai/whl/nightly/ \
+  --extra-index-url https://flashinfer.ai/whl/nightly/cu130/ \
+  --index-strategy unsafe-best-match
+
+hf download poolside/Laguna-S-2.1-NVFP4
+hf download poolside/Laguna-S-2.1-DFlash-NVFP4   # 73 GB total
+```
+
+*Serve:*
+
+```bash
+export CUTE_DSL_ARCH=sm_121a          # arch string for FP4 kernel JIT
+export PATH=/usr/local/cuda/bin:$PATH # nvcc for JIT
+export MAX_JOBS=4                     # cap JIT fan-out; see warning below
+source ~/venvs/vllm025/bin/activate
+
+vllm serve poolside/Laguna-S-2.1-NVFP4 \
+  --speculative-config '{"model":"poolside/Laguna-S-2.1-DFlash-NVFP4","num_speculative_tokens":15}' \
+  --enable-auto-tool-choice \
+  --tool-call-parser poolside_v1 \
+  --reasoning-parser poolside_v1 \
+  --override-generation-config '{"temperature":0.7,"top_p":0.95}' \
+  --max-num-seqs 32 \
+  --max-model-len 262144 \
+  --gpu-memory-utilization 0.85 \
+  --host 0.0.0.0 --port 8000
+```
+
+- You do not need backend flags: auto-selection picks FlashInferCutlass, which
+  runs natively on `sm_121`. Do not set `--linear-backend flashinfer_b12x` on
+  0.25.1; the opt-in is broken there and it is slower anyway.
+- Keep the `--override-generation-config`. Many clients send no sampling
+  parameters, and the raw defaults degrade output on the NVFP4 quantization. The
+  model's `generation_config.json` sets `top_k 20` (eval-certified truncation),
+  which handles that. Do not add `min_p`: vLLM rejects `min_p` and `logit_bias`
+  under speculative decoding, so putting it in the defaults returns a 400 on
+  every sampled request.
+- `--max-num-seqs 32` is required: DFlash crashes vLLM at the default of 256.
+- The first start takes about 15 minutes (weight load from NVMe, JIT, and graph
+  capture).
+
+> [!WARNING]
+> Never drop `MAX_JOBS=4` on a cold `~/.cache/flashinfer`. An uncapped nvcc
+> fan-out can exhaust the 128 GB unified memory and take the whole machine down.
+> A warm cache makes it a no-op, but the cache is cold again after any config or
+> shape change.
+
+Prefill runs 600-800 tok/s, decode is around 15 tok/s on prose and 22-24 on
+code, and DFlash accepts 2.9-3.1 tokens per step. At the 256K setting you get
+830-870K KV tokens. Two requests at once roughly double total throughput.
+Without speculation, decode sits at 13-14 tok/s on every engine we tried on the
+GB10, which is the memory-bandwidth ceiling for this model. Speculation is how
+you get past it.
+
+*pool CLI against the box:*
+
+```bash
+POOLSIDE_STANDALONE_BASE_URL=http://<spark-ip>:8000/v1 \
+POOLSIDE_STANDALONE_MODEL=poolside/Laguna-S-2.1-NVFP4 \
+POOLSIDE_STANDALONE_CONTEXT_LENGTH=262144 \
+POOLSIDE_API_KEY=dummy pool
+```
+
+Use the box's IPv4 address rather than the `.local` mDNS name: Go's resolver can
+pick the link-local IPv6 and fail with "no route to host". On macOS the terminal
+app needs Local Network permission, and the system curl is exempt from that
+check, so if curl works but pool does not, it is the permission and not the
+network.
+
+#### SGLang
+
+Laguna S 2.1 is supported in SGLang via [sgl-project/sglang#24204](https://github.com/sgl-project/sglang/pull/24204). Quantization is detected automatically from `quantization_config`, so no extra flags are required. See the [SGLang cookbook entry](https://docs.sglang.io/cookbook/autoregressive/Poolside/Laguna-S-2.1) and the main [Laguna S 2.1 model card](https://huggingface.co/poolside/Laguna-S-2.1) for a serving recipe.
+
+#### Transformers
+
+The full Transformers recipe is on the main [Laguna S 2.1 model card](https://huggingface.co/poolside/Laguna-S-2.1). Substitute `poolside/Laguna-S-2.1-NVFP4` for the model ID; quantization is detected automatically from `quantization_config`.
+
+#### TRT-LLM
+
+Laguna S 2.1 support ships in TensorRT-LLM `>=1.3.0rc16`; see the install recipe on the main [Laguna S 2.1 model card](https://huggingface.co/poolside/Laguna-S-2.1). Substitute `poolside/Laguna-S-2.1-NVFP4` for the model ID; quantization is detected automatically from `quantization_config`, no extra flags required.
+
+```python
+from tensorrt_llm import LLM
+
+llm = LLM(model="poolside/Laguna-S-2.1-NVFP4", trust_remote_code=True)
+```
+
+#### Ollama
+
+Available on the [Ollama library](https://ollama.com/laguna-s-2.1).
+
+
+## Controlling reasoning
+
+Laguna S 2.1-NVFP4 uses the same reasoning controls (interleaved thinking, preserved reasoning, and the `enable_thinking` flag) as the base model. See the [Controlling reasoning](https://huggingface.co/poolside/Laguna-S-2.1#controlling-reasoning) section of the main Laguna S 2.1 model card.
+
+## License
+
+This model is licensed under the [OpenMDW-1.1 License](https://huggingface.co/poolside/Laguna-S-2.1-NVFP4/blob/main/LICENSE.md).
+
+## Intended and Responsible Use 
+
+Laguna S 2.1-NVFP4 is designed for software engineering and agentic coding use cases, and you are responsible for confirming that it is appropriate for your intended application. Laguna S 2.1-NVFP4 is subject to the [OpenMDW-1.1 License](https://huggingface.co/poolside/Laguna-S-2.1-NVFP4/blob/main/LICENSE.md), and should be used consistently with Poolside's [Acceptable Use Policy](https://poolside.ai/legal/acceptable-use-policy). We advise against circumventing Laguna S 2.1-NVFP4 safety guardrails without implementing substantially equivalent mitigations appropriate for your use case.
+
+Please report security vulnerabilities or safety concerns to [security@poolside.ai](mailto:security@poolside.ai).
+

chat_template.jinja

@@ -0,0 +1,93 @@
+{#- Iteration on laguna_glm_thinking_v8/chat_template.jinja -#}
+{#- No formatting instructions -#}
+{{- "〈|EOS|〉" -}}
+{%- set enable_thinking = enable_thinking | default(false) -%}
+{%- set add_generation_prompt = add_generation_prompt | default(false) -%}
+
+{#- ───── header (system message) ───── -#}
+{#- A caller-supplied system message with empty content opts out of the default below, producing no <system> block — used to train without a 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 -%}
+  {%- set messages = messages[1:] -%}
+{%- endif -%}
+
+{%- set has_sys = system_message and system_message.strip() -%}
+{%- if has_sys or tools or enable_thinking -%}
+  {{- "<system>" -}}
+
+  {%- if has_sys -%}
+    {{- system_message.rstrip() -}}
+    {%- if tools -%}{{- "\n\n" -}}{%- endif -%}
+  {%- endif -%}
+
+  {%- if tools -%}
+    {{- "### Tools\n\n" -}}
+    {{- "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 -%}
+      {{- (tool | tojson) ~ "\n" -}}
+    {%- endfor -%}
+    {{- "</available_tools>" -}}
+  {%- endif -%}
+
+  {{- "</system>\n" -}}
+{%- endif -%}
+
+{#- ───── main loop ───── -#}
+{%- for message in messages -%}
+  {%- set content = message.content if message.content is string else "" -%}
+  {%- if message.role == "user" -%}
+    {{- "<user>" + content + "</user>\n" -}}
+  {%- elif message.role == "assistant" -%}
+    {%- generation -%}
+      {{- "<assistant>" -}}
+      {#- Extract reasoning content from message.reasoning (vLLM field name) or message.reasoning_content -#}
+      {%- 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 -%}
+      {#- Display reasoning content for all messages if enable_thinking -#}
+      {%- if enable_thinking -%}
+        {{- '<think>' + reasoning_content + '</think>' -}}
+      {%- else -%}
+        {{- '</think>' -}}
+      {%- endif -%}
+      {#- Display main content (trailing newline only when no tool_calls follow) -#}
+      {%- if content -%}
+        {{- content -}}
+      {%- 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>" -}}
+            {{- "<arg_value>" -}}{{- v | tojson(ensure_ascii=False) if v is not string else v -}}{{- "</arg_value>" -}}
+          {%- endfor -%}
+          {{- "</tool_call>" -}}
+        {%- endfor -%}
+      {%- endif -%}
+      {{- "</assistant>\n" -}}
+    {%- endgeneration -%}
+  {%- elif message.role == "tool" -%}
+    {{- "<tool_response>" + content + "</tool_response>\n" -}}
+  {%- elif message.role == "system" -%}
+    {#- Render additional system messages (the first one, if any, is handled separately in the header and was sliced off above) -#}
+    {{- "<system>" + content + "</system>\n" -}}
+  {%- endif -%}
+{%- endfor -%}
+{#- ───── generation prompt ───── -#}
+{%- if add_generation_prompt -%}
+  {{- "<assistant>" -}}
+  {#- ───── Include reasoning mode directive ───── -#}
+  {%- if enable_thinking -%}
+    {{- '<think>' -}}
+  {%- else -%}
+    {{- '</think>' -}}
+  {%- endif -%}
+{%- endif -%}

config.json

@@ -0,0 +1,337 @@
+{
+  "architectures": [
+    "LagunaForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_laguna.LagunaConfig",
+    "AutoModelForCausalLM": "modeling_laguna.LagunaForCausalLM"
+  },
+  "model_type": "laguna",
+  "vocab_size": 100352,
+  "hidden_size": 3072,
+  "intermediate_size": 12288,
+  "num_hidden_layers": 48,
+  "num_attention_heads": 48,
+  "num_key_value_heads": 8,
+  "head_dim": 128,
+  "max_position_embeddings": 262144,
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "rms_norm_eps": 1e-06,
+  "num_experts": 256,
+  "num_experts_per_tok": 10,
+  "moe_intermediate_size": 1024,
+  "shared_expert_intermediate_size": 1024,
+  "norm_topk_prob": true,
+  "router_aux_loss_coef": 0.0,
+  "moe_router_logit_softcapping": 0.0,
+  "decoder_sparse_step": 1,
+  "mlp_only_layers": [
+    0
+  ],
+  "bos_token_id": 2,
+  "eos_token_id": [
+    2,
+    24
+  ],
+  "pad_token_id": 9,
+  "tie_word_embeddings": false,
+  "use_cache": true,
+  "torch_dtype": "bfloat16",
+  "gating": "per-head",
+  "sliding_window": 512,
+  "rope_parameters": {
+    "full_attention": {
+      "rope_theta": 500000.0,
+      "rope_type": "yarn",
+      "factor": 32.0,
+      "original_max_position_embeddings": 8192,
+      "beta_slow": 1.0,
+      "beta_fast": 32.0,
+      "attention_factor": 1.3465735902799727,
+      "partial_rotary_factor": 0.5
+    },
+    "sliding_attention": {
+      "rope_type": "default",
+      "rope_theta": 10000.0,
+      "partial_rotary_factor": 1.0
+    }
+  },
+  "layer_types": [
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention"
+  ],
+  "moe_apply_router_weight_on_input": false,
+  "mlp_layer_types": [
+    "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"
+  ],
+  "gating_types": [
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head",
+    "per_head"
+  ],
+  "moe_routed_scaling_factor": 2.5,
+  "num_attention_heads_per_layer": [
+    48,
+    72,
+    72,
+    72,
+    48,
+    72,
+    72,
+    72,
+    48,
+    72,
+    72,
+    72,
+    48,
+    72,
+    72,
+    72,
+    48,
+    72,
+    72,
+    72,
+    48,
+    72,
+    72,
+    72,
+    48,
+    72,
+    72,
+    72,
+    48,
+    72,
+    72,
+    72,
+    48,
+    72,
+    72,
+    72,
+    48,
+    72,
+    72,
+    72,
+    48,
+    72,
+    72,
+    72,
+    48,
+    72,
+    72,
+    72
+  ],
+  "quantization_config": {
+    "config_groups": {
+      "group_0": {
+        "format": "nvfp4-pack-quantized",
+        "input_activations": {
+          "actorder": null,
+          "block_structure": null,
+          "dynamic": "local",
+          "group_size": 16,
+          "num_bits": 4,
+          "observer": "static_minmax",
+          "observer_kwargs": {},
+          "scale_dtype": "torch.float8_e4m3fn",
+          "strategy": "tensor_group",
+          "symmetric": true,
+          "type": "float",
+          "zp_dtype": null
+        },
+        "output_activations": null,
+        "targets": [
+          "re:.*experts\\.[0-9]+\\.(gate_proj|up_proj|down_proj)$"
+        ],
+        "weights": {
+          "actorder": null,
+          "block_structure": null,
+          "dynamic": false,
+          "group_size": 16,
+          "num_bits": 4,
+          "observer": "memoryless_minmax",
+          "observer_kwargs": {},
+          "scale_dtype": "torch.float8_e4m3fn",
+          "strategy": "tensor_group",
+          "symmetric": true,
+          "type": "float",
+          "zp_dtype": null
+        }
+      }
+    },
+    "format": "nvfp4-pack-quantized",
+    "global_compression_ratio": null,
+    "ignore": [
+      "lm_head",
+      "re:.*\\.self_attn\\.q_proj$",
+      "re:.*\\.self_attn\\.k_proj$",
+      "re:.*\\.self_attn\\.v_proj$",
+      "re:.*\\.self_attn\\.o_proj$",
+      "re:.*\\.self_attn\\.g_proj$",
+      "re:.*\\.mlp\\.gate$",
+      "model.layers.0.mlp.gate_proj",
+      "model.layers.0.mlp.up_proj",
+      "model.layers.0.mlp.down_proj",
+      "re:.*\\.mlp\\.shared_expert\\.gate_proj$",
+      "re:.*\\.mlp\\.shared_expert\\.up_proj$",
+      "re:.*\\.mlp\\.shared_expert\\.down_proj$"
+    ],
+    "kv_cache_scheme": {
+      "actorder": null,
+      "block_structure": null,
+      "dynamic": false,
+      "group_size": null,
+      "num_bits": 8,
+      "observer": "minmax",
+      "observer_kwargs": {},
+      "scale_dtype": null,
+      "strategy": "tensor",
+      "symmetric": true,
+      "type": "float",
+      "zp_dtype": null
+    },
+    "quant_method": "compressed-tensors",
+    "quantization_status": "compressed",
+    "sparsity_config": {},
+    "transform_config": {},
+    "version": "0.14.1.dev11+gf2ee47b"
+  }
+}

configuration_laguna.py

@@ -0,0 +1,256 @@
+# 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.
+from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig
+try:
+    from transformers.modeling_rope_utils import RopeParameters
+except ImportError:
+    RopeParameters = dict
+try:
+    from transformers.utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available
+except ImportError:
+    def is_causal_conv1d_available(): return False
+    def is_flash_linear_attention_available(): return False
+
+
+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` or `str`, *optional*, defaults to `True`):
+            Attention output gating mode. When ``True`` or ``"per-element"`` a g_proj
+            linear layer with output size ``num_attention_heads * head_dim`` is added
+            and ``attn_output = attn_output * softplus(g_proj(x))``. When ``"per-head"``
+            g_proj has output size ``num_attention_heads`` and the gate broadcasts across
+            ``head_dim``. When ``False`` no gating is applied.
+        partial_rotary_factor (`float`, *optional*):
+            Fraction of head_dim to apply rotary embeddings to. When set, this value is
+            injected into ``rope_parameters`` (and ``swa_rope_parameters``) if not already
+            specified there. When ``None`` the default behaviour of the rope implementation
+            is used (typically full rotary).
+        num_attention_heads_per_layer (`list[int]`, *optional*):
+            Optional per-layer override for ``num_attention_heads``. When provided the list
+            length must equal ``num_hidden_layers`` and each entry is the head count used by
+            that layer. When ``None`` every layer uses ``num_attention_heads``.
+        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.
+        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
+            ``"full_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.
+            When enabled, a per-head bias parameter is added that allows the model to attend
+            to position 0 even when it falls outside the sliding window.
+        swa_rope_parameters (`RopeParameters`, *optional*):
+            Separate RoPE configuration for sliding-window attention layers. When ``None``,
+            SWA layers use the same RoPE as global 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.
+        moe_routed_scaling_factor (`float`, *optional*, defaults to 1.0):
+            Scalar multiplier applied to the routed-expert output before combining with the
+            shared-expert output.
+        moe_apply_router_weight_on_input (`bool`, *optional*, defaults to `False`):
+            When ``True`` the top-k routing weights are multiplied into each expert's input
+            rather than its output. Matches the numerical form used by the trained checkpoint.
+        moe_router_logit_softcapping (`float`, *optional*, defaults to 0.0):
+            Optional soft-capping value ``c`` applied to router logits as
+            ``x = tanh(x / c) * c`` before sigmoid + top-k. Disabled when ``0``.
+        rope_parameters (`RopeParameters`, *optional*):
+            RoPE configuration. Defaults to rope_theta=500000.0.
+    """
+
+    model_type = "laguna"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    # PreTrainedConfig in transformers v5 no longer auto-declares these; subclasses
+    # opt in by providing class-level annotations with defaults.
+    pad_token_id: int | None = None
+    bos_token_id: int | None = None
+    eos_token_id: int | list[int] | None = None
+    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 | str = 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_parameters: RopeParameters | dict[str, RopeParameters] | None = None,
+        partial_rotary_factor: float | None = None,
+        attention_dropout: float = 0.0,
+        sliding_window: int | None = None,
+        layer_types: list[str] | None = None,
+        num_attention_heads_per_layer: list[int] | None = None,
+        swa_attention_sink_enabled: bool = False,
+        swa_rope_parameters: RopeParameters | None = None,
+        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,
+        moe_routed_scaling_factor: float = 1.0,
+        moe_apply_router_weight_on_input: bool = False,
+        moe_router_logit_softcapping: float = 0.0,
+        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]
+
+        # Default layer_types: all layers use full attention (Laguna-M). Laguna-XS
+        # ships an explicit list with a mix of "full_attention" and "sliding_attention".
+        # Downstream mask builders (``create_masks_for_generate``) iterate
+        # ``layer_types``, so it must be a list — not left as ``None``.
+        if layer_types is None:
+            layer_types = ["full_attention"] * num_hidden_layers
+
+        # Default rope_parameters with Laguna's theta
+        if rope_parameters is None:
+            rope_parameters = {"rope_type": "default", "rope_theta": 500000.0}
+
+        # config.json stores SWA rope nested in rope_parameters["sliding_attention"]
+        # and carries no top-level swa_rope_parameters. Derive it here, else the
+        # sliding-window layers silently reuse the full-attention rope.
+        if swa_rope_parameters is None and isinstance(rope_parameters, dict):
+            swa_rope_parameters = rope_parameters.get("sliding_attention")
+
+        # If ``partial_rotary_factor`` is set at the top level, inject it into any
+        # rope dict that does not already carry one so the rotary embedding picks
+        # it up consistently for both full-attention and SWA layers.
+        if partial_rotary_factor is not None:
+            if isinstance(rope_parameters, dict) and "partial_rotary_factor" not in rope_parameters:
+                rope_parameters = {**rope_parameters, "partial_rotary_factor": partial_rotary_factor}
+            if isinstance(swa_rope_parameters, dict) and "partial_rotary_factor" not in swa_rope_parameters:
+                swa_rope_parameters = {
+                    **swa_rope_parameters,
+                    "partial_rotary_factor": partial_rotary_factor,
+                }
+
+        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_parameters = rope_parameters
+        self.partial_rotary_factor = partial_rotary_factor
+        self.attention_dropout = attention_dropout
+        # Sliding window attention arguments
+        self.sliding_window = sliding_window
+        self.layer_types = layer_types
+        self.num_attention_heads_per_layer = num_attention_heads_per_layer
+        self.swa_attention_sink_enabled = swa_attention_sink_enabled
+        self.swa_rope_parameters = swa_rope_parameters
+        # 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.moe_routed_scaling_factor = moe_routed_scaling_factor
+        self.moe_apply_router_weight_on_input = moe_apply_router_weight_on_input
+        self.moe_router_logit_softcapping = moe_router_logit_softcapping
+        self.output_router_logits = output_router_logits
+
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+
+
+__all__ = ["LagunaConfig"]

generation_config.json

@@ -0,0 +1,25 @@
+{
+  "bos_token_id": 2,
+  "do_sample": true,
+  "eos_token_id": [
+    2,
+    24
+  ],
+  "max_new_tokens": 32768,
+  "pad_token_id": 9,
+  "temperature": 1.0,
+  "top_p": 1.0,
+  "speculative_config": {
+    "method": "dflash",
+    "source": "huggingface",
+    "model": "poolside/Laguna-S-2.1-DFlash-NVFP4",
+    "num_speculative_tokens": 15
+  },
+  "tool_call_parser": "poolside_v1",
+  "reasoning_parser": "poolside_v1",
+  "default_chat_template_kwargs": {
+    "enable_thinking": true
+  },
+  "top_k": 20,
+  "min_p": 0.0
+}

model-00001-of-00014.safetensors

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

model-00002-of-00014.safetensors

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

model-00003-of-00014.safetensors

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

model-00004-of-00014.safetensors

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

model-00005-of-00014.safetensors

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

model-00006-of-00014.safetensors

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

model-00007-of-00014.safetensors

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

model-00008-of-00014.safetensors

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

model-00009-of-00014.safetensors

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

model-00010-of-00014.safetensors

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

model-00011-of-00014.safetensors

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

model-00012-of-00014.safetensors

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

model-00013-of-00014.safetensors

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

model-00014-of-00014.safetensors

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

model.safetensors.index.json

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

modeling_laguna.py

@@ -0,0 +1,886 @@
+# 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.
+
+from collections.abc import Callable
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache
+from transformers.integrations import use_experts_implementation, use_kernelized_func
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_outputs import MoeModelOutputWithPast
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+from transformers.processing_utils import Unpack
+from transformers.utils import auto_docstring, can_return_tuple, is_grouped_mm_available
+from transformers.utils.generic import TransformersKwargs, merge_with_config_defaults
+from transformers.utils.output_capturing import OutputRecorder, capture_outputs
+from transformers.cache_utils import DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.integrations import use_kernel_forward_from_hub
+from transformers.masking_utils import create_causal_mask
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils.generic import maybe_autocast
+from .configuration_laguna import LagunaConfig
+
+from transformers import initialization as init
+from transformers.masking_utils import create_sliding_window_causal_mask
+from transformers.modeling_outputs import MoeCausalLMOutputWithPast
+from transformers.utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class LagunaRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+        """
+        LagunaRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class LagunaRotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor  # fix linting for `register_buffer`
+
+    def __init__(self, config: LagunaConfig, device=None):
+        super().__init__()
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+
+        self.rope_type = self.config.rope_parameters["rope_type"]
+        rope_init_fn: Callable = self.compute_default_rope_parameters
+        if self.rope_type != "default":
+            rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+    @staticmethod
+    def compute_default_rope_parameters(config, device=None, seq_len=None) -> tuple["torch.Tensor", float]:
+        """
+        Computes the inverse frequencies according to the original RoPE implementation
+        Args:
+            config ([`~transformers.PreTrainedConfig`]):
+                The model configuration.
+            device (`torch.device`):
+                The device to use for initialization of the inverse frequencies.
+            seq_len (`int`, *optional*):
+                The current sequence length. Unused for this type of RoPE.
+        Returns:
+            Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+            post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+        """
+        base = config.rope_parameters["rope_theta"]
+        head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+        partial = config.rope_parameters.get("partial_rotary_factor", 1.0)
+        dim = int(head_dim * partial)
+        inv_freq = 1.0 / (
+            base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+        )
+        return inv_freq, 1.0
+
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids):
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with maybe_autocast(device_type=device_type, enabled=False):  # Force float32
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+class LagunaMLP(nn.Module):
+    def __init__(self, config, intermediate_size=None):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+class LagunaTopKRouter(nn.Module):
+    """Laguna MoE router using sigmoid scoring (not softmax).
+
+    Supports optional router-logit soft-capping and auxiliary-loss-free load
+    balancing (arXiv:2408.15664): the per-expert bias ``e_score_correction_bias``
+    is added to selection scores but the returned routing weights remain unbiased.
+    The bias lives on the router so accelerate's per-module hooks can co-locate it
+    with the gate — moving it to the experts module would cross a hook boundary
+    and leave the bias on meta under ``device_map="auto"`` / CPU-offload.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.top_k = config.num_experts_per_tok
+        self.num_experts = config.num_experts
+        self.norm_topk_prob = config.norm_topk_prob
+        self.hidden_dim = config.hidden_size
+        self.weight = nn.Parameter(torch.zeros(self.num_experts, self.hidden_dim))
+        # Zero-initialised so inference on checkpoints that don't ship the bias
+        # is a no-op. ``_checkpoint_conversion_mapping`` below remaps the
+        # ``mlp.experts.e_score_correction_bias`` key from vLLM-trained
+        # checkpoints onto this attribute.
+        self.e_score_correction_bias = nn.Parameter(torch.zeros(config.num_experts), requires_grad=False)
+        self.router_logit_softcapping = float(getattr(config, "moe_router_logit_softcapping", 0.0) or 0.0)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        hidden_states = hidden_states.reshape(-1, self.hidden_dim)
+        router_logits = F.linear(hidden_states, self.weight).float()
+        if self.router_logit_softcapping > 0.0:
+            router_logits = torch.tanh(router_logits / self.router_logit_softcapping) * self.router_logit_softcapping
+        routing_scores = torch.sigmoid(router_logits)
+        scores_for_selection = routing_scores + self.e_score_correction_bias.to(routing_scores.dtype)
+        _, selected_experts = torch.topk(scores_for_selection, self.top_k, dim=-1)
+        routing_weights = routing_scores.gather(-1, selected_experts)
+        if self.norm_topk_prob:
+            routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
+        routing_weights = routing_weights.to(hidden_states.dtype)
+        return router_logits, routing_weights, selected_experts
+
+
+@use_experts_implementation
+class LagunaExperts(nn.Module):
+    """Fused expert weights as 3D tensors for batched execution."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.num_experts = config.num_experts
+        self.hidden_dim = config.hidden_size
+        self.intermediate_dim = config.moe_intermediate_size
+        self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim))
+        self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim))
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        top_k_index: torch.Tensor,
+        top_k_weights: torch.Tensor,
+    ) -> torch.Tensor:
+        final_hidden_states = torch.zeros_like(hidden_states)
+        with torch.no_grad():
+            expert_mask = F.one_hot(top_k_index, num_classes=self.num_experts)
+            expert_mask = expert_mask.permute(2, 1, 0)
+            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+
+        for expert_idx in expert_hit:
+            expert_idx = expert_idx[0]
+            if expert_idx == self.num_experts:
+                continue
+            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
+            current_state = hidden_states[token_idx]
+            gate, up = F.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
+            current_hidden_states = self.act_fn(gate) * up
+            current_hidden_states = F.linear(current_hidden_states, self.down_proj[expert_idx])
+            current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
+            final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))
+
+        return final_hidden_states
+
+
+class LagunaSparseMoeBlock(nn.Module):
+    """Laguna MoE block using sigmoid router, fused expert tensors, and a shared expert."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.num_experts = config.num_experts
+        self.routed_scaling_factor = float(getattr(config, "moe_routed_scaling_factor", 1.0))
+        # ``moe_apply_router_weight_on_input=True`` would require scaling each expert's
+        # input (rather than its output) by the routing weight. Supporting it cleanly
+        # alongside the fused experts kernels (``grouped_mm`` / ``batched_mm``) is future
+        # work; for now we fail loudly so a checkpoint that needs it can't silently
+        # diverge from its numerical form.
+        if getattr(config, "moe_apply_router_weight_on_input", False):
+            raise NotImplementedError(
+                "moe_apply_router_weight_on_input=True is not yet supported in the "
+                "transformers implementation of Laguna."
+            )
+        self.gate = LagunaTopKRouter(config)
+        self.experts = LagunaExperts(config)
+        self.shared_expert = LagunaMLP(config, intermediate_size=config.shared_expert_intermediate_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, sequence_length, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+
+        shared_expert_output = self.shared_expert(hidden_states)
+        _, routing_weights, selected_experts = self.gate(hidden_states)
+        expert_output = self.experts(hidden_states, selected_experts, routing_weights)
+        if self.routed_scaling_factor != 1.0:
+            expert_output = expert_output * self.routed_scaling_factor
+
+        expert_output = expert_output + shared_expert_output
+        expert_output = expert_output.reshape(batch_size, sequence_length, hidden_dim)
+        return expert_output
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+# Adapted from transformers.models.glm.modular_glm.apply_rotary_pos_emb
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Removes the interleaving of cos and sin from GLM
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+
+    # Keep half or full tensor for later concatenation
+    rotary_dim = cos.shape[-1]
+    q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
+    k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
+
+    # Apply rotary embeddings on the first half or full tensor
+    q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
+    k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)
+
+    # Concatenate back to full shape
+    q_embed = torch.cat([q_embed, q_pass], dim=-1)
+    k_embed = torch.cat([k_embed, k_pass], dim=-1)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        attn_weights = attn_weights + attention_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+# Laguna attention is identical to Qwen2MoE attention except:
+# - No QKV bias
+# - Explicit head_dim from config
+# - Output gating: attn_output = attn_output * softplus(g_proj(hidden_states)) (optional)
+# - Per-layer sliding window attention with optional attention sinks
+@use_kernelized_func(apply_rotary_pos_emb)
+class LagunaAttention(nn.Module):
+    def __init__(self, config: LagunaConfig, layer_idx: int, num_heads: int | None = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = config.head_dim
+        # Allow the caller (decoder layer) to supply a per-layer head count; fall back
+        # to config.num_attention_heads when not provided.
+        self.num_heads = num_heads if num_heads is not None else config.num_attention_heads
+        self.num_key_value_groups = self.num_heads // config.num_key_value_heads
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+
+        # Per-layer sliding window (follows Gemma2/Cohere2 convention)
+        layer_types = getattr(config, "layer_types", None)
+        if layer_types is not None:
+            self.is_sliding = layer_types[layer_idx] == "sliding_attention"
+            self.sliding_window = config.sliding_window if self.is_sliding else None
+        else:
+            self.is_sliding = False
+            self.sliding_window = None
+
+        # Laguna: no QKV bias, explicit head_dim
+        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * config.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * config.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * config.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * config.head_dim, config.hidden_size, bias=False)
+
+        # Laguna-specific: optional gating projection.
+        # ``gating`` may be:
+        #   - True / "per-element": one gate per (head, head_dim) channel
+        #   - "per-head":           one gate per head, broadcast across head_dim
+        #   - False:                no gating
+        gating = getattr(config, "gating", True)
+        self.gating = bool(gating)
+        self.gate_per_head = gating == "per-head"
+        if self.gating:
+            g_out = self.num_heads if self.gate_per_head else self.num_heads * config.head_dim
+            self.g_proj = nn.Linear(config.hidden_size, g_out, bias=False)
+
+        # Attention sinks (learnable per-head bias for SWA layers)
+        if self.is_sliding and getattr(config, "swa_attention_sink_enabled", False):
+            self.sink = nn.Parameter(torch.zeros(self.num_heads))
+
+        # QK normalization (RMSNorm applied per-head after reshape, before RoPE)
+        self.q_norm = LagunaRMSNorm(config.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = LagunaRMSNorm(config.head_dim, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        past_key_values: Cache | None = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(hidden_shape).transpose(1, 2)
+        key_states = key_states.view(hidden_shape).transpose(1, 2)
+        value_states = value_states.view(hidden_shape).transpose(1, 2)
+
+        # QK normalization (applied per-head before RoPE)
+        query_states = self.q_norm(query_states)
+        key_states = self.k_norm(key_states)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_values is not None:
+            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+        # ``attention_mask`` here is already the correct mask for this layer type —
+        # ``LagunaModel.forward`` builds separate full-attention and sliding-attention
+        # masks (using ``create_causal_mask`` / ``create_sliding_window_causal_mask``)
+        # and the decoder layer passes the right one in.
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+
+        # Laguna-specific: apply gating BEFORE o_proj (optional)
+        if self.gating:
+            gate = F.softplus(self.g_proj(hidden_states).float()).to(attn_output.dtype)
+            if self.gate_per_head:
+                # gate: [..., num_heads]; broadcast across head_dim
+                attn_shape = attn_output.shape
+                attn_output = (
+                    attn_output.view(*attn_shape[:-1], self.num_heads, self.head_dim) * gate.unsqueeze(-1)
+                ).view(attn_shape)
+            else:
+                attn_output = attn_output * gate
+
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output, attn_weights
+
+
+class LagunaDecoderLayer(GradientCheckpointingLayer):
+    """Laguna decoder layer with gated attention and sigmoid-routed MoE."""
+
+    def __init__(self, config: LagunaConfig, layer_idx: int):
+        super().__init__()
+        per_layer_heads = getattr(config, "num_attention_heads_per_layer", None)
+        layer_num_heads = per_layer_heads[layer_idx] if per_layer_heads is not None else config.num_attention_heads
+        # Layer type drives mask and position-embedding dispatch in ``LagunaModel.forward``.
+        layer_types = getattr(config, "layer_types", None)
+        self.attention_type = layer_types[layer_idx] if layer_types is not None else "full_attention"
+        self.self_attn = LagunaAttention(config, layer_idx, num_heads=layer_num_heads)
+        # Use MoE or dense MLP based on layer configuration
+        if (layer_idx not in config.mlp_only_layers) and (
+            config.num_experts > 0 and (layer_idx + 1) % config.decoder_sparse_step == 0
+        ):
+            self.mlp = LagunaSparseMoeBlock(config)
+        else:
+            self.mlp = LagunaMLP(config, intermediate_size=config.intermediate_size)
+        self.input_layernorm = LagunaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = LagunaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.hidden_size = config.hidden_size
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        use_cache: bool | None = False,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        # Self Attention
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+@auto_docstring
+class LagunaPreTrainedModel(PreTrainedModel):
+    config: LagunaConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["LagunaDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _can_compile_fullgraph = (
+        is_grouped_mm_available()
+    )  # https://huggingface.co/docs/transformers/experts_interface#torchcompile
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "router_logits": OutputRecorder(LagunaTopKRouter, index=0),
+        "hidden_states": LagunaDecoderLayer,
+        "attentions": LagunaAttention,
+    }
+    # vLLM-trained Laguna checkpoints store the aux-loss-free routing bias on the
+    # experts module (``mlp.experts.e_score_correction_bias``). In this impl the
+    # bias lives on the router to stay co-located with its consumer across
+    # accelerate's per-module hooks, so remap the legacy key on load.
+    _checkpoint_conversion_mapping = {
+        r"^(.*)\.mlp\.experts\.e_score_correction_bias$": r"\1.mlp.gate.e_score_correction_bias",
+    }
+
+    @torch.no_grad()
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        std = self.config.initializer_range
+        if isinstance(module, LagunaExperts):
+            init.normal_(module.gate_up_proj, mean=0.0, std=std)
+            init.normal_(module.down_proj, mean=0.0, std=std)
+        elif isinstance(module, LagunaTopKRouter):
+            init.normal_(module.weight, mean=0.0, std=std)
+        # Bare ``nn.Parameter``s that are not covered by the parent's generic
+        # Linear/Embedding/norm handling need their own rules so that the
+        # __init__ and from_pretrained(state_dict={}) paths produce identical
+        # weights under a fixed seed.
+        if isinstance(module, LagunaTopKRouter):
+            torch.nn.init.zeros_(module.e_score_correction_bias)
+        if isinstance(module, LagunaAttention) and hasattr(module, "sink"):
+            torch.nn.init.zeros_(module.sink)
+
+
+class LagunaModel(LagunaPreTrainedModel):
+    def __init__(self, config: LagunaConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [LagunaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = LagunaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        # ``LagunaRotaryEmbedding`` inherits ``Qwen2MoeRotaryEmbedding``'s flat-shape
+        # contract — it reads ``config.rope_parameters["rope_type"]`` at the outer
+        # level. Laguna stores rope nested by layer type (``{"full_attention": {...},
+        # ...}``), so pass a config clone with the full-attention sub-dict flattened.
+        rp = getattr(config, "rope_parameters", None)
+        if isinstance(rp, dict) and isinstance(rp.get("full_attention"), dict):
+            import copy
+
+            full_config = copy.deepcopy(config)
+            full_config.rope_parameters = dict(rp["full_attention"])
+            self.rotary_emb = LagunaRotaryEmbedding(config=full_config)
+        else:
+            self.rotary_emb = LagunaRotaryEmbedding(config=config)
+
+        # Separate RoPE for sliding-window attention layers (when configured).
+        # Be careful with ``partial_rotary_factor`` — ``PreTrainedConfig.standardize_rope_params``
+        # unconditionally overwrites ``rope_parameters["partial_rotary_factor"]`` with
+        # ``self.partial_rotary_factor``, so we must align the top-level field on the
+        # cloned config to the SWA value, otherwise the global partial factor silently
+        # clobbers the SWA one.
+        if getattr(config, "swa_rope_parameters", None) is not None:
+            import copy
+
+            swa_config = copy.deepcopy(config)
+            swa_config.rope_parameters = dict(config.swa_rope_parameters)
+            swa_partial = swa_config.rope_parameters.get("partial_rotary_factor")
+            swa_config.partial_rotary_factor = swa_partial
+            self.swa_rotary_emb = LagunaRotaryEmbedding(config=swa_config)
+        else:
+            self.swa_rotary_emb = None
+
+        self.gradient_checkpointing = False
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @merge_with_config_defaults
+    @capture_outputs
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> MoeModelOutputWithPast:
+        from transformers.cache_utils import DynamicCache
+        from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+
+        if position_ids is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            position_ids = (
+                torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+            ).unsqueeze(0)
+
+        # Build one mask per layer-type so each layer can be dispatched with the right
+        # attention pattern (follows the afmoe / cohere2 v5 convention).
+        layer_types = getattr(self.config, "layer_types", None)
+        has_swa = layer_types is not None and "sliding_attention" in layer_types
+        if not isinstance(causal_mask_mapping := attention_mask, dict):
+            mask_kwargs = {
+                "config": self.config,
+                "inputs_embeds": inputs_embeds,
+                "attention_mask": attention_mask,
+                "past_key_values": past_key_values,
+                "position_ids": position_ids,
+            }
+            causal_mask_mapping = {"full_attention": create_causal_mask(**mask_kwargs)}
+            if has_swa:
+                causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs)
+
+        hidden_states = inputs_embeds
+        global_pe = self.rotary_emb(hidden_states, position_ids)
+        # Per-layer-type position embeddings: Laguna optionally uses a different rope for
+        # sliding layers (``swa_rope_parameters``). When absent, SWA layers share the
+        # global rope.
+        if has_swa:
+            swa_pe = self.swa_rotary_emb(hidden_states, position_ids) if self.swa_rotary_emb is not None else global_pe
+            position_embeddings_mapping = {"full_attention": global_pe, "sliding_attention": swa_pe}
+        else:
+            position_embeddings_mapping = None
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            layer_attn_mask = causal_mask_mapping[decoder_layer.attention_type]
+            layer_pos_emb = (
+                position_embeddings_mapping[decoder_layer.attention_type]
+                if position_embeddings_mapping is not None
+                else global_pe
+            )
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=layer_attn_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                position_embeddings=layer_pos_emb,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+
+        return MoeModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+def load_balancing_loss_func(
+    gate_logits: torch.Tensor | tuple[torch.Tensor] | None,
+    num_experts: int | None = None,
+    top_k=2,
+    attention_mask: torch.Tensor | None = None,
+) -> torch.Tensor | int:
+    r"""
+    Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
+
+    See Switch Transformer (https://huggingface.co/papers/2101.03961) for more details. This function implements the loss
+    function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
+    experts is too unbalanced.
+
+    Args:
+        gate_logits:
+            Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of
+            shape [batch_size X sequence_length, num_experts].
+        num_experts:
+            Number of experts
+        top_k:
+            The number of experts to route per-token, can be also interpreted as the `top-k` routing
+            parameter.
+        attention_mask (`torch.Tensor`, *optional*):
+            The attention_mask used in forward function
+            shape [batch_size X sequence_length] if not None.
+
+    Returns:
+        The auxiliary loss.
+    """
+    if gate_logits is None or not isinstance(gate_logits, tuple):
+        return 0
+
+    if isinstance(gate_logits, tuple):
+        compute_device = gate_logits[0].device
+        concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0)
+
+    routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1)
+
+    _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)
+
+    expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
+
+    if attention_mask is None:
+        # Compute the percentage of tokens routed to each experts
+        tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
+
+        # Compute the average probability of routing to these experts
+        router_prob_per_expert = torch.mean(routing_weights, dim=0)
+    else:
+        batch_size, sequence_length = attention_mask.shape
+        num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length)
+
+        # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask
+        expert_attention_mask = (
+            attention_mask[None, :, :, None, None]
+            .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))
+            .reshape(-1, top_k, num_experts)
+            .to(compute_device)
+        )
+
+        # Compute the percentage of tokens routed to each experts
+        tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(
+            expert_attention_mask, dim=0
+        )
+
+        # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert
+        router_per_expert_attention_mask = (
+            attention_mask[None, :, :, None]
+            .expand((num_hidden_layers, batch_size, sequence_length, num_experts))
+            .reshape(-1, num_experts)
+            .to(compute_device)
+        )
+
+        # Compute the average probability of routing to these experts
+        router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(
+            router_per_expert_attention_mask, dim=0
+        )
+
+    overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
+    return overall_loss * num_experts
+
+
+@auto_docstring
+class LagunaForCausalLM(LagunaPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+    _tp_plan = {"lm_head": "colwise_gather_output"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = LagunaModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.router_aux_loss_coef = config.router_aux_loss_coef
+        self.num_experts = config.num_experts
+        self.num_experts_per_tok = config.num_experts_per_tok
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_router_logits: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> MoeCausalLMOutputWithPast:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+        """
+
+        output_router_logits = (
+            output_router_logits if output_router_logits is not None else self.config.output_router_logits
+        )
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs: MoeModelOutputWithPast = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_router_logits=output_router_logits,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
+
+        aux_loss = None
+        if output_router_logits:
+            aux_loss = load_balancing_loss_func(
+                outputs.router_logits,
+                self.num_experts,
+                self.num_experts_per_tok,
+                attention_mask,
+            )
+            if labels is not None:
+                loss += self.router_aux_loss_coef * aux_loss.to(loss.device)  # make sure to reside in the same device
+
+        return MoeCausalLMOutputWithPast(
+            loss=loss,
+            aux_loss=aux_loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            router_logits=outputs.router_logits,
+        )
+
+
+__all__ = ["LagunaForCausalLM", "LagunaModel", "LagunaPreTrainedModel"]
+
+
+# --- Added: register the native Laguna checkpoint-conversion for trust_remote_code loads.
+# transformers >=5.12 skips checkpoint-conversion mappings for custom (remote) code
+# unless explicitly registered, which broke loading the shipped per-expert MoE weights.
+try:
+    from transformers.conversion_mapping import (
+        get_checkpoint_conversion_mapping as _lg_get,
+        register_checkpoint_conversion_mapping as _lg_reg,
+        USER_REGISTERED_MAPPINGS as _lg_user,
+    )
+
+    if "laguna" not in _lg_user:
+        _lg_m = _lg_get("laguna")
+        if _lg_m is not None:
+            _lg_reg("laguna", _lg_m, overwrite=True)
+except Exception:
+    pass

special_tokens_map.json

@@ -0,0 +1,9 @@
+{
+  "bos_token": "〈|EOS|〉",
+  "cls_token": "〈|CLS|〉",
+  "eos_token": "〈|EOS|〉",
+  "mask_token": "〈|MASK|〉",
+  "pad_token": "〈|PAD|〉",
+  "sep_token": "〈|SEP|〉",
+  "unk_token": "〈|UNK|〉"
+}

tokenizer_config.json

@@ -0,0 +1,576 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "〈|UNK|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "〈|CODE_START|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "〈|EOS|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "〈|CODE_END|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "〈|META_START|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "5": {
+      "content": "〈|META_END|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "6": {
+      "content": "〈|FIM_MIDDLE|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "7": {
+      "content": "〈|FIM_SUFFIX|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "8": {
+      "content": "〈|SEP|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "9": {
+      "content": "〈|PAD|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "10": {
+      "content": "〈|CLS|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "11": {
+      "content": "〈|FIM_START|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "12": {
+      "content": "〈|MASK|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "13": {
+      "content": "|◊|",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "14": {
+      "content": "〈|",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "15": {
+      "content": "|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16": {
+      "content": "〈|/",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "17": {
+      "content": "/|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "20": {
+      "content": "〈|SPECIAL_1|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "21": {
+      "content": "〈|SPECIAL_2|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "22": {
+      "content": "〈|SPECIAL_3|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "27": {
+      "content": "〈|SPECIAL_8|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "28": {
+      "content": "〈|SPECIAL_9|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "29": {
+      "content": "〈|SPECIAL_10|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "30": {
+      "content": "〈|SPECIAL_11|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "31": {
+      "content": "〈|SPECIAL_12|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32": {
+      "content": "〈|SPECIAL_13|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "33": {
+      "content": "〈|SPECIAL_14|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "34": {
+      "content": "〈|SPECIAL_15|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "35": {
+      "content": "〈|SPECIAL_16|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "36": {
+      "content": "〈|SPECIAL_17|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "37": {
+      "content": "〈|SPECIAL_18|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "38": {
+      "content": "〈|SPECIAL_19|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "39": {
+      "content": "〈|SPECIAL_20|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "40": {
+      "content": "〈|SPECIAL_21|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "41": {
+      "content": "〈|SPECIAL_22|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "42": {
+      "content": "〈|SPECIAL_23|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "43": {
+      "content": "〈|SPECIAL_24|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "44": {
+      "content": "〈|SPECIAL_25|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "45": {
+      "content": "〈|SPECIAL_26|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "46": {
+      "content": "〈|SPECIAL_27|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "47": {
+      "content": "〈|SPECIAL_28|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "48": {
+      "content": "〈|SPECIAL_29|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "49": {
+      "content": "〈|SPECIAL_30|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50": {
+      "content": "〈|SPECIAL_31|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "51": {
+      "content": "〈|SPECIAL_32|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "52": {
+      "content": "〈|SPECIAL_33|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "53": {
+      "content": "〈|SPECIAL_34|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "54": {
+      "content": "〈|SPECIAL_35|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "55": {
+      "content": "〈|SPECIAL_36|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "56": {
+      "content": "〈|SPECIAL_37|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "57": {
+      "content": "〈|SPECIAL_38|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "58": {
+      "content": "〈|SPECIAL_39|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "59": {
+      "content": "〈|SPECIAL_40|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "60": {
+      "content": "〈|SPECIAL_41|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "61": {
+      "content": "〈|SPECIAL_42|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "62": {
+      "content": "〈|SPECIAL_43|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "63": {
+      "content": "〈|SPECIAL_44|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "64": {
+      "content": "〈|SPECIAL_45|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "65": {
+      "content": "〈|SPECIAL_46|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "66": {
+      "content": "〈|SPECIAL_47|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "67": {
+      "content": "〈|SPECIAL_48|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "68": {
+      "content": "〈|SPECIAL_49|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "69": {
+      "content": "〈|SPECIAL_50|〉",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "18": {
+      "content": "<think>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": false
+    },
+    "19": {
+      "content": "</think>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": false
+    },
+    "23": {
+      "content": "<assistant>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": false
+    },
+    "24": {
+      "content": "</assistant>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": false
+    },
+    "25": {
+      "content": "<tool_call>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": false
+    },
+    "26": {
+      "content": "</tool_call>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": false
+    }
+  },
+  "bos_token": "〈|EOS|〉",
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "〈|CLS|〉",
+  "eos_token": "〈|EOS|〉",
+  "extra_special_tokens": {},
+  "mask_token": "〈|MASK|〉",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "〈|PAD|〉",
+  "sep_token": "〈|SEP|〉",
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "unk_token": "〈|UNK|〉",
+  "chat_template": "{% include 'chat_template.jinja' %}"
+}