Upload Lfm3VlForConditionalGeneration

README.md

@@ -0,0 +1,199 @@
+---
+library_name: transformers
+tags: []
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]

config.json

@@ -0,0 +1,154 @@
+{
+  "architectures": [
+    "Lfm3VlForConditionalGeneration"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_lfm3_vl.Lfm3VlConfig",
+    "AutoModel": "modeling_lfm3_vl.Lfm3VlModel",
+    "AutoModelForImageTextToText": "modeling_lfm3_vl.Lfm3VlForConditionalGeneration"
+  },
+  "bos_token_id": 124894,
+  "downsample_factor": 2,
+  "dtype": "bfloat16",
+  "eos_token_id": 124900,
+  "image_token_id": 125030,
+  "model_type": "lfm3_vl",
+  "pad_token_id": 124893,
+  "projector_bias": true,
+  "projector_hidden_act": "gelu",
+  "projector_hidden_size": 4096,
+  "projector_layer_norm_eps": 1e-06,
+  "text_config": {
+    "_name_or_path": "",
+    "architectures": null,
+    "bos_token_id": 124894,
+    "chunk_size_feed_forward": 0,
+    "conv_L_cache": 7,
+    "conv_bias": false,
+    "dtype": "bfloat16",
+    "eos_token_id": 124900,
+    "head_dim": 128,
+    "hidden_size": 2048,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "intermediate_size": 7680,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_types": [
+      "conv",
+      "conv",
+      "full_attention",
+      "conv",
+      "sliding_attention",
+      "conv",
+      "full_attention",
+      "conv",
+      "sliding_attention",
+      "conv",
+      "full_attention",
+      "conv",
+      "sliding_attention",
+      "conv",
+      "full_attention",
+      "conv",
+      "sliding_attention",
+      "conv",
+      "full_attention",
+      "conv",
+      "sliding_attention",
+      "conv",
+      "full_attention",
+      "conv"
+    ],
+    "max_position_embeddings": 128000,
+    "model_type": "lfm3",
+    "norm_eps": 1e-05,
+    "num_attention_heads": 16,
+    "num_hidden_layers": 24,
+    "num_key_value_heads": 4,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "pad_token_id": 124893,
+    "problem_type": null,
+    "return_dict": true,
+    "rope_parameters": {
+      "full_attention": {
+        "partial_rotary_factor": 0.5,
+        "rope_theta": 1000000.0,
+        "rope_type": "default"
+      },
+      "sliding_attention": {
+        "rope_theta": 10000,
+        "rope_type": "default"
+      }
+    },
+    "sliding_window": 256,
+    "tie_word_embeddings": true,
+    "use_cache": true,
+    "use_sandwich_norm": true,
+    "vocab_size": 128000
+  },
+  "tie_word_embeddings": true,
+  "transformers_version": "5.21.0",
+  "video_token_id": 125029,
+  "vision_config": {
+    "_name_or_path": "",
+    "architectures": null,
+    "chunk_size_feed_forward": 0,
+    "conv_kernel_size": 3,
+    "dtype": "bfloat16",
+    "hidden_size": 1024,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "in_channels": 3,
+    "initializer_range": 0.02,
+    "intermediate_size": 2560,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_types": [
+      "conv",
+      "conv",
+      "full_attention",
+      "conv",
+      "full_attention",
+      "conv",
+      "full_attention",
+      "conv",
+      "full_attention",
+      "conv",
+      "full_attention",
+      "conv",
+      "full_attention",
+      "conv"
+    ],
+    "model_type": "lfm3_vl_vision",
+    "norm_eps": 1e-05,
+    "num_attention_heads": 16,
+    "num_hidden_layers": 14,
+    "num_key_value_heads": 8,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "patch_size": 16,
+    "problem_type": null,
+    "return_dict": true,
+    "rope_parameters": {
+      "rope_theta": 1000000.0,
+      "rope_type": "axial"
+    },
+    "spatial_merge_size": 2,
+    "temporal_patch_size": 2
+  },
+  "vision_end_token_id": 125010,
+  "vision_start_token_id": 125009
+}

configuration_lfm3.py

