Add tiny random LFM3 model for optimum-intel tests
cf7edf4 verified 10 days ago
35.8 kB
# 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"]