Regenerate tiny random LFM3 from LiquidAI/LFM3-1.7B for transformers 5.20.0

Regenerated with supported_models/lfm3/make_tiny_lfm3.py on transformers 5.20.0 from the released LiquidAI/LFM3-1.7B remote code (adds lfm3_conv.py). The previous remote code predates the release. The architecture of the tiny model is unchanged. Tested with the optimum-intel LFM3 tests on transformers 5.20.0.

chat_template.jinja

@@ -43,6 +43,8 @@
             {%- endif -%}
         {%- endfor -%}
         {{- _ns.result -}}
+    {%- elif content is not none -%}
+        {{- content | tojson -}}
     {%- endif -%}
 {%- endmacro -%}
 
@@ -78,9 +80,25 @@
     {{- "}" -}}
 {%- endmacro -%}
 
+{%- macro is_tool_reference(item) -%}
+    {{- "1" if item is mapping and item.get("type") == "tool_reference" else "" -}}
+{%- endmacro -%}
+
+{%- macro referenced_tool_to_json(name) -%}
+    {%- set _ns = namespace(found=none) -%}
+    {%- for tool in tools or [] -%}
+        {%- set fn = tool.function if tool.function is defined else tool -%}
+        {%- if fn.name == name -%}
+            {%- set _ns.found = fn -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {%- if _ns.found is none -%}
+        {{- raise_exception("Tool reference '" + name | string + "' not found in available tools. Every tool referenced by a tool_reference must be in `tools` on every request.") -}}
+    {%- endif -%}
+    {{- tool_to_json(_ns.found) -}}
+{%- endmacro -%}
+
 {%- macro render_tool_response(content) -%}
-    {#- content that already carries the tag (pre-formatted by the caller) is
-        passed through as-is instead of being wrapped a second time -#}
     {%- if content is string and "<tool_response>" in content -%}
         {{- content -}}
     {%- else -%}
@@ -92,7 +110,15 @@
     {%- elif content is iterable -%}
         {%- for item in content -%}
             {%- if not loop.first -%}{{- "\n" -}}{%- endif -%}
-            {%- if item is mapping and item.result is defined and item.result is mapping and item.result.tools is defined -%}
+            {%- if is_tool_reference(item) -%}
+                {%- if loop.first or not is_tool_reference(loop.previtem) -%}
+                    {{- "<additional_tools>\n" -}}
+                {%- endif -%}
+                {{- referenced_tool_to_json(item.name) -}}
+                {%- if loop.last or not is_tool_reference(loop.nextitem) -%}
+                    {{- "\n</additional_tools>" -}}
+                {%- endif -%}
+            {%- elif item is mapping and item.result is defined and item.result is mapping and item.result.tools is defined -%}
                 {{- "<additional_tools>\n" -}}
                 {%- set _tl = item.result.tools -%}
                 {%- if _tl is string -%}
@@ -110,6 +136,8 @@
                 {{- parse_content(item) -}}
             {%- endif -%}
         {%- endfor -%}
+    {%- elif content is not none -%}
+        {{- content | tojson -}}
     {%- endif -%}
     {{- "\n</tool_response>" -}}
     {%- endif -%}
@@ -143,15 +171,37 @@
         {%- set ns.last_user_index = loop.index0 -%}
     {%- endif -%}
 {%- endfor -%}
+{%- set carry = namespace(pending=[]) -%}
 {%- for message in messages -%}
+    {%- set _next = messages[loop.index0 + 1] if not loop.last else none -%}
+    {%- set _merge = message.role == "tool" and not message.get("additional_tools")
+        and _next is not none and _next.role == "tool" and not _next.get("additional_tools")
+        and message.get("tool_call_id") and _next.get("tool_call_id") == message.get("tool_call_id") -%}
+    {%- set _items = [] -%}
+    {%- if message.role == "tool" and not message.get("additional_tools") -%}
+        {%- set _c = message.get("content") -%}
+        {%- if _c is string -%}
+            {%- set _items = [{"type": "text", "text": _c}] if _c else [] -%}
+        {%- elif _c is mapping -%}
+            {%- set _items = [_c] -%}
+        {%- elif _c is iterable -%}
+            {%- set _items = _c | list -%}
+        {%- elif _c is not none -%}
+            {%- set _items = [_c] -%}
+        {%- endif -%}
+    {%- endif -%}
+    {%- if _merge -%}
+        {%- set carry.pending = carry.pending + _items -%}
+    {%- else -%}
     {{- "<|im_start|>" + message.role + "\n" -}}
     {%- if message.role == "assistant" -%}
         {%- generation -%}
+        {%- set thinking = message.thinking if message.thinking is defined else message.reasoning if message.reasoning is defined else message.reasoning_content -%}
         {%- set keep_thinking = preserve_thinking or loop.index0 > ns.last_user_index -%}
-        {%- set thinking = message.thinking or message.reasoning or message.reasoning_content -%}
-        {%- set thinking = thinking if thinking is string else "" -%}
-        {%- if thinking and keep_thinking -%}
-            {{- "<think>" + thinking + "</think>" -}}
+        {%- if thinking is defined and keep_thinking -%}
+            {{- "<think>" + (thinking or "") + "</think>" -}}
+        {%- elif keep_thinking and "</think>" not in (message.get("content") if message.get("content") is string else "") -%}
+            {{- "<think></think>" -}}
         {%- endif -%}
         {%- set _cfm_tag = "CONTINUE_FINAL_MESSAGE_TAG " -%}
         {%- set _has_cfm = false -%}
@@ -182,15 +232,11 @@
     {%- elif message.role == "tool" -%}
         {%- if message.get("additional_tools") -%}
             {%- set _at = message["additional_tools"] -%}
-            {#- definitions may live in additional_tools itself, or in content
-                with additional_tools used as a plain truthy flag -#}
             {%- if _at is mapping or (_at is sequence and _at is not string) -%}
                 {%- set _src = _at -%}
             {%- else -%}
                 {%- set _src = message.get("content") -%}
             {%- endif -%}
-            {#- nothing to reveal: skip the wrapper entirely rather than
-                emitting an empty <additional_tools> block -#}
             {%- if _src is string and "<tool_response>" in _src -%}
                 {{- _src -}}
             {%- elif _src is string and "<additional_tools>" in _src -%}
@@ -209,8 +255,11 @@
                 {%- endif -%}
                 {{- "\n</additional_tools>\n</tool_response>" -}}
             {%- endif -%}
-        {%- elif message.get("content") -%}
-            {{- render_tool_response(message["content"]) -}}
+        {%- elif carry.pending -%}
+            {{- render_tool_response(carry.pending + _items) -}}
+            {%- set carry.pending = [] -%}
+        {%- else -%}
+            {{- render_tool_response("" if message.get("content") is none else message.get("content")) -}}
         {%- endif -%}
         {{- "<|im_end|>\n" -}}
     {%- else %}
@@ -219,6 +268,7 @@
         {%- endif -%}
         {{- "<|im_end|>\n" -}}
     {%- endif %}
+    {%- endif -%}
 {%- endfor -%}
 {%- if add_generation_prompt -%}
     {{- "<|im_start|>assistant\n" -}}

config.json

@@ -4,6 +4,7 @@
   ],
   "auto_map": {
     "AutoConfig": "configuration_lfm3.Lfm3Config",
+    "AutoModel": "modeling_lfm3.Lfm3Model",
     "AutoModelForCausalLM": "modeling_lfm3.Lfm3ForCausalLM"
   },
   "bos_token_id": 124894,
@@ -13,6 +14,7 @@
   "eos_token_id": 124900,
   "head_dim": 8,
   "hidden_size": 16,
+  "initializer_range": 0.02,
   "intermediate_size": 32,
   "layer_types": [
     "conv",
@@ -23,7 +25,7 @@
     "conv",
     "sliding_attention"
   ],
-  "max_position_embeddings": 196608,
+  "max_position_embeddings": 262144,
   "model_type": "lfm3",
   "norm_eps": 1e-05,
   "num_attention_heads": 4,

configuration_lfm3.py

@@ -1,299 +1,132 @@
-from typing import Optional, Union
-
-from transformers.configuration_utils import PretrainedConfig
-from transformers.modeling_rope_utils import RopeParameters
-
-
-class Lfm3Config(PretrainedConfig):
+# Copyright 2026 The HuggingFace 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 huggingface_hub.dataclasses import strict
+
+from transformers.configuration_utils import PreTrainedConfig
+from transformers.utils import auto_docstring
+
+
+@auto_docstring(checkpoint="LiquidAI/LFM3-1.7B-Preview")
+@strict
+class Lfm3Config(PreTrainedConfig):
     r"""
-    Configuration class for [`Lfm3Model`] / [`Lfm3ForCausalLM`].
-
-    LFM3 dense is a hybrid architecture composed of:
-      - GQA full-attention layers (with headwise gated attention)
-      - sliding-window attention layers (may use a different rope_theta from global)
-      - gated short-conv layers (depthwise conv with B*x -> conv -> C*conv_out)
-
-    Args:
-        vocab_size (`int`, *optional*, defaults to 65536):
-            Vocabulary size.
-        hidden_size (`int`, *optional*, defaults to 2048):
-            Hidden / residual stream dimension.
-        intermediate_size (`int`, *optional*, defaults to 7168):
-            SwiGLU MLP intermediate dimension.
-        num_hidden_layers (`int`, *optional*, defaults to 32):
-            Total number of decoder layers.
-        num_attention_heads (`int`, *optional*, defaults to 32):
-            Number of query heads.
-        num_key_value_heads (`int`, *optional*, defaults to 8):
-            Number of KV heads (GQA) for full-attention layers.
-        head_dim (`int`, *optional*):
-            Per-head dim. When `None`, defaults to `hidden_size // num_attention_heads`.
-            May differ from that quotient (e.g. dim=1536, num_heads=16, head_dim=128).
-        max_position_embeddings (`int`, *optional*, defaults to 128_000):
-            Maximum sequence length the model is designed for.
-        norm_eps (`float`, *optional*, defaults to 1e-5):
-            RMSNorm epsilon.
-        use_cache (`bool`, *optional*, defaults to `True`):
-            Whether to return past_key_values.
-        pad_token_id, bos_token_id, eos_token_id (`int`, *optional*):
-            Special token ids.
-        tie_word_embeddings (`bool`, *optional*, defaults to `True`):
-            Tie input/output embeddings.
-        rope_parameters (`RopeParameters` or `dict[str, RopeParameters]`, *optional*):
-            Either a single RopeParameters dict (legacy form — broadcast to every
-            attention layer type) or a Gemma3-style dict keyed by `layer_type`:
-            `{"full_attention": {"rope_type": "default", "rope_theta": 1_000_000,
-            "partial_rotary_factor": 0.5}, "sliding_attention":
-            {"rope_type": "default", "rope_theta": 10_000}}`. Partial RoPE is a
-            global-attention-only knob (Gemma 4 convention); SWA layers always
-            use full RoPE, so the SWA entry never carries `partial_rotary_factor`.
-            For backward compatibility, callers may also pass a legacy flat
-            dict together with a separate `swa_rope_parameters` kwarg; both forms
-            are migrated to the nested representation in `__init__`.
-        conv_bias (`bool`, *optional*, defaults to `False`):
-            Whether the depthwise conv layers use a bias term.
-        conv_L_cache (`int`, *optional*, defaults to 3):
-            Convolution kernel size and the number of conv-state slots cached per layer.
-        sliding_window (`int`, *optional*):
-            Sliding-window size used by sliding-window-attention layers.
-        full_attn_idxs (`list[int]`, *optional*):
-            Convenience: layer indices that use full attention. Only used to build the
-            default `layer_types` when `layer_types` is not explicitly provided.
-        layer_types (`list[str]`, *optional*):
-            Per-layer operator type. Each entry is one of
-            `"full_attention"`, `"sliding_attention"`, `"conv"`.
-            When `None`, defaults to `"full_attention"` for indices in `full_attn_idxs`
-            and `"conv"` for the rest.
-        use_sandwich_norm (`bool`, *optional*, defaults to `False`):
-            Whether to apply an additional RMSNorm AFTER the operator and AFTER the
-            FFN (sandwich norm; quantization-friendly variant).
+    rope_parameters (`dict[str, RopeParameters]`, *optional*):
+        RoPE parameters keyed by attention layer type, for example `{"full_attention": {"rope_type": "default",
+        "rope_theta": 1000000.0, "partial_rotary_factor": 0.5}, "sliding_attention": {"rope_type": "default",
+        "rope_theta": 10000.0}}`. `conv` layers do not use RoPE and have no entry. Full-attention layers rotate the
+        first `partial_rotary_factor * head_dim` dimensions and leave the rest unrotated.
+    conv_bias (`bool`, *optional*, defaults to `False`):
+        Whether to use bias in the conv layers.
+    conv_L_cache (`int`, *optional*, defaults to 7):
+        L_cache dim in the conv layers.
+    layer_types (`list[str]`, *optional*):
+        Per-layer operator type, each one of `"full_attention"`, `"sliding_attention"` or `"conv"`. Defaults to the
+        LFM3 pattern of a full-attention layer every fourth layer starting from the third, a sliding-attention layer
+        in between, and gated short convolutions elsewhere.
+    sliding_window (`int`, *optional*, defaults to 256):
+        Sliding-window size of the `sliding_attention` layers.
+    use_sandwich_norm (`bool`, *optional*, defaults to `True`):
+        Whether to apply an RMSNorm to the operator and feed-forward outputs before they enter the residual stream.
+
+    ```python
+    >>> from transformers import Lfm3Model, Lfm3Config
+
+    >>> # Initializing a LFM3 model
+    >>> configuration = Lfm3Config()
+
+    >>> # Initializing a model from the LFM3-1.7B style configuration
+    >>> model = Lfm3Model(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```
     """
 
     model_type = "lfm3"
     keys_to_ignore_at_inference = ["past_key_values"]
-    # Attention layer types that consume a per-layer-type entry in
-    # `rope_parameters`. `"conv"` layers don't use RoPE.
+    default_theta = 1000000.0
+
+    vocab_size: int = 128_000
+    hidden_size: int = 2048
+    intermediate_size: int = 7680
+    num_hidden_layers: int = 24
+    num_attention_heads: int = 16
+    num_key_value_heads: int = 4
+    max_position_embeddings: int = 128_000
+    initializer_range: float = 0.02
+    norm_eps: float = 0.00001
+    use_cache: bool = True
+    pad_token_id: int | None = 124_893
+    bos_token_id: int | None = 124_894
+    eos_token_id: int | list[int] | None = 124_900
+    tie_word_embeddings: bool = True
+    rope_parameters: dict | None = None
+    conv_bias: bool = False
+    conv_L_cache: int = 7
+    layer_types: list[str] | None = None
     _attention_layer_types = ("full_attention", "sliding_attention")
+    head_dim: int | None = None
+    sliding_window: int | None = 256
+    use_sandwich_norm: bool = True
 
-    def __init__(
-        self,
-        vocab_size: int = 65536,
-        hidden_size: int = 2048,
-        intermediate_size: int = 7168,
-        num_hidden_layers: int = 32,
-        num_attention_heads: int = 32,
-        num_key_value_heads: int = 8,
-        head_dim: Optional[int] = None,
-        max_position_embeddings: int = 128_000,
-        norm_eps: float = 1e-5,
-        use_cache: bool = True,
-        pad_token_id: int = 0,
-        bos_token_id: int = 1,
-        eos_token_id: int = 2,
-        tie_word_embeddings: bool = True,
-        rope_parameters: Optional[Union[RopeParameters, dict[str, RopeParameters]]] = None,
-        conv_bias: bool = False,
-        conv_L_cache: int = 3,
-        sliding_window: Optional[int] = None,
-        full_attn_idxs: Optional[list[int]] = None,
-        layer_types: Optional[list[str]] = None,
-        use_sandwich_norm: bool = False,
-        **kwargs,
-    ):
-        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.max_position_embeddings = max_position_embeddings
-        self.norm_eps = norm_eps
-        self.use_cache = use_cache
-
-        self.conv_bias = conv_bias
-        self.conv_L_cache = conv_L_cache
-
-        self.sliding_window = sliding_window
-
-        self.use_sandwich_norm = use_sandwich_norm
-
-        self.layer_types = layer_types
+    def __post_init__(self, **kwargs):
+        if self.head_dim is None:
+            self.head_dim = self.hidden_size // self.num_attention_heads
+        if self.rope_parameters is None:
+            self.rope_parameters = {
+                "full_attention": {"rope_type": "default", "rope_theta": 1_000_000.0, "partial_rotary_factor": 0.5},
+                "sliding_attention": {"rope_type": "default", "rope_theta": 10_000.0},
+            }
         if self.layer_types is None:
-            full_attn_idxs = (
-                full_attn_idxs if full_attn_idxs is not None else list(range(num_hidden_layers))
-            )
             self.layer_types = [
-                "full_attention" if i in full_attn_idxs else "conv"
-                for i in range(num_hidden_layers)
+                "full_attention" if i % 4 == 2 else "sliding_attention" if i % 4 == 0 and i > 0 else "conv"
+                for i in range(self.num_hidden_layers)
             ]
 
-        # Catch the inconsistent combination at construction time so users see
-        # a clear error here instead of an opaque one from
-        # `create_sliding_window_causal_mask` deep in the forward pass.
-        if "sliding_attention" in self.layer_types and sliding_window is None:
-            raise ValueError(
-                "Lfm3Config has `sliding_attention` in `layer_types` but "
-                "`sliding_window` is None. Set `sliding_window` to the window size "
-                "(in tokens) you trained with."
-            )
-
-        # rope_parameters is the Gemma3-style nested dict keyed by attention
-        # layer_type (e.g. {"full_attention": {...}, "sliding_attention": {...}}).
-        # The trainer→HF conversion writes this nested form directly.
-        self.rope_parameters = self._validate_rope_parameters(rope_parameters)
-
-        tie_word_embeddings = kwargs.get("tie_embedding", tie_word_embeddings)
-        super().__init__(
-            pad_token_id=pad_token_id,
-            bos_token_id=bos_token_id,
-            eos_token_id=eos_token_id,
-            tie_word_embeddings=tie_word_embeddings,
-            **kwargs,
-        )
-
-    def _validate_rope_parameters(
-        self,
-        rope_parameters: Optional[dict[str, dict]],
-    ) -> Optional[dict[str, dict]]:
-        """Keep the per-attention-layer-type rope entries present in the model.
-
-        `rope_parameters` must be the Gemma3-style nested form keyed by attention
-        `layer_type`, e.g. ``{"full_attention": {"rope_type": "default",
-        "rope_theta": 1e6, "partial_rotary_factor": 0.5}, "sliding_attention":
-        {"rope_type": "default", "rope_theta": 1e4}}``. Conv layers have no rope.
-        Partial RoPE is a full-attention-only knob — the SWA entry never carries
-        `partial_rotary_factor`.
-        """
-        if rope_parameters is None:
-            return None
-
-        attention_types_in_model = [
-            lt for lt in (self.layer_types or []) if lt in self._attention_layer_types
-        ]
-        # Conv-only models don't use RoPE.
-        if not attention_types_in_model:
-            return None
+        self.tie_word_embeddings = kwargs.pop("tie_embedding", self.tie_word_embeddings)
+        super().__post_init__(**kwargs)
 
-        if not (
-            isinstance(rope_parameters, dict)
-            and all(k in self._attention_layer_types for k in rope_parameters)
-        ):
+    def validate_architecture(self):
+        """Part of ``@strict``-powered validation."""
+        if len(self.layer_types) != self.num_hidden_layers:
             raise ValueError(
-                "Lfm3 expects nested per-layer-type `rope_parameters` keyed by "
-                "attention layer_type (e.g. {'full_attention': {...}, "
-                "'sliding_attention': {...}}); got "
-                f"{list(rope_parameters) if isinstance(rope_parameters, dict) else type(rope_parameters).__name__}."
+                f"`layer_types` has {len(self.layer_types)} entries but `num_hidden_layers` is {self.num_hidden_layers}."
             )
-
-        return {
-            k: dict(v) for k, v in rope_parameters.items()
-            if k in attention_types_in_model
-        } or None
-
-    def convert_rope_params_to_dict(self, **kwargs):
-        """Override the base PretrainedConfig hook so it doesn't pollute our
-        per-layer-type `rope_parameters` dict with top-level `rope_theta` /
-        `rope_type` keys.
-
-        The base implementation (modeling_rope_utils.py
-        `RotaryEmbeddingConfigMixin.convert_rope_params_to_dict`) does
-        `self.rope_parameters.setdefault("rope_theta", ...)` directly on
-        whatever we set — fine for legacy flat configs, but it puts top-level
-        keys onto our nested dict, which then makes `validate_rope` treat the
-        whole thing as flat and warn about the layer-type keys as
-        "Unrecognized". This override mirrors Gemma3Config: when the user
-        already supplied a nested dict, we trust it and only run validation.
-        """
-        rope_scaling = kwargs.pop("rope_scaling", None)
-        if rope_scaling is not None and self.rope_parameters is None:
-            self.rope_parameters = rope_scaling
-
-        if self.rope_parameters is None:
-            # Conv-only models don't carry RoPE. Skipping the base helpers
-            # avoids a noisy "no RoPE parameters were found" log on every
-            # config construction.
-            return kwargs
-
-        is_nested = isinstance(self.rope_parameters, dict) and all(
-            k in self._attention_layer_types for k in self.rope_parameters.keys()
-        )
-        if is_nested:
-            # Ensure each entry has rope_type; default to "default" if missing.
-            for entry in self.rope_parameters.values():
-                entry.setdefault("rope_type", "default")
-            # validate_rope reads self.ignore_keys_at_rope_validation internally.
-            self.validate_rope()
-            return kwargs
-
-        # Flat / legacy form — fall back to the base behavior.
-        return super().convert_rope_params_to_dict(**kwargs)
-
-    def standardize_rope_params(self):
-        """Override the base PretrainedConfig hook to skip non-attention
-        layer types (e.g. ``"conv"``).
-
-        Callers like vLLM's ``patch_rope_parameters`` (in
-        ``vllm/transformers_utils/config.py``) invoke
-        ``config.standardize_rope_params()`` directly after loading our
-        config, bypassing our ``convert_rope_params_to_dict`` override.
-
-        The base implementation in
-        ``transformers/modeling_rope_utils.py:standardize_rope_params``
-        Case 2 (nested rope) does:
-
-            for layer_type in set(layer_types):
-                rope_parameters[layer_type].setdefault("rope_type", ...)
-
-        This iterates EVERY entry in ``layer_types`` and indexes
-        ``rope_parameters[layer_type]`` blindly. Lfm3 has ``"conv"`` in
-        ``layer_types`` (conv layers don't use RoPE), and our nested
-        ``rope_parameters`` correctly omits a ``"conv"`` entry — so the
-        blind index raises ``KeyError: 'conv'`` and vLLM fails to load.
-
-        Gemma3 dodges this because its ``layer_types`` only contains
-        attention types. Upstream LFM2 dodges it because its
-        ``rope_parameters`` is flat (Case 1), not nested (Case 2).
-
-        Our override: iterate only the layer types that actually have an
-        entry in ``rope_parameters``, applying the same defaulting logic
-        the base does. For the flat / legacy form, defer to base.
-        """
-        rope_parameters = getattr(self, "rope_parameters", None)
-        if not rope_parameters:
-            return
-
-        is_nested = isinstance(rope_parameters, dict) and all(
-            k in self._attention_layer_types for k in rope_parameters.keys()
-        )
-        if not is_nested:
-            return super().standardize_rope_params()
-
-        rope_theta = getattr(self, "rope_theta", None)
-        partial_rotary_factor = getattr(self, "partial_rotary_factor", None)
-        for layer_type in rope_parameters:
-            entry = rope_parameters[layer_type]
-            entry.setdefault("rope_type", entry.get("type", "default"))
-            if rope_theta is not None:
-                entry.setdefault("rope_theta", rope_theta)
-            # Partial RoPE is a global-attention-only knob (Gemma 4 invariant);
-            # never apply it to SWA, even if a top-level
-            # `partial_rotary_factor` is set as a legacy attribute.
-            if partial_rotary_factor is not None and layer_type == "full_attention":
-                entry["partial_rotary_factor"] = partial_rotary_factor
-            if entry["rope_type"] in ("llama3", "yarn", "longrope"):
-                entry.setdefault(
-                    "original_max_position_embeddings", self.max_position_embeddings
+        allowed_layer_types = (*self._attention_layer_types, "conv")
+        invalid_layer_types = [layer_type for layer_type in self.layer_types if layer_type not in allowed_layer_types]
+        if invalid_layer_types:
+            raise ValueError(f"`layer_types` entries must be one of {allowed_layer_types}; got {invalid_layer_types}.")
+        if "sliding_attention" in self.layer_types and self.sliding_window is None:
+            raise ValueError("`layer_types` has `sliding_attention` layers but `sliding_window` is None.")
+        missing_rope = [
+            layer_type
+            for layer_type in self._attention_layer_types
+            if layer_type in self.layer_types and layer_type not in self.rope_parameters
+        ]
+        if missing_rope:
+            raise ValueError(f"`rope_parameters` has no entry for the {missing_rope} layers in `layer_types`.")
+        for layer_type, rope_params in self.rope_parameters.items():
+            if not isinstance(rope_params, dict):
+                continue
+            if (
+                rope_params.get("rope_type", "default") in ("dynamic", "longrope")
+                and rope_params.get("partial_rotary_factor", 1.0) < 1
+            ):
+                raise ValueError(
+                    f"`{rope_params['rope_type']}` RoPE recomputes its frequencies during the forward pass over the "
+                    f"rotated dimension only and cannot keep LFM3's partial rotation on `{layer_type}` layers."
                 )
 
-    @property
-    def rope_scaling(self):
-        """Backward compatibility property: rope_scaling maps to rope_parameters."""
-        return self.rope_parameters
-
-    @rope_scaling.setter
-    def rope_scaling(self, value):
-        self.rope_parameters = value
-
 
 __all__ = ["Lfm3Config"]

lfm3_conv.py

@@ -0,0 +1,202 @@
+# Copyright 2026 The HuggingFace 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.
+"""Inference kernels for LFM3's width-7 depthwise convolution."""
+
+import torch
+
+from transformers.utils import is_triton_available
+
+
+if is_triton_available():
+    import triton
+    import triton.language as tl
+
+    # the length and strides are runtime arguments, or every prompt length would compile its own kernel
+    @triton.jit(do_not_specialize=["T", "X0", "X1", "X2", "S0", "S1", "ST"])
+    def _prefill(
+        X,
+        W,
+        BIAS,
+        SEQ,
+        Y,
+        T,
+        D: tl.constexpr,
+        X0,
+        X1,
+        X2,
+        S0,
+        S1,
+        ST,
+        HAS_BIAS: tl.constexpr,
+        PACKED: tl.constexpr,
+        BLOCK: tl.constexpr,
+    ):
+        # int64 offsets: batch * channels * length can exceed 2**31 elements
+        b, d = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
+        t = tl.program_id(2).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
+        acc = tl.full((BLOCK,), 0, tl.float32)
+        if PACKED:
+            doc = tl.load(SEQ + b * S0 + (t - T + ST) * S1, (t < T) & (t >= T - ST), other=-1)
+        for k in tl.static_range(7):
+            pos = t + k - 6
+            valid = (t < T) & (pos >= 0)
+            if PACKED:
+                prev = tl.load(SEQ + b * S0 + (pos - T + ST) * S1, valid & (pos >= T - ST), other=-1)
+                valid = valid & (doc == prev)
+            x = tl.load(X + b * X0 + d * X1 + pos * X2, valid, other=0).to(tl.float32)
+            w = tl.load(W + d * 7 + k).to(tl.float32)
+            acc += x * w
+        if HAS_BIAS:
+            acc += tl.load(BIAS + d).to(tl.float32)
+        tl.store(Y + (b * D + d) * T + t, acc, t < T)
+
+    @triton.jit
+    def _decode(
+        X,
+        W,
+        BIAS,
+        STATE,
+        Y,
+        D: tl.constexpr,
+        K: tl.constexpr,
+        X0: tl.constexpr,
+        X1: tl.constexpr,
+        S0: tl.constexpr,
+        S1: tl.constexpr,
+        S2: tl.constexpr,
+        HAS_BIAS: tl.constexpr,
+        BLOCK: tl.constexpr,
+    ):
+        b = tl.program_id(0).to(tl.int64)
+        d = tl.program_id(1).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
+        k = tl.arange(0, 8)
+        # Load the complete state before writing the shifted window.
+        old = tl.load(STATE + b * S0 + d[:, None] * S1 + k[None, :] * S2, (d[:, None] < D) & (k[None, :] < K), other=0)
+        x = tl.load(X + b * X0 + d * X1, d < D, other=0)
+        shifted = tl.gather(old, tl.broadcast_to(tl.minimum(k + 1, 7)[None, :], (BLOCK, 8)), 1)
+        shifted = tl.where(k[None, :] == K - 1, x[:, None], shifted)
+        # A cache may store either K=6 previous inputs or K=7 including the oldest discarded input.
+        positions = k + K - 6
+        values = tl.gather(old, tl.broadcast_to(tl.minimum(positions, 7)[None, :], (BLOCK, 8)), 1)
+        values = tl.where(k[None, :] == 6, x[:, None], values)
+        weights = tl.load(W + d[:, None] * 7 + k[None, :], (d[:, None] < D) & (k[None, :] < 7), other=0)
+        acc = tl.sum(values.to(tl.float32) * weights.to(tl.float32), 1)
+        if HAS_BIAS:
+            acc += tl.load(BIAS + d, d < D, other=0).to(tl.float32)
+        tl.store(STATE + b * S0 + d[:, None] * S1 + k[None, :] * S2, shifted, (d[:, None] < D) & (k[None, :] < K))
+        tl.store(Y + b * D + d, acc, d < D)
+
+
+def can_use_lfm3_conv(hidden_states, weight):
+    return (
+        is_triton_available()
+        and hasattr(tl, "gather")  # the decode kernel needs Triton >= 3.3
+        and hidden_states.is_cuda
+        and not torch.is_grad_enabled()
+        and hidden_states.dtype == weight.dtype
+        and hidden_states.dtype in (torch.float16, torch.bfloat16, torch.float32)
+        and weight.shape[-1] == 7
+    )
+
+
+def lfm3_causal_conv1d(hidden_states, weight, bias=None, seq_idx=None):
+    batch, channels, length = hidden_states.shape
+    # the kernels index every tensor by the batch and channel of `hidden_states` and cannot bound the others
+    if weight.shape != (channels, 7) or (bias is not None and bias.shape != (channels,)):
+        raise ValueError(
+            f"`weight` must have shape ({channels}, 7) and `bias` ({channels},), got {tuple(weight.shape)}"
+        )
+    if seq_idx is not None and (
+        seq_idx.ndim != 2 or seq_idx.shape[0] not in (1, batch) or seq_idx.dtype not in (torch.int32, torch.int64)
+    ):
+        raise ValueError(
+            f"`seq_idx` must be an int32 or int64 tensor of shape (1, length) or ({batch}, length), got "
+            f"{seq_idx.dtype} {tuple(seq_idx.shape)}"
+        )
+    if any(t is not None and t.device != hidden_states.device for t in (weight, bias, seq_idx)):
+        raise ValueError(
+            f"`weight`, `bias` and `seq_idx` must be on the device of `hidden_states` ({hidden_states.device})"
+        )
+    output = torch.empty_like(hidden_states, memory_format=torch.contiguous_format)
+    # Triton launches on the current device; the model may live on another one (`device_map`, `.to("cuda:1")`).
+    with torch.cuda.device(hidden_states.device):
+        _prefill[(batch, channels, triton.cdiv(length, 128))](
+            hidden_states,
+            # the kernels index `bias` and `weight` densely
+            weight.contiguous(),
+            bias.contiguous() if bias is not None else None,
+            seq_idx,
+            output,
+            length,
+            channels,
+            *hidden_states.stride(),
+            *(
+                (0 if seq_idx.shape[0] == 1 else seq_idx.stride(0), seq_idx.stride(1))
+                if seq_idx is not None
+                else (0, 0)
+            ),
+            seq_idx.shape[-1] if seq_idx is not None else 0,
+            bias is not None,
+            seq_idx is not None,
+            128,
+            enable_fp_fusion=False,
+        )
+    return output
+
+
+def lfm3_causal_conv1d_update(hidden_states, conv_state, weight, bias=None):
+    batch, channels, length = hidden_states.shape
+    # the kernel addresses one position per row and bounds nothing by the other tensors' shapes
+    if length != 1:
+        raise ValueError(f"`hidden_states` must hold one position per row, got {tuple(hidden_states.shape)}")
+    # the 8-wide tile holds a 6- or 7-wide state (see `positions` in `_decode`)
+    if (
+        conv_state.ndim != 3
+        or conv_state.shape[:2] != (batch, channels)
+        or conv_state.shape[-1] not in (6, 7)
+        or weight.shape != (channels, 7)
+    ):
+        raise ValueError(
+            f"`conv_state` must have shape ({batch}, {channels}, 6 or 7) and `weight` ({channels}, 7), got "
+            f"{tuple(conv_state.shape)} and {tuple(weight.shape)}"
+        )
+    if bias is not None and bias.shape != (channels,):
+        raise ValueError(f"`bias` must have shape ({channels},), got {tuple(bias.shape)}")
+    # the kernel writes the state in place: rows or channels sharing memory (an `expand`ed cache) would race
+    extent = 0
+    for stride, size in sorted(zip(conv_state.stride(), conv_state.shape)):
+        if size > 1 and stride <= extent:
+            raise ValueError(f"`conv_state` must not overlap itself, got strides {conv_state.stride()}")
+        extent += stride * (size - 1)
+    if any(t is not None and t.device != hidden_states.device for t in (conv_state, weight, bias)):
+        raise ValueError(
+            f"`conv_state`, `weight` and `bias` must be on the device of `hidden_states` ({hidden_states.device})"
+        )
+    output = torch.empty_like(hidden_states, memory_format=torch.contiguous_format)
+    with torch.cuda.device(hidden_states.device):
+        _decode[(batch, triton.cdiv(channels, 32))](
+            hidden_states,
+            weight.contiguous(),
+            bias.contiguous() if bias is not None else None,
+            conv_state,
+            output,
+            channels,
+            conv_state.shape[-1],
+            *hidden_states.stride()[:2],
+            *conv_state.stride(),
+            bias is not None,
+            32,
+            enable_fp_fusion=False,
+        )
+    return output

modeling_lfm3.py

@@ -1,44 +1,64 @@
-# LFM3 dense model.
-from typing import Any, Callable, Optional, Union
+# Copyright 2026 The HuggingFace 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.
+import copy
+from collections.abc import Callable
 
 import torch
 import torch.nn.functional as F
 from torch import nn
 
+from transformers import initialization as init
+from transformers.activations import ACT2FN
 from transformers.cache_utils import Cache, DynamicCache
+from transformers.configuration_utils import PreTrainedConfig
 from transformers.generation import GenerationMixin
 from transformers.integrations import (
     use_kernel_forward_from_hub,
-    use_kernel_func_from_hub,
+    use_kernel_func_from_hub_with_fallback,
     use_kernelized_func,
 )
-from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
+from transformers.integrations.accelerate import force_accelerate_hooks
+from transformers.masking_utils import (
+    create_causal_mask,
+    create_recurrent_attention_mask,
+    create_sliding_window_causal_mask,
+)
 from transformers.modeling_layers import GradientCheckpointingLayer
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
 from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from transformers.processing_utils import Unpack
-from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
-from transformers.utils.generic import check_model_inputs, maybe_autocast
-from transformers.utils.import_utils import is_causal_conv1d_available, is_torchdynamo_compiling
+from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check
+from transformers.utils.generic import merge_with_config_defaults
+from transformers.utils.import_utils import is_tracing
+from transformers.utils.output_capturing import capture_outputs
 
 from .configuration_lfm3 import Lfm3Config
-
-
-if is_causal_conv1d_available():
-    from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
-else:
-    causal_conv1d_fn, causal_conv1d_update = None, None
+from .lfm3_conv import can_use_lfm3_conv, lfm3_causal_conv1d, lfm3_causal_conv1d_update
 
 
 @use_kernel_forward_from_hub("RMSNorm")
 class Lfm3RMSNorm(nn.Module):
-    def __init__(self, hidden_size, eps=1e-6):
+    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+        """
+        Lfm3RMSNorm is equivalent to T5LayerNorm
+        """
         super().__init__()
         self.weight = nn.Parameter(torch.ones(hidden_size))
         self.variance_epsilon = eps
 
-    def forward(self, hidden_states):
+    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)
@@ -50,18 +70,9 @@ class Lfm3RMSNorm(nn.Module):
 
 
 class Lfm3RotaryEmbedding(nn.Module):
-    """Per-layer-type rotary embedding (Gemma 3-style).
-
-    `config.rope_parameters` is a dict keyed by `layer_type`
-    (`"full_attention"`, `"sliding_attention"`). For each layer type
-    present in `config.layer_types`, an `{layer_type}_inv_freq` buffer is
-    registered. `forward(x, position_ids, layer_type=...)` selects the right
-    one. Partial RoPE (`partial_rotary_factor < 1.0`) follows the Gemma 4
-    convention: the first `rope_angles = int(partial_rotary_factor * dim) // 2`
-    frequency pairs receive the standard rotation, the remaining
-    `dim//2 - rope_angles` entries are zero (cos=1, sin=0 → identity). The
-    SWA entry must never carry `partial_rotary_factor` — the Lfm3Config
-    normalizer strips it during migration of legacy configs.
+    """
+    Rotary embedding per attention layer type, like Gemma 3; `conv` layers register no buffers. Partial RoPE keeps the
+    frequency progression of the full head dimension and zeroes the tail instead of shrinking the rotated dimension.
     """
 
     def __init__(self, config: Lfm3Config, device=None):
@@ -69,151 +80,99 @@ class Lfm3RotaryEmbedding(nn.Module):
         self.max_seq_len_cached = config.max_position_embeddings
         self.original_max_seq_len = config.max_position_embeddings
         self.config = config
-
-        # Deterministic dedupe of attention-layer types present in this model.
-        # `dict.fromkeys` preserves insertion order so the buffer-registration
-        # order is stable across processes / Python releases — `set()` would
-        # not, which matters for reproducible test output.
-        layer_types_in_model = list(dict.fromkeys(config.layer_types or []))
-        self.attention_layer_types = [
-            lt for lt in layer_types_in_model if lt in config._attention_layer_types
-        ]
-        self.rope_type: dict[str, str] = {}
-
-        rope_parameters = config.rope_parameters or {}
-        for layer_type in self.attention_layer_types:
-            entry = rope_parameters.get(layer_type)
-            if entry is None:
-                continue
-            self.rope_type[layer_type] = entry["rope_type"]
-            if entry["rope_type"] != "default":
-                rope_init_fn: Callable = ROPE_INIT_FUNCTIONS[entry["rope_type"]]
-                inv_freq, attention_scaling = rope_init_fn(self.config, device)
-            else:
-                inv_freq, attention_scaling = self.compute_default_rope_parameters(
-                    self.config, device, layer_type=layer_type
-                )
-            self.register_buffer(f"{layer_type}_inv_freq", inv_freq, persistent=False)
-            self.register_buffer(
-                f"{layer_type}_original_inv_freq", inv_freq.clone(), persistent=False
-            )
-            setattr(self, f"{layer_type}_attention_scaling", attention_scaling)
+        self.layer_types = sorted(set(config.layer_types) & set(config.rope_parameters))
+        self.rope_type = {}
+        for layer_type in self.layer_types:
+            rope_params = self.config.rope_parameters[layer_type]
+            self.rope_type[layer_type] = rope_params["rope_type"]
+            curr_inv_freq, curr_attention_scaling = self.compute_rope_parameters(config, device, layer_type=layer_type)
+            setattr(self, f"{layer_type}_inv_freq", nn.Buffer(curr_inv_freq, persistent=False))
+            setattr(self, f"{layer_type}_original_inv_freq", nn.Buffer(curr_inv_freq.clone(), persistent=False))
+            setattr(self, f"{layer_type}_attention_scaling", curr_attention_scaling)
 
     @staticmethod
     def compute_default_rope_parameters(
-        config: Optional[Lfm3Config] = None,
-        device: Optional["torch.device"] = None,
-        seq_len: Optional[int] = None,
-        layer_type: Optional[str] = None,
-    ) -> tuple["torch.Tensor", float]:
-        entry = config.rope_parameters[layer_type]
-        base = entry["rope_theta"]
-        partial_rotary_factor = entry.get("partial_rotary_factor", 1.0)
-        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
-        attention_factor = 1.0
+        config: Lfm3Config, device=None, layer_type: str | None = None, **kwargs
+    ) -> tuple[torch.Tensor, float]:
+        """
+        Inverse frequencies of the default RoPE with LFM3's zero-tail partial rotation: the first
+        `partial_rotary_factor * head_dim // 2` frequencies follow the progression of the full head dimension and the
+        rest are zero, leaving those dimensions unrotated.
 
+        Returns:
+            Tuple of (`torch.Tensor`, `float`): the inverse frequencies and the attention scaling factor (always 1).
+        """
+        base = config.rope_parameters[layer_type]["rope_theta"]
+        partial_rotary_factor = config.rope_parameters[layer_type].get("partial_rotary_factor", 1.0)
+        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
         rope_angles = int(partial_rotary_factor * dim) // 2
-        assert 0 < rope_angles <= dim // 2, (
-            f"partial_rotary_factor={partial_rotary_factor} with dim={dim} produced "
-            f"rope_angles={rope_angles}, must be in (0, dim//2={dim // 2}]."
-        )
-        inv_freq_rotated = 1.0 / (
-            base
-            ** (
-                torch.arange(0, 2 * rope_angles, 2, dtype=torch.int64).to(
-                    device=device, dtype=torch.float
-                )
-                / dim
-            )
-        )
-        nope_angles = dim // 2 - rope_angles
-        if nope_angles > 0:
-            inv_freq = torch.cat(
-                (
-                    inv_freq_rotated,
-                    torch.zeros(nope_angles, dtype=torch.float, device=device),
-                ),
-                dim=0,
-            )
-        else:
-            inv_freq = inv_freq_rotated
-        return inv_freq, attention_factor
+        inv_freq = 1.0 / (base ** (torch.arange(0, 2 * rope_angles, 2, dtype=torch.float) / dim))
+        inv_freq = torch.cat((inv_freq, torch.zeros(dim // 2 - rope_angles, dtype=torch.float)))
+        return inv_freq.to(device), 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, layer_type: Optional[str] = None):
+    def forward(self, x, position_ids, layer_type):
         inv_freq = getattr(self, f"{layer_type}_inv_freq")
         attention_scaling = getattr(self, f"{layer_type}_attention_scaling")
 
-        inv_freq_expanded = (
-            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() * attention_scaling
-            sin = emb.sin() * attention_scaling
+        freqs = position_ids[..., None].float() * inv_freq.to(device=x.device, dtype=torch.float)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        cos = emb.cos() * attention_scaling
+        sin = emb.sin() * attention_scaling
 
         return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
 
+    @classmethod
+    def compute_rope_parameters(
+        cls, config: Lfm3Config, device=None, layer_type: str | None = None
+    ) -> tuple[torch.Tensor, float]:
+        """Inverse frequencies of any RoPE type in the zero-tail layout: scaled types are evaluated over the full head
+        dimension and truncated to the rotated prefix."""
+        rope_params = config.rope_parameters[layer_type]
+        if rope_params["rope_type"] == "default":
+            return cls.compute_default_rope_parameters(config, device, layer_type=layer_type)
+        full_config = copy.copy(config)
+        full_config.rope_parameters = {
+            **config.rope_parameters,
+            layer_type: {**rope_params, "partial_rotary_factor": 1.0},
+        }
+        inv_freq, attention_factor = ROPE_INIT_FUNCTIONS[rope_params["rope_type"]](
+            full_config, device, layer_type=layer_type
+        )
+        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+        rope_angles = int(rope_params.get("partial_rotary_factor", 1.0) * dim) // 2
+        inv_freq = torch.cat((inv_freq[:rope_angles], inv_freq.new_zeros(dim // 2 - rope_angles)))
+        return inv_freq, attention_factor
+
 
 class Lfm3MLP(nn.Module):
-    def __init__(self, config: Lfm3Config, intermediate_size: Optional[int] = None):
+    def __init__(self, config: Lfm3Config):
         super().__init__()
-        self.hidden_size = config.hidden_size
-        self.intermediate_size = (
-            config.intermediate_size if intermediate_size is None else intermediate_size
-        )
-        self.w1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.w3 = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.w2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.w1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.w3 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.w2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
 
     def forward(self, x):
         return self.w2(F.silu(self.w1(x)) * self.w3(x))
 
 
-
 def rotate_half(x):
     """Rotates half the hidden dims of the input."""
     x1 = x[..., : x.shape[-1] // 2]
-    x2 = x[..., x.shape[-1] // 2:]
+    x2 = x[..., x.shape[-1] // 2 :]
     return torch.cat((-x2, x1), dim=-1)
 
 
-@use_kernel_func_from_hub("rotary_pos_emb")
-def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
-    """Apply Rotary Position Embedding to q / k.
-
-    The rotation is computed in fp32 to match the native (Llama-style) RoPE
-    path: native does a complex multiply in fp32 and casts to bf16 once at the
-    end. The upstream default form (multiplying bf16 q/k by bf16 cos/sin)
-    accrues ~3× more ULP-level rounding error per element, which propagates
-    through softmax to ~5% per-layer native↔HF divergence in bf16. When the
-    kernel hub provides a fused rotary kernel, that takes over via the
-    decorator; this body is the Python fallback.
-    """
-    in_dtype = q.dtype
-    cos = cos.unsqueeze(unsqueeze_dim).float()
-    sin = sin.unsqueeze(unsqueeze_dim).float()
-    q_f = q.float()
-    k_f = k.float()
-    q_embed = (q_f * cos) + (rotate_half(q_f) * sin)
-    k_embed = (k_f * cos) + (rotate_half(k_f) * sin)
-    return q_embed.to(in_dtype), k_embed.to(in_dtype)
-
-
 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
-    )
+    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)
 
 
@@ -222,7 +181,7 @@ def eager_attention_forward(
     query: torch.Tensor,
     key: torch.Tensor,
     value: torch.Tensor,
-    attention_mask: Optional[torch.Tensor],
+    attention_mask: torch.Tensor | None,
     scaling: float,
     dropout: float = 0.0,
     **kwargs: Unpack[TransformersKwargs],
@@ -232,8 +191,7 @@ def eager_attention_forward(
 
     attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
     if attention_mask is not None:
-        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
-        attn_weights = attn_weights + causal_mask
+        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)
@@ -243,58 +201,44 @@ def eager_attention_forward(
     return attn_output, attn_weights
 
 
-@use_kernelized_func(apply_rotary_pos_emb)
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+    """Applies RoPE to the query and key tensors in float32, like the reference implementation."""
+    cos = cos.unsqueeze(unsqueeze_dim).float()
+    sin = sin.unsqueeze(unsqueeze_dim).float()
+    q_embed = (q.float() * cos) + (rotate_half(q.float()) * sin)
+    k_embed = (k.float() * cos) + (rotate_half(k.float()) * sin)
+    return q_embed.to(q.dtype), k_embed.to(k.dtype)
+
+
 class Lfm3Attention(nn.Module):
-    """Multi-headed attention. Supports GQA, optional sliding window, and a
-    headwise sigmoid gate on the attention output.
-
-    The caller (``Lfm3Model.forward``) is responsible for building the right
-    mask for this layer's type (sliding-window-causal for SWA layers, full
-    causal otherwise). This forward only also passes ``sliding_window=...`` to
-    the attention backend for FA2's native fast path; SDPA / eager rely on the
-    precomputed mask alone.
-    """
+    """Llama attention with QK RMSNorm, a sigmoid output gate per head, a sliding window on `sliding_attention`
+    layers and RoPE in float32."""
 
-    def __init__(self, config: Lfm3Config, layer_idx: int, sliding_window: Optional[int] = None):
+    def __init__(self, config: Lfm3Config, layer_idx: int):
         super().__init__()
         self.config = config
         self.layer_idx = layer_idx
-        self.head_dim = getattr(config, "head_dim", None) or (
-            config.hidden_size // config.num_attention_heads
-        )
-        self.sliding_window = sliding_window
-        self.is_causal = True
-
-        self.num_key_value_heads = config.num_key_value_heads
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
         self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
-        self.scaling = self.head_dim ** -0.5
-        self.q_proj = nn.Linear(
-            config.hidden_size, config.num_attention_heads * self.head_dim, bias=False
-        )
-        self.k_proj = nn.Linear(
-            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False
-        )
-        self.v_proj = nn.Linear(
-            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False
-        )
-        self.o_proj = nn.Linear(
-            config.num_attention_heads * self.head_dim, config.hidden_size, bias=False
-        )
+        self.scaling = self.head_dim**-0.5
+        self.is_causal = True
+        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
         self.q_norm = Lfm3RMSNorm(self.head_dim, eps=config.norm_eps)
         self.k_norm = Lfm3RMSNorm(self.head_dim, eps=config.norm_eps)
-        # Headwise sigmoid gate on the attention output (one scalar per query head).
-        self.gate_proj = nn.Linear(
-            config.hidden_size, config.num_attention_heads, bias=False
-        )
+        self.gate_proj = nn.Linear(config.hidden_size, config.num_attention_heads, bias=False)
+        self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None
 
     def forward(
         self,
         hidden_states: torch.Tensor,
         position_embeddings: tuple[torch.Tensor, torch.Tensor],
-        attention_mask: Optional[torch.Tensor],
-        past_key_values: Optional[Cache] = None,
-        **kwargs,
-    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        attention_mask: torch.Tensor | None,
+        past_key_values: Cache | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
         input_shape = hidden_states.shape[:-1]
         hidden_shape = (*input_shape, -1, self.head_dim)
 
@@ -306,13 +250,11 @@ class Lfm3Attention(nn.Module):
         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
-            )
+            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
 
-        attention_interface: Callable = eager_attention_forward
-        if self.config._attn_implementation != "eager":
-            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
 
         attn_output, attn_weights = attention_interface(
             self,
@@ -325,43 +267,71 @@ class Lfm3Attention(nn.Module):
             sliding_window=self.sliding_window,
             **kwargs,
         )
-        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
-        # Headwise sigmoid gate: one scalar per query head, broadcast over head_dim.
-        num_heads = self.config.num_attention_heads
-        gate = torch.sigmoid(self.gate_proj(hidden_states)).view(*input_shape, num_heads, 1)
-        attn_output = attn_output.view(*input_shape, num_heads, self.head_dim)
-        attn_output = (gate * attn_output).view(*input_shape, -1)
-        output = self.o_proj(attn_output)
-        return output, attn_weights
+        gate = torch.sigmoid(self.gate_proj(hidden_states)).unsqueeze(-1)
+        attn_output = (attn_output * gate).reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
 
 
 def apply_mask_to_padding_states(hidden_states, attention_mask):
-    """Zero out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66 .
-
-    NOTE (upstream-faithful, mamba/Lfm2): the `shape[0] > 1` (batch > 1) guard means
-    padding is NOT zeroed for a single (batch=1) sequence — so the conv can bleed
-    pad positions for a left-padded batch-of-1. Matches upstream; single-sequence
-    decode isn't padded in practice, but it's a latent footgun for padded bsz-1 prefill.
     """
-    if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+    Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66
+    """
+    # NOTE: attention mask is a 2D boolean tensor
+    if attention_mask is not None:
         dtype = hidden_states.dtype
         hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
-    return hidden_states
-
 
-kernel_modules = (causal_conv1d_fn, causal_conv1d_update)
-is_fast_path_available = all(kernel_modules)
+    return hidden_states
 
 
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
+def causal_conv1d_update(
+    hidden_states: torch.Tensor,
+    conv_state: torch.Tensor,
+    weight: nn.Parameter,
+    bias: nn.Parameter | None = None,
+    activation: str | None = None,
+):
+    _, hidden_size, seq_len = hidden_states.shape
+    state_len = conv_state.shape[-1]
+
+    hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype)
+    conv_state.copy_(hidden_states_new[:, :, -state_len:])
+    out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size)
+    out = out[:, :, -seq_len:]
+    if activation is not None:
+        out = ACT2FN[activation](out)
+    return out.to(hidden_states.dtype)
+
+
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
+def causal_conv1d_fn(
+    hidden_states: torch.Tensor,
+    weight: nn.Parameter,
+    bias: nn.Parameter | None = None,
+    activation: str | None = None,
+    **kwargs,
+):
+    _, hidden_size, seq_len = hidden_states.shape
+    padding = weight.shape[-1] - 1
+
+    out = F.conv1d(
+        hidden_states.to(weight.dtype),
+        weight=weight.unsqueeze(1),
+        bias=bias,
+        padding=padding,
+        groups=hidden_size,
+    )[:, :, :seq_len]
+    if activation is not None:
+        out = ACT2FN[activation](out)
+    return out.to(hidden_states.dtype)
+
+
+@use_kernelized_func([causal_conv1d_update, causal_conv1d_fn])
 class Lfm3ShortConv(nn.Module):
