diff --git a/scripts/detect_affected_models.py b/scripts/detect_affected_models.py index 0b4b9e6a..55494f0c 100644 --- a/scripts/detect_affected_models.py +++ b/scripts/detect_affected_models.py @@ -58,17 +58,19 @@ def classify_file(path: str) -> str: """Classify a changed file path. Returns one of: 'model', 'traceable', 'shared_infra', - 'test', 'other'. + 'test_config', 'test', 'other'. """ normalized = path.replace("\\", "/") if not normalized.startswith("src/mobius/"): # Test infrastructure files that affect all models - if normalized in ( - "tests/conftest.py", - "tests/_test_configs.py", - ): + if normalized == "tests/conftest.py": return "shared_infra" + if normalized == "tests/_test_configs.py": + # A config-only change is broad, but new model implementations also + # add an entry here. Defer the run-all decision until model files + # have been collected so those PRs can use import-graph scoping. + return "test_config" if normalized.endswith("_test.py") or normalized.startswith("tests/"): return "test" return "other" @@ -427,6 +429,7 @@ def detect_affected_models( """ affected: set[str] = set() run_all = False + test_config_changed = False # Classify files model_files: list[str] = [] @@ -436,6 +439,8 @@ def detect_affected_models( if category == "shared_infra": run_all = True break + elif category == "test_config": + test_config_changed = True elif category == "model": # Deleted model files could break dependents — run all full_path = _PROJECT_ROOT / path @@ -453,6 +458,9 @@ def detect_affected_models( if run_all: return {"affected": [], "run_all": True} + if test_config_changed and not model_files: + return {"affected": [], "run_all": True} + if not model_files and not traceable_files: return {"affected": [], "run_all": False} diff --git a/scripts/detect_affected_models_test.py b/scripts/detect_affected_models_test.py index 64a170b8..a472f056 100644 --- a/scripts/detect_affected_models_test.py +++ b/scripts/detect_affected_models_test.py @@ -72,7 +72,7 @@ def test_test_infra_conftest(self): assert classify_file("tests/conftest.py") == "shared_infra" def test_test_infra_configs(self): - assert classify_file("tests/_test_configs.py") == "shared_infra" + assert classify_file("tests/_test_configs.py") == "test_config" def test_readme(self): assert classify_file("README.md") == "other" @@ -200,6 +200,29 @@ def test_configs_change_no_run_all(self): assert result["run_all"] is False assert result["affected"] == [] + def test_test_configs_only_runs_all(self): + result = detect_affected_models(["tests/_test_configs.py"]) + assert result == {"affected": [], "run_all": True} + + def test_test_configs_with_model_uses_model_scope(self): + result = detect_affected_models( + [ + "tests/_test_configs.py", + "src/mobius/models/lfm2.py", + ] + ) + assert result["run_all"] is False + assert result["affected"] == ["lfm2"] + + def test_test_configs_with_unmapped_task_still_runs_all(self): + result = detect_affected_models( + [ + "tests/_test_configs.py", + "src/mobius/tasks/_causal_lm.py", + ] + ) + assert result == {"affected": [], "run_all": True} + def test_unrelated_file_no_affected(self): result = detect_affected_models(["README.md"]) assert result["run_all"] is False diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 475b69b9..86a0cb31 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -65,7 +65,7 @@ def _resolve_hidden_act(config, model_type: str) -> str | None: or getattr(config, "afn", None) # LLaDA/OLMo expose the activation as ``activation_type`` (e.g. "silu"). or getattr(config, "activation_type", None) - or ("silu" if model_type in ("qwen", "chatglm") else None) + or ("silu" if model_type in ("qwen", "chatglm", "lfm2") else None) # gelu_activation is a boolean (XLM) — must be after all string # attrs so it cannot override an explicit hidden_act. or ("gelu" if getattr(config, "gelu_activation", False) else None) @@ -411,6 +411,10 @@ class ArchitectureConfig(BaseModelConfig): linear_num_key_heads: int | None = None linear_num_value_heads: int | None = None + # Double-gated short-convolution config (LFM2-style hybrid layers). + short_conv_kernel: int = 3 + short_conv_bias: bool = False + rms_norm_eps: float = 1e-6 # Rotary embedding config. @@ -706,6 +710,8 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig: linear_value_head_dim=(getattr(config, "linear_value_head_dim", None)), linear_num_key_heads=(getattr(config, "linear_num_key_heads", None)), linear_num_value_heads=(getattr(config, "linear_num_value_heads", None)), + short_conv_kernel=getattr(config, "conv_L_cache", 3), + short_conv_bias=getattr(config, "conv_bias", False), pad_token_id=(getattr(config, "pad_token_id", 0)), model_type=model_type, bos_token_id=getattr(config, "bos_token_id", None), @@ -786,6 +792,7 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig: model_type in ( "gemma3_text", + "lfm2", "flex_olmo", "olmoe", "olmo2", diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index 1c611e97..0b7750fc 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -68,6 +68,7 @@ HunYuanVLMoTModel, InternLM2CausalLMModel, LayerNormCausalLMModel, + Lfm2CausalLMModel, LLaDAModel, Llama4CausalLMModel, MoECausalLMModel, @@ -463,6 +464,7 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None: "hunyuan_v1_dense": ModelRegistration(HunYuanV1DenseCausalLMModel), "internlm2": ModelRegistration(InternLM2CausalLMModel), "llama4_text": ModelRegistration(Llama4CausalLMModel), + "lfm2": ModelRegistration(Lfm2CausalLMModel), "llada": ModelRegistration(LLaDAModel, task="masked-diffusion"), "modernbert-decoder": ModelRegistration(ModernBertDecoderModel), "mpt": ModelRegistration(MPTCausalLMModel), @@ -951,6 +953,7 @@ def _create_default_registry() -> ModelRegistry: # --- Hybrid SSM+Attention --- "jamba": "ai21labs/Jamba-v0.1", "bamba": "ibm-fms/Bamba-9B", + "lfm2": "LiquidAI/LFM2.5-230M", # --- Multimodal --- "qwen2_vl": "Qwen/Qwen2-VL-2B-Instruct", @@ -1199,6 +1202,7 @@ def _create_default_registry() -> ModelRegistry: "llama": "llama", "code_llama": "llama", "llama4_text": "llama", + "lfm2": "lfm", "mllama": "llama", "mistral": "mistral", "mistral3": "mistral", @@ -1270,6 +1274,7 @@ def _create_default_registry() -> ModelRegistry: "falcon_mamba": "ssm", "jamba": "hybrid-ssm+attn", "bamba": "hybrid-mamba2+attn", + "lfm2": "hybrid-conv+attn", "qwen3_next": "moe+linear-attn", } diff --git a/src/mobius/_testing/torch_reference.py b/src/mobius/_testing/torch_reference.py index 095e3211..d72d3ab9 100644 --- a/src/mobius/_testing/torch_reference.py +++ b/src/mobius/_testing/torch_reference.py @@ -288,8 +288,8 @@ def torch_forward( input_ids: np.ndarray, attention_mask: np.ndarray, position_ids: np.ndarray, - past_key_values: list[tuple[np.ndarray, np.ndarray]] | None = None, -) -> tuple[np.ndarray, list[tuple[np.ndarray, np.ndarray]]]: + past_key_values: object | None = None, +) -> tuple[np.ndarray, object]: """Run a single forward pass on a HuggingFace causal LM model. Args: @@ -297,10 +297,12 @@ def torch_forward( input_ids: [batch, seq_len] int64 numpy array. attention_mask: [batch, total_seq_len] int64 numpy array. position_ids: [batch, seq_len] int64 numpy array. - past_key_values: Optional list of (key, value) numpy array tuples. + past_key_values: Optional list of (key, value) numpy array tuples, or + an opaque HuggingFace Cache for hybrid recurrent models. Returns: - Tuple of (logits as numpy, list of (key, value) numpy tuples). + Tuple of logits and either a list of KV numpy tuples or an opaque + HuggingFace Cache when model-specific recurrent state must be retained. """ import inspect @@ -323,16 +325,21 @@ def torch_forward( kwargs["position_ids"] = pos_t if past_key_values is not None: - from transformers.cache_utils import DynamicCache - - cache = DynamicCache() - for layer_idx, (k, v) in enumerate(past_key_values): - cache.update( - torch.from_numpy(k).to(device=device, dtype=dtype), - torch.from_numpy(v).to(device=device, dtype=dtype), - layer_idx, - ) - kwargs["past_key_values"] = cache + from transformers.cache_utils import Cache, DynamicCache + + if isinstance(past_key_values, Cache): + # Hybrid caches carry model-specific recurrent state (for example + # LFM2 conv windows) that cannot be reconstructed from KV pairs. + kwargs["past_key_values"] = past_key_values + else: + cache = DynamicCache() + for layer_idx, (k, v) in enumerate(past_key_values): + cache.update( + torch.from_numpy(k).to(device=device, dtype=dtype), + torch.from_numpy(v).to(device=device, dtype=dtype), + layer_idx, + ) + kwargs["past_key_values"] = cache try: outputs = model(**kwargs) @@ -351,6 +358,10 @@ def torch_forward( # Extract KV cache if available (Mamba models don't have it) present_kv: list[tuple[np.ndarray, np.ndarray]] = [] cache = getattr(outputs, "past_key_values", None) + layer_types = getattr(model.config, "layer_types", None) or [] + if "conv" in layer_types: + return logits, cache + if cache is not None and hasattr(cache, "layers"): for layer_idx in range(len(cache.layers)): layer_cache = cache.layers[layer_idx] diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py index 3baeb6e7..90f0cafb 100644 --- a/src/mobius/components/__init__.py +++ b/src/mobius/components/__init__.py @@ -32,6 +32,7 @@ "GatedDeltaNet", "GatedMLP", "GatedRMSNorm", + "GatedShortConv", "ClippableLinear", "GroupNorm", "GQAContext", @@ -268,6 +269,7 @@ from mobius.components._sanm_attention import ( SANMEncoderLayer as SANMEncoderLayer, ) +from mobius.components._short_conv import GatedShortConv from mobius.components._ssm import ( JambaSelectiveScan as JambaSelectiveScan, ) diff --git a/src/mobius/components/_short_conv.py b/src/mobius/components/_short_conv.py new file mode 100644 index 00000000..bc2f2ed8 --- /dev/null +++ b/src/mobius/components/_short_conv.py @@ -0,0 +1,117 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Double-gated causal short convolution with recurrent state.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from onnxscript import OpBuilder, nn + +from mobius.components._common import INT64_MAX, Linear + +if TYPE_CHECKING: + import onnx_ir as ir + + +class _DepthwiseShortConv1d(nn.Module): + """Depthwise Conv1D with a full-kernel recurrent cache.""" + + def __init__(self, channels: int, kernel_size: int, *, bias: bool): + super().__init__() + self.weight = nn.Parameter([channels, 1, kernel_size]) + self.bias = nn.Parameter([channels]) if bias else None + self._channels = channels + self._kernel_size = kernel_size + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + conv_state: ir.Value, + ) -> tuple[ir.Value, ir.Value]: + # The cache contract is K-wide, but causal Conv only needs its newest + # K-1 values. Appending T inputs then produces exactly T outputs. + past = op.Slice( + conv_state, + op.Constant(value_ints=[1]), + op.Constant(value_ints=[INT64_MAX]), + op.Constant(value_ints=[2]), + ) + conv_input = op.Concat(past, hidden_states, axis=2) + if self.bias is None: + conv_output = op.Conv( + conv_input, + self.weight, + kernel_shape=[self._kernel_size], + group=self._channels, + ) + else: + conv_output = op.Conv( + conv_input, + self.weight, + self.bias, + kernel_shape=[self._kernel_size], + group=self._channels, + ) + + present_state = op.Slice( + conv_input, + op.Constant(value_ints=[-self._kernel_size]), + op.Constant(value_ints=[INT64_MAX]), + op.Constant(value_ints=[2]), + ) + return conv_output, present_state + + +class GatedShortConv(nn.Module): + """Double-gated depthwise causal convolution used by hybrid language models. + + The input projection produces ``B``, ``C``, and ``x`` branches. The + depthwise convolution consumes ``B * x``, then the second gate forms + ``C * conv(B * x)`` before the output projection. + """ + + def __init__(self, hidden_size: int, kernel_size: int, *, bias: bool = False): + super().__init__() + self.in_proj = Linear(hidden_size, 3 * hidden_size, bias=bias) + self.conv = _DepthwiseShortConv1d(hidden_size, kernel_size, bias=bias) + self.out_proj = Linear(hidden_size, hidden_size, bias=bias) + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + conv_state: ir.Value, + attention_mask: ir.Value | None = None, + ) -> tuple[ir.Value, ir.Value]: + """Apply the gated convolution and return output plus updated state. + + Args: + hidden_states: Input activations shaped ``(batch, seq, hidden)``. + conv_state: Previous full convolution window shaped + ``(batch, hidden, kernel_size)``. + attention_mask: Optional padding mask shaped + ``(batch, past_seq + seq)``. + """ + if attention_mask is not None: + # Recurrent layers only consume the mask for the current token span. + seq_len = op.Shape(hidden_states, start=1, end=2) + current_mask = op.Slice( + attention_mask, + op.Neg(seq_len), + op.Constant(value_ints=[INT64_MAX]), + op.Constant(value_ints=[1]), + ) + current_mask = op.Unsqueeze(current_mask, op.Constant(value_ints=[-1])) + hidden_states = op.Mul(hidden_states, op.CastLike(current_mask, hidden_states)) + + # (B, T, H) -> (B, 3H, T), split into the two gates and conv input. + projected = op.Transpose(self.in_proj(op, hidden_states), perm=[0, 2, 1]) + gate_b, gate_c, conv_input = op.Split(projected, num_outputs=3, axis=1, _outputs=3) + conv_input = op.Mul(gate_b, conv_input) # (B, H, T) + conv_output, present_state = self.conv(op, conv_input, conv_state) + output = op.Mul(gate_c, conv_output) # (B, H, T) + output = self.out_proj(op, op.Transpose(output, perm=[0, 2, 1])) + return output, present_state diff --git a/src/mobius/components/_short_conv_test.py b/src/mobius/components/_short_conv_test.py new file mode 100644 index 00000000..96802806 --- /dev/null +++ b/src/mobius/components/_short_conv_test.py @@ -0,0 +1,25 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +from mobius._testing import create_test_builder, create_test_input +from mobius.components import GatedShortConv + + +def test_gated_short_conv_builds_stateful_graph(): + component = GatedShortConv(hidden_size=16, kernel_size=3) + builder, op, graph = create_test_builder() + hidden_states = create_test_input(builder, "hidden_states", [1, 4, 16]) + conv_state = create_test_input(builder, "conv_state", [1, 16, 3]) + attention_mask = create_test_input(builder, "attention_mask", [1, 4]) + + output, present_state = component(op, hidden_states, conv_state, attention_mask) + builder._adapt_outputs([output, present_state], "") + + assert any(node.op_type == "Conv" for node in graph) + assert {name for name, _ in component.named_parameters()} == { + "in_proj.weight", + "conv.weight", + "out_proj.weight", + } diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 3b4786b5..c7be2f9b 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -83,6 +83,7 @@ "gemma4_unified_text": "gemma4_text", "mistral": "mistral", "mistral3": "mistral3", + "lfm2": "lfm2", # HunYuan-V1 dense / Hy-MT1.5 — generic decoder LLM type accepted by # ORT GenAI (see onnxruntime-genai/src/models/model_type.h LLM list). "hunyuan_v1_dense": "decoder", diff --git a/src/mobius/integrations/ort_genai/genai_config.py b/src/mobius/integrations/ort_genai/genai_config.py index 36c6761e..4c7a948c 100644 --- a/src/mobius/integrations/ort_genai/genai_config.py +++ b/src/mobius/integrations/ort_genai/genai_config.py @@ -169,6 +169,8 @@ def __init__( decoder_inputs: dict[str, str] | None = None, decoder_filename: str | None = None, supports_in_place_kv_cache: bool | None = None, + layer_types: list[str] | None = None, + conv_cache_size: int | None = None, ): self.model_type = model_type self.vocab_size = vocab_size @@ -193,6 +195,8 @@ def __init__( # to the EP capability flag, preserving existing behaviour for callers # that don't introspect the graph. self._supports_in_place_kv_cache = supports_in_place_kv_cache + self._layer_types = layer_types + self._conv_cache_size = conv_cache_size # Optional VLM fields (set via with_vision()) self._vision: dict[str, Any] | None = None @@ -253,6 +257,8 @@ def from_config( decoder_inputs=decoder_inputs, decoder_filename=decoder_filename, supports_in_place_kv_cache=supports_in_place_kv_cache, + layer_types=getattr(config, "layer_types", None), + conv_cache_size=getattr(config, "short_conv_kernel", None), ) def with_vision( @@ -423,6 +429,11 @@ def generate(self) -> dict[str, Any]: "num_hidden_layers": self.num_hidden_layers, "num_key_value_heads": self.num_key_value_heads, } + if self.model_type == "lfm2": + decoder["layer_types"] = self._layer_types or [] + decoder["conv_cache_size"] = self._conv_cache_size or 3 + decoder["inputs"]["past_conv_names"] = "past_key_values.%d.conv_state" + decoder["outputs"]["present_conv_names"] = "present.%d.conv_state" # Model section model: dict[str, Any] = { @@ -458,6 +469,10 @@ def generate(self) -> dict[str, Any]: context_length=self.context_length, supports_in_place_kv_cache=self._supports_in_place_kv_cache, ) + if self.model_type == "lfm2": + # ORT GenAI's LFM2 cache mixes fixed convolution windows with + # dynamic attention KV; shared in-place KV buffers are unsupported. + search["past_present_share_buffer"] = False search.update(self._search_overrides) return { diff --git a/src/mobius/integrations/ort_genai/genai_config_test.py b/src/mobius/integrations/ort_genai/genai_config_test.py index 8ee45d18..b4ae1fd9 100644 --- a/src/mobius/integrations/ort_genai/genai_config_test.py +++ b/src/mobius/integrations/ort_genai/genai_config_test.py @@ -79,6 +79,27 @@ def test_llm_decoder_outputs(self): assert outputs["present_key_names"] == "present.%d.key" assert outputs["present_value_names"] == "present.%d.value" + def test_lfm2_decoder_declares_hybrid_cache(self): + gen = GenaiConfigGenerator( + "lfm2", + vocab_size=65536, + hidden_size=1024, + num_hidden_layers=4, + num_attention_heads=16, + num_key_value_heads=8, + head_dim=64, + layer_types=["conv", "conv", "full_attention", "conv"], + conv_cache_size=3, + ) + + decoder = gen.generate()["model"]["decoder"] + + assert decoder["layer_types"] == ["conv", "conv", "full_attention", "conv"] + assert decoder["conv_cache_size"] == 3 + assert decoder["inputs"]["past_conv_names"] == "past_key_values.%d.conv_state" + assert decoder["outputs"]["present_conv_names"] == "present.%d.conv_state" + assert gen.generate()["search"]["past_present_share_buffer"] is False + def test_token_ids_included_when_set(self): """Token IDs are included in the model section.""" gen = GenaiConfigGenerator( diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 9ca9e2ce..b3d61257 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -75,6 +75,7 @@ "LLaDAModel", "LLaVAModel", "LayerNormCausalLMModel", + "Lfm2CausalLMModel", "LongcatFlashCausalLMModel", "MPTCausalLMModel", "Mamba2CausalLMModel", @@ -211,6 +212,7 @@ from mobius.models.internvl import InternVL2Model from mobius.models.jamba import JambaCausalLMModel from mobius.models.jetmoe import JetMoeCausalLMModel +from mobius.models.lfm2 import Lfm2CausalLMModel from mobius.models.llada import LLaDAModel from mobius.models.llama4 import Llama4CausalLMModel from mobius.models.llava import LLaVAModel diff --git a/src/mobius/models/lfm2.py b/src/mobius/models/lfm2.py new file mode 100644 index 00000000..cc7333d4 --- /dev/null +++ b/src/mobius/models/lfm2.py @@ -0,0 +1,188 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""LiquidAI LFM2 hybrid short-convolution and GQA causal language model.""" + +from __future__ import annotations + +import dataclasses + +import onnx_ir as ir +import torch +from onnxscript import OpBuilder, nn + +from mobius._configs import ArchitectureConfig +from mobius.components import ( + MLP, + Attention, + Embedding, + GatedShortConv, + RMSNorm, + create_padding_mask, + initialize_rope, +) +from mobius.models.base import CausalLMModel + + +class Lfm2RMSNorm(RMSNorm): + """LFM2 RMSNorm with fp32 variance accumulation, matching Transformers.""" + + def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: + hidden_states_f32 = op.Cast(hidden_states, to=ir.DataType.FLOAT) + variance = op.ReduceMean( + op.Mul(hidden_states_f32, hidden_states_f32), + [-1], + keepdims=1, + ) + normalized_f32 = op.Mul( + hidden_states_f32, + op.Reciprocal(op.Sqrt(op.Add(variance, self.variance_epsilon))), + ) + # Transformers casts the normalized activation back before applying gamma. + return op.Mul(op.CastLike(normalized_f32, hidden_states), self.weight) + + +class Lfm2DecoderLayer(nn.Module): + """LFM2 pre-norm decoder layer with either short convolution or full GQA.""" + + def __init__(self, config: ArchitectureConfig, layer_idx: int): + super().__init__() + layer_types = config.layer_types or [] + self.layer_type = ( + layer_types[layer_idx] if layer_idx < len(layer_types) else "full_attention" + ) + if self.layer_type == "conv": + self.conv = GatedShortConv( + config.hidden_size, + config.short_conv_kernel, + bias=config.short_conv_bias, + ) + else: + self.self_attn = Attention(config, rms_norm_class=Lfm2RMSNorm) + + self.feed_forward = MLP(config) + self.operator_norm = Lfm2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.ffn_norm = Lfm2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + attention_mask: ir.Value, + attention_bias: ir.Value, + position_embeddings: tuple[ir.Value, ir.Value], + past_key_value: tuple[ir.Value, ...], + ) -> tuple[ir.Value, tuple[ir.Value, ...]]: + residual = hidden_states + operator_input = self.operator_norm(op, hidden_states) + + if self.layer_type == "conv": + (conv_state,) = past_key_value + operator_output, present_state = self.conv( + op, + operator_input, + conv_state, + attention_mask, + ) + present_key_value = (present_state,) + else: + operator_output, present_key_value = self.self_attn( + op, + hidden_states=operator_input, + attention_bias=attention_bias, + position_embeddings=position_embeddings, + past_key_value=past_key_value, + ) + + # Both operators are pre-normalized and feed separate residual branches. + hidden_states = op.Add(residual, operator_output) # (B, T, H) + feed_forward = self.feed_forward(op, self.ffn_norm(op, hidden_states)) + hidden_states = op.Add(hidden_states, feed_forward) # (B, T, H) + return hidden_states, present_key_value + + +class Lfm2TextModel(nn.Module): + """LFM2 decoder backbone with mixed convolution and full-attention layers.""" + + def __init__(self, config: ArchitectureConfig): + super().__init__() + self.embed_tokens = Embedding( + config.vocab_size, + config.hidden_size, + config.pad_token_id, + ) + self.layers = nn.ModuleList( + [Lfm2DecoderLayer(config, i) for i in range(config.num_hidden_layers)] + ) + self.embedding_norm = Lfm2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = initialize_rope(config) + + def forward( + self, + op: OpBuilder, + input_ids: ir.Value, + attention_mask: ir.Value, + position_ids: ir.Value, + past_key_values: list[tuple[ir.Value, ...]] | None = None, + ) -> tuple[ir.Value, list[tuple[ir.Value, ...]]]: + hidden_states = self.embed_tokens(op, input_ids) # (B, T) -> (B, T, H) + position_embeddings = self.rotary_emb(op, position_ids) + # ONNX Attention applies causality internally; this mask carries padding only. + attention_bias = create_padding_mask( + op, + input_ids=input_ids, + attention_mask=attention_mask, + ) + + present_key_values = [] + past_kvs = past_key_values or [None] * len(self.layers) + for layer, past_kv in zip(self.layers, past_kvs): + hidden_states, present_kv = layer( + op, + hidden_states=hidden_states, + attention_mask=attention_mask, + attention_bias=attention_bias, + position_embeddings=position_embeddings, + past_key_value=past_kv, + ) + present_key_values.append(present_kv) + + hidden_states = self.embedding_norm(op, hidden_states) # (B, T, H) + return hidden_states, present_key_values + + +class Lfm2CausalLMModel(CausalLMModel): + """LiquidAI LFM2 causal LM with double-gated short convolutions and QK-norm GQA.""" + + default_task: str = "hybrid-text-generation" + category: str = "Hybrid Convolution+Attention" + + def __init__(self, config: ArchitectureConfig): + # LFM2 hardcodes per-head Q/K RMSNorm and SiLU-gated feed-forward blocks. + config = dataclasses.replace( + config, + attn_qk_norm=True, + hidden_act=config.hidden_act or "silu", + ) + super().__init__(config) + self.model = Lfm2TextModel(config) + if config.tie_word_embeddings: + self.lm_head.weight = self.model.embed_tokens.weight + + def preprocess_weights( + self, + state_dict: dict[str, torch.Tensor], + ) -> dict[str, torch.Tensor]: + """Map upstream LFM2 projection names to shared mobius components.""" + renamed: dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + new_key = ( + key.replace(".self_attn.out_proj.", ".self_attn.o_proj.") + .replace(".self_attn.q_layernorm.", ".self_attn.q_norm.") + .replace(".self_attn.k_layernorm.", ".self_attn.k_norm.") + .replace(".feed_forward.w1.", ".feed_forward.gate_proj.") + .replace(".feed_forward.w3.", ".feed_forward.up_proj.") + .replace(".feed_forward.w2.", ".feed_forward.down_proj.") + ) + renamed[new_key] = value + return super().preprocess_weights(renamed) diff --git a/src/mobius/tasks/_cache_utils.py b/src/mobius/tasks/_cache_utils.py index c69787e4..edc79a5c 100644 --- a/src/mobius/tasks/_cache_utils.py +++ b/src/mobius/tasks/_cache_utils.py @@ -279,12 +279,12 @@ def _make_hybrid_cache_inputs( pairs.append((conv_state, rec_state)) elif ltype == "conv": # ShortConv layers: conv_state only (no SSM state) - # State: (batch, hidden_size, short_conv_kernel - 1) + # LFM2/ORT GenAI state: full K-wide pre-convolution window. short_conv_kernel = getattr(config, "short_conv_kernel", 3) conv_state = builder.input( f"{prefix}.{i}.conv_state", dtype=dtype, - shape=[batch, config.hidden_size, short_conv_kernel - 1], + shape=[batch, config.hidden_size, short_conv_kernel], ) pairs.append((conv_state,)) # 1-tuple: conv has no second state elif ltype in ("mlp", "moe"): diff --git a/testdata/cases/causal-lm/lfm2_5-230m.yaml b/testdata/cases/causal-lm/lfm2_5-230m.yaml new file mode 100644 index 00000000..a6091b70 --- /dev/null +++ b/testdata/cases/causal-lm/lfm2_5-230m.yaml @@ -0,0 +1,17 @@ +model_id: "LiquidAI/LFM2.5-230M" +model_type: "lfm2" +revision: "13a53837c4906b4f7405932532ba85d182bb013b" +task_type: "text-generation" +dtype: "float32" + +inputs: + prompts: + - "Here is my poem:" + +level: "L4+L5" + +generation: + max_new_tokens: 20 + do_sample: false + +notes: "LFM2.5 230M. Hybrid double-gated short-convolution and QK-norm GQA architecture." diff --git a/testdata/golden/causal-lm/lfm2_5-230m.json b/testdata/golden/causal-lm/lfm2_5-230m.json new file mode 100644 index 00000000..774640ca --- /dev/null +++ b/testdata/golden/causal-lm/lfm2_5-230m.json @@ -0,0 +1,42 @@ +{ + "top1_id": 509, + "top2_id": 730, + "top10_ids": [ + 509, + 730, + 997, + 767, + 508, + 859, + 941, + 3604, + 835, + 1371 + ], + "top10_logits": [ + "0x1.3570240000000p+4", + "0x1.ea47bc0000000p+3", + "0x1.e976380000000p+3", + "0x1.d3e9580000000p+3", + "0x1.b0d2200000000p+3", + "0x1.afc7180000000p+3", + "0x1.a2a1f20000000p+3", + "0x1.a0470c0000000p+3", + "0x1.995ed40000000p+3", + "0x1.8a34020000000p+3" + ], + "logits_summary": [ + "0x1.3570240000000p+4", + "-0x1.87f8da0000000p+3", + "-0x1.9756a73e10300p+0", + "0x1.87fe4869c1794p+1" + ], + "input_ids": [ + 1, + 9151, + 856, + 1727, + 15543, + 535 + ] +} diff --git a/testdata/golden/causal-lm/lfm2_5-230m_generation.json b/testdata/golden/causal-lm/lfm2_5-230m_generation.json new file mode 100644 index 00000000..e9c3c0ba --- /dev/null +++ b/testdata/golden/causal-lm/lfm2_5-230m_generation.json @@ -0,0 +1,27 @@ +{ + "model_id": "LiquidAI/LFM2.5-230M", + "prompt": "Here is my poem:", + "generated_tokens": [ + 509, + 1098, + 4979, + 27870, + 521, + 768, + 20480, + 14559, + 521, + 3604, + 41429, + 2793, + 53306, + 884, + 779, + 2031, + 521, + 3604, + 542, + 2829 + ], + "generated_text": "\n\nThe sun rises, a golden orb, \nIts light dances on the world, \nA sym" +} diff --git a/tests/_test_configs.py b/tests/_test_configs.py index 19ea1668..800884f8 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -107,6 +107,16 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: CAUSAL_LM_CONFIGS: list[tuple[str, dict, bool]] = [ # === Text Generation (Llama-compatible) === ("llama", {}, True), + ( + "lfm2", + { + "layer_types": ["conv", "full_attention"], + "attn_qk_norm": True, + "short_conv_kernel": 3, + "short_conv_bias": False, + }, + True, + ), ("mistral", {}, False), ("qwen2", {}, True), ("cohere", {"tie_word_embeddings": True, "logit_scale": 0.0625}, True), diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index 8c986cdc..bb7a6e4c 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -283,6 +283,10 @@ def test_graph_builds_without_weights(self, model_type: str, config_overrides: d assert f"present.{i}.ssm_state" in output_names, ( f"Missing present.{i}.ssm_state" ) + elif ltype == "conv": + assert f"present.{i}.conv_state" in output_names, ( + f"Missing present.{i}.conv_state" + ) else: assert f"present.{i}.key" in output_names, f"Missing present.{i}.key" assert f"present.{i}.value" in output_names, f"Missing present.{i}.value" diff --git a/tests/integration_test.py b/tests/integration_test.py index 99218321..24ed48d2 100644 --- a/tests/integration_test.py +++ b/tests/integration_test.py @@ -85,6 +85,7 @@ def _model_accessible(model_id: str) -> bool: _TEXT_MODELS = [ # CausalLMModel (base: llama/mistral/qwen2 architecture) pytest.param("Qwen/Qwen2.5-0.5B", False, id="qwen2.5-0.5b"), + pytest.param("LiquidAI/LFM2.5-230M", False, id="lfm2.5-230m"), pytest.param("HuggingFaceTB/SmolLM-135M", False, id="smollm-135m"), # SmolLM3 (per-layer RoPE gating via no_rope_layers) pytest.param("HuggingFaceTB/SmolLM3-3B", False, id="smollm3-3b"), @@ -194,6 +195,14 @@ def _make_prefill_feeds(config, input_ids, attention_mask, position_ids): "position_ids": position_ids, } for i in range(config.num_hidden_layers): + layer_types = config.layer_types or [] + layer_type = layer_types[i] if i < len(layer_types) else "full_attention" + if layer_type == "conv": + feeds[f"past_key_values.{i}.conv_state"] = np.zeros( + (1, config.hidden_size, config.short_conv_kernel), + dtype=np.float32, + ) + continue feeds[f"past_key_values.{i}.key"] = np.zeros( (1, config.num_key_value_heads, 0, config.head_dim), dtype=np.float32, @@ -215,6 +224,13 @@ def _make_decode_feeds( "position_ids": decode_position_ids, } for i in range(config.num_hidden_layers): + layer_types = config.layer_types or [] + layer_type = layer_types[i] if i < len(layer_types) else "full_attention" + if layer_type == "conv": + feeds[f"past_key_values.{i}.conv_state"] = onnx_prefill_out[ + f"present.{i}.conv_state" + ] + continue feeds[f"past_key_values.{i}.key"] = onnx_prefill_out[f"present.{i}.key"] feeds[f"past_key_values.{i}.value"] = onnx_prefill_out[f"present.{i}.value"] return feeds diff --git a/tests/ort_genai_test.py b/tests/ort_genai_test.py index 162c9853..74fa97bd 100644 --- a/tests/ort_genai_test.py +++ b/tests/ort_genai_test.py @@ -219,6 +219,68 @@ def _write_processor_config(processor, output_dir: str) -> None: ] +@pytest.mark.integration +@pytest.mark.integration_slow +def test_lfm2_text_generation(tmp_path): + """Build LFM2.5, load it in ORT GenAI, and run deterministic generation.""" + from mobius import build + from mobius.integrations.ort_genai.auto_export import write_ort_genai_config + + model_id = "LiquidAI/LFM2.5-230M" + pkg = build(model_id, dtype="f32", load_weights=True) + output_dir = str(tmp_path / "lfm2") + pkg.save(output_dir) + write_ort_genai_config(pkg, output_dir, hf_model_id=model_id) + + with open(os.path.join(output_dir, "genai_config.json"), encoding="utf-8") as f: + config = json.load(f) + decoder = config["model"]["decoder"] + assert decoder["layer_types"] == pkg.config.layer_types + assert decoder["inputs"]["past_conv_names"] == "past_key_values.%d.conv_state" + + model = ort_genai.Model(output_dir) + tokenizer = ort_genai.Tokenizer(model) + input_ids = [pkg.config.bos_token_id, *tokenizer.encode("Here is my poem:")] + params = ort_genai.GeneratorParams(model) + params.set_search_options( + max_length=len(input_ids) + 20, + do_sample=False, + ) + generator = ort_genai.Generator(model, params) + generator.append_tokens(input_ids) + + generated_tokens = list(input_ids) + for _ in range(20): + if generator.is_done(): + break + generator.generate_next_token() + generated_tokens.append(generator.get_next_tokens()[0]) + + assert generated_tokens[len(input_ids) :] == [ + 509, + 1098, + 4979, + 27870, + 521, + 768, + 20480, + 14559, + 521, + 3604, + 41429, + 2793, + 53306, + 884, + 779, + 2031, + 521, + 3604, + 542, + 2829, + ] + assert tokenizer.decode(generated_tokens).strip() + + @pytest.mark.integration @pytest.mark.integration_slow @pytest.mark.parametrize("model_id", _MODELS) diff --git a/tests/synthetic_parity_test.py b/tests/synthetic_parity_test.py index 306406ed..d5e928f9 100644 --- a/tests/synthetic_parity_test.py +++ b/tests/synthetic_parity_test.py @@ -510,6 +510,16 @@ def _create_hf_config(model_type: str, config_overrides: dict): i for i, lt in enumerate(layer_types) if lt in ("full_attention", "attention") ] + if hf_model_type == "lfm2": + hf_kwargs["conv_L_cache"] = hf_kwargs.pop("short_conv_kernel", 3) + hf_kwargs["conv_bias"] = hf_kwargs.pop("short_conv_bias", False) + hf_kwargs["block_auto_adjust_ff_dim"] = False + hf_kwargs["norm_eps"] = hf_kwargs.pop("rms_norm_eps") + hf_kwargs["rope_parameters"] = { + "rope_type": "default", + "rope_theta": 10_000.0, + } + # Jamba uses attn_layer_offset/attn_layer_period if hf_model_type in ("jamba",) and "layer_types" in hf_kwargs: layer_types = hf_kwargs.pop("layer_types") @@ -648,6 +658,8 @@ def _create_hf_config(model_type: str, config_overrides: dict): "attn_qk_norm", "attn_qk_norm_full", "post_feedforward_norm", + "short_conv_kernel", + "short_conv_bias", # dual_ln is a mobius-only flag for Falcon/Bloom parallel attention; # HF controls this behavior via new_decoder_architecture=True. "dual_ln",