Laguna-XS-2.1

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

.eval_results/swe-bench_pro.yaml

@@ -0,0 +1,7 @@
+- dataset:
+    id: ScaleAI/SWE-bench_Pro
+    task_id: SWE_Bench_Pro
+  value: 47.6
+  source:
+    url: https://huggingface.co/poolside/Laguna-XS-2.1
+    name: Model Card

.eval_results/swe-bench_verified.yaml

@@ -0,0 +1,7 @@
+- dataset:
+    id: SWE-bench/SWE-bench_Verified
+    task_id: swe_bench_%_resolved
+  value: 70.9
+  source:
+    url: https://huggingface.co/poolside/Laguna-XS-2.1
+    name: Model Card

.eval_results/terminal-bench-2.0.yaml

@@ -0,0 +1,7 @@
+- dataset:
+    id: harborframework/terminal-bench-2.0
+    task_id: terminalbench_2
+  value: 37.5
+  source:
+    url: https://huggingface.co/poolside/Laguna-XS-2.1
+    name: Model Card

.gitattributes

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

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,429 @@
+---
+library_name: transformers
+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-xs-2.1
+- vllm
+license: openmdw-1.1
+pipeline_tag: text-generation
+---
+
+<p align="center">
+  <img alt="poolside-banner" src="https://poolside.ai/assets/laguna/laguna-xs2-1-banner.svg" width="800px">
+</p>
+
+<p align="center">
+  <a href="https://openrouter.ai/poolside/laguna-xs-2.1"><strong>Use on OpenRouter</strong></a> ·
+  <a href="https://poolside.ai/blog/introducing-laguna-xs-2-1"><strong>Release blog post</strong></a>
+</p>
+
+<br>
+
+# Laguna XS 2.1
+Laguna XS 2.1 is a 33B total parameter Mixture-of-Experts model with 3B activated parameters per token designed for agentic coding and long-horizon work on a local machine. This model is an upgraded version of our [Laguna XS.2](https://huggingface.co/poolside/Laguna-XS.2) model with a +5.4% jump on SWE-bench Multilingual as well as stronger performance on terminal-style tasks.
+
+> [!NOTE]
+> For more details on how we train, including on data automixing and async off-policy agent RL, check out our recent [technical report](https://poolside.ai/assets/laguna/laguna-m1-xs2-technical-report.pdf).
+
+## Highlights
+- **Mixed SWA and global attention layout**: Laguna XS 2.1 uses sigmoid gating with per-layer rotary scales, enabling mixed SWA (Sliding Window Attention) and global attention layers in a 3:1 ratio (across 40 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 33B total parameters and 3B activated, Laguna XS 2.1 is compact enough to run on a Mac with 36 GB of RAM. Available on [Ollama](https://ollama.com/library/laguna-xs-2.1) and [llama.cpp](https://github.com/ggml-org/llama.cpp/pull/25165). High-quality FP8, NVFP4 and INT4 quantized variants available ([see the collection](https://huggingface.co/collections/poolside/laguna-xs-21))
+- **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: 33B total with 3B activated per token
+- Optimizer: Muon
+- Layers: 40 layers (10 layers with global attention, 30 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-xs2-1-chart.svg" width="800px">
+</p>
+
+| Model                     | Size (total params.) | SWE-bench Verified | SWE-bench Multilingual | SWE-Bench Pro (Public Dataset) | Terminal-Bench 2.0 |
+|---------------------------|----------------------|--------------------|------------------------|--------------------------------|--------------------|
+| **Laguna XS 2.1**         | 33B                  | 70.9%              | 63.1%                  | 47.6%                          | 37.5%              |
+| **Laguna XS.2**           | 33B                  | 69.9%              | 57.7%                  | 46.3%                          | 35.7%              |
+| Qwen3.6-35B-A3B           | 35B                  | 73.4%              | 67.2%                  | 49.5%                          | 51.5%              |
+| North Mini Code           | 30B                  | 67.6%              | -                      | 40.2%                          | 36.0%              |
+| MAI-Code-1-Flash          | 137B                 | 71.6%              | 65.5%                  | 51.2%                          | 54.8%              |
+| gpt-oss-120B              | 120B                 | -                  | -                      | 16.2%                          | 18.7%              |
+| Claude Haiku 4.5          | -                    | 73.3%              | -                      | 39.5%                          | 29.8%              |
+| GPT-5.4 Nano              | -                    | -                  | -                      | 52.4%                          | 46.3%              |
+
+*We used the highest publicly-referenced scores for all comparison models across each benchmark. In all cases these were official scores published in release blog posts or equivalent, with the exception of gpt-oss-120b and Claude Haiku 4.5 where the highest published (verified) scores for SWE-Bench Pro and Terminal-Bench 2.0 are from their respective official leaderboards.*
+
+<details>
+<summary>Expand for benchmarking methodology</summary>
+
+All benchmarking for Laguna XS 2.1 was completed using Laude Institute’s Harbor Framework with our [agent harness](https://github.com/poolsideai/pool), with a maximum of 500 steps and sandboxed execution. The same sampling parameters were used for all Laguna XS 2.1 benchmarking: temperature=1.0, top_k=20 and top_p=1, with thinking mode enabled and a context length of 256K tokens. All tasks were run in their own sandbox using 8 GB RAM/2 CPUs, with the exception of Terminal-Bench 2.0, which used 48 GB RAM/32 CPUs.
+
+Some base task images and verifiers were patched to fix infrastructure reliability issues inherent in task setup, such as rate limits on third-party dependencies in external registries used by the verifier. All four agentic benchmarks were run with patched images. We also ran a reward-hack judge post-hoc on Laguna XS 2.1 evaluation runs and did not find significant reward hacking after joint judge review and manual review.
+
+- SWE-bench Verified: mean pass@1 averaged over 4 attempts per task
+- SWE-bench Multilingual: mean pass@1 averaged over 4 attempts per task
+- SWE-Bench Pro: mean pass@1 averaged over 2 attempts per task
+- Terminal-Bench 2.0: mean pass@1 averaged over 5 attempts per task; 48 GB RAM/32 CPUs
+
+</details>
+
+## Usage
+
+Laguna XS 2.1 has launch-day support in vLLM, SGLang, Transformers and Llama.cpp, and TRT-LLM thanks to the support of the team at NVIDIA.
+
+The fastest way to get started is using OpenRouter.
+
+> [!NOTE]
+> We are providing free inference for a limited time for Laguna XS 2.1, as well as our larger 225B model, Laguna M.1. Visit our [provider page on OpenRouter](https://openrouter.ai/poolside/) to get started.
+
+### pool
+
+**pool** is a lightweight terminal-based coding agent and a dual [Agent Client Protocol](https://agentclientprotocol.com/get-started) client-server.
+
+Download and install for macOS and Linux:
+
+```shell
+curl -fsSL https://downloads.poolside.ai/pool/install.sh | bash
+```
+
+Launch and `> Log in with Poolside` to get a free, limited-use API key.
+
+Alternatively, use Poolside models alongside the wide catalog of models available on OpenRouter by selecting `> Log in with OpenRouter`.
+
+```shell
+pool login
+```
+
+Use in any [ACP client](https://agentclientprotocol.com/get-started/clients). Configure Zed and JetBrains automatically:
+
+```shell
+pool acp setup --editor zed|jetbrains
+```
+
+Use pool with Ollama with one-command setup:
+
+```shell
+ollama pull laguna-xs-2.1
+ollama launch pool --model laguna-xs-2.1
+```
+
+#### Feedback and issues
+
+Submit feedback with `/feedback` and read the [full documentation on GitHub](https://github.com/poolsideai/pool).
+
+### Local deployment
+
+Laguna XS 2.1 is supported in vLLM, SGLang, Transformers and Llama.cpp, and TRT-LLM thanks to the support of the team at NVIDIA. Use Laguna-XS 2.1 with Ollama (with MLX support) and the mlx-lm framework for the best experience on your local machine.
+
+#### vLLM
+
+Serve Laguna XS 2.1 locally with vLLM and query it from any OpenAI-compatible client (see [Controlling reasoning](#controlling-reasoning) for tool calls, streaming, and reasoning extraction):
+
+> [!NOTE]
+> Laguna XS 2.1 support is available in vLLM 0.21.0 and later ([vllm-project/vllm#41129](https://github.com/vllm-project/vllm/pull/41129)).
+
+```shell
+pip install 'vllm>=0.21.0'
+
+vllm serve \
+    --model poolside/Laguna-XS-2.1 \
+    --tool-call-parser poolside_v1 \
+    --reasoning-parser poolside_v1 \
+    --enable-auto-tool-choice \
+    --served-model-name laguna \
+    --default-chat-template-kwargs '{"enable_thinking": true}'
+```
+
+See the [vLLM recipes page](https://recipes.vllm.ai/poolside/Laguna-XS-2.1) for additional deployment guidance.
+
+> [!NOTE]
+> **Optional: speculative decoding with DFlash.** For lower latency, pair Laguna XS 2.1 with the [DFlash speculator](https://huggingface.co/poolside/Laguna-XS-2.1-DFlash), a 5-layer Llama-style draft model that proposes up to 7 tokens per step at ~70% per-position acceptance on coding tasks. vLLM support is in progress in [vllm-project/vllm#46853](https://github.com/vllm-project/vllm/pull/46853); once it lands, add `--speculative-config '{"model":"poolside/Laguna-XS-2.1-DFlash","num_speculative_tokens":7,"method":"dflash"}'` to the serve command above.
+
+#### SGLang
+
+Laguna XS 2.1 is supported in SGLang via [sgl-project/sglang#24204](https://github.com/sgl-project/sglang/pull/24204). See the [SGLang cookbook entry](https://docs.sglang.io/cookbook/autoregressive/Poolside/Laguna-XS-2.1) for a serving recipe.
+
+```shell
+git clone https://github.com/sgl-project/sglang.git
+cd sglang
+pip install -e "python[all]"
+
+python -m sglang.launch_server \
+  --model-path poolside/Laguna-XS-2.1 \
+  --tp-size 8 \
+  --mem-fraction-static 0.7 \
+  --reasoning-parser poolside_v1 \
+  --trust-remote-code
+```
+
+> [!NOTE]
+> **Optional: speculative decoding with DFlash.** The [DFlash speculator](https://huggingface.co/poolside/Laguna-XS-2.1-DFlash) can be paired with Laguna XS 2.1 for lower latency. SGLang support was added in [sgl-project/sglang#29446](https://github.com/sgl-project/sglang/pull/29446). Add `--speculative-algorithm DFLASH \ --speculative-draft-model-path poolside/Laguna-XS-2.1-DFlash-FP8` to the serve command above.
+
+#### Transformers
+
+Laguna XS 2.1 is supported in Transformers `v5.7.0` and later ([huggingface/transformers#45673](https://github.com/huggingface/transformers/pull/45673)).
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_id = "poolside/Laguna-XS-2.1"
+
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    dtype=torch.bfloat16,
+    device_map="auto",
+)
+
+messages = [
+    {"role": "user", "content": "Write a Python retry wrapper with exponential backoff."},
+]
+
+# Reasoning is on by default; pass enable_thinking=False to skip the <think> block.
+inputs = tokenizer.apply_chat_template(
+    messages,
+    add_generation_prompt=True,
+    return_tensors="pt",
+    enable_thinking=True,
+).to(model.device)
+
+outputs = model.generate(
+    inputs,
+    max_new_tokens=1024,
+    do_sample=True,
+    temperature=1.0,
+    top_k=20,
+)
+
+response = tokenizer.decode(outputs[0][inputs.shape[-1]:], skip_special_tokens=True)
+print(response)
+```
+
+#### TRT-LLM
+
+Laguna XS 2.1 support is merged into TensorRT-LLM ([NVIDIA/TensorRT-LLM#13559](https://github.com/NVIDIA/TensorRT-LLM/pull/13559)) and ships in the `v1.3.0rc16` pre-release wheels and later.
+
+> [!NOTE]
+> A stable `v1.3.0` has not been released yet, so install a pre-release wheel from NVIDIA's package index (the latest stable, `v1.2.x`, does not include Laguna XS 2.1 support). To build from source instead, see the TensorRT-LLM build documentation. Laguna XS 2.1 support is on `main`.
+
+Install the **CUDA-13 `torch` build first**, then TensorRT-LLM. The default PyPI `torch` is a CUDA-12 build whose `cuda-bindings` pin conflicts with TRT-LLM's `cuda-python 13.x`, so a bare `pip install tensorrt-llm` fails to resolve; installing the cu130 `torch` up front avoids it.
+
+```shell
+# 1. CUDA-13 torch build (pins cuda-bindings 13.x, matching TRT-LLM's cuda-python)
+pip install 'torch==2.10.0' torchvision --index-url https://download.pytorch.org/whl/cu130
+
+# 2. TRT-LLM from NVIDIA's index (torch already satisfied, so it is not replaced)
+pip install --pre 'tensorrt-llm>=1.3.0rc16' \
+    --extra-index-url https://pypi.nvidia.com \
+    --extra-index-url https://download.pytorch.org/whl/cu130
+```
+
+This resolves to `tensorrt-llm 1.3.0rc20` with `torch 2.10.0+cu130`, `cuda-python 13.0.3`, and `transformers 5.5.4`.
+
+Load the checkpoint directly with `trust_remote_code=True`. No `transformers` compatibility overlay is required: `v1.3.0rc16+` pins `transformers 5.5.4`, which provides the symbols Laguna XS 2.1's config needs (earlier TRT-LLM releases pinned `transformers 4.57`, which did not).
+
+```python
+from tensorrt_llm import LLM, SamplingParams
+
+llm = LLM(
+    model="poolside/Laguna-XS-2.1",
+    trust_remote_code=True,
+    tensor_parallel_size=1,
+)
+
+sampling = SamplingParams(max_tokens=1024, temperature=1.0, top_k=20)
+out = llm.generate(["Write a Python retry wrapper with exponential backoff."], sampling)
+print(out[0].outputs[0].text)
+```
+
+Or serve with an OpenAI-compatible endpoint:
+
+```shell
+trtllm-serve poolside/Laguna-XS-2.1 --port 8000 --trust-remote-code --tool_parser poolside_v1 --reasoning_parser laguna
+```
+
+The Laguna XS 2.1 tool-call and reasoning parsers are built into TRT-LLM `>=1.3.0rc16` (shipped with [#13559](https://github.com/NVIDIA/TensorRT-LLM/pull/13559)), so no extra install is needed. Note that the flag names differ from vLLM's (`--tool_parser`, and the reasoning parser is `laguna`, not `poolside_v1`).
+
+The same recipe works for the [FP8](https://huggingface.co/poolside/Laguna-XS-2.1-FP8) and [NVFP4](https://huggingface.co/poolside/Laguna-XS-2.1-NVFP4) variants: quantization is detected automatically from `quantization_config`, no extra flags required.
+
+> [!NOTE]
+> **Optional: speculative decoding with DFlash.** The [DFlash speculator](https://huggingface.co/poolside/Laguna-XS-2.1-DFlash) can be paired with Laguna XS 2.1 for lower latency. TRT-LLM support is in progress in [NVIDIA/TensorRT-LLM#15666](https://github.com/NVIDIA/TensorRT-LLM/pull/15666).
+
+#### llama.cpp
+
+> [!NOTE]
+> Requires building llama.cpp from the upstream PR that adds Laguna XS 2.1 support until it lands ([ggml-org/llama.cpp#25165](https://github.com/ggml-org/llama.cpp/pull/25165)).
+
+Official GGUF conversions (BF16 and Q4\_K\_M) are available at [poolside/Laguna-XS-2.1-GGUF](https://huggingface.co/poolside/Laguna-XS-2.1-GGUF).
+
+```shell
+# Build llama.cpp from the PR branch
+git clone https://github.com/ggml-org/llama.cpp && cd llama.cpp
+git fetch origin pull/25165/head:laguna && git checkout laguna
+cmake -B build && cmake --build build -j
+
+# Download a GGUF and serve an OpenAI-compatible endpoint
+huggingface-cli download poolside/Laguna-XS-2.1-GGUF Laguna-XS-2.1-Q4_K_M.gguf --local-dir ~/models/Laguna-XS-2.1-GGUF
+./build/bin/llama-server -m ~/models/Laguna-XS-2.1-GGUF/Laguna-XS-2.1-Q4_K_M.gguf --jinja --port 8000
+```
+
+#### Ollama
+
+Available on the [Ollama library](https://ollama.com/library/laguna-xs-2.1).
+
+    ollama run laguna-xs-2.1          # default — Q4_K_M (imatrix)
+    ollama run laguna-xs-2.1:q8_0     # higher precision
+    ollama run laguna-xs-2.1:bf16     # full precision
+
+Reasoning and tool-calling work out of the box via the built-in `laguna` template.
+
+> [!NOTE]
+> **macOS (Metal) users:** Chat (`ollama run` / `/api/chat`) works as expected on Linux/CUDA. On macOS/Metal it may currently return empty output; the root cause is not yet fully understood and we're investigating it with the Ollama team. On a Mac, use a Linux/CUDA host, or the `/api/generate` endpoint with `"raw": true`.
+
+## Controlling reasoning
+
+Laguna XS 2.1 has native reasoning support and is designed to work best with *preserved thinking*, where `reasoning` content from prior assistant messages is preserved in the message history. This model will generally reason before calling tools and between tool calls.
+
+Reasoning may not be generated in follow-up steps if prior thinking blocks are dropped (i.e., thinking is not preserved) when messages are reconstructed over multiple steps.
+
+<details>
+<summary>Expand for example</summary>
+
+```python
+import json
+from openai import OpenAI
+
+client = OpenAI(
+  base_url="https://openrouter.ai/api/v1",
+  api_key="...",
+)
+
+model = "poolside/laguna-xs-2.1"
+
+tools = [{"type": "function", "function": {
+  "name": "shell",
+  "description": "Execute a bash command and return the output.",
+  "parameters": {"type": "object", "properties": {"cmd": {"type": "string"}}, "required": ["cmd"]},
+}}]
+
+messages = [
+  {"role": "system", "content": "You are a coding agent with access to a shell tool."},
+  {"role": "user", "content": "Run uname -a"},
+]
+
+# Thinking is enabled by default when the server sets --default-chat-template-kwargs {"enable_thinking": True}
+# When using OpenRouter's Chat API (https://openrouter.ai/api/v1), this flag is set by default
+response = client.chat.completions.create(
+  model=model,
+  messages=messages,
+  tools=tools,
+  stream=True,
+)
+
+reasoning, content, tool_calls = "", "", []
+for chunk in response:
+  delta = chunk.choices[0].delta
+  if hasattr(delta, "reasoning_content") and delta.reasoning_content:
+    reasoning += delta.reasoning_content
+  if hasattr(delta, "content") and delta.content:
+    content += delta.content
+  if hasattr(delta, "tool_calls") and delta.tool_calls:
+    for tc in delta.tool_calls:
+      if tc.index >= len(tool_calls):
+        tool_calls.append({"id": tc.id, "function": {"name": "", "arguments": ""}})
+      if tc.function.name:
+        tool_calls[tc.index]["function"]["name"] = tc.function.name
+      if tc.function.arguments:
+        tool_calls[tc.index]["function"]["arguments"] += tc.function.arguments
+
+print(f"Reasoning: {reasoning}\nContent: {content}\nTool calls: {tool_calls}\n")
+
+# Return reasoning in the next request for best performance
+messages.append({
+  "role": "assistant",
+  "content": content,
+  "reasoning_content": reasoning,
+  "tool_calls": [{"id": tc["id"], "type": "function", "function": tc["function"]} for tc in tool_calls]
+})
+
+messages.append({
+  "role": "tool",
+  "tool_call_id": tool_calls[0]["id"],
+  "content": json.dumps({"stdout": "Darwin arm64", "exit_code": "0"})
+})
+
+response = client.chat.completions.create(
+  model=model,
+  messages=messages,
+  tools=tools,
+  stream=True,
+)
+
+reasoning, content = "", ""
+for chunk in response:
+  delta = chunk.choices[0].delta
+  if hasattr(delta, "reasoning_content") and delta.reasoning_content:
+    reasoning += delta.reasoning_content
+  if hasattr(delta, "content") and delta.content:
+    content += delta.content
+
+print(f"Reasoning: {reasoning}\nContent: {content}")
+```
+
+</details>
+
+### Disabling reasoning
+
+You can disable thinking by setting `enable_thinking` to `False` in a request or by not providing `--default-chat-template-kwargs {"enable_thinking": True}` or equivalent when starting the server.
+
+<details>
+<summary>Expand for example</summary>
+
+```python
+from openai import OpenAI
+client = OpenAI()
+
+completion = client.chat.completions.create(
+  model="poolside/laguna-xs-2.1",
+  messages=[
+    {"role": "user", "content": "Write a retry wrapper with exponential backoff."}
+  ],
+  extra_body={
+    "chat_template_kwargs": { "enable_thinking": False },
+  },
+  stream=True
+)
+
+for chunk in completion:
+    print(chunk.choices[0].delta)
+```
+
+</details>
+
+For agentic coding use cases, we recommend enabling thinking and preserving reasoning in message history as outlined in the [Controlling reasoning](#controlling-reasoning) section.
+
+## License
+
+This model is licensed under the [OpenMDW-1.1 License](https://huggingface.co/poolside/Laguna-XS-2.1/blob/main/LICENSE.md).
+
+## Intended and Responsible Use 
+
+Laguna XS 2.1 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 XS 2.1 is subject to the [OpenMDW-1.1 License](https://huggingface.co/poolside/Laguna-XS-2.1/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 XS 2.1 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,132 @@
+{#- Copied from laguna_glm_thinking_v4/chat_template.jinja -#}
+{#- Removes prefix that references <think> token, and replaces message.reasoning_content reference with message.reasoning -#}
+{{- "〈|EOS|〉" -}}
+{%- set enable_thinking = enable_thinking | default(false) -%}
+{%- set render_assistant_messages_raw = render_assistant_messages_raw | default(false) -%}
+{%- set add_generation_prompt = add_generation_prompt | default(false) -%}
+
+{#- ───── header (system message) ───── -#}
+{%- set system_message = "" -%}
+{%- if messages and messages[0].role == "system" -%}
+  {%- set system_message = messages[0].content -%}
+{%- endif -%}
+
+{%- if (system_message and system_message.strip()) or tools -%}
+  {{- "<system>\n" -}}
+
+  {%- if system_message and system_message.strip() -%}
+    {{- "\n" -}}
+    {{- system_message.rstrip() -}}
+  {%- endif -%}
+
+  {%- if tools -%}
+    {{- "\n\n### Tools\n\n" -}}
+    {%- set ns = namespace(tool_string="You may call functions to assist with the user query.\n"
+          ~ "All available function signatures are listed below:\n"
+          ~ "<available_tools>\n") -%}
+    {%- for tool in tools -%}
+      {%- set ns.tool_string = ns.tool_string ~ (tool | tojson) ~ "\n" -%}
+    {%- endfor -%}
+    {%- if enable_thinking -%}
+      {%- set tool_string = ns.tool_string + "</available_tools>\n\n" ~
+            "Wrap your thinking in '<think>', '</think>' tags, followed by a function call. For each function call, return an unescaped XML-like object with function name and arguments within '<tool_call>' and '</tool_call>' tags, like here:\n" ~
+            "<think> your thoughts here </think>\n" ~
+            "<tool_call>function-name\n<arg_key>argument-key</arg_key>\n<arg_value>value-of-argument-key</arg_value>\n" ~
+            "</tool_call>" -%}
+    {%- else -%}
+      {%- set tool_string = ns.tool_string + "</available_tools>\n\n" ~
+            "For each function call, return an unescaped XML-like object " ~
+            "with function name and arguments within '<tool_call>' and '</tool_call>' tags, like here:\n" ~
+            "<tool_call>function-name\n<arg_key>argument-key</arg_key>\n<arg_value>value-of-argument-key</arg_value>\n" ~
+            "</tool_call>" -%}
+    {%- endif -%}
+    {{- tool_string -}}
+  {%- endif -%}
+
+  {{- "\n</system>\n" -}}
+{%- endif -%}
+
+{#- ───── main loop ───── -#}
+{%- for message in messages -%}
+  {%- set content = message.content if message.content is string else "" -%}
+  {%- if message.role == "user" -%}
+    {{- "<user>\n" + content + "\n</user>\n" -}}
+  {%- elif message.role == "assistant" -%}
+    {%- generation -%}
+      {{- "<assistant>\n" -}}
+      {%- if render_assistant_messages_raw -%}
+        {#- Raw mode: prepend the generation prompt token, then dump content verbatim. -#}
+        {#- The generation prompt is <think> when enable_thinking, </think> otherwise. -#}
+        {#- Only prepend if content doesn't already start with it. -#}
+        {%- if enable_thinking -%}
+          {%- if not content.startswith('<think>') -%}
+            {{- '<think>' -}}
+          {%- endif -%}
+        {%- else -%}
+          {%- if not content.startswith('</think>') -%}
+            {{- '</think>' -}}
+          {%- endif -%}
+        {%- endif -%}
+        {{- content -}}
+        {#- Append closing tag if content doesn't already end with it. -#}
+        {%- if not content.endswith('</assistant>\n') and not content.endswith('</assistant>') -%}
+          {{- '\n</assistant>' -}}
+        {%- endif -%}
+        {{- "\n" -}}
+      {%- else -%}
+        {#- Extract reasoning content from message.reasoning (vLLM field name) or message.reasoning_content, or from <think> tags -#}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning is string %}
+          {%- set reasoning_content = message.reasoning %}
+        {%- elif message.reasoning_content is string %}
+          {%- set reasoning_content = message.reasoning_content %}
+        {%- endif %}
+        {#- Always strip <think> tags from content if present to avoid duplication -#}
+        {%- if '</think>' in content %}
+          {%- if not reasoning_content %}
+            {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+          {%- endif %}
+          {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+        {%- endif %}
+        {#- Display reasoning content for all messages -#}
+        {%- if reasoning_content -%}
+          {{- '<think>\n' + reasoning_content.strip() + '\n</think>\n' -}}
+        {%- else -%}
+          {{- '</think>\n' -}}
+        {%- endif -%}
+        {#- Display main content -#}
+        {%- if content.strip() -%}
+          {{- content.strip() ~ "\n" -}}
+        {%- endif -%}
+        {%- if message.tool_calls -%}
+          {%- for tool_call in message.tool_calls -%}
+            {%- set function_data = tool_call.function -%}
+            {{- '<tool_call>' + function_data.name }}
+            {% set _args = function_data.arguments %}
+            {%- for k, v in _args.items() -%}
+              {{- "<arg_key>" ~ k ~ "</arg_key>\n" -}}
+              {{- "<arg_value>"}}{{ v | tojson(ensure_ascii=False) if v is not string else v }}{{ "</arg_value>\n" -}}
+            {%- endfor -%}
+            {{- "</tool_call>\n" -}}
+          {%- endfor -%}
+        {%- endif -%}
+        {{- "</assistant>\n" -}}
+      {%- endif -%}
+    {%- endgeneration -%}
+  {%- elif message.role == "tool" -%}
+    {{- "<tool_response>\n" + content + "\n</tool_response>\n" -}}
+  {%- elif message.role == "system" and loop.index0 != 0 -%}
+    {#- Render additional system messages (skip the first one which is handled separately in the header) -#}
+    {{- "<system>\n" + content + "\n</system>\n" -}}
+  {%- endif -%}
+{%- endfor -%}
+{#- ───── generation prompt ───── -#}
+{%- if add_generation_prompt -%}
+  {{- "<assistant>\n" -}}
+  {#- ───── Include reasoning mode directive ───── -#}
+  {%- if not enable_thinking %}
+    {{- '</think>' -}}
+  {%- else %}
+    {{- '<think>' -}}
+  {%- endif %}
+{%- endif -%}

config.json

@@ -0,0 +1,229 @@
+{
+  "architectures": [
+    "LagunaForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_laguna.LagunaConfig",
+    "AutoModelForCausalLM": "modeling_laguna.LagunaForCausalLM"
+  },
+  "model_type": "laguna",
+  "vocab_size": 100352,
+  "hidden_size": 2048,
+  "intermediate_size": 8192,
+  "num_hidden_layers": 40,
+  "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": 8,
+  "moe_intermediate_size": 512,
+  "shared_expert_intermediate_size": 512,
+  "norm_topk_prob": true,
+  "router_aux_loss_coef": 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": 64.0,
+      "attention_factor": 1.0,
+      "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"
+  ],
+  "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"
+  ],
+  "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"
+  ],
+  "moe_routed_scaling_factor": 2.5,
+  "num_attention_heads_per_layer": [
+    48,
+    64,
+    64,
+    64,
+    48,
+    64,
+    64,
+    64,
+    48,
+    64,
+    64,
+    64,
+    48,
+    64,
+    64,
+    64,
+    48,
+    64,
+    64,
+    64,
+    48,
+    64,
+    64,
+    64,
+    48,
+    64,
+    64,
+    64,
+    48,
+    64,
+    64,
+    64,
+    48,
+    64,
+    64,
+    64,
+    48,
+    64,
+    64,
+    64
+  ]
+}

configuration_laguna.py

@@ -0,0 +1,245 @@
+# 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
+from transformers.modeling_rope_utils import RopeParameters
+from transformers.utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available
+
+
+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}
+
+        # 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,24 @@
+{
+ "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,
+ "min_p": 0.0,
+ "speculative_config": {
+  "method": "dflash",
+  "source": "huggingface",
+  "model": "poolside/Laguna-XS-2.1-DFlash",
+  "num_speculative_tokens": 15
+ },
+ "tool_call_parser": "poolside_v1",
+ "reasoning_parser": "poolside_v1",
+ "default_chat_template_kwargs": {
+  "enable_thinking": true
+ }
+}

model-00001-of-00014.safetensors

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

model-00002-of-00014.safetensors

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

model-00003-of-00014.safetensors

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

model-00004-of-00014.safetensors

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

model-00005-of-00014.safetensors

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

model-00006-of-00014.safetensors

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

model-00007-of-00014.safetensors

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

model-00008-of-00014.safetensors

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

model-00009-of-00014.safetensors

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

model-00010-of-00014.safetensors

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

model-00011-of-00014.safetensors

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

model-00012-of-00014.safetensors

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

model-00013-of-00014.safetensors

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

model-00014-of-00014.safetensors

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

modeling_laguna.py

@@ -0,0 +1,879 @@
+# 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 ...cache_utils import DynamicCache
+        from ...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"]

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,575 @@
+{
+  "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|〉"
+}