-    """Standard gated short conv block:
-
-        proj = in_proj(x)               # 3*D
-        B, C, x' = proj.chunk(3, dim=-1)
-        conv_in = B * x'
-        y = C * depthwise_conv1d(conv_in)
-        out = out_proj(y)
-    """
+    """LFM2 gated short convolution with packed-document boundaries (`seq_idx`), the padding mask re-applied after a
+    biased `in_proj`, and width-7 Triton inference kernels with a PyTorch fallback."""
 
     def __init__(
         self,
@@ -371,165 +341,118 @@ class Lfm3ShortConv(nn.Module):
         super().__init__()
         self.config = config
         self.layer_idx = layer_idx
-        self.L_cache = config.conv_L_cache
+        self.conv_kernel_size = config.conv_L_cache
         self.bias = config.conv_bias
 
-        D = config.hidden_size
         self.conv = nn.Conv1d(
-            in_channels=D,
-            out_channels=D,
-            kernel_size=self.L_cache,
-            groups=D,
+            in_channels=config.hidden_size,
+            out_channels=config.hidden_size,
+            kernel_size=self.conv_kernel_size,
+            groups=config.hidden_size,
             bias=self.bias,
-            padding=self.L_cache - 1,
+            padding=self.conv_kernel_size - 1,
         )
-        self.in_proj = nn.Linear(D, 3 * D, bias=self.bias)
-        self.out_proj = nn.Linear(D, D, bias=self.bias)
-
-    def cuda_kernels_forward(
-        self,
-        x: torch.Tensor,
-        past_key_values: Optional[Cache] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        seq_idx: Optional[torch.Tensor] = None,
-    ):
-        x = apply_mask_to_padding_states(x, attention_mask)
-        D = self.config.hidden_size
-        proj = self.in_proj(x).transpose(-1, -2)
-        # Split order matches Lfm3MoeShortConv("both") chunk(3): B first, C second, x third
-        B, C, x = proj.split([D, D, D], dim=-2)
-        conv_in = B * x
-
-        conv_weights = self.conv.weight.view(self.conv.weight.size(0), self.conv.weight.size(2))
-        if past_key_values is not None and past_key_values.has_previous_state(self.layer_idx):
-            conv_states = past_key_values.layers[self.layer_idx].conv_states
-            if isinstance(conv_states, dict):
-                conv_states = conv_states[0]
-            conv_out = causal_conv1d_update(
-                conv_in.squeeze(-1),
-                conv_states,
-                conv_weights,
-                self.conv.bias,
-                None,
-            )
-            conv_out = conv_out.unsqueeze(-1)
-        else:
-            if past_key_values is not None:
-                conv_state = nn.functional.pad(conv_in, (self.L_cache - conv_in.shape[-1], 0))
-                past_key_values.update_conv_state(conv_state, self.layer_idx)
+        self.in_proj = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=self.bias)
+        self.out_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=self.bias)
 