@@ -0,0 +1,132 @@
+# 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")
+@strict
+class Lfm3Config(PreTrainedConfig):
+    r"""
+    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"]
+    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 __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:
+            self.layer_types = [
+                "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)
+            ]
+
+        self.tie_word_embeddings = kwargs.pop("tie_embedding", self.tie_word_embeddings)
+        super().__post_init__(**kwargs)
+
+    def validate_architecture(self):
+        """Part of ``@strict``-powered validation."""
+        if len(self.layer_types) != self.num_hidden_layers:
+            raise ValueError(
+                f"`layer_types` has {len(self.layer_types)} entries but `num_hidden_layers` is {self.num_hidden_layers}."
+            )
+        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."
+                )
+
+
+__all__ = ["Lfm3Config"]

configuration_lfm3_vl.py

@@ -0,0 +1,211 @@
+# 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.models.auto import CONFIG_MAPPING, AutoConfig
+from transformers.utils import auto_docstring
+
+from .configuration_lfm3 import Lfm3Config
+
+
+def _config_class(model_type):
+    # the LFM3 configs ship with the checkpoint instead of the transformers registry
+    remote = {"lfm3": Lfm3Config, "lfm3_vl_vision": Lfm3VlVisionConfig}
+    return remote.get(model_type) or CONFIG_MAPPING[model_type]
+
+
+@auto_docstring(checkpoint="LiquidAI/LFM3-VL-2B")
+@strict
+class Lfm3VlVisionConfig(PreTrainedConfig):
+    r"""
+    layer_types (`list[str]`, *optional*):
+        Per-layer operator type, each one of `"full_attention"` or `"conv"`. Defaults to two gated convolution
+        layers followed by alternating attention and convolution layers.
+    patch_size (`int`, *optional*, defaults to 16):
+        The spatial size of each image patch.
+    temporal_patch_size (`int`, *optional*, defaults to 2):
+        The number of video frames packed into one patch.
+    spatial_merge_size (`int`, *optional*, defaults to 2):
+        The side of the square block of patches merged into one language token.
+    conv_kernel_size (`int`, *optional*, defaults to 3):
+        The kernel size of the depthwise 2D convolution layers.
+    norm_eps (`float`, *optional*, defaults to 1e-05):
+        The epsilon used by the RMS normalization layers.
+    """
+
+    model_type = "lfm3_vl_vision"
+    base_config_key = "vision_config"
+    default_theta = 1_000_000.0
+    default_rope_type = "axial"
+    _attention_layer_types = ("full_attention",)
+
+    hidden_size: int = 1024
+    intermediate_size: int = 2560
+    num_hidden_layers: int = 14
+    num_attention_heads: int = 16
+    num_key_value_heads: int = 8
+    layer_types: list[str] | None = None
+    patch_size: int = 16
+    temporal_patch_size: int = 2
+    in_channels: int = 3
+    spatial_merge_size: int = 2
+    conv_kernel_size: int = 3
+    norm_eps: float = 1e-5
+    initializer_range: float = 0.02
+    rope_parameters: dict | None = None
+
+    def __post_init__(self, **kwargs):
+        if self.layer_types is None:
+            self.layer_types = ["conv" if i < 2 or i % 2 else "full_attention" for i in range(self.num_hidden_layers)]
+        super().__post_init__(**kwargs)
+
+    def validate_architecture(self):
+        """Part of ``@strict``-powered validation."""
+        if len(self.layer_types) != self.num_hidden_layers:
+            raise ValueError(
+                f"`layer_types` has {len(self.layer_types)} entries but `num_hidden_layers` is {self.num_hidden_layers}."
+            )
+        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 self.hidden_size % self.num_attention_heads:
+            raise ValueError("`hidden_size` must be a multiple of `num_attention_heads`.")
+
+
+# The first checkpoints spell every text and `vision_*` field at the top level of the config
+_LEGACY_TEXT_KEYS = (
+    "vocab_size",
+    "hidden_size",
+    "intermediate_size",
+    "num_hidden_layers",
+    "num_attention_heads",
+    "num_key_value_heads",
+    "head_dim",
+    "layer_types",
+    "conv_bias",
+    "conv_L_cache",
+    "sliding_window",
+    "max_position_embeddings",
+    "use_sandwich_norm",
+    "use_cache",
+)
+_LEGACY_VISION_KEYS = {
+    "vision_hidden_size": "hidden_size",
+    "vision_intermediate_size": "intermediate_size",
+    "vision_num_hidden_layers": "num_hidden_layers",
+    "vision_num_attention_heads": "num_attention_heads",
+    "vision_num_key_value_heads": "num_key_value_heads",
+    "vision_layer_types": "layer_types",
+    "vision_patch_size": "patch_size",
+    "vision_temporal_patch_size": "temporal_patch_size",
+    "vision_num_channels": "in_channels",
+    "vision_spatial_merge_size": "spatial_merge_size",
+}
+_LEGACY_DROPPED_KEYS = ("vision_patch_dim", "vision_max_position_embeddings")
+
+
+@auto_docstring(checkpoint="LiquidAI/LFM3-VL-2B")
+@strict
+class Lfm3VlConfig(PreTrainedConfig):
+    r"""
+    video_token_id (`int`, *optional*, defaults to 125029):
+        The id of the placeholder token for video patches.
+    vision_start_token_id (`int`, *optional*, defaults to 125009):
+        The id of the token opening an image or a video frame.
+    vision_end_token_id (`int`, *optional*, defaults to 125010):
+        The id of the token closing an image or a video frame.
+    projector_layer_norm_eps (`float`, *optional*, defaults to 1e-06):
+        The epsilon of the layer normalization applied to the vision features before projection.
+    downsample_factor (`int`, *optional*, defaults to 2):
+        The side of the square block of vision tokens merged by the projector, equal to the vision encoder's
+        `spatial_merge_size`.
+    """
+
+    model_type = "lfm3_vl"
+    sub_configs = {"text_config": AutoConfig, "vision_config": Lfm3VlVisionConfig}
+
+    text_config: dict | PreTrainedConfig | None = None
+    vision_config: dict | PreTrainedConfig | None = None
+    image_token_id: int = 125_030
+    video_token_id: int = 125_029
+    vision_start_token_id: int = 125_009
+    vision_end_token_id: int = 125_010
+    projector_hidden_act: str = "gelu"
+    projector_hidden_size: int = 4096
+    projector_bias: bool = True
+    projector_layer_norm_eps: float = 1e-6
+    downsample_factor: int = 2
+    tie_word_embeddings: bool = True
+
+    def __post_init__(self, **kwargs):
+        if "vision_hidden_size" in kwargs:
+            kwargs = self._convert_legacy_flat_config(kwargs)
+
+        if isinstance(self.vision_config, dict):
+            self.vision_config = self.sub_configs["vision_config"](**self.vision_config)
+        elif self.vision_config is None:
+            self.vision_config = self.sub_configs["vision_config"]()
+
+        if isinstance(self.text_config, dict):
+            self.text_config = {"model_type": "lfm3", **self.text_config}
+            self.text_config = _config_class(self.text_config["model_type"])(**self.text_config)
+        elif self.text_config is None:
+            self.text_config = _config_class("lfm3")()
+        self.validate_vision_fields()
+
+        super().__post_init__(**kwargs)
+
+    def validate_vision_fields(self):
+        """Part of `@strict`-powered validation, so `validate()` and `save_pretrained` refuse later edits too."""
+        if self.downsample_factor != self.vision_config.spatial_merge_size:
+            raise ValueError("downsample_factor must match vision_config.spatial_merge_size")
+        # an `inputs_embeds` forward embeds both placeholder ids to find them
+        vocab_size = self.text_config.vocab_size
+        for name in ("image_token_id", "video_token_id"):
+            if not 0 <= getattr(self, name) < vocab_size:
+                raise ValueError(f"{name}={getattr(self, name)} must lie in [0, text_config.vocab_size={vocab_size})")
+
+    def _convert_legacy_flat_config(self, kwargs: dict) -> dict:
+        text_config = {key: kwargs.pop(key) for key in _LEGACY_TEXT_KEYS if key in kwargs}
+        text_config["norm_eps"] = kwargs.pop("rms_norm_eps", 1e-5)
+        text_config["rope_parameters"] = {
+            "full_attention": {
+                "rope_type": "default",
+                "rope_theta": kwargs.pop("full_attention_rope_theta", 1_000_000.0),
+                "partial_rotary_factor": kwargs.pop("partial_rotary_factor", 0.5),
+            },
+            "sliding_attention": {
+                "rope_type": "default",
+                "rope_theta": kwargs.pop("sliding_attention_rope_theta", 10_000.0),
+            },
+        }
+        for key in ("pad_token_id", "bos_token_id", "eos_token_id"):
+            if kwargs.get(key) is not None:
+                text_config[key] = kwargs[key]
+        text_config["tie_word_embeddings"] = self.tie_word_embeddings
+
+        vision_config = {new: kwargs.pop(old) for old, new in _LEGACY_VISION_KEYS.items() if old in kwargs}
+        vision_config["norm_eps"] = text_config["norm_eps"]
+        vision_config["rope_parameters"] = {"rope_type": "axial", "rope_theta": kwargs.pop("vision_rope_theta", 1e6)}
+        for key in _LEGACY_DROPPED_KEYS:
+            kwargs.pop(key, None)
+
+        self.text_config = text_config
+        self.vision_config = vision_config
+        return kwargs
+
+
+__all__ = ["Lfm3VlConfig", "Lfm3VlVisionConfig"]

generation_config.json

@@ -0,0 +1,14 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 124894,
+  "do_sample": true,
+  "eos_token_id": 124900,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 124893,
+  "repetition_penalty": 1.1,
+  "temperature": 0.1,
+  "top_k": 50,
+  "transformers_version": "5.21.0",
+  "use_cache": true
+}

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

model.safetensors

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

modeling_lfm3.py

@@ -0,0 +1,881 @@
+# 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_with_fallback,
+    use_kernelized_func,
+)
+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, 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
+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: 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: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class Lfm3RotaryEmbedding(nn.Module):
+    """
+    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):
+        super().__init__()
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+        self.config = config
+        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: 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
+        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):
+        inv_freq = getattr(self, f"{layer_type}_inv_freq")
+        attention_scaling = getattr(self, f"{layer_type}_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):
+        super().__init__()
+        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 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        attn_weights = attn_weights + attention_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+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):
+    """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):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        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.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)
+        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: 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)
+
+        query_states = self.q_norm(self.q_proj(hidden_states).view(*hidden_shape)).transpose(1, 2)
+        key_states = self.k_norm(self.k_proj(hidden_states).view(*hidden_shape)).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(*hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_values is not None:
+            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            **kwargs,
+        )
+        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):
+    """
+    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
+
+
+@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):
+    """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,
+        config: Lfm3Config,
+        layer_idx: int,
+    ):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.conv_kernel_size = config.conv_L_cache
+        self.bias = config.conv_bias
+
+        self.conv = nn.Conv1d(
+            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.conv_kernel_size - 1,
+        )
+        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)
+
+        self.layer_type = config.layer_types[layer_idx]
+
+    @force_accelerate_hooks("conv")
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        past_key_values: Cache | None = None,
+        attention_mask: torch.Tensor | None = None,
+        seq_idx: torch.IntTensor | None = None,
+    ):
+        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 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 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:
+                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 * hidden_states
+        y = y.transpose(-1, -2).contiguous()
+        y = self.out_proj(y)
+        return y
+
+    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.is_attention_layer = config.layer_types[layer_idx] in config._attention_layer_types
+        self.use_sandwich_norm = config.use_sandwich_norm
+
+        if self.is_attention_layer:
+            self.self_attn = Lfm3Attention(config, layer_idx)
+        else:
+            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:
+            self.post_operator_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
+            self.post_ffn_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        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
+        if self.is_attention_layer:
+            hidden_states, _ = self.self_attn(
+                hidden_states=self.operator_norm(hidden_states),
+                position_embeddings=position_embeddings,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                **kwargs,
+            )
+        else:
+            hidden_states = self.conv(
+                hidden_states=self.operator_norm(hidden_states),
+                past_key_values=past_key_values,
+                attention_mask=attention_mask,
+                seq_idx=seq_idx,
+            )
+        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)
+        hidden_states = hidden_states + ffn_out
+
+        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
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Lfm3DecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+
+    _can_compile_fullgraph = True
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "hidden_states": Lfm3DecoderLayer,
+        "attentions": Lfm3Attention,
+    }
+
+    @torch.no_grad()
+    def _init_weights(self, module):
+        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
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        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
+        self.embedding_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @merge_with_config_defaults
+    @capture_outputs
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPast:
+        r"""
+        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:
+            past_key_values = DynamicCache(config=self.config)
+        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
+
+        # 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 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
+        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 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=causal_mask_mapping[layer_type],
+                position_embeddings=position_embeddings.get(layer_type),
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                seq_idx=seq_idx,
+                **kwargs,
+            )
+
+        hidden_states = self.embedding_norm(hidden_states)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+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_gather_output"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+    _fsdp_plan = {"lm_head": "keep_full_weight"}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = Lfm3Model(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> CausalLMOutputWithPast:
+        r"""
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, Lfm3ForCausalLM
+
+        >>> model = Lfm3ForCausalLM.from_pretrained("LiquidAI/LFM3-1.7B")
+        >>> tokenizer = AutoTokenizer.from_pretrained("LiquidAI/LFM3-1.7B")
+
+        >>> 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,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            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"]

modeling_lfm3_vl.py

@@ -0,0 +1,965 @@
+# 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 collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache
+from transformers.configuration_utils import PreTrainedConfig
+from transformers.generation import GenerationMixin
+from transformers.integrations import use_kernel_forward_from_hub
+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, BaseModelOutputWithPooling, CausalLMOutputWithPast
+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,
+    torch_compilable_check,
+)
+from transformers.utils.generic import (
+    accepts_precomputed_kwargs,
+    is_flash_attention_requested,
+    maybe_autocast,
+    merge_with_config_defaults,
+)
+from transformers.utils.output_capturing import OutputRecorder, capture_outputs
+from transformers.vision_utils import get_max_seqlen, get_vision_attention_seqlens, get_vision_position_ids
+
+from .configuration_lfm3_vl import Lfm3VlConfig, Lfm3VlVisionConfig
+from .modeling_lfm3 import Lfm3Model
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class Lfm3VlRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+        """
+        Lfm3VlRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class Lfm3VlMLP(nn.Module):
+    def __init__(self, config: Lfm3VlConfig):
+        super().__init__()
+        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))
+
+
+class Lfm3VlVisionRotaryEmbedding(nn.Module):
+    """
+    Simple axial 2D rope with same freqs used for H and W grids. The freqs are
+    pre-computed using `head-dim//4` which is later used to concat H and W positions.
+    The final angles rotate over the whole head dim, no partial rotation involved.
+    """
+
+    def __init__(self, config: Lfm3VlVisionConfig):
+        super().__init__()
+        self.config = config
+
+        self.rope_type = self.config.rope_parameters["rope_type"]
+        rope_init_fn: Callable = self.compute_axial_rope_parameters
+        if self.rope_type != "axial":
+            raise ValueError(f"{self.__class__.__name__} supports only axial rope, but requested {self.rope_type}")
+        inv_freq, self.attention_scaling = rope_init_fn(self.config)
+
+        self.inv_freq = nn.Buffer(inv_freq, persistent=False)
+        self.original_inv_freq = nn.Buffer(inv_freq.clone(), persistent=False)
+
+    @staticmethod
+    def compute_axial_rope_parameters(config: Lfm3VlVisionConfig, **kwargs) -> tuple[torch.Tensor, float]:
+        """Inverse frequencies of the axial RoPE over half the head dimension per axis, with an attention scaling of 1."""
+        base = config.rope_parameters["rope_theta"]
+        spatial_dim = config.hidden_size // config.num_attention_heads // 2
+        inv_freq = 1.0 / (base ** (torch.arange(0, spatial_dim, 2, dtype=torch.float) / spatial_dim))
+        return inv_freq, 1.0
+
+    @torch.no_grad()
+    def forward(self, x, position_ids):
+        # position_ids: (2, N) — row 0 = h coords, row 1 = w coords
+        position_ids_expanded = position_ids[..., None].float()
+        device_type = x.device.type if isinstance(x.device.type, str) else "cpu"
+        with maybe_autocast(device_type=device_type, enabled=False):
+            freqs = position_ids_expanded * self.inv_freq.float()
+            cos = freqs.cos() * self.attention_scaling
+            sin = freqs.sin() * self.attention_scaling
+
+        cos = self.recomposition_frequencies(cos)
+        sin = self.recomposition_frequencies(sin)
+        return cos, sin
+
+    def recomposition_frequencies(self, freq):
+        """
+        Recompose the frequencies into the final spatial layout used per each grid.
+        """
+        freq_h, freq_w = freq[:, 0], freq[:, 1]
+        freq_hw = torch.cat([freq_h, freq_w], dim=-1)
+        return torch.cat([freq_hw, freq_hw], dim=-1)
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb_vision(
+    q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    orig_q_dtype = q.dtype
+    orig_k_dtype = k.dtype
+    q, k = q.float(), k.float()
+    cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float()
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    q_embed = q_embed.to(orig_q_dtype)
+    k_embed = k_embed.to(orig_k_dtype)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        attn_weights = attn_weights + attention_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+class Lfm3VlVisionAttention(nn.Module):
+    """Bidirectional grouped-query attention over the patches of one frame, with RMSNorm on the query and key heads."""
+
+    def __init__(self, config: Lfm3VlVisionConfig):
+        super().__init__()
+        self.config = config
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = 0.0
+        self.is_causal = False
+        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.out_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
+        self.q_layernorm = Lfm3VlRMSNorm(self.head_dim, eps=config.norm_eps)
+        self.k_layernorm = Lfm3VlRMSNorm(self.head_dim, eps=config.norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        max_seqlen: int | None = None,
+        frame_shapes: list[list[int]] | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        seq_length = hidden_states.shape[0]
+        query_states = self.q_layernorm(self.q_proj(hidden_states).view(seq_length, self.num_heads, self.head_dim))
+        key_states = self.k_layernorm(
+            self.k_proj(hidden_states).view(seq_length, self.num_key_value_heads, self.head_dim)
+        )
+        value_states = self.v_proj(hidden_states).view(seq_length, self.num_key_value_heads, self.head_dim)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin)
+
+        query_states = query_states.transpose(0, 1).unsqueeze(0)
+        key_states = key_states.transpose(0, 1).unsqueeze(0)
+        value_states = value_states.transpose(0, 1).unsqueeze(0)
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        if is_flash_attention_requested(self.config):
+            max_seqlen = get_max_seqlen(cu_seqlens, self.config, kwargs={"max_seqlen": max_seqlen})
+            attn_output, attn_weights = attention_interface(
+                self,
+                query_states,
+                key_states,
+                value_states,
+                attention_mask=None,
+                scaling=self.scaling,
+                dropout=0.0,
+                cu_seq_lens_q=cu_seqlens,
+                cu_seq_lens_k=cu_seqlens,
+                max_length_q=max_seqlen,
+                max_length_k=max_seqlen,
+                is_causal=False,
+                **kwargs,
+            )
+        else:
+            # Other implementations: frames of one shape (every video, equal-size images) attend as one batch, mixed
+            # shapes frame by frame. `frame_shapes` gives the lengths without a device sync.
+            if frame_shapes is not None:
+                lengths = [height * width for height, width in frame_shapes]
+            else:
+                lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
+            if len(set(lengths)) == 1:
+                query_states, key_states, value_states = (
+                    tensor.reshape(-1, len(lengths), lengths[0], self.head_dim).transpose(0, 1)
+                    for tensor in (query_states, key_states, value_states)
+                )
+                attn_output, attn_weights = attention_interface(
+                    self,
+                    query_states,
+                    key_states,
+                    value_states,
+                    attention_mask=None,
+                    scaling=self.scaling,
+                    dropout=0.0,
+                    is_causal=False,
+                    **kwargs,
+                )
+            else:
+                splits = [torch.split(tensor, lengths, dim=2) for tensor in (query_states, key_states, value_states)]
+                attn_outputs, attn_weights = [], []
+                for q, k, v in zip(*splits):
+                    output, weights = attention_interface(
+                        self,
+                        q,
+                        k,
+                        v,
+                        attention_mask=None,
+                        scaling=self.scaling,
+                        dropout=0.0,
+                        is_causal=False,
+                        **kwargs,
+                    )
+                    attn_outputs.append(output)
+                    attn_weights.append(weights)
+                attn_output = torch.cat(attn_outputs, dim=1)
+                # frames of different shapes have differently sized maps: one entry per frame, or None under sdpa
+                attn_weights = None if attn_weights[0] is None else tuple(attn_weights)
+
+        attn_output = attn_output.reshape(seq_length, -1).contiguous()
+        attn_output = self.out_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class Lfm3VlVisionConv2d(nn.Module):
+    """Gated depthwise 2D convolution over the patch grid of each frame."""
+
+    def __init__(self, config: Lfm3VlVisionConfig):
+        super().__init__()
+        self.spatial_merge_size = config.spatial_merge_size
+        self.in_proj = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=False)
+        self.spatial_conv = nn.Conv2d(
+            config.hidden_size,
+            config.hidden_size,
+            kernel_size=config.conv_kernel_size,
+            padding=config.conv_kernel_size // 2,
+            groups=config.hidden_size,
+            bias=False,
+        )
+        self.out_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+
+    def forward(self, hidden_states: torch.Tensor, frame_shapes: list[list[int]]):
+        B, C, x = self.in_proj(hidden_states).chunk(3, dim=-1)
+        gated = B * x
+        merge = self.spatial_merge_size
+        num_frames = len(frame_shapes)
+        if len(set(map(tuple, frame_shapes))) == 1:
+            # frames of one shape (every video, equal-size images) convolve as one batch: merge-block order -> raster
+            # grid, convolve, back
+            height, width = frame_shapes[0]
+            frames = gated.view(num_frames, height // merge, width // merge, merge, merge, -1)
+            frames = frames.permute(0, 5, 1, 3, 2, 4).reshape(num_frames, -1, height, width)
+            frames = self.spatial_conv(frames)
+            frames = frames.view(num_frames, -1, height // merge, merge, width // merge, merge).permute(
+                0, 2, 4, 3, 5, 1
+            )
+            conv_output = frames.reshape(gated.shape[0], -1)
+        else:
+            conv_outputs = []
+            lengths = [height * width for height, width in frame_shapes]
+            for frame, (height, width) in zip(torch.split(gated, lengths, dim=0), frame_shapes):
+                frame = frame.view(height // merge, width // merge, merge, merge, -1)
+                frame = frame.permute(4, 0, 2, 1, 3).reshape(1, -1, height, width)
+                frame = self.spatial_conv(frame)
+                frame = frame.view(-1, height // merge, merge, width // merge, merge).permute(1, 3, 2, 4, 0)
+                conv_outputs.append(frame.reshape(height * width, -1))
+            conv_output = torch.cat(conv_outputs, dim=0)
+        return self.out_proj(C * conv_output)
+
+
+class Lfm3VlVisionLayer(GradientCheckpointingLayer):
+    def __init__(self, config: Lfm3VlVisionConfig, layer_idx: int):
+        super().__init__()
+        self.is_attention_layer = config.layer_types[layer_idx] == "full_attention"
+        if self.is_attention_layer:
+            self.self_attn = Lfm3VlVisionAttention(config)
+        else:
+            self.conv = Lfm3VlVisionConv2d(config)
+        self.feed_forward = Lfm3VlMLP(config)
+        self.operator_norm = Lfm3VlRMSNorm(config.hidden_size, eps=config.norm_eps)
+        self.ffn_norm = Lfm3VlRMSNorm(config.hidden_size, eps=config.norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        frame_shapes: list[list[int]],
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        if self.is_attention_layer:
+            operator_output, _ = self.self_attn(
+                self.operator_norm(hidden_states),
+                cu_seqlens=cu_seqlens,
+                position_embeddings=position_embeddings,
+                frame_shapes=frame_shapes,
+                **kwargs,
+            )
+        else:
+            operator_output = self.conv(self.operator_norm(hidden_states), frame_shapes)
+        hidden_states = hidden_states + operator_output
+        hidden_states = hidden_states + self.feed_forward(self.ffn_norm(hidden_states))
+        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 Lfm3VlPreTrainedModel(PreTrainedModel):
+    config: Lfm3VlConfig
+    base_model_prefix = "model"
+    input_modalities = ("image", "video", "text")
+    supports_gradient_checkpointing = True
+    _skip_keys_device_placement = ["past_key_values"]
+
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    # the cached decode step traces without graph breaks, like the text model
+    _can_compile_fullgraph = True
+    _supports_flex_attn = True
+    _supports_attention_backend = True
+    _no_split_modules = ["Lfm3VlVisionLayer"]
+
+    @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:
+        # the static-cache masks of the text model are sized by layer type, not by the generic first attention layer
+        return create_masks_by_layer_type(config, inputs_embeds, attention_mask, past_key_values, position_ids)
+
+
+def get_vision_frame_shapes(grid_thw: torch.Tensor) -> list[list[int]]:
+    return grid_thw[:, 1:].repeat_interleave(grid_thw[:, 0], dim=0).tolist()
+
+
+@auto_docstring(
+    custom_intro="""
+    The LFM3-VL vision encoder: a stack of gated convolution and bidirectional attention layers over the patches of
+    each image or video frame, embedded by a linear patchifier.
+    """
+)
+class Lfm3VlVisionModel(Lfm3VlPreTrainedModel):
+    config: Lfm3VlVisionConfig
+    input_modalities = ("image", "video")
+    _input_embed_layer = "patchifier"
+    _can_record_outputs = {
+        "hidden_states": Lfm3VlVisionLayer,
+        "attentions": OutputRecorder(Lfm3VlVisionAttention, index=1),
+    }
+
+    def __init__(self, config: Lfm3VlVisionConfig):
+        super().__init__(config)
+        self.spatial_merge_size = config.spatial_merge_size
+        patch_dim = config.in_channels * config.temporal_patch_size * config.patch_size**2
+        self.patchifier = nn.Linear(patch_dim, config.hidden_size, bias=False)
+        self.embedding_norm = Lfm3VlRMSNorm(config.hidden_size, eps=config.norm_eps)
+        self.rotary_pos_emb = Lfm3VlVisionRotaryEmbedding(config)
+        self.layers = nn.ModuleList(
+            [Lfm3VlVisionLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    @merge_with_config_defaults
+    @capture_outputs
+    @auto_docstring
+    def forward(
+        self, pixel_values: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs]
+    ) -> BaseModelOutputWithPooling:
+        r"""
+        pixel_values (`torch.Tensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`):
+            The flattened patches of every image or video, grouped by merge block.
+        grid_thw (`torch.LongTensor` of shape `(num_images_or_videos, 3)`):
+            The temporal, height and width dimensions of the patch grid of each image or video.
+        """
+        position_ids = get_vision_position_ids(grid_thw, self.spatial_merge_size, kwargs=kwargs)
+        cu_seqlens, max_seqlen = get_vision_attention_seqlens(grid_thw, self.config, kwargs=kwargs)
+        frame_shapes = get_vision_frame_shapes(grid_thw)
+
+        hidden_states = self.patchifier(pixel_values.to(self.patchifier.weight.dtype))
+        hidden_states = self.embedding_norm(hidden_states)
+        position_embeddings = self.rotary_pos_emb(hidden_states, position_ids)
+
+        for layer in self.layers:
+            hidden_states = layer(
+                hidden_states,
+                cu_seqlens=cu_seqlens,
+                frame_shapes=frame_shapes,
+                max_seqlen=max_seqlen,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+        return BaseModelOutputWithPooling(last_hidden_state=hidden_states)
+
+
+class Lfm3VlMultiModalProjector(nn.Module):
+    """Merges each block of `downsample_factor**2` vision tokens into one language token."""
+
+    def __init__(self, config: Lfm3VlConfig):
+        super().__init__()
+        self.downsample_factor = config.downsample_factor
+        self.layer_norm = nn.LayerNorm(config.vision_config.hidden_size, eps=config.projector_layer_norm_eps)
+        self.linear_1 = nn.Linear(
+            config.vision_config.hidden_size * config.downsample_factor**2,
+            config.projector_hidden_size,
+            bias=config.projector_bias,
+        )
+        self.act = ACT2FN[config.projector_hidden_act]
+        self.linear_2 = nn.Linear(
+            config.projector_hidden_size, config.text_config.hidden_size, bias=config.projector_bias
+        )
+
+    def forward(self, image_features: torch.Tensor):
+        hidden_states = self.layer_norm(image_features).reshape(-1, self.linear_1.in_features)
+        hidden_states = self.linear_1(hidden_states)
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.linear_2(hidden_states)
+        return hidden_states
+
+
+@auto_docstring(
+    custom_intro="""
+    Base class for Lfm3Vl outputs, with hidden states and attentions.
+    """
+)
+@dataclass
+class Lfm3VlModelOutputWithPast(BaseModelOutputWithPast):
+    r"""
+    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+        It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+        `past_key_values` input) to speed up sequential decoding.
+    image_hidden_states (`torch.FloatTensor`, *optional*):
+        A `torch.FloatTensor` of size `(num_image_tokens, hidden_size)`.
+        image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
+    video_hidden_states (`torch.FloatTensor`, *optional*):
+        A `torch.FloatTensor` of size `(num_video_tokens, hidden_size)`.
+        video_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
+    """
+
+    image_hidden_states: torch.FloatTensor | None = None
+    video_hidden_states: torch.FloatTensor | None = None
+
+
+@auto_docstring(
+    custom_intro="""
+    Base class for Lfm3Vl causal language model (or autoregressive) outputs.
+    """
+)
+@dataclass
+class Lfm3VlCausalLMOutputWithPast(CausalLMOutputWithPast):
+    r"""
+    image_hidden_states (`torch.FloatTensor`, *optional*):
+        A `torch.FloatTensor` of size `(num_image_tokens, hidden_size)`.
+        image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
+    video_hidden_states (`torch.FloatTensor`, *optional*):
+        A `torch.FloatTensor` of size `(num_video_tokens, hidden_size)`.
+        video_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
+    """
+
+    image_hidden_states: torch.FloatTensor | None = None
+    video_hidden_states: torch.FloatTensor | None = None
+
+
+@auto_docstring(
+    custom_intro="""
+    The Lfm3Vl model which consists of a vision backbone and a language model, without a language modeling head.
+    """
+)
+class Lfm3VlModel(Lfm3VlPreTrainedModel):
+    def __init__(self, config: Lfm3VlConfig):
+        super().__init__(config)
+        self.vision_tower = Lfm3VlVisionModel(config.vision_config)
+        self.multi_modal_projector = Lfm3VlMultiModalProjector(config)
+        self.language_model = Lfm3Model(config.text_config)
+        self.post_init()
+
+    @accepts_precomputed_kwargs(modality="image")
+    @can_return_tuple
+    @auto_docstring(
+        custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection."
+    )
+    def get_image_features(
+        self,
+        pixel_values: torch.FloatTensor,
+        image_grid_thw: torch.LongTensor,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | BaseModelOutputWithPooling:
+        r"""
+        pixel_values (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`):
+            The flattened patches of every image, grouped by merge block.
+        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`):
+            The temporal, height and width dimensions of the patch grid of each image.
+        """
+        image_outputs = self.vision_tower(pixel_values, grid_thw=image_grid_thw, return_dict=True, **kwargs)
+        image_features = self.multi_modal_projector(image_outputs.last_hidden_state)
+        split_sizes = (image_grid_thw.prod(-1) // self.vision_tower.spatial_merge_size**2).tolist()
+        image_outputs.pooler_output = torch.split(image_features, split_sizes)
+        return image_outputs
+
+    def get_placeholder_mask(
+        self,
+        input_ids: torch.LongTensor,
+        inputs_embeds: torch.FloatTensor,
+        image_features: torch.FloatTensor | None = None,
+        video_features: torch.FloatTensor | None = None,
+    ):
+        """Placeholder masks from `input_ids` or `inputs_embeds`, checked against the feature counts."""
+        if input_ids is None:
+            special_image_mask = inputs_embeds == self.get_input_embeddings()(
+                torch.full((), self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+            )
+            special_image_mask = special_image_mask.all(-1)
+            special_video_mask = inputs_embeds == self.get_input_embeddings()(
+                torch.full((), self.config.video_token_id, dtype=torch.long, device=inputs_embeds.device)
+            )
+            special_video_mask = special_video_mask.all(-1)
+        else:
+            special_image_mask = input_ids == self.config.image_token_id
+            special_video_mask = input_ids == self.config.video_token_id
+
+        n_image_tokens = special_image_mask.sum()
+        special_image_mask = special_image_mask.unsqueeze(-1).to(inputs_embeds.device)
+        if image_features is not None:
+            torch_compilable_check(
+                n_image_tokens * inputs_embeds.shape[-1] == image_features.numel(),
+                f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {image_features.shape[0]}",
+            )
+
+        n_video_tokens = special_video_mask.sum()
+        special_video_mask = special_video_mask.unsqueeze(-1).to(inputs_embeds.device)
+        if video_features is not None:
+            torch_compilable_check(
+                n_video_tokens * inputs_embeds.shape[-1] == video_features.numel(),
+                f"Video features and video tokens do not match, tokens: {n_video_tokens}, features: {video_features.shape[0]}",
+            )
+        return special_image_mask, special_video_mask
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        pixel_values: torch.FloatTensor | None = None,
+        image_grid_thw: torch.LongTensor | None = None,
+        pixel_values_videos: torch.FloatTensor | None = None,
+        video_grid_thw: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        mm_encoder_outputs: dict[str, BaseModelOutputWithPooling] | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | Lfm3VlModelOutputWithPast:
+        r"""
+        pixel_values (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`, *optional*):
+            The flattened patches of every image, grouped by merge block.
+        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+            The temporal, height and width dimensions of the patch grid of each image.
+        pixel_values_videos (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`, *optional*):
+            The flattened patches of every video, grouped by merge block.
+        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+            The temporal, height and width dimensions of the patch grid of each video.
+        """
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if (pixel_values is not None or pixel_values_videos is not None) and mm_encoder_outputs is not None:
+            raise ValueError(
+                "You cannot specify both pixel_values/pixel_values_videos and mm_encoder_outputs at the same time"
+            )
+
+        if inputs_embeds is None:
+            inputs_embeds = self.get_input_embeddings()(input_ids)
+
+        mm_encoder_outputs = mm_encoder_outputs if mm_encoder_outputs is not None else {}
+        if (pixel_values is None) != (image_grid_thw is None):
+            raise ValueError("pixel_values and image_grid_thw must be supplied together")
+        if (pixel_values_videos is None) != (video_grid_thw is None):
+            raise ValueError("pixel_values_videos and video_grid_thw must be supplied together")
+        if mm_encoder_outputs.get("image") is None and pixel_values is not None:
+            mm_encoder_outputs["image"] = self.get_image_features(
+                pixel_values, image_grid_thw, return_dict=True, **kwargs
+            )
+        if mm_encoder_outputs.get("video") is None and pixel_values_videos is not None:
+            mm_encoder_outputs["video"] = self.get_video_features(
+                pixel_values_videos, video_grid_thw, return_dict=True, **kwargs
+            )
+
+        image_features = None
+        if mm_encoder_outputs.get("image") is not None:
+            image_features = torch.cat(mm_encoder_outputs["image"].pooler_output, dim=0).to(
+                inputs_embeds.device, inputs_embeds.dtype
+            )
+            image_mask, _ = self.get_placeholder_mask(
+                input_ids, inputs_embeds=inputs_embeds, image_features=image_features
+            )
+            inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_features)
+
+        video_features = None
+        if mm_encoder_outputs.get("video") is not None:
+            video_features = torch.cat(mm_encoder_outputs["video"].pooler_output, dim=0).to(
+                inputs_embeds.device, inputs_embeds.dtype
+            )
+            _, video_mask = self.get_placeholder_mask(
+                input_ids, inputs_embeds=inputs_embeds, video_features=video_features
+            )
+            inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_features)
+
+        outputs = self.language_model(
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            **kwargs,
+        )
+
+        return Lfm3VlModelOutputWithPast(
+            last_hidden_state=outputs.last_hidden_state,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            image_hidden_states=image_features,
+            video_hidden_states=video_features,
+        )
+
+    @accepts_precomputed_kwargs(modality="video")
+    @can_return_tuple
+    @auto_docstring(
+        custom_intro="Obtains video last hidden states from the vision tower and apply multimodal projection."
+    )
+    def get_video_features(
+        self,
+        pixel_values_videos: torch.FloatTensor,
+        video_grid_thw: torch.LongTensor,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | BaseModelOutputWithPooling:
+        r"""
+        pixel_values_videos (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`):
+            The flattened patches of every video, grouped by merge block.
+        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`):
+            The temporal, height and width dimensions of the patch grid of each video.
+        """
+        return self.get_image_features(pixel_values_videos, video_grid_thw, **kwargs)
+
+
+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."
+        )
+
+
+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(
+    custom_intro="""
+    The LFM3_VL model which consists of a vision backbone and a language model.
+    """
+)
+class Lfm3VlForConditionalGeneration(Lfm3VlPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
+
+    def __init__(self, config: Lfm3VlConfig):
+        super().__init__(config)
+        self.model = Lfm3VlModel(config)
+        self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_output_embeddings(self) -> nn.Module:
+        return self.lm_head
+
+    @auto_docstring
+    def get_image_features(
+        self,
+        pixel_values: torch.FloatTensor,
+        image_grid_thw: torch.LongTensor,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | BaseModelOutputWithPooling:
+        r"""
+        pixel_values (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`):
+            The flattened patches of every image, grouped by merge block.
+        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`):
+            The temporal, height and width dimensions of the patch grid of each image.
+        """
+        return self.model.get_image_features(pixel_values, image_grid_thw, **kwargs)
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        pixel_values: torch.FloatTensor | None = None,
+        image_grid_thw: torch.LongTensor | None = None,
+        pixel_values_videos: torch.FloatTensor | None = None,
+        video_grid_thw: torch.LongTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        mm_encoder_outputs: dict[str, BaseModelOutputWithPooling] | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | Lfm3VlCausalLMOutputWithPast:
+        r"""
+        pixel_values (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`, *optional*):
+            The flattened patches of every image, grouped by merge block.
+        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+            The temporal, height and width dimensions of the patch grid of each image.
+        pixel_values_videos (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`, *optional*):
+            The flattened patches of every video, grouped by merge block.
+        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+            The temporal, height and width dimensions of the patch grid of each video.
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoProcessor, AutoModelForImageTextToText
+
+        >>> model = AutoModelForImageTextToText.from_pretrained("LiquidAI/LFM3-VL-2B")
+        >>> processor = AutoProcessor.from_pretrained("LiquidAI/LFM3-VL-2B")
+
+        >>> conversation = [
+        ...     {
+        ...         "role": "user",
+        ...         "content": [
+        ...             {"type": "image", "url": "https://www.ilankelman.org/stopsigns/australia.jpg"},
+        ...             {"type": "text", "text": "What is in this image?"},
+        ...         ],
+        ...     },
+        ... ]
+
+        >>> inputs = processor.apply_chat_template(
+        ...     conversation, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt"
+        ... )
+
+        >>> outputs = model.generate(**inputs, max_new_tokens=45)
+        >>> processor.batch_decode(outputs, skip_special_tokens=True)[0]
+        ```"""
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            pixel_values=pixel_values,
+            image_grid_thw=image_grid_thw,
+            pixel_values_videos=pixel_values_videos,
+            video_grid_thw=video_grid_thw,
+            use_cache=use_cache,
+            mm_encoder_outputs=mm_encoder_outputs,
+            **kwargs,
+        )
+
+        hidden_states = outputs[0]
+        # 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.text_config.vocab_size, **kwargs
+            )
+
+        return Lfm3VlCausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            image_hidden_states=outputs.image_hidden_states,
+            video_hidden_states=outputs.video_hidden_states,
+        )
+
+    def _prefill(self, input_ids, generation_config, model_kwargs, is_first_iteration=True):
+        # the text model's refusal, before the prefill writes to the cache
+        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)
+
+    @auto_docstring
+    def get_video_features(
+        self,
+        pixel_values_videos: torch.FloatTensor,
+        video_grid_thw: torch.LongTensor,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | BaseModelOutputWithPooling:
+        r"""
+        pixel_values_videos (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`):
+            The flattened patches of every video, grouped by merge block.
+        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`):
+            The temporal, height and width dimensions of the patch grid of each video.
+        """
+        return self.model.get_video_features(pixel_values_videos, video_grid_thw, **kwargs)
+
+
+__all__ = ["Lfm3VlForConditionalGeneration", "Lfm3VlModel", "Lfm3VlPreTrainedModel", "Lfm3VlVisionModel"]