Add tiny random LFM3 model for optimum-intel tests (#1)

- Add tiny random LFM3 model for optimum-intel tests (cf7edf4c3dd449e2ca23f93b6ce8e9ae99b097d4)

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

chat_template.jinja

@@ -0,0 +1,230 @@
+{{- bos_token -}}
+{%- set preserve_thinking = preserve_thinking | default(false) -%}
+{%- set enable_thinking = enable_thinking | default(true) -%}
+
+{%- macro render_value(v) -%}
+    {%- if v is string -%}
+        {{- v -}}
+    {%- elif v is boolean -%}
+        {{- "true" if v else "false" -}}
+    {%- elif v is none -%}
+        {{- "null" -}}
+    {%- elif v is mapping or (v is sequence and v is not string) -%}
+        {{- v | tojson -}}
+    {%- else -%}
+        {{- v -}}
+    {%- endif -%}
+{%- endmacro -%}
+
+{%- macro parse_content(content) -%}
+    {%- if content is string -%}
+        {{- content -}}
+    {%- elif content is mapping and content.get("type") == "image" -%}
+        {{- "<|vision_start|><|image|><|vision_end|>" -}}
+    {%- elif content is mapping and content.get("type") == "video" -%}
+        {{- "<|vision_start|><|video|><|vision_end|>" -}}
+    {%- elif content is mapping and content.get("type") == "text" -%}
+        {{- (content.get("text") or "") | string -}}
+    {%- elif content is mapping -%}
+        {{- content | tojson -}}
+    {%- elif content is iterable -%}
+        {%- set _ns = namespace(result="") -%}
+        {%- for item in content -%}
+            {%- if item is string -%}
+                {%- set _ns.result = _ns.result + item -%}
+            {%- elif item is mapping and item.get("type") == "image" -%}
+                {%- set _ns.result = _ns.result + "<|vision_start|><|image|><|vision_end|>" -%}
+            {%- elif item is mapping and item.get("type") == "video" -%}
+                {%- set _ns.result = _ns.result + "<|vision_start|><|video|><|vision_end|>" -%}
+            {%- elif item is mapping and item.get("type") == "text" -%}
+                {%- set _ns.result = _ns.result + ((item.get("text") or "") | string) -%}
+            {%- else -%}
+                {%- set _ns.result = _ns.result + (item | tojson) -%}
+            {%- endif -%}
+        {%- endfor -%}
+        {{- _ns.result -}}
+    {%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_tool_calls(tool_calls) -%}
+    {%- for tool_call in tool_calls -%}
+        {%- if not loop.first -%}{{- "\n\n" -}}{%- endif -%}
+        {%- set func = tool_call["function"] if "function" in tool_call else tool_call -%}
+        {%- set func_args = func.get("arguments") -%}
+        {%- if func_args is string and (func_args | trim) not in ["", "{}", "null"] -%}
+            {{- raise_exception("Tool call arguments must be a mapping, got a JSON-encoded string: parse arguments with json.loads() before applying the chat template") -}}
+        {%- endif -%}
+        {{- "<tool_call>\n<function=" + func["name"] + ">\n" -}}
+        {%- if func_args is mapping -%}
+            {%- for arg_name, arg_value in func_args.items() -%}
+                {{- "<parameter=" + arg_name + ">\n" + render_value(arg_value) + "\n</parameter>\n" -}}
+            {%- endfor -%}
+        {%- endif -%}
+        {{- "</function>\n</tool_call>" -}}
+    {%- endfor -%}
+{%- endmacro -%}
+
+{%- macro tool_to_json(tool) -%}
+    {%- set fn = tool.function if tool.function is defined else tool -%}
+    {%- set ns_t = namespace(first=true) -%}
+    {{- "{" -}}
+    {%- for k, v in fn.items() -%}
+        {%- if k != "defer_loading" -%}
+            {%- if not ns_t.first -%}{{- ", " -}}{%- endif -%}
+            {%- set ns_t.first = false -%}
+            {{- k | tojson -}}{{- ": " -}}{{- v | tojson -}}
+        {%- endif -%}
+    {%- endfor -%}
+    {{- "}" -}}
+{%- endmacro -%}
+
+{%- macro render_tool_response(content) -%}
+    {#- content that already carries the tag (pre-formatted by the caller) is
+        passed through as-is instead of being wrapped a second time -#}
+    {%- if content is string and "<tool_response>" in content -%}
+        {{- content -}}
+    {%- else -%}
+    {{- "<tool_response>\n" -}}
+    {%- if content is string -%}
+        {{- content -}}
+    {%- elif content is mapping -%}
+        {{- parse_content(content) -}}
+    {%- elif content is iterable -%}
+        {%- for item in content -%}
+            {%- if not loop.first -%}{{- "\n" -}}{%- endif -%}
+            {%- if item is mapping and item.result is defined and item.result is mapping and item.result.tools is defined -%}
+                {{- "<additional_tools>\n" -}}
+                {%- set _tl = item.result.tools -%}
+                {%- if _tl is string -%}
+                    {{- _tl -}}
+                {%- elif _tl is mapping -%}
+                    {{- tool_to_json(_tl) -}}
+                {%- else -%}
+                    {%- for t in _tl -%}
+                        {%- if not loop.first -%}{{- "\n" -}}{%- endif -%}
+                        {{- tool_to_json(t) -}}
+                    {%- endfor -%}
+                {%- endif -%}
+                {{- "\n</additional_tools>" -}}
+            {%- else -%}
+                {{- parse_content(item) -}}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- endif -%}
+    {{- "\n</tool_response>" -}}
+    {%- endif -%}
+{%- endmacro -%}
+
+{%- set ns = namespace(system_prompt="", last_user_index=-1) -%}
+{%- if messages and messages[0]["role"] == "system" -%}
+    {%- if messages[0].get("content") -%}
+        {%- set ns.system_prompt = parse_content(messages[0]["content"]) -%}
+    {%- endif -%}
+    {%- set messages = messages[1:] -%}
+{%- endif -%}
+{%- if tools -%}
+    {%- set tools_ns = namespace(lines=[]) -%}
+    {%- for tool in tools -%}
+        {%- set fn = tool.function if tool.function is defined else tool -%}
+        {%- set _dyn = (fn.defer_loading is defined and fn.defer_loading) or (tool.defer_loading is defined and tool.defer_loading) -%}
+        {%- if not _dyn -%}
+            {%- set tools_ns.lines = tools_ns.lines + [tool_to_json(fn)] -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {%- if tools_ns.lines -%}
+        {%- set ns.system_prompt = ns.system_prompt + ("\n\n" if ns.system_prompt else "") + "<tools>\n" + (tools_ns.lines | join("\n")) + "\n</tools>" -%}
+    {%- endif -%}
+{%- endif -%}
+{%- if ns.system_prompt -%}
+    {{- "<|im_start|>system\n" + ns.system_prompt + "<|im_end|>\n" -}}
+{%- endif -%}
+{%- for message in messages -%}
+    {%- if message["role"] == "user" -%}
+        {%- set ns.last_user_index = loop.index0 -%}
+    {%- endif -%}
+{%- endfor -%}
+{%- for message in messages -%}
+    {{- "<|im_start|>" + message.role + "\n" -}}
+    {%- if message.role == "assistant" -%}
+        {%- generation -%}
+        {%- set keep_thinking = preserve_thinking or loop.index0 > ns.last_user_index -%}
+        {%- set thinking = message.thinking or message.reasoning or message.reasoning_content -%}
+        {%- set thinking = thinking if thinking is string else "" -%}
+        {%- if thinking and keep_thinking -%}
+            {{- "<think>" + thinking + "</think>" -}}
+        {%- endif -%}
+        {%- set _cfm_tag = "CONTINUE_FINAL_MESSAGE_TAG " -%}
+        {%- set _has_cfm = false -%}
+        {%- set content = "" -%}
+        {%- if message.get("content") -%}
+            {%- set content = parse_content(message.content) -%}
+        {%- endif -%}
+        {%- if not keep_thinking -%}
+            {%- if "</think>" in content -%}
+                {%- set content = content.split("</think>")[-1] | trim -%}
+            {%- endif -%}
+            {%- set content = "<think></think>" + content -%}
+        {%- endif -%}
+        {%- if content.endswith(_cfm_tag) -%}
+            {%- set _has_cfm = true -%}
+            {%- set _trunc_len = (content | length) - (_cfm_tag | length) -%}
+            {%- set content = content[:_trunc_len] -%}
+        {%- endif -%}
+        {{- content -}}
+        {%- if message.tool_calls -%}
+            {{- render_tool_calls(message.tool_calls) -}}
+        {%- endif -%}
+        {%- if _has_cfm -%}
+            {{- _cfm_tag -}}
+        {%- endif -%}
+        {{- "<|im_end|>\n" -}}
+        {%- endgeneration -%}
+    {%- elif message.role == "tool" -%}
+        {%- if message.get("additional_tools") -%}
+            {%- set _at = message["additional_tools"] -%}
+            {#- definitions may live in additional_tools itself, or in content
+                with additional_tools used as a plain truthy flag -#}
+            {%- if _at is mapping or (_at is sequence and _at is not string) -%}
+                {%- set _src = _at -%}
+            {%- else -%}
+                {%- set _src = message.get("content") -%}
+            {%- endif -%}
+            {#- nothing to reveal: skip the wrapper entirely rather than
+                emitting an empty <additional_tools> block -#}
+            {%- if _src is string and "<tool_response>" in _src -%}
+                {{- _src -}}
+            {%- elif _src is string and "<additional_tools>" in _src -%}
+                {{- "<tool_response>\n" + _src + "\n</tool_response>" -}}
+            {%- elif _src -%}
+                {{- "<tool_response>\n<additional_tools>\n" -}}
+                {%- if _src is string -%}
+                    {{- _src -}}
+                {%- elif _src is mapping -%}
+                    {{- tool_to_json(_src) -}}
+                {%- else -%}
+                    {%- for t in _src -%}
+                        {%- if not loop.first -%}{{- "\n" -}}{%- endif -%}
+                        {{- tool_to_json(t) -}}
+                    {%- endfor -%}
+                {%- endif -%}
+                {{- "\n</additional_tools>\n</tool_response>" -}}
+            {%- endif -%}
+        {%- elif message.get("content") -%}
+            {{- render_tool_response(message["content"]) -}}
+        {%- endif -%}
+        {{- "<|im_end|>\n" -}}
+    {%- else %}
+        {%- if message.get("content") -%}
+            {{- parse_content(message["content"]) -}}
+        {%- endif -%}
+        {{- "<|im_end|>\n" -}}
+    {%- endif %}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+    {{- "<|im_start|>assistant\n" -}}
+    {%- if enable_thinking -%}
+        {{- "<think>" -}}
+    {%- else -%}
+        {{- "<think></think>" -}}
+    {%- endif -%}
+{%- endif -%}

config.json

@@ -0,0 +1,50 @@
+{
+  "architectures": [
+    "Lfm3ForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_lfm3.Lfm3Config",
+    "AutoModelForCausalLM": "modeling_lfm3.Lfm3ForCausalLM"
+  },
+  "bos_token_id": 124894,
+  "conv_L_cache": 7,
+  "conv_bias": false,
+  "dtype": "float32",
+  "eos_token_id": 124900,
+  "head_dim": 8,
+  "hidden_size": 16,
+  "intermediate_size": 32,
+  "layer_types": [
+    "conv",
+    "conv",
+    "full_attention",
+    "conv",
+    "sliding_attention",
+    "conv",
+    "sliding_attention"
+  ],
+  "max_position_embeddings": 196608,
+  "model_type": "lfm3",
+  "norm_eps": 1e-05,
+  "num_attention_heads": 4,
+  "num_hidden_layers": 7,
+  "num_key_value_heads": 2,
+  "pad_token_id": 124893,
+  "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": 4,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.20.0",
+  "use_cache": true,
+  "use_sandwich_norm": true,
+  "vocab_size": 128000
+}

configuration_lfm3.py

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

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 124894,
+  "eos_token_id": 124900,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 124893,
+  "transformers_version": "5.20.0",
+  "use_cache": true
+}

model.safetensors

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

modeling_lfm3.py

@@ -0,0 +1,830 @@
+# LFM3 dense model.
+from typing import Any, Callable, Optional, Union
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.integrations import (
+    use_kernel_forward_from_hub,
+    use_kernel_func_from_hub,
+    use_kernelized_func,
+)
+from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
+from transformers.utils.generic import check_model_inputs, maybe_autocast
+from transformers.utils.import_utils import is_causal_conv1d_available, is_torchdynamo_compiling
+
+from .configuration_lfm3 import Lfm3Config
+
+
+if is_causal_conv1d_available():
+    from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
+else:
+    causal_conv1d_fn, causal_conv1d_update = None, None
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class Lfm3RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        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):
+    """Per-layer-type rotary embedding (Gemma 3-style).
+
+    `config.rope_parameters` is a dict keyed by `layer_type`
+    (`"full_attention"`, `"sliding_attention"`). For each layer type
+    present in `config.layer_types`, an `{layer_type}_inv_freq` buffer is
+    registered. `forward(x, position_ids, layer_type=...)` selects the right
+    one. Partial RoPE (`partial_rotary_factor < 1.0`) follows the Gemma 4
+    convention: the first `rope_angles = int(partial_rotary_factor * dim) // 2`
+    frequency pairs receive the standard rotation, the remaining
+    `dim//2 - rope_angles` entries are zero (cos=1, sin=0 → identity). The
+    SWA entry must never carry `partial_rotary_factor` — the Lfm3Config
+    normalizer strips it during migration of legacy configs.
+    """
+
+    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
+
+        # Deterministic dedupe of attention-layer types present in this model.
+        # `dict.fromkeys` preserves insertion order so the buffer-registration
+        # order is stable across processes / Python releases — `set()` would
+        # not, which matters for reproducible test output.
+        layer_types_in_model = list(dict.fromkeys(config.layer_types or []))
+        self.attention_layer_types = [
+            lt for lt in layer_types_in_model if lt in config._attention_layer_types
+        ]
+        self.rope_type: dict[str, str] = {}
+
+        rope_parameters = config.rope_parameters or {}
+        for layer_type in self.attention_layer_types:
+            entry = rope_parameters.get(layer_type)
+            if entry is None:
+                continue
+            self.rope_type[layer_type] = entry["rope_type"]
+            if entry["rope_type"] != "default":
+                rope_init_fn: Callable = ROPE_INIT_FUNCTIONS[entry["rope_type"]]
+                inv_freq, attention_scaling = rope_init_fn(self.config, device)
+            else:
+                inv_freq, attention_scaling = self.compute_default_rope_parameters(
+                    self.config, device, layer_type=layer_type
+                )
+            self.register_buffer(f"{layer_type}_inv_freq", inv_freq, persistent=False)
+            self.register_buffer(
+                f"{layer_type}_original_inv_freq", inv_freq.clone(), persistent=False
+            )
+            setattr(self, f"{layer_type}_attention_scaling", attention_scaling)
+
+    @staticmethod
+    def compute_default_rope_parameters(
+        config: Optional[Lfm3Config] = None,
+        device: Optional["torch.device"] = None,
+        seq_len: Optional[int] = None,
+        layer_type: Optional[str] = None,
+    ) -> tuple["torch.Tensor", float]:
+        entry = config.rope_parameters[layer_type]
+        base = entry["rope_theta"]
+        partial_rotary_factor = entry.get("partial_rotary_factor", 1.0)
+        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+        attention_factor = 1.0
+
+        rope_angles = int(partial_rotary_factor * dim) // 2
+        assert 0 < rope_angles <= dim // 2, (
+            f"partial_rotary_factor={partial_rotary_factor} with dim={dim} produced "
+            f"rope_angles={rope_angles}, must be in (0, dim//2={dim // 2}]."
+        )
+        inv_freq_rotated = 1.0 / (
+            base
+            ** (
+                torch.arange(0, 2 * rope_angles, 2, dtype=torch.int64).to(
+                    device=device, dtype=torch.float
+                )
+                / dim
+            )
+        )
+        nope_angles = dim // 2 - rope_angles
+        if nope_angles > 0:
+            inv_freq = torch.cat(
+                (
+                    inv_freq_rotated,
+                    torch.zeros(nope_angles, dtype=torch.float, device=device),
+                ),
+                dim=0,
+            )
+        else:
+            inv_freq = inv_freq_rotated
+        return inv_freq, attention_factor
+
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids, layer_type: Optional[str] = None):
+        inv_freq = getattr(self, f"{layer_type}_inv_freq")
+        attention_scaling = getattr(self, f"{layer_type}_attention_scaling")
+
+        inv_freq_expanded = (
+            inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        )
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = (
+            x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        )
+        with maybe_autocast(device_type=device_type, enabled=False):  # Force float32
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * attention_scaling
+            sin = emb.sin() * attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+class Lfm3MLP(nn.Module):
+    def __init__(self, config: Lfm3Config, intermediate_size: Optional[int] = None):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = (
+            config.intermediate_size if intermediate_size is None else intermediate_size
+        )
+        self.w1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.w3 = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.w2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+
+    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)
+
+
+@use_kernel_func_from_hub("rotary_pos_emb")
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Apply Rotary Position Embedding to q / k.
+
+    The rotation is computed in fp32 to match the native (Llama-style) RoPE
+    path: native does a complex multiply in fp32 and casts to bf16 once at the
+    end. The upstream default form (multiplying bf16 q/k by bf16 cos/sin)
+    accrues ~3× more ULP-level rounding error per element, which propagates
+    through softmax to ~5% per-layer native↔HF divergence in bf16. When the
+    kernel hub provides a fused rotary kernel, that takes over via the
+    decorator; this body is the Python fallback.
+    """
+    in_dtype = q.dtype
+    cos = cos.unsqueeze(unsqueeze_dim).float()
+    sin = sin.unsqueeze(unsqueeze_dim).float()
+    q_f = q.float()
+    k_f = k.float()
+    q_embed = (q_f * cos) + (rotate_half(q_f) * sin)
+    k_embed = (k_f * cos) + (rotate_half(k_f) * sin)
+    return q_embed.to(in_dtype), k_embed.to(in_dtype)
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    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: Optional[torch.Tensor],
+    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:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_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
+
+
+@use_kernelized_func(apply_rotary_pos_emb)
+class Lfm3Attention(nn.Module):
+    """Multi-headed attention. Supports GQA, optional sliding window, and a
+    headwise sigmoid gate on the attention output.
+
+    The caller (``Lfm3Model.forward``) is responsible for building the right
+    mask for this layer's type (sliding-window-causal for SWA layers, full
+    causal otherwise). This forward only also passes ``sliding_window=...`` to
+    the attention backend for FA2's native fast path; SDPA / eager rely on the
+    precomputed mask alone.
+    """
+
+    def __init__(self, config: Lfm3Config, layer_idx: int, sliding_window: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", None) or (
+            config.hidden_size // config.num_attention_heads
+        )
+        self.sliding_window = sliding_window
+        self.is_causal = True
+
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.scaling = self.head_dim ** -0.5
+        self.q_proj = nn.Linear(
+            config.hidden_size, config.num_attention_heads * self.head_dim, bias=False
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim, config.hidden_size, bias=False
+        )
+        self.q_norm = Lfm3RMSNorm(self.head_dim, eps=config.norm_eps)
+        self.k_norm = Lfm3RMSNorm(self.head_dim, eps=config.norm_eps)
+        # Headwise sigmoid gate on the attention output (one scalar per query head).
+        self.gate_proj = nn.Linear(
+            config.hidden_size, config.num_attention_heads, bias=False
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        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 = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            **kwargs,
+        )
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        # Headwise sigmoid gate: one scalar per query head, broadcast over head_dim.
+        num_heads = self.config.num_attention_heads
+        gate = torch.sigmoid(self.gate_proj(hidden_states)).view(*input_shape, num_heads, 1)
+        attn_output = attn_output.view(*input_shape, num_heads, self.head_dim)
+        attn_output = (gate * attn_output).view(*input_shape, -1)
+        output = self.o_proj(attn_output)
+        return output, attn_weights
+
+
+def apply_mask_to_padding_states(hidden_states, attention_mask):
+    """Zero out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66 .
+
+    NOTE (upstream-faithful, mamba/Lfm2): the `shape[0] > 1` (batch > 1) guard means
+    padding is NOT zeroed for a single (batch=1) sequence — so the conv can bleed
+    pad positions for a left-padded batch-of-1. Matches upstream; single-sequence
+    decode isn't padded in practice, but it's a latent footgun for padded bsz-1 prefill.
+    """
+    if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+        dtype = hidden_states.dtype
+        hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
+    return hidden_states
+
+
+kernel_modules = (causal_conv1d_fn, causal_conv1d_update)
+is_fast_path_available = all(kernel_modules)
+
+
+class Lfm3ShortConv(nn.Module):
+    """Standard gated short conv block:
+
+        proj = in_proj(x)               # 3*D
+        B, C, x' = proj.chunk(3, dim=-1)
+        conv_in = B * x'
+        y = C * depthwise_conv1d(conv_in)
+        out = out_proj(y)
+    """
+
+    def __init__(
+        self,
+        config: Lfm3Config,
+        layer_idx: int,
+    ):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.L_cache = config.conv_L_cache
+        self.bias = config.conv_bias
+
+        D = config.hidden_size
+        self.conv = nn.Conv1d(
+            in_channels=D,
+            out_channels=D,
+            kernel_size=self.L_cache,
+            groups=D,
+            bias=self.bias,
+            padding=self.L_cache - 1,
+        )
+        self.in_proj = nn.Linear(D, 3 * D, bias=self.bias)
+        self.out_proj = nn.Linear(D, D, bias=self.bias)
+
+    def cuda_kernels_forward(
+        self,
+        x: torch.Tensor,
+        past_key_values: Optional[Cache] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        seq_idx: Optional[torch.Tensor] = None,
+    ):
+        x = apply_mask_to_padding_states(x, attention_mask)
+        D = self.config.hidden_size
+        proj = self.in_proj(x).transpose(-1, -2)
+        # Split order matches Lfm3MoeShortConv("both") chunk(3): B first, C second, x third
+        B, C, x = proj.split([D, D, D], dim=-2)
+        conv_in = B * x
+
+        conv_weights = self.conv.weight.view(self.conv.weight.size(0), self.conv.weight.size(2))
+        if past_key_values is not None and past_key_values.has_previous_state(self.layer_idx):
+            conv_states = past_key_values.layers[self.layer_idx].conv_states
+            if isinstance(conv_states, dict):
+                conv_states = conv_states[0]
+            conv_out = causal_conv1d_update(
+                conv_in.squeeze(-1),
+                conv_states,
+                conv_weights,
+                self.conv.bias,
+                None,
+            )
+            conv_out = conv_out.unsqueeze(-1)
+        else:
+            if past_key_values is not None:
+                conv_state = nn.functional.pad(conv_in, (self.L_cache - conv_in.shape[-1], 0))
+                past_key_values.update_conv_state(conv_state, self.layer_idx)
+
+            # `seq_idx` (per-token document index, int32) resets the conv at packed
+            # document boundaries so it never bleeds across documents.
+            conv_out = causal_conv1d_fn(
+                conv_in, conv_weights, self.conv.bias, seq_idx=seq_idx, activation=None
+            )
+
+        y = C * conv_out
+        y = self.out_proj(y.transpose(-1, -2).contiguous())
+        return y
+
+    def _segmented_conv(self, conv_in: torch.Tensor, seq_idx: torch.Tensor) -> torch.Tensor:
+        """Causal depthwise conv that resets at packed document boundaries.
+
+        Pure-torch equivalent of `causal_conv1d_fn(..., seq_idx=...)`: out[t] only
+        sums conv_in[t-(L-1)..t] entries that share t's document id. `conv_in` is
+        (b, D, seqlen); `seq_idx` is (b, seqlen) int. Used only on the packed
+        path — the unpacked path keeps the plain `nn.Conv1d` call below.
+        """
+        L = self.L_cache
+        cin_p = nn.functional.pad(conv_in, (L - 1, 0))             # (b, D, L-1+seqlen)
+        windows = cin_p.unfold(-1, L, 1)                            # (b, D, seqlen, L)
+        # window element k of position t is conv_in[t-(L-1)+k]; current token is k=L-1.
+        sidx_p = nn.functional.pad(seq_idx, (L - 1, 0), value=-1)   # (b, L-1+seqlen)
+        sidx_win = sidx_p.unfold(-1, L, 1)                          # (b, seqlen, L)
+        same_doc = (sidx_win == sidx_win[..., -1:]).unsqueeze(1)    # (b, 1, seqlen, L)
+        windows = windows * same_doc.to(windows.dtype)
+        w = self.conv.weight.squeeze(1)                            # (D, L)
+        conv_out = torch.einsum("bdtk,dk->bdt", windows, w)
+        if self.bias:
+            conv_out = conv_out + self.conv.bias[None, :, None]
+        return conv_out
+
+    def slow_forward(
+        self,
+        x: torch.Tensor,
+        past_key_values: Optional[Cache] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        seq_idx: Optional[torch.Tensor] = None,
+    ):
+        seqlen = x.shape[1]
+
+        x = apply_mask_to_padding_states(x, attention_mask)
+        D = self.config.hidden_size
+        proj = self.in_proj(x).transpose(-1, -2)
+        B, C, x = proj.split([D, D, D], dim=-2)
+        conv_in = B * x
+
+        if past_key_values is not None and past_key_values.has_previous_state(self.layer_idx):
+            # update_conv_state rolls the cached window left and appends conv_in.
+            conv_state = past_key_values.update_conv_state(conv_in, self.layer_idx)[..., -self.L_cache :]
+            conv_out = torch.sum(conv_state.to(conv_in.device) * self.conv.weight[:, 0, :], dim=-1)
+            if self.bias:
+                conv_out += self.conv.bias
+
+            conv_out = conv_out.unsqueeze(-1)
+        else:
+            if past_key_values is not None:
+                conv_state = nn.functional.pad(conv_in, (self.L_cache - conv_in.shape[-1], 0))
+                past_key_values.update_conv_state(conv_state, self.layer_idx)
+
+            if seq_idx is not None:
+                conv_out = self._segmented_conv(conv_in, seq_idx)
+            else:
+                conv_out = self.conv(conv_in)[..., :seqlen]
+
+        y = C * conv_out
+        y = y.transpose(-1, -2).contiguous()
+        y = self.out_proj(y)
+        return y
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        past_key_values: Optional[Cache] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        seq_idx: Optional[torch.Tensor] = None,
+    ):
+        # causal_conv1d's stock build only supports conv widths 2-4 and raises
+        # otherwise; build it for wider kernels (e.g. L_cache=7) or omit it for
+        # the pure-torch slow path.
+        if (
+            is_fast_path_available
+            and "cuda" in hidden_states.device.type
+            and not is_torchdynamo_compiling()
+        ):
+            return self.cuda_kernels_forward(
+                hidden_states, past_key_values, attention_mask, seq_idx
+            )
+        return self.slow_forward(
+            hidden_states, past_key_values, attention_mask, seq_idx
+        )
+
+
+class Lfm3DecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: Lfm3Config, layer_idx: int):
+        super().__init__()
+        self.layer_type = config.layer_types[layer_idx]
+        self.is_attention_layer = self.layer_type in ("full_attention", "sliding_attention")
+        self.use_sandwich_norm = getattr(config, "use_sandwich_norm", False)
+
+        if self.layer_type == "full_attention":
+            self.self_attn = Lfm3Attention(config, layer_idx)
+        elif self.layer_type == "sliding_attention":
+            self.self_attn = Lfm3Attention(config, layer_idx, sliding_window=config.sliding_window)
+        elif self.layer_type == "conv":
+            self.conv = Lfm3ShortConv(config, layer_idx)
+        else:
+            raise ValueError(
+                f"Unsupported layer_type {self.layer_type!r} at layer {layer_idx}. "
+                "Lfm3 supports only {full_attention, sliding_attention, conv}."
+            )
+        self.feed_forward = Lfm3MLP(config, intermediate_size=config.intermediate_size)
+        self.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],
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        seq_idx: Optional[torch.Tensor] = None,
+        **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
+
+
+@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 = False
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "hidden_states": Lfm3DecoderLayer,
+        "attentions": Lfm3Attention,
+    }
+
+    def _init_weights(self, module):
+        # Re-compute RotaryEmbedding inv_freq buffers. Required because
+        # `from_pretrained` constructs the model on a meta device and
+        # materializes persistent=False buffers as uninitialized memory after
+        # the fact. We can't rely on values registered during meta-device
+        # __init__ — they aren't preserved through materialization. Re-run
+        # `compute_default_rope_parameters` for each attention layer type so
+        # the per-layer-type inv_freq buffers (`full_attention_inv_freq`,
+        # `sliding_attention_inv_freq`) get the correct values for
+        # their configured theta / partial_rotary_factor, instead of HF's
+        # default branch which assumes a single global table.
+        if isinstance(module, Lfm3RotaryEmbedding):
+            for layer_type in module.attention_layer_types:
+                if layer_type not in module.rope_type:
+                    continue
+                inv_freq, attention_scaling = module.compute_default_rope_parameters(
+                    module.config,
+                    device=getattr(module, f"{layer_type}_inv_freq").device,
+                    layer_type=layer_type,
+                )
+                with torch.no_grad():
+                    getattr(module, f"{layer_type}_inv_freq").copy_(inv_freq)
+                    getattr(module, f"{layer_type}_original_inv_freq").copy_(inv_freq)
+                setattr(module, f"{layer_type}_attention_scaling", attention_scaling)
+            return
+        super()._init_weights(module)
+
+
+@auto_docstring
+class Lfm3Model(Lfm3PreTrainedModel):
+    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.gradient_checkpointing = False
+        # Single Gemma3-style rotary module. Per-attention-layer-type inv_freq
+        # buffers (`full_attention_inv_freq`, `sliding_attention_inv_freq`)
+        # are registered inside the module based on `config.rope_parameters`.
+        # `forward` is dispatched per layer below via `layer_type=...`.
+        self.rotary_emb = Lfm3RotaryEmbedding(config)
+        self.embedding_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @check_model_inputs
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPast:
+        r"""
+        cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+            Indices depicting the position of the input sequence tokens in the sequence. Used to update the
+            cache in the correct position and to infer the complete sequence length.
+        """
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            # DynamicCache(config) builds the right per-layer cache from
+            # config.layer_types: DynamicLayer (full_attention),
+            # DynamicSlidingWindowLayer (sliding_attention, auto-bounded), and
+            # LinearAttentionLayer (conv, conv-state buffer).
+            past_key_values = DynamicCache(config=self.config)
+
+        if cache_position is None:
+            past_seen_tokens = (
+                past_key_values.get_seq_length() if past_key_values is not None else 0
+            )
+            cache_position = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        # Packed-sequence document boundaries for the conv layers: when several
+        # documents are concatenated into one row, `position_ids` restarts at 0
+        # per document, so a cumulative count of those resets gives a per-token
+        # document id. Conv layers use it to avoid bleeding across documents.
+        # Only meaningful during prefill (multi-token); a single decode token has
+        # nothing to segment.
+        seq_idx = None
+        if inputs_embeds.shape[1] > 1:
+            doc_id = (position_ids == 0).cumsum(-1) - 1
+            if doc_id[..., -1].max() > 0:  # >1 document present
+                seq_idx = doc_id.to(torch.int32)
+
+        # Build one causal mask per attention layer type present in the model.
+        # SWA layers get a sliding-window causal mask; full-attention layers
+        # get a plain causal mask. Conv layers receive the RAW 2D padding mask
+        # (matching upstream `Lfm2Model.forward`'s `linear_attention` route)
+        # so `Lfm3ShortConv.apply_mask_to_padding_states` can zero pad
+        # positions during prefill. Conv during decode (single token) doesn't
+        # need a mask — None is fine.
+        mask_kwargs = {
+            "config": self.config,
+            "inputs_embeds": inputs_embeds,
+            "attention_mask": attention_mask,
+            "past_key_values": past_key_values,
+            "position_ids": position_ids,
+        }
+        causal_masks_by_type: dict[str, Optional[torch.Tensor]] = {}
+        attention_types_in_model = {
+            lt for lt in self.config.layer_types if lt in self.config._attention_layer_types
+        }
+        if "full_attention" in attention_types_in_model:
+            causal_masks_by_type["full_attention"] = create_causal_mask(**mask_kwargs)
+        if "sliding_attention" in attention_types_in_model:
+            causal_masks_by_type["sliding_attention"] = create_sliding_window_causal_mask(
+                **mask_kwargs
+            )
+        # Conv layers want the 2D padding mask only during prefill. For
+        # single-token decode (`shape[1] == 1`) there's nothing to pad-zero.
+        conv_mask = attention_mask if inputs_embeds.shape[1] != 1 else None
+
+        hidden_states = inputs_embeds
+        # Pre-compute one (cos, sin) per attention layer type. Conv-only models
+        # leave this dict empty; the per-layer dispatch below only reads it for
+        # attention layers.
+        position_embeddings_by_type = {
+            layer_type: self.rotary_emb(hidden_states, position_ids, layer_type=layer_type)
+            for layer_type in self.rotary_emb.rope_type
+        }
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            if decoder_layer.is_attention_layer:
+                layer_mask = causal_masks_by_type.get(decoder_layer.layer_type)
+            else:
+                layer_mask = conv_mask
+            layer_pos_embeddings = position_embeddings_by_type.get(decoder_layer.layer_type)
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=layer_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                cache_position=cache_position,
+                position_embeddings=layer_pos_embeddings,
+                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,
+        )
+
+
+@auto_docstring
+class Lfm3ForCausalLM(Lfm3PreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+    _tp_plan = {"lm_head": "colwise_rep"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    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: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> CausalLMOutputWithPast:
+        r"""
+        cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+            Indices depicting the position of the input sequence tokens in the sequence. Used to update the
+            cache in the correct position and to infer the complete sequence length.
+        """
+        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,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        slice_indices = (
+            slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        )
+        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,
+        )
+
+
+__all__ = ["Lfm3ForCausalLM", "Lfm3Model", "Lfm3PreTrainedModel"]

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,22 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|startoftext|>",
+  "eos_token": "<|im_end|>",
+  "extra_special_tokens": [],
+  "image_token": "<|image|>",
+  "is_local": true,
+  "legacy": false,
+  "local_files_only": false,
+  "model_max_length": 1000000000000000019884624838656,
+  "model_specific_special_tokens": {
+    "image_token": "<|image|>",
+    "video_token": "<|video|>",
+    "vision_end_token": "<|vision_end|>",
+    "vision_start_token": "<|vision_start|>"
+  },
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "TokenizersBackend",
+  "video_token": "<|video|>",
+  "vision_end_token": "<|vision_end|>",
+  "vision_start_token": "<|vision_start|>"
+}