-            # `seq_idx` (per-token document index, int32) resets the conv at packed
-            # document boundaries so it never bleeds across documents.
-            conv_out = causal_conv1d_fn(
-                conv_in, conv_weights, self.conv.bias, seq_idx=seq_idx, activation=None
-            )
-
-        y = C * conv_out
-        y = self.out_proj(y.transpose(-1, -2).contiguous())
-        return y
-
-    def _segmented_conv(self, conv_in: torch.Tensor, seq_idx: torch.Tensor) -> torch.Tensor:
-        """Causal depthwise conv that resets at packed document boundaries.
-
-        Pure-torch equivalent of `causal_conv1d_fn(..., seq_idx=...)`: out[t] only
-        sums conv_in[t-(L-1)..t] entries that share t's document id. `conv_in` is
-        (b, D, seqlen); `seq_idx` is (b, seqlen) int. Used only on the packed
-        path — the unpacked path keeps the plain `nn.Conv1d` call below.
-        """
-        L = self.L_cache
-        cin_p = nn.functional.pad(conv_in, (L - 1, 0))             # (b, D, L-1+seqlen)
-        windows = cin_p.unfold(-1, L, 1)                            # (b, D, seqlen, L)
-        # window element k of position t is conv_in[t-(L-1)+k]; current token is k=L-1.
-        sidx_p = nn.functional.pad(seq_idx, (L - 1, 0), value=-1)   # (b, L-1+seqlen)
-        sidx_win = sidx_p.unfold(-1, L, 1)                          # (b, seqlen, L)
-        same_doc = (sidx_win == sidx_win[..., -1:]).unsqueeze(1)    # (b, 1, seqlen, L)
-        windows = windows * same_doc.to(windows.dtype)
-        w = self.conv.weight.squeeze(1)                            # (D, L)
-        conv_out = torch.einsum("bdtk,dk->bdt", windows, w)
-        if self.bias:
-            conv_out = conv_out + self.conv.bias[None, :, None]
-        return conv_out
+        self.layer_type = config.layer_types[layer_idx]
 
