# 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"]