Add LiquidAI LFM2.5 hybrid model support

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",

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",
 }
 

src/mobius/_testing/torch_reference.py

@@ -288,19 +288,21 @@ 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:
         model: HuggingFace model in eval mode.
         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]

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,
 )

src/mobius/components/_short_conv.py

@@ -0,0 +1,134 @@
+# 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 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]:
+        # Keep a full K-wide state, matching ORT GenAI's LFM2 cache contract.
+        # Concatenating it with T current values yields T+1 valid Conv outputs;
+        # the first is all-past, so only the final T outputs are propagated.
+        conv_input = op.Concat(conv_state, 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,
+            )
+
+        seq_len = op.Shape(hidden_states, start=2, end=3)
+        conv_len = op.Shape(conv_output, start=2, end=3)
+        output_start = op.Sub(conv_len, seq_len)
+        output = op.Slice(
+            conv_output,
+            output_start,
+            conv_len,
+            op.Constant(value_ints=[2]),
+        )
+
+        input_len = op.Shape(conv_input, start=2, end=3)
+        state_start = op.Sub(input_len, op.Constant(value_ints=[self._kernel_size]))
+        present_state = op.Slice(
+            conv_input,
+            state_start,
+            input_len,
+            op.Constant(value_ints=[2]),
+        )
+        return 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)
+        self._hidden_size = hidden_size
+
+    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)
+            mask_len = op.Shape(attention_mask, start=1, end=2)
+            mask_start = op.Sub(mask_len, seq_len)
+            current_mask = op.Slice(
+                attention_mask,
+                mask_start,
+                mask_len,
+                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,
+            op.Constant(
+                value_ints=[self._hidden_size, self._hidden_size, self._hidden_size]
+            ),
+            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

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",
+    }

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",

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 {