-    def slow_forward(
+    @force_accelerate_hooks("conv")
+    def forward(
         self,
-        x: torch.Tensor,
-        past_key_values: Optional[Cache] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        seq_idx: Optional[torch.Tensor] = None,
+        hidden_states: torch.Tensor,
+        past_key_values: Cache | None = None,
+        attention_mask: torch.Tensor | None = None,
+        seq_idx: torch.IntTensor | None = None,
     ):
-        seqlen = x.shape[1]
-
-        x = apply_mask_to_padding_states(x, attention_mask)
-        D = self.config.hidden_size
-        proj = self.in_proj(x).transpose(-1, -2)
-        B, C, x = proj.split([D, D, D], dim=-2)
-        conv_in = B * x
-
-        if past_key_values is not None and past_key_values.has_previous_state(self.layer_idx):
-            # update_conv_state rolls the cached window left and appends conv_in.
-            conv_state = past_key_values.update_conv_state(conv_in, self.layer_idx)[..., -self.L_cache :]
-            conv_out = torch.sum(conv_state.to(conv_in.device) * self.conv.weight[:, 0, :], dim=-1)
-            if self.bias:
-                conv_out += self.conv.bias
-
-            conv_out = conv_out.unsqueeze(-1)
+        seq_len = hidden_states.shape[1]
+
+        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
+        BCx = self.in_proj(hidden_states).transpose(-1, -2)
+        B, C, x = BCx.chunk(3, dim=-2)
+        hidden_states = B * x
+        if attention_mask is not None:
+            # with `conv_bias` the projection makes the padded positions non-zero again
+            hidden_states = hidden_states * attention_mask[:, None, :].to(hidden_states.dtype)
+
+        use_precomputed_states = past_key_values is not None and past_key_values.has_previous_state(self.layer_idx)
+        use_kernels = 2 <= self.conv_kernel_size <= 4
+        use_wide_kernel = can_use_lfm3_conv(hidden_states, self.conv.weight)
+
+        if use_precomputed_states and seq_len == 1 and not past_key_values.layers[self.layer_idx].record_past:
+            conv_state = past_key_values.layers[self.layer_idx].conv_states[0]
+            if use_wide_kernel and conv_state.shape[-1] in (6, 7):
+                hidden_states = lfm3_causal_conv1d_update(
+                    hidden_states, conv_state, self.conv.weight.squeeze(1), self.conv.bias
+                )
+            elif use_kernels:
+                hidden_states = causal_conv1d_update(
+                    hidden_states, conv_state, self.conv.weight.squeeze(1), self.conv.bias
+                )
+            else:
+                hidden_states = torch.cat([conv_state, hidden_states], dim=-1)
+                conv_state.copy_(hidden_states[:, :, -conv_state.shape[-1] :])
+                hidden_states = self._conv(hidden_states)[:, :, -seq_len:]
         else:
-            if past_key_values is not None:
-                conv_state = nn.functional.pad(conv_in, (self.L_cache - conv_in.shape[-1], 0))
-                past_key_values.update_conv_state(conv_state, self.layer_idx)
+            if past_key_values is not None and seq_idx is not None:
+                # `seq_idx` is only set on an empty cache, and the decode steps read the cached window without it:
+                # only the last document (after any positional left padding) may enter the window
+                last_document = (seq_idx == seq_idx[:, -1:])[:, None].to(hidden_states.dtype)
+                past_key_values.update_conv_state(
+                    hidden_states * last_document, self.layer_idx, conv_kernel_size=self.conv_kernel_size
+                )
+            elif past_key_values is not None:
+                hidden_states = past_key_values.update_conv_state(
+                    hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size
+                )
 
-            if seq_idx is not None:
-                conv_out = self._segmented_conv(conv_in, seq_idx)
+            if use_wide_kernel:
+                hidden_states = lfm3_causal_conv1d(hidden_states, self.conv.weight.squeeze(1), self.conv.bias, seq_idx)
+            elif seq_idx is not None:
+                hidden_states = self._segmented_conv(hidden_states, seq_idx)
+            elif use_kernels:
+                hidden_states = causal_conv1d_fn(hidden_states, self.conv.weight.squeeze(1), self.conv.bias)
             else:
-                conv_out = self.conv(conv_in)[..., :seqlen]
+                hidden_states = self._conv(hidden_states)
+
+            # Drop the additional previous states
+            if past_key_values is not None:
+                hidden_states = hidden_states[:, :, -seq_len:]
 
-        y = C * conv_out
+        y = C * hidden_states
         y = y.transpose(-1, -2).contiguous()
         y = self.out_proj(y)
         return y
 
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        past_key_values: Optional[Cache] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        seq_idx: Optional[torch.Tensor] = None,
-    ):
-        # causal_conv1d's stock build only supports conv widths 2-4 and raises
-        # otherwise; build it for wider kernels (e.g. L_cache=7) or omit it for
-        # the pure-torch slow path.
-        if (
-            is_fast_path_available
-            and "cuda" in hidden_states.device.type
-            and not is_torchdynamo_compiling()
-        ):
-            return self.cuda_kernels_forward(
-                hidden_states, past_key_values, attention_mask, seq_idx
-            )
-        return self.slow_forward(
-            hidden_states, past_key_values, attention_mask, seq_idx
-        )
+    def _conv(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Causal depthwise conv in torch: the `causal_conv1d` kernels only accept widths 2 to 4."""
+        conv_out = self.conv(hidden_states.to(self.conv.weight.dtype))[..., : hidden_states.shape[-1]]
+        return conv_out.to(hidden_states.dtype)
+
+    def _segmented_conv(self, hidden_states: torch.Tensor, seq_idx: torch.Tensor) -> torch.Tensor:
+        """Causal depthwise conv that only reads positions sharing the current token's document index."""
+        kernel_size = self.conv_kernel_size
+        seq_idx = F.pad(seq_idx, (hidden_states.shape[-1] - seq_idx.shape[-1] + kernel_size - 1, 0), value=-1)
+        windows = F.pad(hidden_states, (kernel_size - 1, 0)).unfold(-1, kernel_size, 1)
+        seq_idx_windows = seq_idx.unfold(-1, kernel_size, 1)
+        same_document = (seq_idx_windows == seq_idx_windows[..., -1:]).unsqueeze(1)
+        conv_out = torch.einsum("bdtk,dk->bdt", windows * same_document.to(windows.dtype), self.conv.weight.squeeze(1))
+        if self.bias:
+            conv_out = conv_out + self.conv.bias[None, :, None]
+        return conv_out.to(hidden_states.dtype)
 
 
 class Lfm3DecoderLayer(GradientCheckpointingLayer):
     def __init__(self, config: Lfm3Config, layer_idx: int):
         super().__init__()
-        self.layer_type = config.layer_types[layer_idx]
-        self.is_attention_layer = self.layer_type in ("full_attention", "sliding_attention")
-        self.use_sandwich_norm = getattr(config, "use_sandwich_norm", False)
+        self.is_attention_layer = config.layer_types[layer_idx] in config._attention_layer_types
+        self.use_sandwich_norm = config.use_sandwich_norm
 
-        if self.layer_type == "full_attention":
+        if self.is_attention_layer:
             self.self_attn = Lfm3Attention(config, layer_idx)
-        elif self.layer_type == "sliding_attention":
-            self.self_attn = Lfm3Attention(config, layer_idx, sliding_window=config.sliding_window)
-        elif self.layer_type == "conv":
-            self.conv = Lfm3ShortConv(config, layer_idx)
         else:
-            raise ValueError(
-                f"Unsupported layer_type {self.layer_type!r} at layer {layer_idx}. "
-                "Lfm3 supports only {full_attention, sliding_attention, conv}."
-            )
-        self.feed_forward = Lfm3MLP(config, intermediate_size=config.intermediate_size)
+            self.conv = Lfm3ShortConv(config, layer_idx)
+        self.feed_forward = Lfm3MLP(config)
         self.operator_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
         self.ffn_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
         if self.use_sandwich_norm:
@@ -539,12 +462,11 @@ class Lfm3DecoderLayer(GradientCheckpointingLayer):
     def forward(
         self,
         hidden_states: torch.Tensor,
-        position_embeddings: tuple[torch.Tensor, torch.Tensor],
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[Cache] = None,
-        cache_position: Optional[torch.LongTensor] = None,
-        seq_idx: Optional[torch.Tensor] = None,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        seq_idx: torch.IntTensor | None = None,
         **kwargs,
     ) -> torch.Tensor:
         residual = hidden_states
@@ -567,6 +489,7 @@ class Lfm3DecoderLayer(GradientCheckpointingLayer):
         if self.use_sandwich_norm:
             hidden_states = self.post_operator_norm(hidden_states)
         hidden_states = hidden_states + residual
+
         ffn_out = self.feed_forward(self.ffn_norm(hidden_states))
         if self.use_sandwich_norm:
             ffn_out = self.post_ffn_norm(ffn_out)
@@ -575,6 +498,38 @@ class Lfm3DecoderLayer(GradientCheckpointingLayer):
         return hidden_states
 
 
+def create_masks_by_layer_type(
+    config: PreTrainedConfig,
+    inputs_embeds: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    past_key_values: Cache | None,
+    position_ids: torch.Tensor | None = None,
+) -> dict:
+    """The cache sizes each mask from the layer it is asked about, and the generic mask utilities ask about the first
+    non-sliding layer, which for a leading `conv` layer redirects to the first attention layer of any type: name the
+    first layer of each mask's own type instead."""
+    config = config.get_text_config()
+    mask_kwargs = {
+        "config": config,
+        "inputs_embeds": inputs_embeds,
+        "attention_mask": attention_mask,
+        "past_key_values": past_key_values,
+        "position_ids": position_ids,
+    }
+    layer_types, causal_mask_mapping = config.layer_types, {}
+    if "full_attention" in layer_types:
+        layer_idx = layer_types.index("full_attention")
+        causal_mask_mapping["full_attention"] = create_causal_mask(**mask_kwargs, layer_idx=layer_idx)
+    if "sliding_attention" in layer_types:
+        layer_idx = layer_types.index("sliding_attention")
+        causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(
+            **mask_kwargs, layer_idx=layer_idx
+        )
+    if "conv" in layer_types:
+        causal_mask_mapping["conv"] = create_recurrent_attention_mask(**mask_kwargs)
+    return causal_mask_mapping
+
+
 @auto_docstring
 class Lfm3PreTrainedModel(PreTrainedModel):
     config: Lfm3Config
@@ -585,43 +540,55 @@ class Lfm3PreTrainedModel(PreTrainedModel):
     _supports_flash_attn = True
     _supports_sdpa = True
     _supports_flex_attn = True
-    _can_compile_fullgraph = False
+
+    _can_compile_fullgraph = True
     _supports_attention_backend = True
     _can_record_outputs = {
         "hidden_states": Lfm3DecoderLayer,
         "attentions": Lfm3Attention,
     }
 
+    @torch.no_grad()
     def _init_weights(self, module):
-        # Re-compute RotaryEmbedding inv_freq buffers. Required because
-        # `from_pretrained` constructs the model on a meta device and
-        # materializes persistent=False buffers as uninitialized memory after
-        # the fact. We can't rely on values registered during meta-device
-        # __init__ — they aren't preserved through materialization. Re-run
-        # `compute_default_rope_parameters` for each attention layer type so
-        # the per-layer-type inv_freq buffers (`full_attention_inv_freq`,
-        # `sliding_attention_inv_freq`) get the correct values for
-        # their configured theta / partial_rotary_factor, instead of HF's
-        # default branch which assumes a single global table.
-        if isinstance(module, Lfm3RotaryEmbedding):
-            for layer_type in module.attention_layer_types:
-                if layer_type not in module.rope_type:
-                    continue
-                inv_freq, attention_scaling = module.compute_default_rope_parameters(
-                    module.config,
-                    device=getattr(module, f"{layer_type}_inv_freq").device,
-                    layer_type=layer_type,
-                )
-                with torch.no_grad():
-                    getattr(module, f"{layer_type}_inv_freq").copy_(inv_freq)
-                    getattr(module, f"{layer_type}_original_inv_freq").copy_(inv_freq)
-                setattr(module, f"{layer_type}_attention_scaling", attention_scaling)
-            return
         super()._init_weights(module)
+        if isinstance(module, Lfm3RotaryEmbedding):
+            for layer_type in module.layer_types:
+                curr_inv_freq, _ = module.compute_rope_parameters(module.config, layer_type=layer_type)
+                init.copy_(getattr(module, f"{layer_type}_inv_freq"), curr_inv_freq)
+                init.copy_(getattr(module, f"{layer_type}_original_inv_freq"), curr_inv_freq)
+
+    @staticmethod
+    def create_masks_for_generate(
+        config: PreTrainedConfig,
+        inputs_embeds: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        past_key_values: Cache | None,
+        position_ids: torch.Tensor | None = None,
+        **kwargs,
+    ) -> dict:
+        return create_masks_by_layer_type(config, inputs_embeds, attention_mask, past_key_values, position_ids)
+
+
+def packed_under_padding(position_ids: torch.Tensor, padding_mask: torch.Tensor, filled) -> torch.Tensor:
+    """Padding repeats position 0, so packing is a restart between two kept tokens, or a kept position 0 on a filled
+    cache whose row holds a kept token; attention under a padding mask cannot isolate it."""
+    restarts = torch.diff(position_ids, prepend=position_ids[:, :1] - 1, dim=-1) != 1
+    valid = padding_mask[:, -restarts.shape[-1] :].bool()
+    packed = restarts[:, 1:] & valid[:, 1:] & valid[:, :-1]
+    if filled:
+        # a row whose cached tokens are all padding (chunked prefill of a left-padded batch) has no document to pack
+        # against; a mask covering only the chunk keeps the strict rule
+        cached = padding_mask[:, : padding_mask.shape[-1] - restarts.shape[-1]]
+        cached_valid = cached.bool().any(-1, keepdim=True) if cached.shape[-1] else torch.ones_like(valid[:, :1])
+        packed = packed | (valid[:, :1] & (position_ids[:, :1] == 0) & cached_valid)
+    return packed
 
 
 @auto_docstring
 class Lfm3Model(Lfm3PreTrainedModel):
+    """LFM2 model with sliding-attention layers as a third layer type: the mask and the rotary embedding are built per
+    layer type, and packed documents (`position_ids` restarting) reset the short-conv window through `seq_idx`."""
+
     def __init__(self, config: Lfm3Config):
         super().__init__(config)
         self.padding_idx = config.pad_token_id
@@ -631,123 +598,161 @@ class Lfm3Model(Lfm3PreTrainedModel):
         self.layers = nn.ModuleList(
             [Lfm3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
         )
+        self.rotary_emb = Lfm3RotaryEmbedding(config=config)
         self.gradient_checkpointing = False
-        # Single Gemma3-style rotary module. Per-attention-layer-type inv_freq
-        # buffers (`full_attention_inv_freq`, `sliding_attention_inv_freq`)
-        # are registered inside the module based on `config.rope_parameters`.
-        # `forward` is dispatched per layer below via `layer_type=...`.
-        self.rotary_emb = Lfm3RotaryEmbedding(config)
         self.embedding_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
 
         # Initialize weights and apply final processing
         self.post_init()
 
-    @check_model_inputs
+    @merge_with_config_defaults
+    @capture_outputs
     @auto_docstring
     def forward(
         self,
-        input_ids: Optional[torch.LongTensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[Cache] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        use_cache: Optional[bool] = None,
-        cache_position: Optional[torch.LongTensor] = None,
+        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],
     ) -> BaseModelOutputWithPast:
         r"""
-        cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
-            Indices depicting the position of the input sequence tokens in the sequence. Used to update the
-            cache in the correct position and to infer the complete sequence length.
+        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)` or `dict`, *optional*):
+            Mask to avoid performing attention on padding token indices (1 keeps a token, 0 masks it), or the mask
+            dict of `create_masks_for_generate`. A mask dict is used as given: `generate` refuses packed
+            `position_ids` before its prefill.
         """
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+        # `seq_idx` is a shared kwarg (`DataCollatorWithFlattening(return_seq_idx=True)`, read by Lfm2); LFM3 derives
+        # the documents from the restarting `position_ids` that come with it
+        if kwargs.pop("seq_idx", None) is not None and position_ids is None:
+            raise ValueError("Pass packed documents as `position_ids` restarting at 0; LFM3 does not read `seq_idx`")
 
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
 
         if use_cache and past_key_values is None:
-            # DynamicCache(config) builds the right per-layer cache from
-            # config.layer_types: DynamicLayer (full_attention),
-            # DynamicSlidingWindowLayer (sliding_attention, auto-bounded), and
-            # LinearAttentionLayer (conv, conv-state buffer).
             past_key_values = DynamicCache(config=self.config)
-
-        if cache_position is None:
-            past_seen_tokens = (
-                past_key_values.get_seq_length() if past_key_values is not None else 0
-            )
-            cache_position = torch.arange(
-                past_seen_tokens,
-                past_seen_tokens + inputs_embeds.shape[1],
-                device=inputs_embeds.device,
+        filled = past_key_values is not None and (
+            past_key_values.get_seq_length() > 0
+            if any(t != "conv" for t in self.config.layer_types)
+            else past_key_values.has_previous_state()
+        )
+        packed_cache = getattr(past_key_values, "lfm3_packed", None)
+        if isinstance(packed_cache, torch.Tensor) or packed_cache:
+            # Reset layers have no history to mix. Ordinary caches carry a Python `False` and skip this, so the decode
+            # steps do no tensor work here; only a compiled packed prefill leaves a tensor flag, and `filled` stays
+            # tensor-valued for static caches.
+            packed_message = (
+                "A cache filled by a packed-document forward cannot be continued: its attention spans the documents"
             )
+            packed_cache = packed_cache & filled if isinstance(packed_cache, torch.Tensor) else filled
+            if isinstance(packed_cache, torch.Tensor):
+                torch_compilable_check(~packed_cache, packed_message)
+                if not is_tracing(packed_cache):
+                    past_key_values.lfm3_packed = bool(packed_cache)
+            elif packed_cache:
+                raise ValueError(packed_message)
+            else:
+                past_key_values.lfm3_packed = False
 
-        if position_ids is None:
-            position_ids = cache_position.unsqueeze(0)
-
-        # Packed-sequence document boundaries for the conv layers: when several
-        # documents are concatenated into one row, `position_ids` restarts at 0
-        # per document, so a cumulative count of those resets gives a per-token
-        # document id. Conv layers use it to avoid bleeding across documents.
-        # Only meaningful during prefill (multi-token); a single decode token has
-        # nothing to segment.
+        # The mask utilities detect packing only with `attention_mask=None`, so a multi-token forward with explicit
+        # positions drops the all-ones mask tokenizers emit (a data-dependent check, hence not when traced).
+        explicit_positions = position_ids is not None and inputs_embeds.shape[1] > 1
+        if (
+            explicit_positions
+            and isinstance(attention_mask, torch.Tensor)
+            and attention_mask.ndim == 2
+            and not is_tracing(position_ids)
+            and bool(attention_mask.all())
+        ):
+            attention_mask = None
         seq_idx = None
-        if inputs_embeds.shape[1] > 1:
-            doc_id = (position_ids == 0).cumsum(-1) - 1
-            if doc_id[..., -1].max() > 0:  # >1 document present
-                seq_idx = doc_id.to(torch.int32)
-
-        # Build one causal mask per attention layer type present in the model.
-        # SWA layers get a sliding-window causal mask; full-attention layers
-        # get a plain causal mask. Conv layers receive the RAW 2D padding mask
-        # (matching upstream `Lfm2Model.forward`'s `linear_attention` route)
-        # so `Lfm3ShortConv.apply_mask_to_padding_states` can zero pad
-        # positions during prefill. Conv during decode (single token) doesn't
-        # need a mask — None is fine.
-        mask_kwargs = {
-            "config": self.config,
-            "inputs_embeds": inputs_embeds,
-            "attention_mask": attention_mask,
-            "past_key_values": past_key_values,
-            "position_ids": position_ids,
-        }
-        causal_masks_by_type: dict[str, Optional[torch.Tensor]] = {}
-        attention_types_in_model = {
-            lt for lt in self.config.layer_types if lt in self.config._attention_layer_types
-        }
-        if "full_attention" in attention_types_in_model:
-            causal_masks_by_type["full_attention"] = create_causal_mask(**mask_kwargs)
-        if "sliding_attention" in attention_types_in_model:
-            causal_masks_by_type["sliding_attention"] = create_sliding_window_causal_mask(
-                **mask_kwargs
-            )
-        # Conv layers want the 2D padding mask only during prefill. For
-        # single-token decode (`shape[1] == 1`) there's nothing to pad-zero.
-        conv_mask = attention_mask if inputs_embeds.shape[1] != 1 else None
+        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
+            position_ids = position_ids.unsqueeze(0)
+        elif explicit_positions and not isinstance(attention_mask, dict):
+            # a mask dict comes from `generate`, which refused packing before its prefill
+            if isinstance(attention_mask, torch.Tensor) and attention_mask.ndim == 2:
+                # data-dependent: a traced forward skips it (`generate` refuses packing eagerly before its prefill)
+                if not is_tracing(position_ids):
+                    torch_compilable_check(
+                        ~packed_under_padding(position_ids, attention_mask, filled),
+                        "Packed documents (`position_ids` restarting) cannot be combined with an attention mask: "
+                        "attention would not isolate the documents. Pass `attention_mask=None` without a static cache.",
+                    )
+            else:
+                restarts = torch.diff(position_ids, prepend=position_ids[:, :1] - 1, dim=-1) != 1
+                document_ids = restarts.cumsum(-1)
+                if is_tracing(position_ids) or bool((document_ids[..., -1] > 0).any()):
+                    seq_idx = document_ids.to(torch.int32)
+                # Packed documents are prefilled in one forward on an empty cache (cached tokens would join the first
+                # document), and their cache is never continued
+                if filled:
+                    torch_compilable_check(
+                        (~restarts[:, 1:]).all() & (position_ids[:, 0] != 0).all(),
+                        "Packed documents must be prefilled in one forward; the cache already holds tokens",
+                    )
+                    # Ordinary continuation must retain the convolution window from the cached document.
+                    seq_idx = None
+                if attention_mask is None:
+                    # without a mask only the cacheless packed-sequence mask isolates repeated zero positions (left
+                    # padding): a cache-sized mask and every later decode step would attend them, so they count too
+                    packed = (document_ids[..., -1] > 0).any()
+                else:
+                    # a prepared 4D mask is the caller's and carries the padding (repeated zero positions), but may
+                    # or may not isolate documents, so the cache records the packing and is never continued
+                    steps = torch.diff(position_ids, dim=-1)
+                    packed = ((steps != 1) & ~((steps == 0) & (position_ids[:, :-1] == 0))).any()
+                if past_key_values is not None:
+                    # resolved here, where the prefill already syncs, so the decode steps carry a Python flag: no
+                    # per-step host sync, data-dependent op or CUDA-graph output on the cache
+                    past_key_values.lfm3_packed = packed if is_tracing(position_ids) else bool(packed)
+
+        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,
+            }
+            # The mask utilities only isolate packed documents without a cache, and an empty cache adds nothing.
+            layer_types = self.config.layer_types
+            if seq_idx is not None and past_key_values is not None:
+                attention_layer = next((i for i, t in enumerate(layer_types) if t != "conv"), None)
+                if attention_layer is not None and past_key_values.get_seq_length(attention_layer) == 0:
+                    if past_key_values.is_compileable:
+                        # static layers always return their full length, so the masks are sized by the cache, which
+                        # isolates neither documents nor positional-only left padding
+                        torch_compilable_check(
+                            ~packed,
+                            "Packed documents, or left padding without an attention mask, cannot be prefilled into a "
+                            "static cache: its masks are sized by the cache. Pass the padding as `attention_mask`, or "
+                            "use a `DynamicCache` for packed documents",
+                        )
+                    else:
+                        mask_kwargs["past_key_values"] = None
+            causal_mask_mapping = create_masks_by_layer_type(**mask_kwargs)
 
         hidden_states = inputs_embeds
-        # Pre-compute one (cos, sin) per attention layer type. Conv-only models
-        # leave this dict empty; the per-layer dispatch below only reads it for
-        # attention layers.
-        position_embeddings_by_type = {
-            layer_type: self.rotary_emb(hidden_states, position_ids, layer_type=layer_type)
-            for layer_type in self.rotary_emb.rope_type
+        position_embeddings = {
+            layer_type: self.rotary_emb(hidden_states, position_ids=position_ids, layer_type=layer_type)
+            for layer_type in self.rotary_emb.layer_types
         }
 
-        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
-            if decoder_layer.is_attention_layer:
-                layer_mask = causal_masks_by_type.get(decoder_layer.layer_type)
-            else:
-                layer_mask = conv_mask
-            layer_pos_embeddings = position_embeddings_by_type.get(decoder_layer.layer_type)
+        for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
+            layer_type = self.config.layer_types[i]
             hidden_states = decoder_layer(
                 hidden_states,
-                attention_mask=layer_mask,
+                attention_mask=causal_mask_mapping[layer_type],
+                position_embeddings=position_embeddings.get(layer_type),
                 position_ids=position_ids,
                 past_key_values=past_key_values,
-                cache_position=cache_position,
-                position_embeddings=layer_pos_embeddings,
                 seq_idx=seq_idx,
                 **kwargs,
             )
@@ -760,11 +765,46 @@ class Lfm3Model(Lfm3PreTrainedModel):
         )
 
 
+def reject_packed_generation(config: PreTrainedConfig, input_ids, generation_config, model_kwargs, is_first_iteration):
+    """`generate` decodes from the cache its prefill fills, and no mask it builds isolates packed documents
+    (`position_ids` restarting): the prompt is refused before the prefill, one forward or every chunk, writes to it."""
+    position_ids, attention_mask = model_kwargs.get("position_ids"), model_kwargs.get("attention_mask")
+    cache = model_kwargs.get("past_key_values")
+    if cache is None or position_ids is None or position_ids.ndim != 2:
+        return
+    if attention_mask is not None and attention_mask.ndim != 2:
+        return
+    has_attention = any(t != "conv" for t in config.get_text_config().layer_types)
+    filled = bool(cache.get_seq_length() > 0 if has_attention else cache.has_previous_state())
+    inputs_embeds = model_kwargs.get("inputs_embeds") if is_first_iteration else None
+    length = input_ids.shape[1] if inputs_embeds is None else inputs_embeds.shape[1]
+    if generation_config.prefill_chunk_size is not None:
+        # the chunks run from the start of the prompt
+        position_ids = position_ids[:, :length]
+    else:
+        if filled and has_attention:
+            # the slice `_prefill` feeds: a prompt that spans the cache is cut to its uncached part
+            mask_length = getattr(generation_config, "_mask_length", -1)
+            if attention_mask is not None:
+                mask_length = attention_mask.shape[1]
+            if inputs_embeds is not None or mask_length == length:
+                length -= int(cache.get_seq_length())
+        position_ids = position_ids[:, position_ids.shape[1] - length :]
+    if attention_mask is None:
+        attention_mask = torch.ones_like(position_ids)
+    if bool(packed_under_padding(position_ids, attention_mask, filled).any()):
+        raise ValueError(
+            "`generate` cannot decode packed documents (`position_ids` restarting, or starting at 0 on a cache that "
+            "holds tokens): attention would span them. Generate each document in its own row of a left-padded batch."
+        )
+
+
 @auto_docstring
 class Lfm3ForCausalLM(Lfm3PreTrainedModel, GenerationMixin):
     _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
-    _tp_plan = {"lm_head": "colwise_rep"}
+    _tp_plan = {"lm_head": "colwise_gather_output"}
     _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+    _fsdp_plan = {"lm_head": "keep_full_weight"}
 
     def __init__(self, config):
         super().__init__(config)
@@ -779,22 +819,31 @@ class Lfm3ForCausalLM(Lfm3PreTrainedModel, GenerationMixin):
     @auto_docstring
     def forward(
         self,
-        input_ids: Optional[torch.LongTensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[Cache] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        use_cache: Optional[bool] = None,
-        cache_position: Optional[torch.LongTensor] = None,
-        logits_to_keep: Union[int, torch.Tensor] = 0,
+        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,
+        logits_to_keep: int | torch.Tensor = 0,
         **kwargs: Unpack[TransformersKwargs],
     ) -> CausalLMOutputWithPast:
         r"""
-        cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
-            Indices depicting the position of the input sequence tokens in the sequence. Used to update the
-            cache in the correct position and to infer the complete sequence length.
-        """
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, Lfm3ForCausalLM
+
+        >>> model = Lfm3ForCausalLM.from_pretrained("LiquidAI/LFM3-1.7B-Preview")
+        >>> tokenizer = AutoTokenizer.from_pretrained("LiquidAI/LFM3-1.7B-Preview")
+
+        >>> inputs = tokenizer.apply_chat_template(
+        ...     [{"role": "user", "content": "What is C. elegans?"}], add_generation_prompt=True, return_tensors="pt"
+        ... )
+        >>> generate_ids = model.generate(**inputs, max_new_tokens=64)
+        >>> tokenizer.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True)
+        ```"""
         outputs: BaseModelOutputWithPast = self.model(
             input_ids=input_ids,
             attention_mask=attention_mask,
@@ -802,21 +851,17 @@ class Lfm3ForCausalLM(Lfm3PreTrainedModel, GenerationMixin):
             past_key_values=past_key_values,
             inputs_embeds=inputs_embeds,
             use_cache=use_cache,
-            cache_position=cache_position,
             **kwargs,
         )
 
         hidden_states = outputs.last_hidden_state
-        slice_indices = (
-            slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
-        )
+        # 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=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs
-            )
+            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
 
         return CausalLMOutputWithPast(
             loss=loss,
@@ -826,5 +871,11 @@ class Lfm3ForCausalLM(Lfm3PreTrainedModel, GenerationMixin):
             attentions=outputs.attentions,
         )
 
+    def _prefill(self, input_ids, generation_config, model_kwargs, is_first_iteration=True):
+        # every prefill passes here once before its first forward: a check per chunk would leave the earlier chunks
+        # in the cache, and a dynamic cache would only refuse packing at the first decode step
+        reject_packed_generation(self.config, input_ids, generation_config, model_kwargs, is_first_iteration)
+        return super()._prefill(input_ids, generation_config, model_kwargs, is_first_iteration=is_first_iteration)
+
 
 __all__ = ["Lfm3ForCausalLM", "Lfm3Model", "Lfm3PreTrainedModel"]

tokenizer.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a39eed58008cbcbc09e5b629b81338a654269c64cd13348aaa6090f4db75c7e7
-size 17908258
+oid sha256:41927bd314d17d7d7fadb44e10cd0919f7b71528d7741596d8cbd15c6945ec44
+size 17908634

tokenizer_config.json

@@ -1,20 +1,62 @@
 {
+  "audio_out_token": "<|audio_out|>",
+  "audio_token": "<|audio|>",
   "backend": "tokenizers",
   "bos_token": "<|startoftext|>",
   "eos_token": "<|im_end|>",
   "extra_special_tokens": [],
   "image_token": "<|image|>",
-  "is_local": true,
-  "legacy": false,
+  "is_local": false,
   "local_files_only": false,
   "model_max_length": 1000000000000000019884624838656,
   "model_specific_special_tokens": {
+    "audio_out_token": "<|audio_out|>",
+    "audio_token": "<|audio|>",
     "image_token": "<|image|>",
     "video_token": "<|video|>",
     "vision_end_token": "<|vision_end|>",
     "vision_start_token": "<|vision_start|>"
   },
   "pad_token": "<|pad|>",
+  "response_template": {
+    "defaults": {
+      "role": "assistant"
+    },
+    "fields": {
+      "content": {
+        "close_pattern": "<\\|im_end\\|>\\s*",
+        "content": "text"
+      },
+      "reasoning_content": {
+        "close_pattern": "</think>\\s*",
+        "content": "text",
+        "open_pattern": "<think>\\n?"
+      },
+      "tool_calls": {
+        "close_pattern": "</tool_call>\\s*",
+        "content": "xml-inline",
+        "content_args": {
+          "tag_pattern": "<parameter=(?P<key>[^>\\n]+)>\\n?(?P<value>.*?)\\n?</parameter>(?=\\s*(?:<parameter=|</function>\\s*\\Z|\\Z))",
+          "value_parser": {
+            "args": {
+              "strip": false
+            },
+            "name": "text"
+          }
+        },
+        "open_pattern": "<tool_call>\\s*<function=(?P<name>[^\\n>]+)>",
+        "repeats": true,
+        "transform": {
+          "function": {
+            "arguments": "{content}",
+            "name": "{name}"
+          },
+          "type": "function"
+        }
+      }
+    },
+    "start_anchor_pattern": "<\\|im_start\\|>assistant\\n"
+  },
   "tokenizer_class": "TokenizersBackend",
   "video_token": "<|video|>",
   "vision_end_token": "<|vision_end|>",