Laguna S 2.1 initial release

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

.gitattributes

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

README.md

@@ -0,0 +1,54 @@
+---
+library_name: speculators
+base_model:
+- poolside/Laguna-S-2.1-FP8
+tags:
+- speculative-decoding
+- dflash
+- speculators
+---
+
+<p align="center">
+  <img alt="poolside-banner" src="https://poolside.ai/assets/laguna/laguna-s-2-1-banner.svg" width="800px">
+</p>
+
+<p align="center">
+  <a href="https://openrouter.ai/poolside/laguna-s-2.1"><strong>Use on OpenRouter</strong></a> ·
+  <a href="https://vercel.com/ai-gateway/models/laguna-s-2.1"><strong>Use on Vercel AI Gateway</strong></a> ·
+  <a href="https://poolside.ai/blog/introducing-laguna-s-2-1"><strong>Release blog post</strong></a>
+</p>
+
+<br>
+
+# poolside/Laguna-S-2.1-DFlash-FP8
+
+DFlash speculator for the FP8 target [poolside/Laguna-S-2.1-FP8](https://huggingface.co/poolside/Laguna-S-2.1-FP8). The
+speculator is a 6-layer Laguna-style draft model (BF16); pair it with the FP8 base for
+lower-latency serving via speculative decoding.
+
+Trained: `e0630_rhiemann_baseline` SFT, DFlash Stage-2, 15k steps. Recommended
+serving setting: `num_speculative_tokens=7`.
+DFlash upstream support is in progress (vLLM #46853, SGLang #29446, TRT-LLM #15666). Use
+`poolside/Laguna-S-2.1-FP8` as the target model.
+
+## Benchmarks
+
+Measured with TP=2, `temperature=0`, and `num_speculative_tokens=15`.
+
+### Throughput speedup
+
+| Concurrency | GSM8K | MATH-500 | HumanEval | MBPP | MT-Bench |
+|---:|---:|---:|---:|---:|---:|
+| 1 | 3.179x | 2.938x | 3.269x | 2.380x | 2.603x |
+| 4 | 2.614x | 2.423x | 2.691x | 1.963x | 2.090x |
+| 8 | 2.666x | 2.410x | 2.803x | 1.962x | 2.230x |
+| 16 | 2.618x | 2.364x | 2.866x | 2.031x | 2.302x |
+
+### Acceptance length
+
+| Concurrency | GSM8K | MATH-500 | HumanEval | MBPP | MT-Bench |
+|---:|---:|---:|---:|---:|---:|
+| 1 | 5.748 | 5.197 | 5.889 | 4.247 | 4.663 |
+| 4 | 5.765 | 5.212 | 5.882 | 4.218 | 4.411 |
+| 8 | 5.863 | 5.199 | 6.094 | 4.178 | 4.572 |
+| 16 | 5.787 | 5.161 | 6.144 | 4.291 | 4.600 |

config.json

@@ -0,0 +1,61 @@
+{
+  "attention_bias": false,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 3072,
+  "intermediate_size": 12288,
+  "max_position_embeddings": 262144,
+  "model_type": "laguna",
+  "num_attention_heads": 72,
+  "num_hidden_layers": 6,
+  "num_key_value_heads": 8,
+  "rms_norm_eps": 1e-06,
+  "sliding_window": 512,
+  "vocab_size": 100352,
+  "layer_types": [
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention"
+  ],
+  "rope_theta": 500000.0,
+  "gating": "per-head",
+  "architectures": [
+    "DFlashLagunaForCausalLM"
+  ],
+  "num_experts": 0,
+  "sliding_windows": [
+    512,
+    512,
+    512,
+    512,
+    512,
+    512
+  ],
+  "draft_vocab_size": 100352,
+  "torch_dtype": "bfloat16",
+  "eagle_aux_hidden_state_layer_ids": [
+    2,
+    11,
+    20,
+    30,
+    39,
+    48
+  ],
+  "dflash_config": {
+    "block_size": 16,
+    "mask_token_id": 12,
+    "num_target_layers": 48,
+    "target_layer_ids": [
+      1,
+      10,
+      19,
+      29,
+      38,
+      47
+    ],
+    "causal": true
+  }
+}

config.py

@@ -0,0 +1,188 @@
+from typing import Any, Literal
+
+from pydantic import Field, field_serializer, field_validator
+from transformers import AutoConfig, PretrainedConfig
+from transformers.models.qwen3.modeling_qwen3 import (
+    Qwen3Config,
+)
+
+from speculators import SpeculatorModelConfig
+
+__all__ = [
+    "DFlashSpeculatorConfig",
+]
+
+
+@SpeculatorModelConfig.register("dflash")
+class DFlashSpeculatorConfig(SpeculatorModelConfig):
+    """
+    Configuration for DFlash speculator with vocabulary mapping.
+
+    DFlash features vocabulary mapping between draft (64K) and target (128K)
+    vocabularies, enabling cross-tokenizer speculation.
+
+    :param transformer_layer_config: Configuration for the transformer decoder layer
+    :param draft_vocab_size: Size of draft model vocabulary for speculation
+    """
+
+    speculators_model_type: Literal["dflash"] = "dflash"
+    architectures: list[str] = Field(
+        default_factory=lambda: ["DFlashSpeculator"],
+        description="Model architectures that can load these weights",
+    )
+
+    transformer_layer_config: PretrainedConfig = Field(
+        default_factory=Qwen3Config,
+        description="Configuration for the transformer decoder layer",
+    )
+
+    draft_vocab_size: int = Field(
+        default=32000,
+        description="Size of draft model vocabulary for speculation",
+    )
+
+    block_size: int = Field(
+        default=8,
+        description=(
+            "Default size of the draft block predicted with a forward pass of the model"
+        ),
+    )
+
+    max_anchors: int = Field(
+        default=256,
+        description=(
+            "Maximum number of anchor positions to sample during training "
+            "(controls memory usage and training efficiency)"
+        ),
+    )
+
+    target_hidden_size: int | None = Field(
+        default=None,
+        description="Hidden size of the target model (if different from draft model)",
+    )
+
+    aux_hidden_state_layer_ids: list[int] | None = Field(
+        default=None,
+        description="Layer IDs of the DFlash auxiliary hidden state layers",
+    )
+
+    decoder_layer_type: Literal["qwen3", "laguna_xs"] = Field(
+        default="qwen3",
+        description="Decoder layer implementation used by the DFlash drafter.",
+    )
+
+    mask_token_id: int | None = Field(
+        default=None,
+        description="Token ID used for masking",
+    )
+
+    sliding_window_non_causal: bool = Field(
+        default=False,
+        description="Use non-causal synthetic block attention for sliding-window layers.",
+    )
+
+    sliding_window_base: Literal["fixed_anchor", "moving_query"] = Field(
+        default="moving_query",
+        description=(
+            "Base-token sliding-window lower-bound policy. 'moving_query' matches "
+            "FlashAttention-style SWA during inference; 'fixed_anchor' preserves "
+            "the legacy DFlash training mask."
+        ),
+    )
+
+    loss_type: Literal["distill", "dflash", "lk", "tv"] = Field(
+        default="distill",
+        description="DFlash objective. 'lk' uses hard-label LK loss.",
+    )
+
+    ce_weight: float | None = Field(
+        default=None,
+        description="Additive weight for hard-label DFlash CE.",
+    )
+
+    tv_weight: float | None = Field(
+        default=None,
+        description="Additive weight for full-distribution TV loss.",
+    )
+
+    kl_weight: float | None = Field(
+        default=None,
+        description="Additive weight for full-distribution KL distillation.",
+    )
+
+    lk_lambda: float = Field(
+        default=0.5,
+        description="Blend coefficient for hard-label LK loss.",
+    )
+
+    tv_temperature: float = Field(
+        default=1.0,
+        description="Teacher softmax temperature for TV/KL terms.",
+    )
+
+    cumacc_weight: bool = Field(
+        default=False,
+        description="Weight hard-label DFlash CE by draft cumulative acceptance.",
+    )
+
+    veri_cum_acc: bool = Field(
+        default=False,
+        description="Weight DFlash loss by verifier cumulative acceptance.",
+    )
+
+    veri_acc_temperature: float = Field(
+        default=1.0,
+        description="Temperature for verifier cumulative acceptance weighting.",
+    )
+
+    static_decay_weight: bool = Field(
+        default=True,
+        description="Apply DFlash position decay to hard-label CE.",
+    )
+
+    kl_distill_weight: float = Field(
+        default=0.0,
+        description="Back-compatible alias for kl_weight when kl_weight is unset.",
+    )
+
+    compile_decoder_layers: bool = Field(
+        default=True,
+        description=(
+            "If True, torch.compile each decoder layer forward during training. "
+            "The DFlash loss remains eager."
+        ),
+    )
+
+    @field_serializer("transformer_layer_config")
+    def serialize_transformer_config(self, value: PretrainedConfig) -> dict:
+        """Serialize transformer config to dict."""
+        return value.to_diff_dict()
+
+    @field_validator("transformer_layer_config", mode="before")
+    @classmethod
+    def validate_transformer_config(cls, value: Any) -> PretrainedConfig:
+        """Validate and convert transformer config."""
+        if isinstance(value, dict):
+            config_class: type[PretrainedConfig] = Qwen3Config
+            if "model_type" in value:
+                config_class = AutoConfig.for_model(
+                    model_type=value["model_type"]
+                ).__class__
+            return config_class(**value)
+        return value
+
+    @property
+    def target_vocab_size(self) -> int:
+        """Get target vocabulary size from transformer config."""
+        return self.transformer_layer_config.vocab_size
+
+    def resolve_loss_weights(self) -> tuple[float, float, float]:
+        if self.loss_type == "tv":
+            ce_default, tv_default = 0.0, 1.0
+        else:
+            ce_default, tv_default = 1.0, 0.0
+
+        ce = ce_default if self.ce_weight is None else self.ce_weight
+        tv = tv_default if self.tv_weight is None else self.tv_weight
+        kl = self.kl_distill_weight if self.kl_weight is None else self.kl_weight
+        return float(ce), float(tv), float(kl)

model.safetensors

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